diff --git a/agent-manager/agent/agent_imp.go b/agent-manager/agent/agent_imp.go index b1fd9f25d..2c54ca89b 100644 --- a/agent-manager/agent/agent_imp.go +++ b/agent-manager/agent/agent_imp.go @@ -296,6 +296,17 @@ func (s *AgentService) GetAgentAuth(ctx context.Context, req *ConnectorAuthReque return &ConnectorAuthResponse{Key: agent.AgentKey, TenantId: agent.TenantID}, nil } +// evictIfOwner deletes the AgentStreamMap entry for agentID only if it still +// points to stream. Prevents a slow-exiting prior AgentStream goroutine from +// clobbering the fresh entry a newly-reconnected agent installed. +func (s *AgentService) evictIfOwner(agentID uint, stream AgentService_AgentStreamServer) { + s.AgentStreamMutex.Lock() + if s.AgentStreamMap[agentID] == stream { + delete(s.AgentStreamMap, agentID) + } + s.AgentStreamMutex.Unlock() +} + func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error { id, _, _, err := utils.GetItemsFromContext(stream.Context()) if err != nil { @@ -307,11 +318,12 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error } idUint := uint(idInt) + // Replace any prior entry rather than rejecting the reconnect. A dead + // prior stream's goroutine may still be looping on Recv (see + // utils.WaitForReconnect) and would otherwise block the agent from + // re-registering for minutes. evictIfOwner guards the map so the old + // goroutine's eventual delete does not clobber the fresh entry. s.AgentStreamMutex.Lock() - if _, ok := s.AgentStreamMap[idUint]; ok { - s.AgentStreamMutex.Unlock() - return status.Error(codes.AlreadyExists, "stream already exists") - } s.AgentStreamMap[idUint] = stream s.AgentStreamMutex.Unlock() @@ -324,18 +336,17 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error if err == io.EOF { err = utils.WaitForReconnect(stream.Context(), stream) if err != nil { - s.AgentStreamMutex.Lock() - delete(s.AgentStreamMap, idUint) - s.AgentStreamMutex.Unlock() - + catcher.Info("AgentStream: WaitForReconnect failed, evicting stream", + map[string]any{"agent_id": idUint, "err": err.Error(), "process": "agent-manager"}) + s.evictIfOwner(idUint, stream) return status.Error(codes.Internal, fmt.Sprintf("failed to reconnect: %v", err)) } continue } if err != nil { - s.AgentStreamMutex.Lock() - delete(s.AgentStreamMap, idUint) - s.AgentStreamMutex.Unlock() + catcher.Info("AgentStream: Recv errored, evicting stream", + map[string]any{"agent_id": idUint, "err": err.Error(), "process": "agent-manager"}) + s.evictIfOwner(idUint, stream) return status.Error(codes.Internal, fmt.Sprintf("failed to receive message: %v", err)) } diff --git a/agent-manager/agent/agent_imp_test.go b/agent-manager/agent/agent_imp_test.go new file mode 100644 index 000000000..b63adc5d3 --- /dev/null +++ b/agent-manager/agent/agent_imp_test.go @@ -0,0 +1,44 @@ +package agent + +import ( + "context" + "testing" + + "google.golang.org/grpc/metadata" +) + +// fakeAgentStream is the smallest thing that satisfies +// AgentService_AgentStreamServer for identity-comparison tests. +type fakeAgentStream struct{ id int } + +func (fakeAgentStream) Send(*BidirectionalStream) error { return nil } +func (fakeAgentStream) Recv() (*BidirectionalStream, error) { return nil, nil } +func (fakeAgentStream) SetHeader(metadata.MD) error { return nil } +func (fakeAgentStream) SendHeader(metadata.MD) error { return nil } +func (fakeAgentStream) SetTrailer(metadata.MD) {} +func (fakeAgentStream) Context() context.Context { return context.Background() } +func (fakeAgentStream) SendMsg(any) error { return nil } +func (fakeAgentStream) RecvMsg(any) error { return nil } + +// TestEvictIfOwner_LeavesForeignStream: an old goroutine returning long after +// a fresh reconnect must NOT clobber the fresh entry. +// TestEvictIfOwner_RemovesOwnedStream: the current owner cleans up on exit. +func TestEvictIfOwner(t *testing.T) { + s := &AgentService{AgentStreamMap: map[uint]AgentService_AgentStreamServer{}} + old := &fakeAgentStream{id: 1} + fresh := &fakeAgentStream{id: 2} + + s.AgentStreamMap[42] = fresh + s.evictIfOwner(42, old) + if _, ok := s.AgentStreamMap[42]; !ok { + t.Fatal("evictIfOwner clobbered a fresh stream owned by a different goroutine") + } + if s.AgentStreamMap[42] != fresh { + t.Fatal("evictIfOwner replaced the fresh entry with something else") + } + + s.evictIfOwner(42, fresh) + if _, ok := s.AgentStreamMap[42]; ok { + t.Fatal("evictIfOwner did not remove the owned entry") + } +} diff --git a/backend/pkg/agentmanager/client.go b/backend/pkg/agentmanager/client.go index 7e3f9a69e..a0336c6a9 100644 --- a/backend/pkg/agentmanager/client.go +++ b/backend/pkg/agentmanager/client.go @@ -352,7 +352,18 @@ func (c *AgentManagerClient) GetCollectorIntegrationState(ctx context.Context, c return resp, nil } +// ProcessCommand opens a bidi stream, sends one command, and reads one result. +// It does NOT call CloseSend — parity with ProcessCommandStream / Java. Half- +// closing the panel-side stream races the agent-manager's ProcessCommand +// handler (agent-manager/agent/agent_imp.go) into an EOF path that has been +// observed to leave AgentStreamMap[agentID] empty, after which every +// subsequent panel call (SOAR + console) returns codes.NotFound "agent not +// found or is disconnected". The ctx cancellation on function return is what +// tears the stream down cleanly. func (c *AgentManagerClient) ProcessCommand(ctx context.Context, cmd *agent.UtmCommand) (*agent.CommandResult, error) { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + stream, err := c.panelService.ProcessCommand(ctx) if err != nil { return nil, fmt.Errorf("agentmanager: ProcessCommand open stream: %w", err) @@ -361,9 +372,6 @@ func (c *AgentManagerClient) ProcessCommand(ctx context.Context, cmd *agent.UtmC if err := stream.Send(cmd); err != nil { return nil, fmt.Errorf("agentmanager: ProcessCommand send: %w", err) } - if err := stream.CloseSend(); err != nil { - return nil, fmt.Errorf("agentmanager: ProcessCommand close send: %w", err) - } result, err := stream.Recv() if err != nil && err != io.EOF {