fix(agent): prevent possible deadlock when stopping SSH server (#2280)

agent/connection_manager.go

@@ -155,7 +155,9 @@ func (c *ConnectionManager) handleEvent(event ConnectionEvent) {
 	case WebSocketConnect:
 		c.handleStateChange(WebSocketConnected)
 	case SSHConnect:
-		c.handleStateChange(SSHConnected)
+		if c.State == Disconnected {
+			c.handleStateChange(SSHConnected)
+		}
 	case WebSocketDisconnect:
 		if c.State == WebSocketConnected {
 			c.handleStateChange(Disconnected)

agent/connection_manager_test.go

@@ -114,6 +114,12 @@ func TestConnectionManager_EventHandling(t *testing.T) {
 			event:         SSHConnect,
 			expectedState: SSHConnected,
 		},
+		{
+			name:          "SSH connect from WebSocket connected (no change)",
+			initialState:  WebSocketConnected,
+			event:         SSHConnect,
+			expectedState: WebSocketConnected,
+		},
 		{
 			name:          "WebSocket disconnect from connected",
 			initialState:  WebSocketConnected,

agent/server.go

@@ -265,6 +265,5 @@ func (a *Agent) StopServer() error {
 	slog.Info("Stopping SSH server")
 	_ = a.server.Close()
 	a.server = nil
-	a.connectionManager.eventChan <- SSHDisconnect
 	return nil
 }

agent/server_test.go

@@ -198,6 +198,28 @@ func TestStartServerDisableSSH(t *testing.T) {
 	assert.Contains(t, err.Error(), "SSH disabled")
 }
 
+func TestStopServerDoesNotBlockWhenEventQueueFull(t *testing.T) {
+	agent := createTestAgent(t)
+	agent.server = &ssh.Server{}
+	agent.connectionManager.eventChan = make(chan ConnectionEvent, 1)
+	agent.connectionManager.eventChan <- WebSocketConnect
+
+	done := make(chan error, 1)
+	go func() {
+		done <- agent.StopServer()
+	}()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("StopServer blocked on the connection event queue")
+	}
+
+	assert.Nil(t, agent.server)
+	assert.Equal(t, WebSocketConnect, <-agent.connectionManager.eventChan)
+}
+
 /////////////////////////////////////////////////////////////////
 //////////////////// ParseKeys Tests ////////////////////////////
 /////////////////////////////////////////////////////////////////