diff --git a/packages/pam/local/base-proxy.go b/packages/pam/local/base-proxy.go index 9f90d8e1..4195c030 100644 --- a/packages/pam/local/base-proxy.go +++ b/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: - } } } diff --git a/packages/pam/local/database-proxy.go b/packages/pam/local/database-proxy.go index 34bd7e7d..76361258 100644 --- a/packages/pam/local/database-proxy.go +++ b/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()) } diff --git a/packages/pam/local/kubernetes-proxy.go b/packages/pam/local/kubernetes-proxy.go index 7f942080..4c8680ee 100644 --- a/packages/pam/local/kubernetes-proxy.go +++ b/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()) } diff --git a/packages/pam/local/rdp-proxy.go b/packages/pam/local/rdp-proxy.go index ce4b1ef1..47db28fe 100644 --- a/packages/pam/local/rdp-proxy.go +++ b/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()) } diff --git a/packages/pam/local/redis-proxy.go b/packages/pam/local/redis-proxy.go index 0d9d9075..43168bd3 100644 --- a/packages/pam/local/redis-proxy.go +++ b/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()) } diff --git a/packages/pam/local/ssh-proxy.go b/packages/pam/local/ssh-proxy.go index 93263e8c..de36795f 100644 --- a/packages/pam/local/ssh-proxy.go +++ b/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()) }