Skip to content

Commit 1610206

Browse files
Backlog/v12 soar agent disconnection (#2515)
* i18n[frontend](soar): add WAITING, EXECUTING, DEAD execution statuses * fix[backend,agent-manager](soar): stop SOAR runs ghosting the agent stream
1 parent 4eb16d3 commit 1610206

3 files changed

Lines changed: 77 additions & 14 deletions

File tree

agent-manager/agent/agent_imp.go

Lines changed: 22 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,17 @@ func (s *AgentService) GetAgentAuth(ctx context.Context, req *ConnectorAuthReque
296296
return &ConnectorAuthResponse{Key: agent.AgentKey, TenantId: agent.TenantID}, nil
297297
}
298298

299+
// evictIfOwner deletes the AgentStreamMap entry for agentID only if it still
300+
// points to stream. Prevents a slow-exiting prior AgentStream goroutine from
301+
// clobbering the fresh entry a newly-reconnected agent installed.
302+
func (s *AgentService) evictIfOwner(agentID uint, stream AgentService_AgentStreamServer) {
303+
s.AgentStreamMutex.Lock()
304+
if s.AgentStreamMap[agentID] == stream {
305+
delete(s.AgentStreamMap, agentID)
306+
}
307+
s.AgentStreamMutex.Unlock()
308+
}
309+
299310
func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error {
300311
id, _, _, err := utils.GetItemsFromContext(stream.Context())
301312
if err != nil {
@@ -307,11 +318,12 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error
307318
}
308319
idUint := uint(idInt)
309320

321+
// Replace any prior entry rather than rejecting the reconnect. A dead
322+
// prior stream's goroutine may still be looping on Recv (see
323+
// utils.WaitForReconnect) and would otherwise block the agent from
324+
// re-registering for minutes. evictIfOwner guards the map so the old
325+
// goroutine's eventual delete does not clobber the fresh entry.
310326
s.AgentStreamMutex.Lock()
311-
if _, ok := s.AgentStreamMap[idUint]; ok {
312-
s.AgentStreamMutex.Unlock()
313-
return status.Error(codes.AlreadyExists, "stream already exists")
314-
}
315327
s.AgentStreamMap[idUint] = stream
316328
s.AgentStreamMutex.Unlock()
317329

@@ -324,18 +336,17 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error
324336
if err == io.EOF {
325337
err = utils.WaitForReconnect(stream.Context(), stream)
326338
if err != nil {
327-
s.AgentStreamMutex.Lock()
328-
delete(s.AgentStreamMap, idUint)
329-
s.AgentStreamMutex.Unlock()
330-
339+
catcher.Info("AgentStream: WaitForReconnect failed, evicting stream",
340+
map[string]any{"agent_id": idUint, "err": err.Error(), "process": "agent-manager"})
341+
s.evictIfOwner(idUint, stream)
331342
return status.Error(codes.Internal, fmt.Sprintf("failed to reconnect: %v", err))
332343
}
333344
continue
334345
}
335346
if err != nil {
336-
s.AgentStreamMutex.Lock()
337-
delete(s.AgentStreamMap, idUint)
338-
s.AgentStreamMutex.Unlock()
347+
catcher.Info("AgentStream: Recv errored, evicting stream",
348+
map[string]any{"agent_id": idUint, "err": err.Error(), "process": "agent-manager"})
349+
s.evictIfOwner(idUint, stream)
339350
return status.Error(codes.Internal, fmt.Sprintf("failed to receive message: %v", err))
340351
}
341352

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package agent
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"google.golang.org/grpc/metadata"
8+
)
9+
10+
// fakeAgentStream is the smallest thing that satisfies
11+
// AgentService_AgentStreamServer for identity-comparison tests.
12+
type fakeAgentStream struct{ id int }
13+
14+
func (fakeAgentStream) Send(*BidirectionalStream) error { return nil }
15+
func (fakeAgentStream) Recv() (*BidirectionalStream, error) { return nil, nil }
16+
func (fakeAgentStream) SetHeader(metadata.MD) error { return nil }
17+
func (fakeAgentStream) SendHeader(metadata.MD) error { return nil }
18+
func (fakeAgentStream) SetTrailer(metadata.MD) {}
19+
func (fakeAgentStream) Context() context.Context { return context.Background() }
20+
func (fakeAgentStream) SendMsg(any) error { return nil }
21+
func (fakeAgentStream) RecvMsg(any) error { return nil }
22+
23+
// TestEvictIfOwner_LeavesForeignStream: an old goroutine returning long after
24+
// a fresh reconnect must NOT clobber the fresh entry.
25+
// TestEvictIfOwner_RemovesOwnedStream: the current owner cleans up on exit.
26+
func TestEvictIfOwner(t *testing.T) {
27+
s := &AgentService{AgentStreamMap: map[uint]AgentService_AgentStreamServer{}}
28+
old := &fakeAgentStream{id: 1}
29+
fresh := &fakeAgentStream{id: 2}
30+
31+
s.AgentStreamMap[42] = fresh
32+
s.evictIfOwner(42, old)
33+
if _, ok := s.AgentStreamMap[42]; !ok {
34+
t.Fatal("evictIfOwner clobbered a fresh stream owned by a different goroutine")
35+
}
36+
if s.AgentStreamMap[42] != fresh {
37+
t.Fatal("evictIfOwner replaced the fresh entry with something else")
38+
}
39+
40+
s.evictIfOwner(42, fresh)
41+
if _, ok := s.AgentStreamMap[42]; ok {
42+
t.Fatal("evictIfOwner did not remove the owned entry")
43+
}
44+
}

backend/pkg/agentmanager/client.go

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,18 @@ func (c *AgentManagerClient) GetCollectorIntegrationState(ctx context.Context, c
352352
return resp, nil
353353
}
354354

355+
// ProcessCommand opens a bidi stream, sends one command, and reads one result.
356+
// It does NOT call CloseSend — parity with ProcessCommandStream / Java. Half-
357+
// closing the panel-side stream races the agent-manager's ProcessCommand
358+
// handler (agent-manager/agent/agent_imp.go) into an EOF path that has been
359+
// observed to leave AgentStreamMap[agentID] empty, after which every
360+
// subsequent panel call (SOAR + console) returns codes.NotFound "agent not
361+
// found or is disconnected". The ctx cancellation on function return is what
362+
// tears the stream down cleanly.
355363
func (c *AgentManagerClient) ProcessCommand(ctx context.Context, cmd *agent.UtmCommand) (*agent.CommandResult, error) {
364+
ctx, cancel := context.WithCancel(ctx)
365+
defer cancel()
366+
356367
stream, err := c.panelService.ProcessCommand(ctx)
357368
if err != nil {
358369
return nil, fmt.Errorf("agentmanager: ProcessCommand open stream: %w", err)
@@ -361,9 +372,6 @@ func (c *AgentManagerClient) ProcessCommand(ctx context.Context, cmd *agent.UtmC
361372
if err := stream.Send(cmd); err != nil {
362373
return nil, fmt.Errorf("agentmanager: ProcessCommand send: %w", err)
363374
}
364-
if err := stream.CloseSend(); err != nil {
365-
return nil, fmt.Errorf("agentmanager: ProcessCommand close send: %w", err)
366-
}
367375

368376
result, err := stream.Recv()
369377
if err != nil && err != io.EOF {

0 commit comments

Comments
 (0)