diff --git a/client/internal/receivedprocessor.go b/client/internal/receivedprocessor.go index 942505a0..3c917251 100644 --- a/client/internal/receivedprocessor.go +++ b/client/internal/receivedprocessor.go @@ -51,10 +51,17 @@ func newReceivedProcessor( // This function will call any relevant callbacks. func (r *receivedProcessor) ProcessReceivedMessage(ctx context.Context, msg *protobufs.ServerToAgent) { if r.callbacks != nil { + // Note that anytime we add a new command capabilities we need to add a check here. + // This is because we want to ignore commands that the agent does not have the capability + // to process. if msg.Command != nil { - r.rcvCommand(msg.Command) - // If a command message exists, other messages will be ignored - return + if r.hasCapability(protobufs.AgentCapabilities_AgentCapabilities_AcceptsRestartCommand) { + r.rcvCommand(msg.Command) + // If a command message exists, other messages will be ignored + return + } else { + r.logger.Debugf("Ignoring Command, agent does not have AcceptsCommands capability") + } } scheduled, err := r.rcvFlags(ctx, protobufs.ServerToAgentFlags(msg.Flags)) diff --git a/client/internal/wsreceiver_test.go b/client/internal/wsreceiver_test.go index 44cdfdbc..c9ac5ca7 100644 --- a/client/internal/wsreceiver_test.go +++ b/client/internal/wsreceiver_test.go @@ -77,7 +77,8 @@ func TestServerToAgentCommand(t *testing.T) { remoteConfigStatus: &protobufs.RemoteConfigStatus{}, } sender := WSSender{} - receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, &sender, &clientSyncedState, nil, 0) + capabilities := protobufs.AgentCapabilities_AgentCapabilities_AcceptsRestartCommand + receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, &sender, &clientSyncedState, nil, capabilities) receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{ Command: test.command, }) @@ -100,7 +101,8 @@ func TestServerToAgentCommandExclusive(t *testing.T) { }, } clientSyncedState := ClientSyncedState{} - receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, nil, &clientSyncedState, nil, 0) + capabilities := protobufs.AgentCapabilities_AgentCapabilities_AcceptsRestartCommand + receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, nil, &clientSyncedState, nil, capabilities) receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{ Command: &protobufs.ServerToAgentCommand{ Type: protobufs.CommandType_CommandType_Restart,