feat(pam): support terminating active sessions (#167)

saifsmailbox98 · web-flow · commit b98e2a490245 · 2026-04-02T04:17:43.000+05:30
* feat: support killing active PAM sessions via gateway session registry

Add a session registry to the gateway so that ALPN cancellation signals
can find and close active proxy connections. When an admin terminates a
session from the UI, the gateway now kills the proxy immediately instead
of waiting for expiry.

- Add pamSessions registry (map + mutex) to Gateway struct
- Register/deregister sessions around HandlePAMProxy dispatch
- HandlePAMCancellation now calls cancelSession to close the proxy conn
- Per-session context so expiry timer exits cleanly on cancellation
- HandleGatewayDisconnect on BaseProxyServer to exit the CLI proxy
  cleanly when the backend connection drops

* fix: track multiple connections per session in gateway PAM registry

The registry previously stored a single entry per session ID, so each
new client connection (e.g., multiple psql windows) overwrote the
previous one. On termination, only the last connection was killed.

Change the registry to a slice per session ID so CancelPAMSession
closes all active connections for the session.

* fix: only trigger proxy shutdown on actual gateway errors

HandleGatewayDisconnect was called unconditionally when errCh received
a value, but io.Copy returns nil on clean EOF (normal client disconnect).
This caused the entire proxy to shut down when a user simply exited psql.

Only call HandleGatewayDisconnect when err != nil, so normal client
disconnects don't kill the proxy.

Also remove redundant shutdownCh select guard before shutdownOnce.Do
since sync.Once already provides the once-only guarantee.

* fix: distinguish gateway disconnect from normal client disconnect

Split the shared errCh into separate gatewayErrCh and clientErrCh so we
can detect which side closed the connection first. Only call
HandleGatewayDisconnect when the gateway side drops — not when the
client (e.g., psql) exits normally.

Previously, HandleGatewayDisconnect was called unconditionally on any
errCh value, which shut down the entire proxy when a user simply exited
their client. With the err != nil guard, it was never called because
both normal and admin-terminated closes produce nil from io.Copy.

Also remove redundant shutdownCh select guard in HandleGatewayDisconnect
since sync.Once already provides the once-only guarantee.

* refactor: extract disconnect channel logic into base proxy helpers
diff --git a/packages/gateway-v2/gateway.go b/packages/gateway-v2/gateway.go
@@ -83,6 +83,11 @@ type GatewayConfig struct {
 	ReconnectDelay time.Duration
 }
 
+type pamSessionEntry struct {
+	cancel context.CancelFunc
+	conn   *tls.Conn
+}
+
 type Gateway struct {
 	GatewayID string
 
@@ -110,6 +115,10 @@ type Gateway struct {
 	heartbeatStarted bool
 	heartbeatMu      sync.Mutex
 	notifyOnce       sync.Once
+
+	// PAM session registry for active proxy connections (multiple connections per session)
+	pamSessions   map[string][]*pamSessionEntry
+	pamSessionsMu sync.Mutex
 }
 
 // NewGateway creates a new gateway instance
@@ -137,9 +146,51 @@ func NewGateway(config *GatewayConfig) (*Gateway, error) {
 		cancel:                cancel,
 		pamCredentialsManager: pamCredentialsManager,
 		pamSessionUploader:    session.NewSessionUploader(httpClient, pamCredentialsManager),
+		pamSessions:           make(map[string][]*pamSessionEntry),
 	}, nil
 }
 
+// RegisterPAMSession registers an active PAM proxy connection for cancellation support
+func (g *Gateway) RegisterPAMSession(sessionID string, cancel context.CancelFunc, conn *tls.Conn) {
+	g.pamSessionsMu.Lock()
+	defer g.pamSessionsMu.Unlock()
+	g.pamSessions[sessionID] = append(g.pamSessions[sessionID], &pamSessionEntry{cancel: cancel, conn: conn})
+}
+
+// DeregisterPAMSession removes a specific connection from the session registry
+func (g *Gateway) DeregisterPAMSession(sessionID string, conn *tls.Conn) {
+	g.pamSessionsMu.Lock()
+	defer g.pamSessionsMu.Unlock()
+	entries := g.pamSessions[sessionID]
+	for i, e := range entries {
+		if e.conn == conn {
+			g.pamSessions[sessionID] = append(entries[:i], entries[i+1:]...)
+			break
+		}
+	}
+	if len(g.pamSessions[sessionID]) == 0 {
+		delete(g.pamSessions, sessionID)
+	}
+}
+
+// CancelPAMSession kills all active connections for a PAM session
+func (g *Gateway) CancelPAMSession(sessionID string) bool {
+	g.pamSessionsMu.Lock()
+	entries, ok := g.pamSessions[sessionID]
+	if ok {
+		delete(g.pamSessions, sessionID)
+	}
+	g.pamSessionsMu.Unlock()
+	if !ok {
+		return false
+	}
+	for _, e := range entries {
+		e.conn.Close()
+		e.cancel()
+	}
+	return true
+}
+
 func (g *Gateway) registerHeartBeat(ctx context.Context, errCh chan error) {
 	sendHeartbeat := func() error {
 		if err := api.CallGatewayHeartBeatV2(g.httpClient); err != nil {
@@ -608,7 +659,10 @@ func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) {
 		}
 		return
 	} else if forwardConfig.Mode == ForwardModePAM {
-		if err := pam.HandlePAMProxy(g.ctx, tlsConn, &forwardConfig.PAMConfig, g.httpClient); err != nil {
+		sessionCtx, sessionCancel := context.WithCancel(g.ctx)
+		g.RegisterPAMSession(forwardConfig.PAMConfig.SessionId, sessionCancel, tlsConn)
+		defer g.DeregisterPAMSession(forwardConfig.PAMConfig.SessionId, tlsConn)
+		if err := pam.HandlePAMProxy(sessionCtx, tlsConn, &forwardConfig.PAMConfig, g.httpClient); err != nil {
 			if err.Error() == "unexpected EOF" {
 				log.Debug().Err(err).Msg("PAM proxy handler ended with unexpected connection termination")
 			} else {
@@ -617,7 +671,7 @@ func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) {
 		}
 		return
 	} else if forwardConfig.Mode == ForwardModePAMCancellation {
-		if err := pam.HandlePAMCancellation(g.ctx, tlsConn, &forwardConfig.PAMConfig, g.httpClient); err != nil {
+		if err := pam.HandlePAMCancellation(g.ctx, tlsConn, &forwardConfig.PAMConfig, g.httpClient, g.CancelPAMSession); err != nil {
 			log.Error().Err(err).Msg("PAM cancellation proxy handler ended with error")
 		}
 		return
diff --git a/packages/pam/local/base-proxy.go b/packages/pam/local/base-proxy.go
@@ -262,6 +262,41 @@ func (b *BaseProxyServer) FallbackToAPITermination() {
 	}
 }
 
+// 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)
+		b.cancel()
+	})
+}
+
+// NewDisconnectChannels creates the error channels used to distinguish gateway
+// disconnects from normal client disconnects.
+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) {
+	select {
+	case <-gatewayErrCh:
+		b.HandleGatewayDisconnect()
+	case <-clientErrCh:
+		// Normal client disconnect, proxy stays running
+	case <-connCtx.Done():
+		log.Info().Msg("Connection cancelled by context")
+	}
+}
+
 // WaitForConnectionsWithTimeout waits for active connections to close with a timeout
 func (b *BaseProxyServer) WaitForConnectionsWithTimeout(timeout time.Duration) {
 	done := make(chan struct{})
diff --git a/packages/pam/local/database-proxy.go b/packages/pam/local/database-proxy.go
@@ -265,9 +265,9 @@ func (p *DatabaseProxyServer) handleConnection(clientConn net.Conn) {
 	connCtx, connCancel := context.WithCancel(p.ctx)
 	defer connCancel()
 
-	errCh := make(chan error, 2)
+	gatewayErrCh, clientErrCh := p.NewDisconnectChannels()
 
-	// Bidirectional data forwarding with context cancellation
+	// Gateway → Client: if this side closes first, the gateway dropped the connection
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(clientConn, gatewayConn)
@@ -278,9 +278,10 @@ func (p *DatabaseProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Gateway to client copy ended")
 			}
 		}
-		errCh <- err
+		gatewayErrCh <- err
 	}()
 
+	// Client → Gateway: if this side closes first, the client disconnected normally
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(gatewayConn, clientConn)
@@ -291,14 +292,10 @@ func (p *DatabaseProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Client to gateway copy ended")
 			}
 		}
-		errCh <- err
+		clientErrCh <- err
 	}()
 
-	select {
-	case <-errCh:
-	case <-connCtx.Done():
-		log.Info().Msg("Connection cancelled by context")
-	}
+	p.WaitForDisconnect(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
@@ -304,9 +304,9 @@ func (p *KubernetesProxyServer) handleConnection(clientConn net.Conn) {
 	connCtx, connCancel := context.WithCancel(p.ctx)
 	defer connCancel()
 
-	errCh := make(chan error, 2)
+	gatewayErrCh, clientErrCh := p.NewDisconnectChannels()
 
-	// Bidirectional data forwarding with context cancellation
+	// Gateway → Client: if this side closes first, the gateway dropped the connection
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(clientConn, gatewayConn)
@@ -317,9 +317,10 @@ func (p *KubernetesProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Gateway to client copy ended")
 			}
 		}
-		errCh <- err
+		gatewayErrCh <- err
 	}()
 
+	// Client → Gateway: if this side closes first, the client disconnected normally
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(gatewayConn, clientConn)
@@ -330,14 +331,10 @@ func (p *KubernetesProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Client to gateway copy ended")
 			}
 		}
-		errCh <- err
+		clientErrCh <- err
 	}()
 
-	select {
-	case <-errCh:
-	case <-connCtx.Done():
-		log.Info().Msg("Connection cancelled by context")
-	}
+	p.WaitForDisconnect(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Info().Msgf("Connection closed for client: %s", clientConn.RemoteAddr().String())
 }
diff --git a/packages/pam/local/redis-proxy.go b/packages/pam/local/redis-proxy.go
@@ -250,9 +250,9 @@ func (p *RedisProxyServer) handleConnection(clientConn net.Conn) {
 	connCtx, connCancel := context.WithCancel(p.ctx)
 	defer connCancel()
 
-	errCh := make(chan error, 2)
+	gatewayErrCh, clientErrCh := p.NewDisconnectChannels()
 
-	// Bidirectional data forwarding with context cancellation
+	// Gateway → Client: if this side closes first, the gateway dropped the connection
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(clientConn, gatewayConn)
@@ -263,9 +263,10 @@ func (p *RedisProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Gateway to client copy ended")
 			}
 		}
-		errCh <- err
+		gatewayErrCh <- err
 	}()
 
+	// Client → Gateway: if this side closes first, the client disconnected normally
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(gatewayConn, clientConn)
@@ -276,14 +277,10 @@ func (p *RedisProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Client to gateway copy ended")
 			}
 		}
-		errCh <- err
+		clientErrCh <- err
 	}()
 
-	select {
-	case <-errCh:
-	case <-connCtx.Done():
-		log.Info().Msg("Connection cancelled by context")
-	}
+	p.WaitForDisconnect(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
@@ -361,10 +361,9 @@ func (p *SSHProxyServer) handleConnection(clientConn net.Conn) {
 	connCtx, connCancel := context.WithCancel(p.ctx)
 	defer connCancel()
 
-	errCh := make(chan error, 2)
+	gatewayErrCh, clientErrCh := p.NewDisconnectChannels()
 
-	// Bidirectional data forwarding with context cancellation
-	// Client (local SSH) → Gateway (SSH proxy)
+	// Client (local SSH) → Gateway (SSH proxy): if this side closes first, the client disconnected normally
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(gatewayConn, clientConn)
@@ -375,10 +374,10 @@ func (p *SSHProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Client to gateway copy ended")
 			}
 		}
-		errCh <- err
+		clientErrCh <- err
 	}()
 
-	// Gateway (SSH proxy) → Client (local SSH)
+	// Gateway (SSH proxy) → Client (local SSH): if this side closes first, the gateway dropped the connection
 	go func() {
 		defer connCancel()
 		_, err := io.Copy(clientConn, gatewayConn)
@@ -389,14 +388,10 @@ func (p *SSHProxyServer) handleConnection(clientConn net.Conn) {
 				log.Debug().Err(err).Msg("Gateway to client copy ended")
 			}
 		}
-		errCh <- err
+		gatewayErrCh <- err
 	}()
 
-	select {
-	case <-errCh:
-	case <-connCtx.Done():
-		log.Debug().Msg("Connection cancelled by context")
-	}
+	p.WaitForDisconnect(gatewayErrCh, clientErrCh, connCtx)
 
 	log.Debug().Msgf("SSH connection closed for client: %s", clientConn.RemoteAddr().String())
 }
diff --git a/packages/pam/pam-proxy.go b/packages/pam/pam-proxy.go
@@ -78,9 +78,16 @@ func HandlePAMCapabilities(ctx context.Context, conn *tls.Conn, gatewayName stri
 	return nil
 }
 
-func HandlePAMCancellation(ctx context.Context, conn *tls.Conn, pamConfig *GatewayPAMConfig, httpClient *resty.Client) error {
+func HandlePAMCancellation(ctx context.Context, conn *tls.Conn, pamConfig *GatewayPAMConfig, httpClient *resty.Client, cancelSession func(string) bool) error {
 	log.Info().Str("sessionId", pamConfig.SessionId).Msg("Received session termination message")
 
+	// Kill the active proxy connection if it exists in the registry
+	if cancelled := cancelSession(pamConfig.SessionId); cancelled {
+		log.Info().Str("sessionId", pamConfig.SessionId).Msg("Active proxy session cancelled via registry")
+	} else {
+		log.Info().Str("sessionId", pamConfig.SessionId).Msg("No active proxy session found in registry (may have already ended)")
+	}
+
 	if err := pamConfig.SessionUploader.CleanupPAMSession(pamConfig.SessionId, "cancellation"); err != nil {
 		log.Error().Err(err).Str("sessionId", pamConfig.SessionId).Msg("Failed to cleanup PAM session")
 	}
