Backlog/v12 soar agent disconnection

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))
 		}
 

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")
+	}
+}

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 {