From 731c85506df0a8cadbf9c659e22ef751c29b05dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20S=C3=A1nchez?= Date: Thu, 3 Sep 2026 10:41:31 -0600 Subject: [PATCH 1/2] fix[agent-manager](grpc): prevent command system freeze when a panel disconnects mid-command --- agent-manager/agent/agent_imp.go | 243 +++++++++++++++----------- agent-manager/agent/agent_imp_test.go | 124 +++++++++++++ 2 files changed, 265 insertions(+), 102 deletions(-) diff --git a/agent-manager/agent/agent_imp.go b/agent-manager/agent/agent_imp.go index 2c54ca89b..6f57ed60b 100644 --- a/agent-manager/agent/agent_imp.go +++ b/agent-manager/agent/agent_imp.go @@ -264,8 +264,8 @@ func (s *AgentService) ListAgents(ctx context.Context, req *ListRequest) (*ListA if req.GetTenantId() != "" { filter = append(filter, utils.Filter{ Field: "tenant_id", - Op: utils.Is, - Value:sanitizeTenant(req.GetTenantId()), + Op: utils.Is, + Value: sanitizeTenant(req.GetTenantId()), }) } @@ -353,24 +353,49 @@ func (s *AgentService) AgentStream(stream AgentService_AgentStreamServer) error switch msg := in.StreamMessage.(type) { case *BidirectionalStream_Result: catcher.Info("Received command result from agent", map[string]any{"agent_id": msg.Result.AgentId, "result": msg.Result.Result, "process": "agent-manager"}) - cmdID := msg.Result.GetCmdId() - - s.CommandResultChannelM.Lock() - if resultChan, ok := s.CommandResultChannel[cmdID]; ok { - resultChan <- &CommandResult{ - AgentId: msg.Result.AgentId, - Result: msg.Result.Result, - CmdId: cmdID, - ExecutedAt: msg.Result.ExecutedAt, - } - } else if OnCommandResultHook == nil || !OnCommandResultHook(msg.Result) { - catcher.Error("failed to find result channel for CmdID", nil, map[string]any{"cmdID": cmdID, "process": "agent-manager"}) + if !s.tryDeliverResult(msg.Result) && + (OnCommandResultHook == nil || !OnCommandResultHook(msg.Result)) { + catcher.Error("failed to find result channel for CmdID", nil, map[string]any{"cmdID": msg.Result.GetCmdId(), "process": "agent-manager"}) } - s.CommandResultChannelM.Unlock() } } } +// tryDeliverResult hands an agent command result to the panel waiting on it. +// It must never block: the AgentStream goroutine is the single consumer of the +// agent's socket, and a stuck send here freezes every result for that agent +// (and the global CommandResultChannelM for the whole service). A missing +// slot, or a slot whose panel already went away (result unclaimed), returns +// false without blocking. +func (s *AgentService) tryDeliverResult(result *CommandResult) bool { + cmdID := result.GetCmdId() + s.CommandResultChannelM.Lock() + defer s.CommandResultChannelM.Unlock() + resultChan, ok := s.CommandResultChannel[cmdID] + if !ok { + return false + } + select { + case resultChan <- result: + return true + default: + // Slot is full: the panel that owned it left without claiming the + // result. Drop the stale slot so the stream keeps flowing. + delete(s.CommandResultChannel, cmdID) + catcher.Error("dropping agent result: panel disconnected before claiming it", nil, map[string]any{"cmdID": cmdID, "process": "agent-manager"}) + return false + } +} + +// reclaimResultSlot removes the panel command slot, releasing any AgentStream +// goroutine blocked on delivering a result for it. Called via defer on every +// exit path of handlePanelCommand. +func (s *AgentService) reclaimResultSlot(cmdID string) { + s.CommandResultChannelM.Lock() + delete(s.CommandResultChannel, cmdID) + s.CommandResultChannelM.Unlock() +} + func (s *AgentService) ProcessCommand(stream PanelService_ProcessCommandServer) error { for { cmd, err := stream.Recv() @@ -380,103 +405,117 @@ func (s *AgentService) ProcessCommand(stream PanelService_ProcessCommandServer) if err != nil { return status.Error(codes.Internal, fmt.Sprintf("failed to receive message: %v", err)) } - streamId, err := strconv.Atoi(cmd.AgentId) - if err != nil { - return status.Error(codes.InvalidArgument, "invalid agent ID") - } - agentStream, ok := s.AgentStreamMap[uint(streamId)] - if !ok { - return status.Errorf(codes.NotFound, "agent not found or is disconnected") + if err := s.handlePanelCommand(stream, cmd); err != nil { + return err } + } +} - target := &models.Agent{} - if dErr := s.DBConnection.GetFirst(target, "id = ?", streamId); dErr == nil && target.NoRemoteControl { - return status.Errorf(codes.PermissionDenied, - "agent %d was installed with remote control disabled; it can only be changed on the machine itself", streamId) - } - if cmd.GetOriginId() == "" { - return status.Errorf(codes.NotFound, "agent origin ID not provided") - } - if cmd.GetOriginType() == "" { - return status.Errorf(codes.NotFound, "agent origin TYPE not provided") - } - if cmd.GetReason() == "" { - return status.Errorf(codes.NotFound, "agent command reason not provided") - } +// handlePanelCommand forwards one panel command to the agent stream and waits +// for the result (max 5 minutes). The result slot is created before the send +// and reclaimed via defer on EVERY exit — success, timeout, and panel +// disconnect — so a result that arrives after the panel is gone can never +// block the AgentStream goroutine. +func (s *AgentService) handlePanelCommand(stream PanelService_ProcessCommandServer, cmd *UtmCommand) error { + streamId, err := strconv.Atoi(cmd.AgentId) + if err != nil { + return status.Error(codes.InvalidArgument, "invalid agent ID") + } + agentStream, ok := s.AgentStreamMap[uint(streamId)] + if !ok { + return status.Errorf(codes.NotFound, "agent not found or is disconnected") + } - cmdID := cmd.GetCmdId() - if cmdID == "" { - cmdID = uuid.New().String() - } + target := &models.Agent{} + if dErr := s.DBConnection.GetFirst(target, "id = ?", streamId); dErr == nil && target.NoRemoteControl { + return status.Errorf(codes.PermissionDenied, + "agent %d was installed with remote control disabled; it can only be changed on the machine itself", streamId) + } + if cmd.GetOriginId() == "" { + return status.Errorf(codes.NotFound, "agent origin ID not provided") + } + if cmd.GetOriginType() == "" { + return status.Errorf(codes.NotFound, "agent origin TYPE not provided") + } + if cmd.GetReason() == "" { + return status.Errorf(codes.NotFound, "agent command reason not provided") + } - s.CommandResultChannelM.Lock() - s.CommandResultChannel[cmdID] = make(chan *CommandResult) - s.CommandResultChannelM.Unlock() + cmdID := cmd.GetCmdId() + if cmdID == "" { + cmdID = uuid.New().String() + } - histCommand := createHistoryCommand(cmd, cmdID, uint(streamId)) - err = s.DBConnection.Create(&histCommand) - if err != nil { - catcher.Error("unable to create a new command history", err, map[string]any{"process": "agent-manager"}) - } + s.CommandResultChannelM.Lock() + // Buffered by 1: a result can always be absorbed even when the panel is no + // longer waiting, so delivery in AgentStream never blocks on it. + s.CommandResultChannel[cmdID] = make(chan *CommandResult, 1) + s.CommandResultChannelM.Unlock() + defer s.reclaimResultSlot(cmdID) - var lock sync.Locker - if LockStreamHook != nil { - lock = LockStreamHook(uint(streamId)) + histCommand := createHistoryCommand(cmd, cmdID, uint(streamId)) + if cErr := s.DBConnection.Create(&histCommand); cErr != nil { + catcher.Error("unable to create a new command history", cErr, map[string]any{"process": "agent-manager"}) + } + + var lock sync.Locker + if LockStreamHook != nil { + lock = LockStreamHook(uint(streamId)) + } + func() { + if lock != nil { + lock.Lock() + defer lock.Unlock() } - func() { - if lock != nil { - lock.Lock() - defer lock.Unlock() - } - err = agentStream.Send(&BidirectionalStream{ - StreamMessage: &BidirectionalStream_Command{ - Command: &UtmCommand{ - AgentId: cmd.AgentId, - Command: replaceSecretValues(cmd.Command), - CmdId: cmdID, - Shell: cmd.Shell, - }, + err = agentStream.Send(&BidirectionalStream{ + StreamMessage: &BidirectionalStream_Command{ + Command: &UtmCommand{ + AgentId: cmd.AgentId, + Command: replaceSecretValues(cmd.Command), + CmdId: cmdID, + Shell: cmd.Shell, }, - }) - }() - if err != nil { - return status.Errorf(codes.Internal, "failed to send command to agent: %v", err) - } - - select { - case result := <-s.CommandResultChannel[cmdID]: - err = s.DBConnection.Upsert( - &models.AgentCommand{}, - "agent_id = ? AND cmd_id = ?", - map[string]interface{}{"command_status": models.Executed, "result": result.Result}, - cmd.AgentId, cmdID, - ) - if err != nil { - catcher.Error("failed to update command status", err, map[string]any{"process": "agent-manager"}) - } - - err = stream.Send(result) - if err != nil { - return err - } - case <-time.After(5 * time.Minute): - s.CommandResultChannelM.Lock() - delete(s.CommandResultChannel, cmdID) - s.CommandResultChannelM.Unlock() - - _ = s.DBConnection.Upsert( - &models.AgentCommand{}, - "agent_id = ? AND cmd_id = ?", - map[string]interface{}{"command_status": models.Error, "result": "command timed out after 5 minutes"}, - cmd.AgentId, cmdID, - ) - - return status.Errorf(codes.DeadlineExceeded, "agent did not respond within 5 minutes") + }, + }) + }() + if err != nil { + return status.Errorf(codes.Internal, "failed to send command to agent: %v", err) + } + + select { + case result := <-s.CommandResultChannel[cmdID]: + if uErr := s.DBConnection.Upsert( + &models.AgentCommand{}, + "agent_id = ? AND cmd_id = ?", + map[string]interface{}{"command_status": models.Executed, "result": result.Result}, + cmd.AgentId, cmdID, + ); uErr != nil { + catcher.Error("failed to update command status", uErr, map[string]any{"process": "agent-manager"}) } - s.CommandResultChannelM.Lock() - delete(s.CommandResultChannel, cmdID) - s.CommandResultChannelM.Unlock() + return stream.Send(result) + case <-time.After(5 * time.Minute): + _ = s.DBConnection.Upsert( + &models.AgentCommand{}, + "agent_id = ? AND cmd_id = ?", + map[string]interface{}{"command_status": models.Error, "result": "command timed out after 5 minutes"}, + cmd.AgentId, cmdID, + ) + + return status.Errorf(codes.DeadlineExceeded, "agent did not respond within 5 minutes") + case <-stream.Context().Done(): + // The panel went away (viewer disconnected, backend request canceled). + // Mark the history row, then exit without waiting: the deferred slot + // reclamation keeps the AgentStream goroutine unblocked when the agent + // finally responds. + _ = s.DBConnection.Upsert( + &models.AgentCommand{}, + "agent_id = ? AND cmd_id = ?", + map[string]interface{}{"command_status": models.Error, "result": "panel disconnected before the agent responded"}, + cmd.AgentId, cmdID, + ) + + return stream.Context().Err() } } diff --git a/agent-manager/agent/agent_imp_test.go b/agent-manager/agent/agent_imp_test.go index b63adc5d3..80a4ca99b 100644 --- a/agent-manager/agent/agent_imp_test.go +++ b/agent-manager/agent/agent_imp_test.go @@ -3,6 +3,7 @@ package agent import ( "context" "testing" + "time" "google.golang.org/grpc/metadata" ) @@ -42,3 +43,126 @@ func TestEvictIfOwner(t *testing.T) { t.Fatal("evictIfOwner did not remove the owned entry") } } + +// runWithTimeout reports whether fn returns within d. A blocking delivery +// (the pre-fix deadlock) fails the test instead of hanging it for minutes. +func runWithTimeout(t *testing.T, d time.Duration, fn func()) bool { + t.Helper() + done := make(chan struct{}) + go func() { + defer close(done) + fn() + }() + select { + case <-done: + return true + case <-time.After(d): + t.Fatalf("operation blocked for %s — the AgentStream result delivery path deadlocks", d) + return false + } +} + +// TestTryDeliverResult_WaitingPanel: the normal path — a panel is blocked on +// the slot's channel and receives the result. The reader reports what it got +// over an unbuffered channel, so the delivery is what unblocks it. +func TestTryDeliverResult_WaitingPanel(t *testing.T) { + s := &AgentService{ + AgentStreamMap: map[uint]AgentService_AgentStreamServer{}, + CommandResultChannel: map[string]chan *CommandResult{}, + } + s.CommandResultChannel["cmd-1"] = make(chan *CommandResult, 1) + + received := make(chan *CommandResult) // unbuffered: sent only by the reader + go func() { + received <- <-s.CommandResultChannel["cmd-1"] + }() + // Let the reader block in its receive before delivering. + select { + case <-received: + t.Fatal("reader returned without any delivery") + case <-time.After(50 * time.Millisecond): + } + + if !s.tryDeliverResult(&CommandResult{AgentId: "7", CmdId: "cmd-1", Result: "ok"}) { + t.Fatal("tryDeliverResult returned false for a waiting panel") + } + + select { + case got := <-received: + if got == nil || got.Result != "ok" { + t.Fatalf("panel received %v, want the delivered result", got) + } + case <-time.After(time.Second): + t.Fatal("panel did not receive the delivered result") + } +} + +// TestTryDeliverResult_NoSlot: an unknown cmd_id must not block and must +// report undelivered. +func TestTryDeliverResult_NoSlot(t *testing.T) { + s := &AgentService{CommandResultChannel: map[string]chan *CommandResult{}} + if !runWithTimeout(t, time.Second, func() { + if s.tryDeliverResult(&CommandResult{CmdId: "missing", Result: "x"}) { + t.Error("tryDeliverResult returned true with no slot registered") + } + }) { + t.Fatal("delivery with no slot blocked") + } +} + +// TestTryDeliverResult_PanelGone is the regression test for the freeze: the +// panel disconnected without claiming its result. The first result is +// absorbed by the buffer; a second delivery must drop the stale slot and +// return immediately. Before the fix, any delivery into a slot with no +// reader blocked the AgentStream goroutine forever, holding the global +// CommandResultChannelM and the whole command system with it. +func TestTryDeliverResult_PanelGone(t *testing.T) { + s := &AgentService{CommandResultChannel: map[string]chan *CommandResult{}} + s.CommandResultChannel["cmd-2"] = make(chan *CommandResult, 1) + + // First result: absorbed into the buffer, nobody reads it. + if !runWithTimeout(t, time.Second, func() { + if !s.tryDeliverResult(&CommandResult{CmdId: "cmd-2", Result: "first"}) { + t.Error("first result should be absorbed by the buffer") + } + }) { + t.Fatal("first delivery blocked") + } + + // Second result: must not block, and must reclaim the orphaned slot. + if !runWithTimeout(t, 2*time.Second, func() { + if s.tryDeliverResult(&CommandResult{CmdId: "cmd-2", Result: "second"}) { + t.Error("second result delivered into a full stale slot") + } + }) { + t.Fatal("second delivery blocked") + } + + s.CommandResultChannelM.Lock() + _, still := s.CommandResultChannel["cmd-2"] + s.CommandResultChannelM.Unlock() + if still { + t.Fatal("orphaned slot was not reclaimed after a dropped delivery") + } +} + +// TestReclaimResultSlot: every handlePanelCommand exit path runs this via +// defer; it must remove the slot, and be safe to call twice. +func TestReclaimResultSlot(t *testing.T) { + s := &AgentService{CommandResultChannel: map[string]chan *CommandResult{}} + s.CommandResultChannel["cmd-3"] = make(chan *CommandResult, 1) + + s.reclaimResultSlot("cmd-3") + + s.CommandResultChannelM.Lock() + _, still := s.CommandResultChannel["cmd-3"] + s.CommandResultChannelM.Unlock() + if still { + t.Fatal("reclaimResultSlot did not remove the slot") + } + + // Reclaiming an unknown id is a harmless no-op (double-reclaim safety). + if !runWithTimeout(t, time.Second, func() { s.reclaimResultSlot("cmd-3") }) { + t.Fatal("reclaiming a missing slot blocked") + } +} From 8d33e36cfea429332788455e8bf13e816b610fc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alex=20S=C3=A1nchez?= Date: Thu, 3 Sep 2026 10:41:55 -0600 Subject: [PATCH 2/2] fix[installer](nginx): make agent-manager gRPC streams persistent in the front-end proxy --- installer/templates/front-end.go | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/installer/templates/front-end.go b/installer/templates/front-end.go index bdc415062..febd7d689 100644 --- a/installer/templates/front-end.go +++ b/installer/templates/front-end.go @@ -70,24 +70,33 @@ server { proxy_request_buffering off; } + # The agent's persistent gRPC streams to agent-manager. nginx's + # grpc_*_timeout of 0 means "no inactivity timeout" (removing the + # directive would silently fall back to a 60s default), so an idle + # AgentStream / PingService / CollectorService is never torn down by the + # proxy. Liveness is carried by the gRPC keepalive (30s ping / 10s + # timeout on both ends) plus grpc_socket_keepalive, which still reaps a + # genuinely dead peer. location /agent.AgentService/ { grpc_pass grpcs://$utmstack_agent_manager_grpc; - grpc_read_timeout 900; - grpc_send_timeout 900; + grpc_read_timeout 0; + grpc_send_timeout 0; client_body_timeout 1h; grpc_socket_keepalive on; } location /agent.PanelService/ { grpc_pass grpcs://$utmstack_agent_manager_grpc; - grpc_read_timeout 900; - grpc_send_timeout 900; + grpc_read_timeout 0; + grpc_send_timeout 0; + grpc_socket_keepalive on; } location /agent.CollectorService/ { grpc_pass grpcs://$utmstack_agent_manager_grpc; - grpc_read_timeout 900; - grpc_send_timeout 900; + grpc_read_timeout 0; + grpc_send_timeout 0; + grpc_socket_keepalive on; } # log-input's ingest, whose service lives in the SDK's "plugins" package. @@ -99,8 +108,9 @@ server { location /agent.PingService/ { grpc_pass grpcs://$utmstack_agent_manager_grpc; - grpc_read_timeout 900; - grpc_send_timeout 900; + grpc_read_timeout 0; + grpc_send_timeout 0; + grpc_socket_keepalive on; } client_max_body_size 200M;