fix: multi-connection proxy pam

packages/pam/local/base-proxy.go

@@ -347,46 +347,21 @@ func (b *BaseProxyServer) fallbackToAPITerminationWith(session LiveSession) {
 	}
 }
 
-// HandleGatewayDisconnect should be called when a gateway connection drops unexpectedly
-// (i.e., not initiated by the user via Ctrl+C). This happens when:
-//   - An administrator terminates the session from the Infisical UI
-//   - The session expires on the gateway side
-//   - The gateway or relay goes down
-//
-// It prints a message and triggers proxy shutdown so the CLI process exits
-// cleanly instead of hanging with a dead backend connection.
-func (b *BaseProxyServer) HandleGatewayDisconnect() {
-	b.shutdownOnce.Do(func() {
-		fmt.Println("\nConnection to session lost. Shutting down proxy...")
-		close(b.shutdownCh)
-		// Guarded rather than assumed. This is exported and BaseProxyServer is built in several
-		// places, so a constructor that forgets cancel should lose its shutdown signal, not panic on
-		// a gateway drop, which is the one moment this runs.
-		if b.cancel != nil {
-			b.cancel()
-		}
-	})
-}
-
-// NewDisconnectChannels creates the error channels used to distinguish gateway
-// disconnects from normal client disconnects.
+// NewDisconnectChannels creates the error channels a proxied connection uses to report which side
+// finished first.
 func (b *BaseProxyServer) NewDisconnectChannels() (gatewayErrCh, clientErrCh chan error) {
 	return make(chan error, 1), make(chan error, 1)
 }
 
-// WaitForDisconnect blocks until either the gateway or client side of a proxied
-// connection closes. If the gateway disconnects, the proxy shuts down.
-func (b *BaseProxyServer) WaitForDisconnect(gatewayErrCh, clientErrCh <-chan error, connCtx context.Context) {
+// WaitForConnectionClose blocks until either side of one proxied connection finishes, and ends
+// only that connection. A gateway-side close is not a session-level event: the gateway closes a
+// stream whenever the resource does, which is routine for pooled clients and for exchanges the
+// server answers by closing, such as a Postgres CancelRequest.
+func (b *BaseProxyServer) WaitForConnectionClose(gatewayErrCh, clientErrCh <-chan error, connCtx context.Context) {
 	select {
 	case <-gatewayErrCh:
-		b.HandleGatewayDisconnect()
 	case <-clientErrCh:
 	case <-connCtx.Done():
-		select {
-		case <-gatewayErrCh:
-			b.HandleGatewayDisconnect()
-		default:
-		}
 	}
 }
 

packages/pam/local/database-proxy.go

@@ -172,7 +172,7 @@ func (p *DatabaseProxyServer) handleConnection(clientConn net.Conn) {
 		clientErrCh <- err
 	}()
 
-	p.WaitForDisconnect(gatewayErrCh, clientErrCh, connCtx)
+	p.WaitForConnectionClose(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Info().Msgf("Connection closed for client: %s", clientConn.RemoteAddr().String())
 }

packages/pam/local/kubernetes-proxy.go

@@ -22,7 +22,6 @@ type KubernetesProxyServer struct {
 	kubeConfigOriginalContext string
 }
 
-
 func (p *KubernetesProxyServer) SetupKubeconfig(clusterName string) error {
 	configLoader := clientcmd.NewDefaultClientConfigLoadingRules()
 	config, err := configLoader.Load()
@@ -201,12 +200,7 @@ func (p *KubernetesProxyServer) handleConnection(clientConn net.Conn) {
 	connCtx, connCancel := context.WithCancel(p.ctx)
 	defer connCancel()
 
-	// For Kubernetes, each kubectl command opens a separate connection.
-	// Unlike persistent protocols (SSH, databases), the gateway closing after
-	// handling a request is normal — not a session-level disconnect.
-	// So we just wait for either side to finish and return, without triggering
-	// HandleGatewayDisconnect which would shut down the entire proxy.
-	done := make(chan struct{}, 2)
+	gatewayErrCh, clientErrCh := p.NewDisconnectChannels()
 
 	// Gateway → Client
 	go func() {
@@ -219,7 +213,7 @@ func (p *KubernetesProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Gateway to client copy ended")
 			}
 		}
-		done <- struct{}{}
+		gatewayErrCh <- err
 	}()
 
 	// Client → Gateway
@@ -233,15 +227,10 @@ func (p *KubernetesProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Client to gateway copy ended")
 			}
 		}
-		done <- struct{}{}
+		clientErrCh <- err
 	}()
 
-	// Wait for either side to finish — this is a per-connection close, not a session close
-	select {
-	case <-done:
-	case <-connCtx.Done():
-		log.Info().Msg("Connection cancelled by context")
-	}
+	p.WaitForConnectionClose(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Info().Msgf("Connection closed for client: %s", clientConn.RemoteAddr().String())
 }

packages/pam/local/rdp-proxy.go

@@ -174,7 +174,7 @@ func (p *RDPProxyServer) handleConnection(clientConn net.Conn) {
 		clientErrCh <- err
 	}()
 
-	p.WaitForDisconnect(gatewayErrCh, clientErrCh, connCtx)
+	p.WaitForConnectionClose(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Info().Msgf("RDP connection closed for client: %s", clientConn.RemoteAddr().String())
 }

packages/pam/local/redis-proxy.go

@@ -173,7 +173,7 @@ func (p *RedisProxyServer) handleConnection(clientConn net.Conn) {
 		clientErrCh <- err
 	}()
 
-	p.WaitForDisconnect(gatewayErrCh, clientErrCh, connCtx)
+	p.WaitForConnectionClose(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Info().Msgf("Connection closed for client: %s", clientConn.RemoteAddr().String())
 }

packages/pam/local/ssh-proxy.go

@@ -176,7 +176,7 @@ func (p *SSHProxyServer) handleConnection(clientConn net.Conn) {
 		gatewayErrCh <- err
 	}()
 
-	p.WaitForDisconnect(gatewayErrCh, clientErrCh, connCtx)
+	p.WaitForConnectionClose(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Debug().Msgf("SSH connection closed for client: %s", clientConn.RemoteAddr().String())
 }