supported resource types for gateway

x032205 · x032205 · commit 2a3fb2cdd59e · 2025-12-16T15:52:18.000-05:00
diff --git a/packages/gateway-v2/gateway.go b/packages/gateway-v2/gateway.go
@@ -34,6 +34,7 @@ const (
 	ForwardModeTCP             ForwardMode = "TCP"
 	ForwardModePAM             ForwardMode = "PAM"
 	ForwardModePAMCancellation ForwardMode = "PAM_CANCELLATION"
+	ForwardModePAMCapabilities ForwardMode = "PAM_CAPABILITIES"
 	ForwardModePing            ForwardMode = "PING"
 )
 
@@ -443,7 +444,7 @@ func (g *Gateway) setupTLSConfig() error {
 		ClientCAs:  clientCAPool,
 		ClientAuth: tls.RequireAndVerifyClientCert,
 		MinVersion: tls.VersionTLS12,
-		NextProtos: []string{"infisical-http-proxy", "infisical-tcp-proxy", "infisical-ping", "infisical-pam-proxy", "infisical-pam-session-cancellation"},
+		NextProtos: []string{"infisical-http-proxy", "infisical-tcp-proxy", "infisical-ping", "infisical-pam-proxy", "infisical-pam-session-cancellation", "infisical-pam-capabilities"},
 	}
 
 	return nil
@@ -620,6 +621,14 @@ func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) {
 			log.Error().Err(err).Msg("PAM cancellation proxy handler ended with error")
 		}
 		return
+	} else if forwardConfig.Mode == ForwardModePAMCapabilities {
+		log.Info().Msg("Starting PAM capabilities handler")
+		if err := pam.HandlePAMCapabilities(g.ctx, tlsConn); err != nil {
+			log.Error().Err(err).Msg("PAM capabilities handler ended with error")
+		} else {
+			log.Info().Msg("PAM capabilities handler completed")
+		}
+		return
 	} else if forwardConfig.Mode == ForwardModePing {
 		log.Info().Msg("Starting ping handler")
 		if err := handlePing(g.ctx, tlsConn, reader); err != nil {
@@ -666,6 +675,10 @@ func (g *Gateway) parseForwardConfigFromALPN(tlsConn *tls.Conn, reader *bufio.Re
 		config.Mode = ForwardModePAMCancellation
 		return config, nil
 
+	case "infisical-pam-capabilities":
+		config.Mode = ForwardModePAMCapabilities
+		return config, nil
+
 	case "infisical-ping":
 		config.Mode = ForwardModePing
 		return config, nil
diff --git a/packages/pam/local/base-proxy.go b/packages/pam/local/base-proxy.go
@@ -4,14 +4,18 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/json"
 	"fmt"
+	"io"
 	"net"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/Infisical/infisical-merge/packages/api"
+	"github.com/Infisical/infisical-merge/packages/pam"
 	"github.com/Infisical/infisical-merge/packages/util"
 	"github.com/go-resty/resty/v2"
 	"github.com/rs/zerolog/log"
@@ -29,6 +33,7 @@ type BaseProxyServer struct {
 	gatewayServerCertChain string
 	sessionExpiry          time.Time
 	sessionId              string
+	resourceType           string
 	ctx                    context.Context
 	cancel                 context.CancelFunc
 	activeConnections      sync.WaitGroup
@@ -88,6 +93,64 @@ func (b *BaseProxyServer) CreateRelayConnection() (net.Conn, error) {
 	return conn, nil
 }
 
+// FetchGatewayCapabilities fetches the supported resource types from the gateway
+func (b *BaseProxyServer) FetchGatewayCapabilities() (*pam.PAMCapabilitiesResponse, error) {
+	relayConn, err := b.CreateRelayConnection()
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to relay: %w", err)
+	}
+	defer relayConn.Close()
+
+	gatewayConn, err := b.CreateGatewayConnection(relayConn, ALPNInfisicalPAMCapabilities)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to gateway: %w", err)
+	}
+	defer gatewayConn.Close()
+
+	// Read length prefix (4 bytes)
+	lengthBytes := make([]byte, 4)
+	if _, err := io.ReadFull(gatewayConn, lengthBytes); err != nil {
+		return nil, fmt.Errorf("failed to read length prefix: %w", err)
+	}
+
+	length := uint32(lengthBytes[0])<<24 | uint32(lengthBytes[1])<<16 | uint32(lengthBytes[2])<<8 | uint32(lengthBytes[3])
+
+	// Read JSON data
+	data := make([]byte, length)
+	if _, err := io.ReadFull(gatewayConn, data); err != nil {
+		return nil, fmt.Errorf("failed to read capabilities response: %w", err)
+	}
+
+	var response pam.PAMCapabilitiesResponse
+	if err := json.Unmarshal(data, &response); err != nil {
+		return nil, fmt.Errorf("failed to parse capabilities response: %w", err)
+	}
+
+	log.Debug().Strs("supportedTypes", response.SupportedResourceTypes).Msg("Received gateway capabilities")
+	return &response, nil
+}
+
+// ValidateResourceTypeSupported checks if the resource type is supported by the gateway
+func (b *BaseProxyServer) ValidateResourceTypeSupported() error {
+	capabilities, err := b.FetchGatewayCapabilities()
+	if err != nil {
+		return fmt.Errorf("failed to fetch gateway capabilities: %w", err)
+	}
+
+	if slices.Contains(capabilities.SupportedResourceTypes, b.resourceType) {
+		return nil
+	}
+
+	return fmt.Errorf(`Your Gateway does not support '%s' PAM accounts.
+
+To fix this:
+1. Update your Gateway deployment
+2. Restart the Gateway service
+3. Retry your access command
+
+Gateway upgrade guide: https://infisical.com/docs/documentation/platform/gateways/gateway-deployment`, b.resourceType)
+}
+
 // CreateGatewayConnection establishes a mTLS connection to the gateway over the relay
 func (b *BaseProxyServer) CreateGatewayConnection(relayConn net.Conn, alpn ALPN) (net.Conn, error) {
 	// Load gateway certificates
diff --git a/packages/pam/local/database-proxy.go b/packages/pam/local/database-proxy.go
@@ -32,6 +32,7 @@ type ALPN string
 const (
 	ALPNInfisicalPAMProxy        ALPN = "infisical-pam-proxy"
 	ALPNInfisicalPAMCancellation ALPN = "infisical-pam-session-cancellation"
+	ALPNInfisicalPAMCapabilities ALPN = "infisical-pam-capabilities"
 )
 
 func askForApprovalRequestTrigger() (bool, error) {
@@ -130,12 +131,18 @@ func StartDatabaseLocalProxy(accessToken string, accountPath string, projectID s
 			gatewayServerCertChain: pamResponse.GatewayServerCertificateChain,
 			sessionExpiry:          time.Now().Add(duration),
 			sessionId:              pamResponse.SessionId,
+			resourceType:           pamResponse.ResourceType,
 			ctx:                    ctx,
 			cancel:                 cancel,
 			shutdownCh:             make(chan struct{}),
 		},
 	}
 
+	if err := proxy.ValidateResourceTypeSupported(); err != nil {
+		util.HandleError(err, "Gateway version outdated")
+		return
+	}
+
 	err = proxy.Start(port)
 	if err != nil {
 		util.HandleError(err, "Failed to start proxy server")
diff --git a/packages/pam/local/kubernetes-proxy.go b/packages/pam/local/kubernetes-proxy.go
@@ -77,12 +77,18 @@ func StartKubernetesLocalProxy(accessToken string, accountPath string, projectId
 			gatewayServerCertChain: pamResponse.GatewayServerCertificateChain,
 			sessionExpiry:          time.Now().Add(duration),
 			sessionId:              pamResponse.SessionId,
+			resourceType:           pamResponse.ResourceType,
 			ctx:                    ctx,
 			cancel:                 cancel,
 			shutdownCh:             make(chan struct{}),
 		},
 	}
 
+	if err := proxy.ValidateResourceTypeSupported(); err != nil {
+		util.HandleError(err, "Gateway version outdated")
+		return
+	}
+
 	err = proxy.Start(port)
 	if err != nil {
 		util.HandleError(err, "Failed to start proxy server")
diff --git a/packages/pam/local/ssh-proxy.go b/packages/pam/local/ssh-proxy.go
@@ -70,12 +70,18 @@ func StartSSHLocalProxy(accessToken string, accountPath string, projectID string
 			gatewayServerCertChain: pamResponse.GatewayServerCertificateChain,
 			sessionExpiry:          time.Now().Add(duration),
 			sessionId:              pamResponse.SessionId,
+			resourceType:           pamResponse.ResourceType,
 			ctx:                    ctx,
 			cancel:                 cancel,
 			shutdownCh:             make(chan struct{}),
 		},
 	}
 
+	if err := proxy.ValidateResourceTypeSupported(); err != nil {
+		util.HandleError(err, "Gateway version outdated")
+		return
+	}
+
 	// Start the local TCP proxy on a random port
 	err = proxy.Start(0) // 0 = random port
 	if err != nil {
diff --git a/packages/pam/pam-proxy.go b/packages/pam/pam-proxy.go
@@ -4,14 +4,14 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"encoding/json"
 	"fmt"
 	"net/url"
 	"time"
 
 	"github.com/Infisical/infisical-merge/packages/pam/handlers"
 	"github.com/Infisical/infisical-merge/packages/pam/handlers/kubernetes"
 	"github.com/Infisical/infisical-merge/packages/pam/handlers/mysql"
-	"github.com/Infisical/infisical-merge/packages/pam/handlers/ssh"
 	"github.com/Infisical/infisical-merge/packages/pam/session"
 	"github.com/go-resty/resty/v2"
 	"github.com/rs/zerolog/log"
@@ -25,6 +25,52 @@ type GatewayPAMConfig struct {
 	SessionUploader    *session.SessionUploader
 }
 
+type PAMCapabilitiesResponse struct {
+	SupportedResourceTypes []string `json:"supportedResourceTypes"`
+}
+
+func GetSupportedResourceTypes() []string {
+	return []string{
+		session.ResourceTypePostgres,
+		session.ResourceTypeMysql,
+		session.ResourceTypeSSH,
+		session.ResourceTypeKubernetes,
+	}
+}
+
+// HandlePAMCapabilities handles the capabilities request from the client
+func HandlePAMCapabilities(ctx context.Context, conn *tls.Conn) error {
+	response := PAMCapabilitiesResponse{
+		SupportedResourceTypes: GetSupportedResourceTypes(),
+	}
+
+	data, err := json.Marshal(response)
+	if err != nil {
+		log.Error().Err(err).Msg("Failed to marshal capabilities response")
+		return fmt.Errorf("failed to marshal capabilities response: %w", err)
+	}
+
+	// Write length prefix (4 bytes) followed by JSON data
+	length := uint32(len(data))
+	lengthBytes := []byte{
+		byte(length >> 24),
+		byte(length >> 16),
+		byte(length >> 8),
+		byte(length),
+	}
+
+	if _, err := conn.Write(lengthBytes); err != nil {
+		return fmt.Errorf("failed to write length prefix: %w", err)
+	}
+
+	if _, err := conn.Write(data); err != nil {
+		return fmt.Errorf("failed to write capabilities response: %w", err)
+	}
+
+	log.Debug().Strs("supportedTypes", response.SupportedResourceTypes).Msg("Sent PAM capabilities to client")
+	return nil
+}
+
 func HandlePAMCancellation(ctx context.Context, conn *tls.Conn, pamConfig *GatewayPAMConfig, httpClient *resty.Client) error {
 	log.Info().Str("sessionId", pamConfig.SessionId).Msg("Received session termination message")
 
