diff --git a/e2e/go.mod b/e2e/go.mod index 861d9d01..e0cb09a7 100644 --- a/e2e/go.mod +++ b/e2e/go.mod @@ -49,7 +49,7 @@ require ( github.com/asaskevich/govalidator v0.0.0-20200907205600-7a23bdc65eef // indirect github.com/awnumar/memcall v0.4.0 // indirect github.com/awnumar/memguard v0.23.0 // indirect - github.com/aws/aws-sdk-go-v2 v1.30.3 // indirect + github.com/aws/aws-sdk-go-v2 v1.38.1 // indirect github.com/aws/aws-sdk-go-v2/config v1.27.27 // indirect github.com/aws/aws-sdk-go-v2/credentials v1.17.27 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 // indirect @@ -61,7 +61,7 @@ require ( github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 // indirect github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 // indirect github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 // indirect - github.com/aws/smithy-go v1.20.3 // indirect + github.com/aws/smithy-go v1.22.5 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/bodgit/ntlmssp v0.0.0-20240506230425-31973bb52d9b // indirect @@ -102,7 +102,7 @@ require ( github.com/docker/go-metrics v0.0.1 // indirect github.com/docker/go-units v0.5.0 // indirect github.com/dprotaso/go-yit v0.0.0-20220510233725-9ba8df137936 // indirect - github.com/dustin/go-humanize v1.0.0 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect github.com/dvsekhvalnov/jose2go v1.7.0 // indirect github.com/ebitengine/purego v0.8.4 // indirect github.com/eiannone/keyboard v0.0.0-20220611211555-0d226195f203 // indirect @@ -132,7 +132,7 @@ require ( github.com/go-resty/resty/v2 v2.16.5 // indirect github.com/go-sql-driver/mysql v1.7.1 // indirect github.com/go-viper/mapstructure/v2 v2.4.0 // indirect - github.com/goccy/go-json v0.10.2 // indirect + github.com/goccy/go-json v0.10.5 // indirect github.com/godbus/dbus/v5 v5.1.0 // indirect github.com/gofrs/flock v0.12.1 // indirect github.com/gofrs/uuid v4.4.0+incompatible // indirect @@ -143,8 +143,8 @@ require ( github.com/golang/glog v1.2.5 // indirect github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect github.com/golang/protobuf v1.5.4 // indirect - github.com/golang/snappy v0.0.4 // indirect - github.com/google/flatbuffers v1.12.1 // indirect + github.com/golang/snappy v1.0.0 // indirect + github.com/google/flatbuffers v25.2.10+incompatible // indirect github.com/google/gnostic-models v0.6.9 // indirect github.com/google/go-cmp v0.7.0 // indirect github.com/google/gofuzz v1.2.0 // indirect diff --git a/e2e/go.sum b/e2e/go.sum index 72a6ca30..e6fc1bff 100644 --- a/e2e/go.sum +++ b/e2e/go.sum @@ -121,8 +121,8 @@ github.com/awnumar/memcall v0.4.0 h1:B7hgZYdfH6Ot1Goaz8jGne/7i8xD4taZie/PNSFZ29g github.com/awnumar/memcall v0.4.0/go.mod h1:8xOx1YbfyuCg3Fy6TO8DK0kZUua3V42/goA5Ru47E8w= github.com/awnumar/memguard v0.23.0 h1:sJ3a1/SWlcuKIQ7MV+R9p0Pvo9CWsMbGZvcZQtmc68A= github.com/awnumar/memguard v0.23.0/go.mod h1:olVofBrsPdITtJ2HgxQKrEYEMyIBAIciVG4wNnZhW9M= -github.com/aws/aws-sdk-go-v2 v1.30.3 h1:jUeBtG0Ih+ZIFH0F4UkmL9w3cSpaMv9tYYDbzILP8dY= -github.com/aws/aws-sdk-go-v2 v1.30.3/go.mod h1:nIQjQVp5sfpQcTc9mPSr1B0PaWK5ByX9MOoDadSN4lc= +github.com/aws/aws-sdk-go-v2 v1.38.1 h1:j7sc33amE74Rz0M/PoCpsZQ6OunLqys/m5antM0J+Z8= +github.com/aws/aws-sdk-go-v2 v1.38.1/go.mod h1:9Q0OoGQoboYIAJyslFyF1f5K1Ryddop8gqMhWx/n4Wg= github.com/aws/aws-sdk-go-v2/config v1.27.27 h1:HdqgGt1OAP0HkEDDShEl0oSYa9ZZBSOmKpdpsDMdO90= github.com/aws/aws-sdk-go-v2/config v1.27.27/go.mod h1:MVYamCg76dFNINkZFu4n4RjDixhVr51HLj4ErWzrVwg= github.com/aws/aws-sdk-go-v2/credentials v1.17.27 h1:2raNba6gr2IfA0eqqiP2XiQ0UVOpGPgDSi0I9iAP+UI= @@ -145,8 +145,8 @@ github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 h1:yiwVzJW2ZxZTurVbYWA7QOrA github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4/go.mod h1:0oxfLkpz3rQ/CHlx5hB7H69YUpFiI1tql6Q6Ne+1bCw= github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 h1:ZsDKRLXGWHk8WdtyYMoGNO7bTudrvuKpDKgMVRlepGE= github.com/aws/aws-sdk-go-v2/service/sts v1.30.3/go.mod h1:zwySh8fpFyXp9yOr/KVzxOl8SRqgf/IDw5aUt9UKFcQ= -github.com/aws/smithy-go v1.20.3 h1:ryHwveWzPV5BIof6fyDvor6V3iUL7nTfiTKXHiW05nE= -github.com/aws/smithy-go v1.20.3/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E= +github.com/aws/smithy-go v1.22.5 h1:P9ATCXPMb2mPjYBgueqJNCA5S9UfktsW0tTxi+a7eqw= +github.com/aws/smithy-go v1.22.5/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= @@ -309,8 +309,9 @@ github.com/docker/libtrust v0.0.0-20160708172513-aabc10ec26b7/go.mod h1:cyGadeNE github.com/dprotaso/go-yit v0.0.0-20191028211022-135eb7262960/go.mod h1:9HQzr9D/0PGwMEbC3d5AB7oi67+h4TsQqItC1GVYG58= github.com/dprotaso/go-yit v0.0.0-20220510233725-9ba8df137936 h1:PRxIJD8XjimM5aTknUK9w6DHLDox2r2M3DI4i2pnd3w= github.com/dprotaso/go-yit v0.0.0-20220510233725-9ba8df137936/go.mod h1:ttYvX5qlB+mlV1okblJqcSMtR4c52UKxDiX9GRBS8+Q= -github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/dvsekhvalnov/jose2go v0.0.0-20170216131308-f21a8cedbbae/go.mod h1:7BvyPhdbLxMXIYTFPLsyJRFMsKmOZnQmzh6Gb+uquuM= github.com/dvsekhvalnov/jose2go v1.7.0 h1:bnQc8+GMnidJZA8zc6lLEAb4xNrIqHwO+9TzqvtQZPo= github.com/dvsekhvalnov/jose2go v1.7.0/go.mod h1:QsHjhyTlD/lAVqn/NSbVZmSCGeDehTB/mPZadG+mhXU= @@ -408,8 +409,8 @@ github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM= github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs= github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM= -github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= -github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= +github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk= github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= @@ -466,14 +467,15 @@ github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= -github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM= -github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= +github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs= +github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q= github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ= github.com/google/certificate-transparency-go v1.0.10-0.20180222191210-5ab67e519c93 h1:jc2UWq7CbdszqeH6qu1ougXMIUBfSy8Pbh/anURYbGI= github.com/google/certificate-transparency-go v1.0.10-0.20180222191210-5ab67e519c93/go.mod h1:QeJfpSbVSfYc7RgB3gJFj9cbuQMMchQxrWXz8Ruopmg= -github.com/google/flatbuffers v1.12.1 h1:MVlul7pQNoDzWRLTw5imwYsl+usrS1TXG2H4jg6ImGw= github.com/google/flatbuffers v1.12.1/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/flatbuffers v25.2.10+incompatible h1:F3vclr7C3HpB1k9mxCGRMXq6FdUalZ6H/pNX4FP1v0Q= +github.com/google/flatbuffers v25.2.10+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/gnostic-models v0.6.9 h1:MU/8wDLif2qCXZmzncUQ/BOfxWfthHi63KqpoNbWqVw= github.com/google/gnostic-models v0.6.9/go.mod h1:CiWsm0s6BSQd1hRn8/QmxqB6BesYcbSZxsz9b0KuDBw= github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= diff --git a/go.mod b/go.mod index 40b03466..76d74803 100644 --- a/go.mod +++ b/go.mod @@ -50,6 +50,7 @@ require ( github.com/spf13/viper v1.8.1 github.com/stretchr/testify v1.11.1 github.com/wasilibs/go-re2 v1.10.0 + github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 go.mongodb.org/mongo-driver/v2 v2.5.0 go.mozilla.org/pkcs7 v0.9.0 golang.org/x/crypto v0.55.0 @@ -203,7 +204,6 @@ require ( github.com/xdg-go/scram v1.2.0 // indirect github.com/xdg-go/stringprep v1.0.4 // indirect github.com/xtgo/uuid v0.0.0-20140804021211-a0b114877d4c // indirect - github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect go.mongodb.org/mongo-driver v1.17.9 // indirect go.opencensus.io v0.24.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect diff --git a/packages/api/model.go b/packages/api/model.go index e3e7da7a..e5e3ab06 100644 --- a/packages/api/model.go +++ b/packages/api/model.go @@ -1033,6 +1033,11 @@ type PAMSessionCredentials struct { Token string `json:"token,omitempty"` Tokens map[string]string `json:"tokens,omitempty"` ServiceAccountEmail string `json:"serviceAccountEmail,omitempty"` + Account string `json:"account,omitempty"` + Warehouse string `json:"warehouse,omitempty"` + Schema string `json:"schema,omitempty"` + Role string `json:"role,omitempty"` + PrivateKeyPassphrase string `json:"privateKeyPassphrase,omitempty"` } type MFASessionStatus string diff --git a/packages/gateway-v2/test_connection_handler.go b/packages/gateway-v2/test_connection_handler.go index 84e97552..0cc9991a 100644 --- a/packages/gateway-v2/test_connection_handler.go +++ b/packages/gateway-v2/test_connection_handler.go @@ -20,6 +20,7 @@ import ( "time" mssqlhandler "github.com/Infisical/infisical-merge/packages/pam/handlers/mssql" + snowflakehandler "github.com/Infisical/infisical-merge/packages/pam/handlers/snowflake" "github.com/go-ldap/ldap/v3" "github.com/go-sql-driver/mysql" "github.com/jackc/pgx/v5" @@ -55,6 +56,7 @@ const ( testConnModeLDAP = "ldap" testConnModeKubernetes = "kubernetes" testConnModeSSH = "ssh" + testConnModeSnowflake = "snowflake" testConnModeTCP = "tcp" ) @@ -97,6 +99,20 @@ type redisTestParams struct { SslCertificate string `json:"sslCertificate"` } +type snowflakeTestParams struct { + Account string `json:"account"` + AuthMethod string `json:"authMethod"` + Username string `json:"username"` + Password string `json:"password"` + Token string `json:"token"` + PrivateKey string `json:"privateKey"` + PrivateKeyPassphrase string `json:"privateKeyPassphrase"` + Warehouse string `json:"warehouse"` + Database string `json:"database"` + Schema string `json:"schema"` + Role string `json:"role"` +} + type ldapTestParams struct { Username string `json:"username"` Password string `json:"password"` @@ -609,6 +625,31 @@ func handleTestConnection(w http.ResponseWriter, r *http.Request) { return } op = func() error { return doKubernetesConnectionTest(ctx, target.host, target.port, params) } + case testConnModeSnowflake: + var params snowflakeTestParams + if !decode(¶ms) { + return + } + op = func() error { + proxy := snowflakehandler.NewSnowflakeProxy(snowflakehandler.SnowflakeProxyConfig{ + Account: params.Account, + Username: params.Username, + AuthMethod: params.AuthMethod, + Password: params.Password, + Token: params.Token, + PrivateKey: params.PrivateKey, + PrivateKeyPass: params.PrivateKeyPassphrase, + Warehouse: params.Warehouse, + Database: params.Database, + Schema: params.Schema, + Role: params.Role, + }) + if err := proxy.Connect(ctx); err != nil { + return authFailure(err) + } + defer proxy.Close() + return authFailure(proxy.Probe(ctx)) + } case testConnModeSSH: var params sshTestParams if !decode(¶ms) { diff --git a/packages/pam/handlers/snowflake/proxy.go b/packages/pam/handlers/snowflake/proxy.go new file mode 100644 index 00000000..ca1ffd3b --- /dev/null +++ b/packages/pam/handlers/snowflake/proxy.go @@ -0,0 +1,302 @@ +package snowflake + +import ( + "context" + "crypto/subtle" + "encoding/json" + "fmt" + "net" + "net/http" + "regexp" + "strings" + "sync" + "time" + + "github.com/Infisical/infisical-merge/packages/pam/session" + "github.com/rs/zerolog/log" +) + +// Snowflake has no wire protocol, so the gateway answers the REST API its drivers speak and runs each +// statement itself. The client never holds a Snowflake token, so it cannot bypass the recorded session. +type SnowflakeProxyConfig struct { + Account string + Username string + AuthMethod string + Password string + Token string + PrivateKey string + PrivateKeyPass string + Warehouse string + Database string + Schema string + Role string + + SessionID string + SessionExpiry time.Time + SessionLogger session.SessionLogger + BlockedCommands []*regexp.Regexp +} + +const ( + AuthMethodKeyPair = "key-pair" + AuthMethodToken = "programmatic-access-token" + AuthMethodPassword = "password" + + loginTimeout = 30 * time.Second + + statementTimeout = 30 * time.Minute +) + +type SnowflakeProxy struct { + config SnowflakeProxyConfig + upstream *upstream + token string + sessionCtx map[string]string + parameters json.RawMessage +} + +func NewSnowflakeProxy(config SnowflakeProxyConfig) *SnowflakeProxy { + return &SnowflakeProxy{config: config} +} + +func (p *SnowflakeProxy) Connect(ctx context.Context) error { + client := newUpstream(p.config) + loginCtx, cancel := context.WithTimeout(ctx, loginTimeout) + defer cancel() + + if err := client.login(loginCtx); err != nil { + client.close() + return err + } + + token, err := sessionToken(p.config) + if err != nil { + client.close() + return err + } + + p.upstream, p.token, p.sessionCtx, p.parameters = client, token, client.sessionCtx, client.parameters + return nil +} + +// Probe runs one statement and checks what the session actually opened with: Snowflake accepts a +// warehouse, role or database the credential can't use and just leaves it unset. +func (p *SnowflakeProxy) Probe(ctx context.Context) error { + for _, requested := range [][3]string{ + {"warehouse", p.config.Warehouse, p.sessionCtx["warehouseName"]}, + {"role", p.config.Role, p.sessionCtx["roleName"]}, + {"database", p.config.Database, p.sessionCtx["databaseName"]}, + {"schema", p.config.Schema, p.sessionCtx["schemaName"]}, + } { + if requested[1] != "" && !strings.EqualFold(requested[1], requested[2]) { + return fmt.Errorf("the user cannot use the %s %q", requested[0], requested[1]) + } + } + + probeCtx, cancel := context.WithTimeout(ctx, loginTimeout) + defer cancel() + + _, err := p.upstream.query(probeCtx, "SELECT 1", nil) + return err +} + +func (p *SnowflakeProxy) Close() { + if p.upstream == nil { + return + } + ctx, cancel := context.WithTimeout(context.Background(), loginTimeout) + defer cancel() + if err := p.upstream.logout(ctx); err != nil { + log.Debug().Err(err).Str("sessionId", p.config.SessionID).Msg("Failed to close the Snowflake session") + } + + p.upstream.close() +} + +func (p *SnowflakeProxy) HandleConnection(ctx context.Context, clientConn net.Conn) error { + defer clientConn.Close() + + l := log.With().Str("sessionId", p.config.SessionID).Str("resourceType", "snowflake").Logger() + + server := &http.Server{ + Handler: p.router(l), + ReadHeaderTimeout: 30 * time.Second, + } + + listener := newSingleConnListener(clientConn) + + go func() { + <-ctx.Done() + listener.Close() + server.Close() + }() + + if err := server.Serve(listener); err != nil && !isListenerDone(err) { + l.Debug().Err(err).Msg("Snowflake proxy stopped") + } + return nil +} + +func (p *SnowflakeProxy) router(l zeroLogger) http.Handler { + mux := http.NewServeMux() + mux.HandleFunc("/session/v1/login-request", p.handleLogin) + mux.HandleFunc("/session/token-request", p.handleTokenRequest) + mux.HandleFunc("/session/heartbeat", p.handleOK) + mux.HandleFunc("/session", p.handleOK) + mux.HandleFunc("/queries/v1/query-request", p.handleQuery(l)) + mux.HandleFunc("/queries/v1/abort-request", p.handleAbort(l)) + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + l.Debug().Str("path", r.URL.Path).Msg("Snowflake proxy received an unsupported request") + writeFailure(w, errCodeUnsupported, fmt.Sprintf("This session does not support %s.", r.URL.Path)) + }) + return mux +} + +func (p *SnowflakeProxy) handleLogin(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + writeFailure(w, errCodeAuth, "Unsupported method") + return + } + + // The gateway already authenticated, so whatever the client sent is ignored rather than forwarded + writeEnvelope(w, envelope{Success: true, Data: loginData(p.token, p.config, p.sessionCtx, p.parameters)}) +} + +func (p *SnowflakeProxy) handleTokenRequest(w http.ResponseWriter, r *http.Request) { + if !p.authorize(r) { + writeFailure(w, errCodeAuth, "Authentication token has expired") + return + } + writeEnvelope(w, envelope{Success: true, Data: renewData(p.token)}) +} + +func (p *SnowflakeProxy) handleOK(w http.ResponseWriter, r *http.Request) { + writeEnvelope(w, envelope{Success: true, Data: map[string]any{}}) +} + +func (p *SnowflakeProxy) handleQuery(l zeroLogger) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if !p.authorize(r) { + writeFailure(w, errCodeAuth, "Authentication token has expired") + return + } + + var body struct { + SqlText string `json:"sqlText"` + Bindings json.RawMessage `json:"bindings"` + } + if err := decodeRequest(w, r, &body); err != nil { + writeFailure(w, errCodeStatement, statementReadError(err)) + return + } + + statement := strings.TrimSpace(body.SqlText) + if statement == "" { + writeFailure(w, errCodeStatement, "Empty statement") + return + } + + if blocked := p.blockedBy(statement); blocked != nil { + p.logStatement(statement, fmt.Sprintf("BLOCKED: %s", blocked.String())) + l.Info().Str("pattern", blocked.String()).Msg("Blocked a statement by policy") + writeFailure(w, errCodeStatement, "This statement is blocked by the command blocking policy on this account.") + return + } + + queryCtx, cancel := context.WithTimeout(r.Context(), statementTimeout) + defer cancel() + + if key := p.inflightKey(r.URL.Query().Get("requestId")); key != "" { + inflightQueries.Store(key, &inflightQuery{upstream: p.upstream, cancel: cancel}) + defer inflightQueries.Delete(key) + } + + result, err := p.upstream.query(queryCtx, statement, body.Bindings) + if err != nil { + if queryCtx.Err() != nil { + abortUpstream(r.Context(), p.upstream, l) + } + p.logStatement(statement, fmt.Sprintf("ERROR: %s", err)) + writeUpstreamFailure(w, err) + return + } + + p.logStatement(statement, result.summary()) + writeEnvelope(w, envelope{Success: true, Data: result.data()}) + } +} + +// A driver sends its cancellation on a second connection, which the tunnel makes a second proxy, so the +// statement is found by the request id the client gave it rather than through this proxy's own state. +var inflightQueries sync.Map + +type inflightQuery struct { + upstream *upstream + cancel context.CancelFunc +} + +func (p *SnowflakeProxy) inflightKey(requestID string) string { + if requestID == "" { + return "" + } + return p.config.SessionID + "/" + requestID +} + +func (p *SnowflakeProxy) handleAbort(l zeroLogger) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if !p.authorize(r) { + writeFailure(w, errCodeAuth, "Authentication token has expired") + return + } + + var body struct { + RequestID string `json:"requestId"` + } + if err := decodeRequest(w, r, &body); err == nil { + if entry, ok := inflightQueries.Load(p.inflightKey(body.RequestID)); ok { + query := entry.(*inflightQuery) + abortUpstream(r.Context(), query.upstream, l) + query.cancel() + } + } + + writeEnvelope(w, envelope{Success: true, Data: map[string]any{}}) + } +} + +// Detached from the request, which is often already cancelled by the time the abort is worth sending +func abortUpstream(ctx context.Context, client *upstream, l zeroLogger) { + abortCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginTimeout) + defer cancel() + + if err := client.abortQuery(abortCtx); err != nil { + l.Debug().Err(err).Msg("Snowflake refused a cancellation") + } +} + +func (p *SnowflakeProxy) authorize(r *http.Request) bool { + presented := strings.TrimSuffix(strings.TrimPrefix(r.Header.Get("Authorization"), `Snowflake Token="`), `"`) + return subtle.ConstantTimeCompare([]byte(presented), []byte(p.token)) == 1 +} + +func (p *SnowflakeProxy) blockedBy(statement string) *regexp.Regexp { + for _, pattern := range p.config.BlockedCommands { + if pattern.MatchString(statement) { + return pattern + } + } + return nil +} + +func (p *SnowflakeProxy) logStatement(input, output string) { + if p.config.SessionLogger == nil { + return + } + if err := p.config.SessionLogger.LogEntry(session.SessionLogEntry{ + Timestamp: time.Now(), + Input: input, + Output: output, + }); err != nil { + log.Error().Err(err).Str("sessionId", p.config.SessionID).Msg("Failed to log a Snowflake statement") + } +} diff --git a/packages/pam/handlers/snowflake/proxy_test.go b/packages/pam/handlers/snowflake/proxy_test.go new file mode 100644 index 00000000..64e1c54e --- /dev/null +++ b/packages/pam/handlers/snowflake/proxy_test.go @@ -0,0 +1,451 @@ +package snowflake + +import ( + "bytes" + "compress/gzip" + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/json" + "encoding/pem" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "regexp" + "strings" + "testing" + "time" + + "github.com/Infisical/infisical-merge/packages/pam/session" + "github.com/stretchr/testify/require" +) + +type recordingLogger struct { + entries []session.SessionLogEntry +} + +func (r *recordingLogger) LogEntry(entry session.SessionLogEntry) error { + r.entries = append(r.entries, entry) + return nil +} +func (r *recordingLogger) LogSessionEvent(session.SessionEvent) error { return nil } +func (r *recordingLogger) LogHttpEvent(session.HttpEvent) error { return nil } +func (r *recordingLogger) Close() error { return nil } + +func newTestProxy(t *testing.T, blocked ...string) (*SnowflakeProxy, *recordingLogger) { + t.Helper() + + patterns := make([]*regexp.Regexp, 0, len(blocked)) + for _, p := range blocked { + patterns = append(patterns, regexp.MustCompile(p)) + } + + logger := &recordingLogger{} + proxy := NewSnowflakeProxy(SnowflakeProxyConfig{ + Account: "acme-test", + Username: "pam_svc", + SessionID: "session-1", + SessionLogger: logger, + BlockedCommands: patterns, + }) + token, err := sessionToken(proxy.config) + require.NoError(t, err) + proxy.token = token + + return proxy, logger +} + +func postJSON(t *testing.T, handler http.HandlerFunc, token string, sqlText string, gzipped bool) map[string]any { + t.Helper() + + encoded, err := json.Marshal(map[string]any{"sqlText": sqlText}) + require.NoError(t, err) + + var payload bytes.Buffer + if gzipped { + writer := gzip.NewWriter(&payload) + _, err = writer.Write(encoded) + require.NoError(t, err) + require.NoError(t, writer.Close()) + } else { + payload.Write(encoded) + } + + req := httptest.NewRequest(http.MethodPost, "/queries/v1/query-request", bytes.NewReader(payload.Bytes())) + if gzipped { + req.Header.Set("Content-Encoding", "gzip") + } + if token != "" { + req.Header.Set("Authorization", `Snowflake Token="`+token+`"`) + } + + recorder := httptest.NewRecorder() + handler(recorder, req) + + var decoded map[string]any + require.NoError(t, json.NewDecoder(recorder.Body).Decode(&decoded)) + return decoded +} + +func TestLoginHandsOutTheGatewayToken(t *testing.T) { + proxy, _ := newTestProxy(t) + + recorder := httptest.NewRecorder() + proxy.handleLogin(recorder, httptest.NewRequest(http.MethodPost, "/session/v1/login-request", nil)) + + var body map[string]any + require.NoError(t, json.NewDecoder(recorder.Body).Decode(&body)) + require.True(t, body["success"].(bool)) + // The client is never given a Snowflake token, so it cannot reach the account outside the session + require.Equal(t, proxy.token, body["data"].(map[string]any)["token"]) +} + +func TestQueryRequiresTheToken(t *testing.T) { + proxy, logger := newTestProxy(t) + + for _, token := range []string{"wrong", proxy.token[:len(proxy.token)-1], proxy.token + "x"} { + body := postJSON(t, proxy.handleQuery(zeroLogger{}), token, "SELECT 1", false) + require.False(t, body["success"].(bool)) + } + require.Empty(t, logger.entries) +} + +func TestBlockedStatementNeverReachesSnowflake(t *testing.T) { + proxy, logger := newTestProxy(t, `(?i)^\s*drop\b`) + + // A nil driver connection proves the statement was refused before it would have run + body := postJSON(t, proxy.handleQuery(zeroLogger{}), proxy.token, " DROP TABLE orders", false) + + require.False(t, body["success"].(bool)) + require.Contains(t, body["message"].(string), "command blocking policy") + require.Len(t, logger.entries, 1) + require.Equal(t, "DROP TABLE orders", logger.entries[0].Input) + require.Contains(t, logger.entries[0].Output, "BLOCKED:") +} + +func TestBlockedStatementIsReadFromAGzippedBody(t *testing.T) { + proxy, logger := newTestProxy(t, `(?i)\btruncate\b`) + + body := postJSON(t, proxy.handleQuery(zeroLogger{}), proxy.token, "TRUNCATE TABLE orders", true) + + require.False(t, body["success"].(bool)) + require.Len(t, logger.entries, 1) + require.Equal(t, "TRUNCATE TABLE orders", logger.entries[0].Input) +} + +func TestResultTranslation(t *testing.T) { + result := &queryResult{ + columns: []column{{Name: "ID", Type: "text"}, {Name: "NAME", Type: "text"}}, + rows: [][]any{{"1", "alice"}, {"2", nil}}, + queryID: "01b2-real", + elapsed: 5 * time.Millisecond, + } + + data := result.data() + require.Equal(t, result.columns, data["rowtype"]) + require.Equal(t, [][]any{{"1", "alice"}, {"2", nil}}, data["rowset"]) + require.EqualValues(t, 2, data["total"]) + // Snowflake's own id, so a client that cancels or looks the query up names the right one + require.Equal(t, "01b2-real", data["queryId"]) + require.Equal(t, "2 row(s), 5ms", result.summary()) + + truncated := &queryResult{rows: make([][]any, 3), truncated: true} + require.Contains(t, truncated.summary(), "truncated at") +} + +// Snowflake reports the real row count even when only the first slice comes back inline +func TestQueryRecordsTruncationAgainstTheReportedTotal(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"success":true,"data":{"queryResultFormat":"json","rowtype":[],"rowset":[["1"]],"total":42}}`) + })) + t.Cleanup(server.Close) + + client := newUpstream(SnowflakeProxyConfig{Account: "acme-test"}) + client.baseURL = server.URL + + result, err := client.query(t.Context(), "SELECT * FROM events", nil) + + require.NoError(t, err) + require.True(t, result.truncated) + require.Contains(t, result.summary(), "truncated at") +} + +// A client on the official Go driver decodes sessionId as an int64 and fails the login on anything else +func TestLoginReportsANumericSessionID(t *testing.T) { + data := loginData("token", SnowflakeProxyConfig{SessionID: "9f1c7a3e-0e3d-4d5a-8b21-1f2c3d4e5f60"}, nil, nil) + + id, ok := data["sessionId"].(int64) + require.True(t, ok) + require.Positive(t, id) + require.Equal(t, id, loginData("t", SnowflakeProxyConfig{SessionID: "9f1c7a3e-0e3d-4d5a-8b21-1f2c3d4e5f60"}, nil, nil)["sessionId"]) +} + +// Cancelling has to reach Snowflake, or the statement runs on and the warehouse bills for it +func TestAbortForwardsTheStatementToSnowflake(t *testing.T) { + var aborted string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + aborted = string(body) + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"success":true,"data":{}}`) + })) + t.Cleanup(server.Close) + + proxy, _ := newTestProxy(t) + proxy.upstream = newUpstream(proxy.config) + proxy.upstream.baseURL = server.URL + proxy.upstream.requestID = "req-1" + + // The driver cancels on a second connection, so the lookup crosses proxies + inflightQueries.Store(proxy.inflightKey("req-1"), &inflightQuery{upstream: proxy.upstream, cancel: func() {}}) + t.Cleanup(func() { inflightQueries.Delete(proxy.inflightKey("req-1")) }) + + request := httptest.NewRequest(http.MethodPost, "/queries/v1/abort-request", strings.NewReader(`{"requestId":"req-1"}`)) + request.Header.Set("Authorization", `Snowflake Token="`+proxy.token+`"`) + recorder := httptest.NewRecorder() + proxy.handleAbort(zeroLogger{})(recorder, request) + + require.Contains(t, aborted, `"requestId":"req-1"`) +} + +func TestCloseEndsTheSnowflakeSession(t *testing.T) { + var deleted bool + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleted = r.URL.Path == "/session" && r.URL.Query().Get("delete") == "true" + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"success":true,"data":{}}`) + })) + t.Cleanup(server.Close) + + proxy, _ := newTestProxy(t) + proxy.upstream = newUpstream(proxy.config) + proxy.upstream.baseURL = server.URL + + proxy.Close() + + require.True(t, deleted) +} + +// Snowflake accepts a warehouse or role the credential can't use and just leaves it unset +func TestProbeRejectsAnUnusableWarehouse(t *testing.T) { + proxy, _ := newTestProxy(t) + proxy.config.Warehouse = "analytics_wh" + proxy.sessionCtx = map[string]string{"warehouseName": ""} + + require.ErrorContains(t, proxy.Probe(t.Context()), "analytics_wh") + + // Snowflake upper-cases an unquoted identifier + proxy.sessionCtx = map[string]string{"warehouseName": "ANALYTICS_WH"} + proxy.upstream = newUpstream(proxy.config) + proxy.upstream.baseURL = "http://127.0.0.1:1" + require.NotContains(t, proxy.Probe(t.Context()).Error(), "analytics_wh") +} + +// The gateway replaces http.DefaultTransport, so the upstream client must carry its own. +func TestUpstreamCarriesItsOwnTransport(t *testing.T) { + client := newUpstream(SnowflakeProxyConfig{Account: "acme-test"}) + + require.NotNil(t, client.client.Transport) + require.NotEqual(t, http.DefaultTransport, client.client.Transport) + require.Equal(t, "https://acme-test.snowflakecomputing.com", client.baseURL) +} + +// A client hanging up has to end the session, or the recording is never flushed +func TestListenerUnblocksWhenTheConnectionCloses(t *testing.T) { + client, server := net.Pipe() + t.Cleanup(func() { _ = client.Close() }) + + listener := newSingleConnListener(server) + conn, err := listener.Accept() + require.NoError(t, err) + + require.NoError(t, conn.Close()) + + _, err = listener.Accept() + require.ErrorIs(t, err, net.ErrClosed) +} + +// Snowflake expires a session token roughly hourly, inside the length of a long PAM session +func TestQueryRenewsAnExpiredSessionToken(t *testing.T) { + var renewals int + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/session/token-request": + renewals++ + require.Contains(t, r.Header.Get("Authorization"), "master-token") + fmt.Fprint(w, `{"success":true,"data":{"sessionToken":"fresh-token","masterToken":"master-token"}}`) + case "/queries/v1/query-request": + if strings.Contains(r.Header.Get("Authorization"), "stale-token") { + fmt.Fprintf(w, `{"success":false,"code":"%s","message":"session expired"}`, sessionExpiredCode) + return + } + fmt.Fprint(w, `{"success":true,"data":{"queryResultFormat":"json","rowtype":[{"name":"N"}],"rowset":[["1"]]}}`) + default: + t.Errorf("unexpected path %s", r.URL.Path) + } + })) + t.Cleanup(server.Close) + + client := newUpstream(SnowflakeProxyConfig{Account: "acme-test"}) + client.baseURL = server.URL + client.token, client.masterTok = "stale-token", "master-token" + + result, err := client.query(t.Context(), "SELECT 1", nil) + + require.NoError(t, err) + require.Equal(t, 1, renewals) + require.Equal(t, "fresh-token", client.token) + require.Len(t, result.rows, 1) +} + +func TestQueryForwardsBindings(t *testing.T) { + var forwarded string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + forwarded = string(body) + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"success":true,"data":{"queryResultFormat":"json","rowtype":[],"rowset":[]}}`) + })) + t.Cleanup(server.Close) + + client := newUpstream(SnowflakeProxyConfig{Account: "acme-test"}) + client.baseURL = server.URL + + _, err := client.query(t.Context(), "SELECT ?", json.RawMessage(`{"1":{"type":"TEXT","value":"PUBLIC"}}`)) + + require.NoError(t, err) + require.Contains(t, forwarded, `"bindings":{"1":{"type":"TEXT","value":"PUBLIC"}}`) +} + +// A passphrase left over from an earlier credential must not break an unencrypted key +func TestParsePrivateKeyIgnoresAStalePassphrase(t *testing.T) { + key, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + der, err := x509.MarshalPKCS8PrivateKey(key) + require.NoError(t, err) + plain := string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})) + + parsed, err := parsePrivateKey(plain, "left-over-passphrase") + + require.NoError(t, err) + require.Equal(t, key.N, parsed.N) +} + +// A query slower than about 45 seconds comes back "in progress" with a URL to poll +func TestQueryPollsAnInProgressResult(t *testing.T) { + var polls int + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/queries/v1/query-request" { + fmt.Fprintf(w, `{"success":false,"code":"%s","data":{"getResultUrl":"/queries/v1/result"}}`, queryInProgressCode) + return + } + polls++ + if polls < 2 { + fmt.Fprintf(w, `{"success":false,"code":"%s","data":{"getResultUrl":"/queries/v1/result"}}`, queryInProgressCode) + return + } + fmt.Fprint(w, `{"success":true,"data":{"queryResultFormat":"json","rowtype":[{"name":"N"}],"rowset":[["1"]]}}`) + })) + t.Cleanup(server.Close) + + client := newUpstream(SnowflakeProxyConfig{Account: "acme-test"}) + client.baseURL = server.URL + + result, err := client.query(t.Context(), "SELECT long_running()", nil) + + require.NoError(t, err) + require.Equal(t, 2, polls) + require.Len(t, result.rows, 1) +} + +// Each client connection gets its own proxy, so a token minted by one must be accepted by the next +func TestSessionTokenIsStableAcrossConnectionsAndIndependentOfCredentials(t *testing.T) { + expiry := time.Now().Add(time.Hour) + config := SnowflakeProxyConfig{SessionID: "session-1", Password: "hunter2", SessionExpiry: expiry} + + first, err := sessionToken(config) + require.NoError(t, err) + again, err := sessionToken(SnowflakeProxyConfig{SessionID: "session-1", Password: "different", SessionExpiry: expiry}) + require.NoError(t, err) + require.Equal(t, first, again) + require.NotContains(t, first, "hunter2") + + other, err := sessionToken(SnowflakeProxyConfig{SessionID: "session-2", SessionExpiry: expiry}) + require.NoError(t, err) + require.NotEqual(t, first, other) +} + +// An expired session's secret must not linger in the gateway once the session is over +func TestSessionSecretsAreSweptOnceExpired(t *testing.T) { + _, err := sessionToken(SnowflakeProxyConfig{SessionID: "stale", SessionExpiry: time.Now().Add(-time.Minute)}) + require.NoError(t, err) + + _, err = sessionToken(SnowflakeProxyConfig{SessionID: "fresh", SessionExpiry: time.Now().Add(time.Hour)}) + require.NoError(t, err) + + _, present := sessionSecrets.Load("stale") + require.False(t, present) +} + +// The JDBC driver reads AUTOCOMMIT off the login response and panics when it is missing +func TestLoginForwardsSnowflakeSessionParameters(t *testing.T) { + params := json.RawMessage(`[{"name":"AUTOCOMMIT","value":true}]`) + + data := loginData("token", SnowflakeProxyConfig{}, nil, params) + require.Equal(t, params, data["parameters"]) + + require.Equal(t, json.RawMessage("[]"), loginData("token", SnowflakeProxyConfig{}, nil, nil)["parameters"]) +} + +func TestGzippedRequestIsBoundedAfterDecompression(t *testing.T) { + var payload bytes.Buffer + writer := gzip.NewWriter(&payload) + _, err := writer.Write([]byte(`{"sqlText":"` + strings.Repeat("A", 4<<20) + `"}`)) + require.NoError(t, err) + require.NoError(t, writer.Close()) + require.Less(t, payload.Len(), maxRequestBytes) + + req := httptest.NewRequest(http.MethodPost, "/queries/v1/query-request", bytes.NewReader(payload.Bytes())) + req.Header.Set("Content-Encoding", "gzip") + + var out struct { + SqlText string `json:"sqlText"` + } + require.Error(t, decodeRequest(httptest.NewRecorder(), req, &out)) +} + +func TestUpstreamFailureAndTotalArePreserved(t *testing.T) { + recorder := httptest.NewRecorder() + writeUpstreamFailure(recorder, &snowflakeError{code: "002003", sqlState: "42S02", queryID: "q-1", message: "no such table"}) + + var body map[string]any + require.NoError(t, json.NewDecoder(recorder.Body).Decode(&body)) + require.Equal(t, "002003", body["code"]) + require.Equal(t, "no such table", body["message"]) + require.Equal(t, "42S02", body["data"].(map[string]any)["sqlState"]) + + truncated := &queryResult{rows: make([][]any, 10), total: 12000, truncated: true} + require.EqualValues(t, 12000, truncated.data()["total"]) + require.Equal(t, 10, truncated.data()["returned"]) +} + +// An unknown auth method must fail rather than quietly attempting a password login +func TestLoginRejectsAnUnknownAuthMethod(t *testing.T) { + client := newUpstream(SnowflakeProxyConfig{Account: "acme-test", AuthMethod: "saml"}) + + err := client.login(t.Context()) + + require.ErrorContains(t, err, "unsupported Snowflake authentication method") +} diff --git a/packages/pam/handlers/snowflake/rest.go b/packages/pam/handlers/snowflake/rest.go new file mode 100644 index 00000000..6e37aadd --- /dev/null +++ b/packages/pam/handlers/snowflake/rest.go @@ -0,0 +1,301 @@ +package snowflake + +import ( + "compress/gzip" + "crypto/rand" + "crypto/rsa" + "crypto/sha256" + "crypto/x509" + "encoding/base64" + "encoding/binary" + "encoding/json" + "encoding/pem" + "errors" + "fmt" + "io" + "net" + "net/http" + "strings" + "sync" + "time" + + "github.com/rs/zerolog" + "github.com/youmark/pkcs8" +) + +type zeroLogger = zerolog.Logger + +const ( + errCodeAuth = "390114" + errCodeStatement = "000904" + errCodeUnsupported = "000002" + sessionExpiredCode = "390112" + queryInProgressCode = "333333" + queryInProgressAsyncCode = "333334" + + maxRequestBytes = 1 << 20 + maxResponseBytes = 128 << 20 + // A ceiling on what one statement can pull into gateway memory + maxRows = 10000 +) + +type envelope struct { + Data any `json:"data"` + Message *string `json:"message"` + Code *string `json:"code"` + Success bool `json:"success"` +} + +type column struct { + Name string `json:"name"` + Type string `json:"type"` + Length int `json:"length"` + Precision *int `json:"precision"` + Scale *int `json:"scale"` + Nullable bool `json:"nullable"` + ByteLength int `json:"byteLength"` + Database string `json:"database"` + Schema string `json:"schema"` + Table string `json:"table"` + Collation *string `json:"collation"` +} + +func writeEnvelope(w http.ResponseWriter, body envelope) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(body) +} + +// http.MaxBytesReader gives back an opaque read error, which reads like a syntax problem otherwise +func statementReadError(err error) string { + var tooLarge *http.MaxBytesError + if errors.As(err, &tooLarge) { + return fmt.Sprintf("The statement is larger than the %d MB a proxied session accepts.", maxRequestBytes>>20) + } + return "Could not read the statement" +} + +func writeFailure(w http.ResponseWriter, code string, message string) { + writeEnvelope(w, envelope{Data: map[string]any{}, Message: &message, Code: &code}) +} + +func writeUpstreamFailure(w http.ResponseWriter, err error) { + var sfErr *snowflakeError + if !errors.As(err, &sfErr) || sfErr.code == "" { + writeFailure(w, errCodeStatement, err.Error()) + return + } + message := sfErr.message + writeEnvelope(w, envelope{ + Data: map[string]any{"sqlState": sfErr.sqlState, "queryId": sfErr.queryID}, + Message: &message, + Code: &sfErr.code, + }) +} + +// The drivers gzip every request body, so the encoding has to be honoured before parsing. +func decodeRequest(w http.ResponseWriter, r *http.Request, out any) error { + var reader io.Reader = http.MaxBytesReader(w, r.Body, maxRequestBytes) + + if r.Header.Get("Content-Encoding") == "gzip" { + gzipReader, err := gzip.NewReader(reader) + if err != nil { + return err + } + defer gzipReader.Close() + // MaxBytesReader caps the compressed bytes + reader = io.LimitReader(gzipReader, maxRequestBytes) + } + + return json.NewDecoder(reader).Decode(out) +} + +// Built from scratch rather than cloned from http.DefaultTransport, which the gateway has replaced. +func snowflakeTransport() *http.Transport { + return &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext, + MaxIdleConns: 10, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } +} + +type sessionSecret struct { + value string + expires time.Time +} + +// Every client connection opens its own tunnel and so its own proxy, so the token has to be shared across them +var sessionSecrets sync.Map + +func sessionToken(config SnowflakeProxyConfig) (string, error) { + now := time.Now() + sessionSecrets.Range(func(key, value any) bool { + if secret, ok := value.(sessionSecret); ok && secret.expires.Before(now) { + sessionSecrets.Delete(key) + } + return true + }) + + if existing, ok := sessionSecrets.Load(config.SessionID); ok { + return existing.(sessionSecret).value, nil + } + + raw := make([]byte, 32) + if _, err := rand.Read(raw); err != nil { + return "", fmt.Errorf("could not generate a session token: %w", err) + } + stored, _ := sessionSecrets.LoadOrStore(config.SessionID, sessionSecret{ + value: base64.RawURLEncoding.EncodeToString(raw), + expires: config.SessionExpiry, + }) + return stored.(sessionSecret).value, nil +} + +// Clients read session parameters such as AUTOCOMMIT off this response and some panic when they are +// missing, so Snowflake's own list is forwarded rather than dropped. +func loginData(token string, config SnowflakeProxyConfig, sessionCtx map[string]string, parameters json.RawMessage) map[string]any { + if len(parameters) == 0 { + parameters = json.RawMessage("[]") + } + return map[string]any{ + "token": token, + "masterToken": token, + "validityInSeconds": 3600, + "masterValidityInSeconds": 3600, + "displayUserName": config.Username, + "serverVersion": "8.0.0", + "firstLogin": false, + "healthCheckInterval": 45, + "newClientForUpgrade": nil, + // A client on the official Go driver decodes this as an int64, so the PAM session's own UUID + // can't be handed over as-is + "sessionId": numericSessionID(config.SessionID), + "parameters": parameters, + "sessionInfo": sessionCtx, + } +} + +func numericSessionID(sessionID string) int64 { + digest := sha256.Sum256([]byte(sessionID)) + return int64(binary.BigEndian.Uint64(digest[:8]) &^ (1 << 63)) +} + +func renewData(token string) map[string]any { + return map[string]any{ + "sessionToken": token, + "masterToken": token, + "validityInSecondsST": 3600, + "validityInSecondsMT": 3600, + "masterValidityInSeconds": 3600, + } +} + +type queryResult struct { + columns []column + rows [][]any + total int64 + queryID string + truncated bool + elapsed time.Duration +} + +func (r *queryResult) summary() string { + summary := fmt.Sprintf("%d row(s), %dms", len(r.rows), r.elapsed.Milliseconds()) + if r.truncated { + summary = fmt.Sprintf("%s, truncated at %d", summary, maxRows) + } + return summary +} + +func (r *queryResult) data() map[string]any { + return map[string]any{ + "rowtype": r.columns, + "rowset": r.rows, + // total is what the statement produced; returned is what survived the row cap + "total": max(r.total, int64(len(r.rows))), + "returned": len(r.rows), + "queryId": r.queryID, + "queryResultFormat": "json", + "parameters": []any{}, + } +} + +func parsePrivateKey(privateKeyPem string, passphrase string) (*rsa.PrivateKey, error) { + block, _ := pem.Decode([]byte(privateKeyPem)) + if block == nil { + return nil, errors.New("the stored private key is not valid PEM") + } + + // Driven by the key itself, not by whether a passphrase happens to be stored, so a passphrase left + // over from an earlier credential doesn't break a plain key + if strings.Contains(block.Type, "ENCRYPTED") { + key, err := pkcs8.ParsePKCS8PrivateKeyRSA(block.Bytes, []byte(passphrase)) + if err != nil { + return nil, fmt.Errorf("could not decrypt the private key: %w", err) + } + return key, nil + } + + parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes) + if err != nil { + return nil, fmt.Errorf("could not read the private key: %w", err) + } + + key, ok := parsed.(*rsa.PrivateKey) + if !ok { + return nil, errors.New("Snowflake key pair authentication requires an RSA private key") + } + return key, nil +} + +type singleConnListener struct { + conns chan net.Conn + closed chan struct{} + once sync.Once +} + +func newSingleConnListener(conn net.Conn) *singleConnListener { + listener := &singleConnListener{conns: make(chan net.Conn, 1), closed: make(chan struct{})} + // http.Server closes the connection when the client hangs up, which is what ends the session and + // lets the deferred session-logger close flush the recording + listener.conns <- &closeNotifyConn{Conn: conn, onClose: listener.Close} + return listener +} + +type closeNotifyConn struct { + net.Conn + onClose func() error + once sync.Once +} + +func (c *closeNotifyConn) Close() error { + err := c.Conn.Close() + c.once.Do(func() { _ = c.onClose() }) + return err +} + +func (l *singleConnListener) Accept() (net.Conn, error) { + select { + case conn := <-l.conns: + return conn, nil + case <-l.closed: + return nil, net.ErrClosed + } +} + +func (l *singleConnListener) Close() error { + l.once.Do(func() { close(l.closed) }) + return nil +} + +func (l *singleConnListener) Addr() net.Addr { + return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)} +} + +func isListenerDone(err error) bool { + return errors.Is(err, net.ErrClosed) || errors.Is(err, http.ErrServerClosed) +} diff --git a/packages/pam/handlers/snowflake/upstream.go b/packages/pam/handlers/snowflake/upstream.go new file mode 100644 index 00000000..e2467b9f --- /dev/null +++ b/packages/pam/handlers/snowflake/upstream.go @@ -0,0 +1,450 @@ +package snowflake + +import ( + "bytes" + "compress/gzip" + "context" + "crypto/sha256" + "crypto/x509" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "runtime" + "strings" + "sync" + "time" + + "github.com/golang-jwt/jwt/v5" + "github.com/google/uuid" +) + +// The gateway's own client for Snowflake's REST API. Only login and query are implemented, which is all +// a proxied session does; stage file transfer and the Arrow result format are deliberately unsupported. +type upstream struct { + config SnowflakeProxyConfig + client *http.Client + baseURL string + mu sync.Mutex + token string + masterTok string + requestID string + sessionCtx map[string]string + parameters json.RawMessage +} + +type upstreamEnvelope struct { + Data json.RawMessage `json:"data"` + Message string `json:"message"` + Code string `json:"code"` + Success bool `json:"success"` +} + +type loginResponse struct { + Token string `json:"token"` + MasterToken string `json:"masterToken"` + Parameters json.RawMessage `json:"parameters"` + SessionInfo struct { + DatabaseName string `json:"databaseName"` + SchemaName string `json:"schemaName"` + WarehouseName string `json:"warehouseName"` + RoleName string `json:"roleName"` + } `json:"sessionInfo"` +} + +type resultChunk struct { + URL string `json:"url"` + RowCount int `json:"rowCount"` +} + +type queryResponse struct { + GetResultURL string `json:"getResultUrl"` + QueryID string `json:"queryId"` + RowType []column `json:"rowtype"` + RowSet [][]any `json:"rowset"` + Total int64 `json:"total"` + Chunks []resultChunk `json:"chunks"` + ChunkHeaders map[string]string `json:"chunkHeaders"` + QueryResultFormat string `json:"queryResultFormat"` +} + +func newUpstream(config SnowflakeProxyConfig) *upstream { + return &upstream{ + config: config, + baseURL: fmt.Sprintf("https://%s.snowflakecomputing.com", config.Account), + client: &http.Client{Transport: snowflakeTransport(), Timeout: statementTimeout}, + } +} + +func (u *upstream) close() { + if transport, ok := u.client.Transport.(*http.Transport); ok { + transport.CloseIdleConnections() + } +} + +func (u *upstream) login(ctx context.Context) error { + data := map[string]any{ + "CLIENT_APP_ID": "Go", + "CLIENT_APP_VERSION": "1.19.1", + "SVN_REVISION": "", + "ACCOUNT_NAME": u.config.Account, + "LOGIN_NAME": u.config.Username, + "CLIENT_ENVIRONMENT": map[string]any{ + "APPLICATION": "Infisical", + "APPLICATION_PATH": "", + "OS": runtime.GOOS, + "OS_VERSION": runtime.GOARCH, + "OCSP_MODE": "FAIL_OPEN", + "GO_VERSION": runtime.Version(), + }, + // JSON keeps results parseable without the Arrow decoder the official driver pulls in + "SESSION_PARAMETERS": map[string]any{"CLIENT_RESULT_FORMAT": "JSON", "GO_QUERY_RESULT_FORMAT": "json"}, + } + + switch u.config.AuthMethod { + case AuthMethodKeyPair: + assertion, err := u.keyPairAssertion() + if err != nil { + return err + } + data["AUTHENTICATOR"] = "SNOWFLAKE_JWT" + data["TOKEN"] = assertion + case AuthMethodToken: + data["AUTHENTICATOR"] = "PROGRAMMATIC_ACCESS_TOKEN" + data["TOKEN"] = u.config.Token + data["PASSWORD"] = u.config.Token + case AuthMethodPassword: + data["PASSWORD"] = u.config.Password + default: + return fmt.Errorf("unsupported Snowflake authentication method %q", u.config.AuthMethod) + } + + params := url.Values{} + for name, value := range map[string]string{ + "databaseName": u.config.Database, + "schemaName": u.config.Schema, + "warehouse": u.config.Warehouse, + "roleName": u.config.Role, + } { + if value != "" { + params.Set(name, value) + } + } + + envelope, err := u.post(ctx, "/session/v1/login-request", params, map[string]any{"data": data}, false) + if err != nil { + return err + } + + var parsed loginResponse + if err := json.Unmarshal(envelope.Data, &parsed); err != nil { + return fmt.Errorf("could not read the Snowflake login response: %w", err) + } + + u.token, u.masterTok, u.parameters = parsed.Token, parsed.MasterToken, parsed.Parameters + u.sessionCtx = map[string]string{ + "databaseName": parsed.SessionInfo.DatabaseName, + "schemaName": parsed.SessionInfo.SchemaName, + "warehouseName": parsed.SessionInfo.WarehouseName, + "roleName": parsed.SessionInfo.RoleName, + } + return nil +} + +// Snowflake identifies a key pair by the SHA-256 fingerprint of its public key, carried in the issuer. +func (u *upstream) keyPairAssertion() (string, error) { + key, err := parsePrivateKey(u.config.PrivateKey, u.config.PrivateKeyPass) + if err != nil { + return "", err + } + + der, err := x509.MarshalPKIXPublicKey(&key.PublicKey) + if err != nil { + return "", fmt.Errorf("could not read the public key: %w", err) + } + digest := sha256.Sum256(der) + fingerprint := "SHA256:" + base64.StdEncoding.EncodeToString(digest[:]) + + account := strings.ToUpper(strings.Split(u.config.Account, ".")[0]) + subject := fmt.Sprintf("%s.%s", account, strings.ToUpper(u.config.Username)) + + return jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.RegisteredClaims{ + Issuer: fmt.Sprintf("%s.%s", subject, fingerprint), + Subject: subject, + IssuedAt: jwt.NewNumericDate(time.Now()), + ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour)), + }).SignedString(key) +} + +// Bindings are passed through untouched: a prepared statement from a JDBC client carries its parameters +// here, and dropping them makes every parameterised query fail. +func (u *upstream) query(ctx context.Context, statement string, bindings json.RawMessage) (*queryResult, error) { + started := time.Now() + body := map[string]any{"sqlText": statement, "asyncExec": false, "sequenceId": 1, "isInternal": false} + if len(bindings) > 0 && string(bindings) != "null" { + body["bindings"] = bindings + } + + // Kept so a cancellation can name this statement to Snowflake + requestID := uuid.NewString() + params := url.Values{} + params.Set("requestId", requestID) + u.mu.Lock() + u.requestID = requestID + u.mu.Unlock() + + envelope, err := u.post(ctx, "/queries/v1/query-request", params, body, true) + // Snowflake expires a session token roughly hourly, well inside a long PAM session + if isSessionExpired(err) { + if renewErr := u.renew(ctx); renewErr != nil { + return nil, renewErr + } + envelope, err = u.post(ctx, "/queries/v1/query-request", params, body, true) + } + if err != nil { + return nil, err + } + + parsed, err := decodeQueryResponse(envelope) + if err != nil { + return nil, err + } + + // A query slower than about 45 seconds comes back "in progress" with a URL to poll until it lands + for envelope.Code == queryInProgressCode || envelope.Code == queryInProgressAsyncCode { + if parsed.GetResultURL == "" { + return nil, fmt.Errorf("Snowflake reported the query as running but gave no result URL") + } + if envelope, err = u.get(ctx, parsed.GetResultURL); err != nil { + return nil, err + } + if parsed, err = decodeQueryResponse(envelope); err != nil { + return nil, err + } + } + + if format := strings.ToLower(parsed.QueryResultFormat); format != "" && format != "json" { + return nil, fmt.Errorf("Snowflake returned an unsupported result format (%s)", format) + } + + result := &queryResult{columns: parsed.RowType, rows: parsed.RowSet, queryID: parsed.QueryID} + if result.rows == nil { + result.rows = [][]any{} + } + + // Rows past the first slice live in pre-signed object storage. The gateway fetches them so the data + // stays on its side of the tunnel rather than flowing straight to the client. + for _, chunk := range parsed.Chunks { + if len(result.rows) >= maxRows { + break + } + rows, err := u.fetchChunk(ctx, chunk, parsed.ChunkHeaders) + if err != nil { + return nil, err + } + result.rows = append(result.rows, rows...) + } + + if len(result.rows) > maxRows { + result.rows = result.rows[:maxRows] + } + result.total = parsed.Total + result.truncated = int64(len(result.rows)) < parsed.Total + result.elapsed = time.Since(started) + return result, nil +} + +// Snowflake keeps running a statement the client walked away from, and the warehouse bills for it. +func (u *upstream) abortQuery(ctx context.Context) error { + u.mu.Lock() + requestID := u.requestID + u.mu.Unlock() + if requestID == "" { + return nil + } + + _, err := u.post(ctx, "/queries/v1/abort-request", url.Values{}, map[string]string{"requestId": requestID}, true) + return err +} + +func (u *upstream) logout(ctx context.Context) error { + params := url.Values{} + params.Set("delete", "true") + _, err := u.post(ctx, "/session", params, map[string]any{}, true) + return err +} + +func decodeQueryResponse(envelope *upstreamEnvelope) (*queryResponse, error) { + var parsed queryResponse + if err := json.Unmarshal(envelope.Data, &parsed); err != nil { + return nil, fmt.Errorf("could not read the Snowflake response: %w", err) + } + return &parsed, nil +} + +func (u *upstream) get(ctx context.Context, path string) (*upstreamEnvelope, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.baseURL+path, nil) + if err != nil { + return nil, err + } + u.mu.Lock() + token := u.token + u.mu.Unlock() + req.Header.Set("Authorization", fmt.Sprintf(`Snowflake Token="%s"`, token)) + req.Header.Set("Accept", "application/json") + + return u.send(req) +} + +func (u *upstream) fetchChunk(ctx context.Context, chunk resultChunk, headers map[string]string) ([][]any, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, chunk.URL, nil) + if err != nil { + return nil, err + } + for name, value := range headers { + req.Header.Set(name, value) + } + + resp, err := u.client.Do(req) + if err != nil { + return nil, fmt.Errorf("could not fetch a result chunk: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("could not fetch a result chunk (HTTP %d)", resp.StatusCode) + } + + body, err := decompress(resp) + if err != nil { + return nil, err + } + + // A chunk is a bare list of rows rather than a whole envelope + var rows [][]any + if err := json.Unmarshal([]byte("["+strings.TrimSuffix(strings.TrimSpace(string(body)), ",")+"]"), &rows); err != nil { + return nil, fmt.Errorf("could not read a result chunk: %w", err) + } + return rows, nil +} + +func (u *upstream) renew(ctx context.Context) error { + u.mu.Lock() + master, old := u.masterTok, u.token + u.mu.Unlock() + + envelope, err := u.post(ctx, "/session/token-request", url.Values{}, + map[string]string{"oldSessionToken": old, "requestType": "RENEW"}, false, master) + if err != nil { + return fmt.Errorf("the Snowflake session expired and could not be renewed: %w", err) + } + + var parsed struct { + SessionToken string `json:"sessionToken"` + MasterToken string `json:"masterToken"` + } + if err := json.Unmarshal(envelope.Data, &parsed); err != nil { + return fmt.Errorf("could not read the Snowflake renewal response: %w", err) + } + + u.mu.Lock() + u.token, u.masterTok = parsed.SessionToken, parsed.MasterToken + u.mu.Unlock() + return nil +} + +type snowflakeError struct{ code, sqlState, queryID, message string } + +func (e *snowflakeError) Error() string { return e.message } + +func isSessionExpired(err error) bool { + var sfErr *snowflakeError + return errors.As(err, &sfErr) && sfErr.code == sessionExpiredCode +} + +// Snowflake rejects a request without its own identifiers, so they are set on every call. +func (u *upstream) post(ctx context.Context, path string, params url.Values, body any, authenticated bool, asToken ...string) (*upstreamEnvelope, error) { + encoded, err := json.Marshal(body) + if err != nil { + return nil, err + } + + if params.Get("requestId") == "" { + params.Set("requestId", uuid.NewString()) + } + params.Set("request_guid", uuid.NewString()) + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.baseURL+path+"?"+params.Encode(), bytes.NewReader(encoded)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + switch { + case len(asToken) > 0: + req.Header.Set("Authorization", fmt.Sprintf(`Snowflake Token="%s"`, asToken[0])) + case authenticated: + u.mu.Lock() + token := u.token + u.mu.Unlock() + req.Header.Set("Authorization", fmt.Sprintf(`Snowflake Token="%s"`, token)) + } + + return u.send(req) +} + +func (u *upstream) send(req *http.Request) (*upstreamEnvelope, error) { + resp, err := u.client.Do(req) + if err != nil { + return nil, fmt.Errorf("could not reach Snowflake: %w", err) + } + defer resp.Body.Close() + + raw, err := decompress(resp) + if err != nil { + return nil, err + } + + var envelope upstreamEnvelope + if err := json.Unmarshal(raw, &envelope); err != nil { + return nil, fmt.Errorf("Snowflake returned an unreadable response (HTTP %d)", resp.StatusCode) + } + // An in-progress response is not a failure; the caller polls on it + if !envelope.Success && envelope.Code != queryInProgressCode && envelope.Code != queryInProgressAsyncCode { + message := strings.TrimSpace(envelope.Message) + if message == "" { + message = fmt.Sprintf("Snowflake refused the request (HTTP %d)", resp.StatusCode) + } + var detail struct { + SQLState string `json:"sqlState"` + QueryID string `json:"queryId"` + } + _ = json.Unmarshal(envelope.Data, &detail) + return nil, &snowflakeError{code: envelope.Code, sqlState: detail.SQLState, queryID: detail.QueryID, message: message} + } + return &envelope, nil +} + +func decompress(resp *http.Response) ([]byte, error) { + var reader io.Reader = resp.Body + if resp.Header.Get("Content-Encoding") == "gzip" { + gzipReader, err := gzip.NewReader(resp.Body) + if err != nil { + return nil, err + } + defer gzipReader.Close() + reader = gzipReader + } + + body, err := io.ReadAll(io.LimitReader(reader, maxResponseBytes+1)) + if err != nil { + return nil, err + } + if len(body) > maxResponseBytes { + return nil, errors.New("the Snowflake response was too large to read safely") + } + return body, nil +} diff --git a/packages/pam/local/access.go b/packages/pam/local/access.go index 1a681d1e..f7557f39 100644 --- a/packages/pam/local/access.go +++ b/packages/pam/local/access.go @@ -27,6 +27,7 @@ const ( AccountTypeMongoDB = "mongodb" AccountTypeOracleDB = "oracledb" AccountTypeRedis = "redis" + AccountTypeSnowflake = "snowflake" AccountTypeKubernetes = "kubernetes" AccountTypeAwsIam = "aws-iam" AccountTypeGcpServiceAccount = "gcp-service-account" @@ -133,6 +134,8 @@ func StartPAMAccess(accessToken string, opts AccessOptions) { startSSHAccess(httpClient, &pamResponse, displayPath, opts) case AccountTypeRedis: startRedisProxy(httpClient, &pamResponse, displayPath, durationStr, port) + case AccountTypeSnowflake: + startSnowflakeProxy(httpClient, &pamResponse, displayPath, durationStr, port) case AccountTypeKubernetes: startKubernetesProxy(httpClient, &pamResponse, displayPath, durationStr, port) case AccountTypeAwsIam: diff --git a/packages/pam/local/snowflake-proxy.go b/packages/pam/local/snowflake-proxy.go new file mode 100644 index 00000000..5f049252 --- /dev/null +++ b/packages/pam/local/snowflake-proxy.go @@ -0,0 +1,126 @@ +package pam + +import ( + "context" + "fmt" + "net/url" + "os" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/Infisical/infisical-merge/packages/api" + "github.com/Infisical/infisical-merge/packages/util" + "github.com/go-resty/resty/v2" + "github.com/rs/zerolog/log" +) + +// The local end is a plain pipe to the gateway, which authenticates, runs each statement and records it. +// Clients connect with TLS off, since the hop to the gateway is already an encrypted tunnel. +func startSnowflakeProxy(httpClient *resty.Client, response *api.PAMAccessResponse, path, durationStr string, port int) { + duration, err := time.ParseDuration(durationStr) + if err != nil { + util.HandleError(err, "Failed to parse duration") + return + } + + ctx, cancel := context.WithCancel(context.Background()) + + proxy := &DatabaseProxyServer{ + BaseProxyServer: BaseProxyServer{ + httpClient: httpClient, + relayHost: response.RelayHost, + relayClientCert: response.RelayClientCertificate, + relayClientKey: response.RelayClientPrivateKey, + relayServerCertChain: response.RelayServerCertificateChain, + gatewayClientCert: response.GatewayClientCertificate, + gatewayClientKey: response.GatewayClientPrivateKey, + gatewayServerCertChain: response.GatewayServerCertificateChain, + sessionExpiry: time.Now().Add(duration), + sessionId: response.SessionId, + resourceType: response.AccountType, + ctx: ctx, + cancel: cancel, + shutdownCh: make(chan struct{}), + }, + } + + if err := proxy.ValidateResourceTypeSupported(); err != nil { + util.HandleError(err, "Gateway version outdated") + return + } + + if err := proxy.Start(port); err != nil { + util.HandleError(err, "Failed to start proxy server") + return + } + + folder, account := parsePath(path) + log.Info().Msgf("Snowflake proxy server listening on port %d", proxy.port) + printSnowflakeSessionInfo(folder, account, duration, response.Metadata, proxy.port) + + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM) + go func() { + sig := <-sigChan + log.Info().Msgf("Received signal %v, initiating graceful shutdown...", sig) + proxy.gracefulShutdown() + }() + + proxy.Run() +} + +func printSnowflakeSessionInfo(folder, account string, duration time.Duration, metadata map[string]string, port int) { + target := metadata["account"] + if target == "" { + target = "" + } + + rule := "**********************************************************************\n" + divider := "----------------------------------------------------------------------\n" + + fmt.Print("\n" + rule) + fmt.Printf(" Snowflake Proxy Session Started! \n") + fmt.Print(rule + "\n") + if folder != "" { + fmt.Printf(" Folder: %s\n", folder) + } + fmt.Printf(" Account: %s\n", account) + fmt.Printf(" Duration: %s\n\n", duration.String()) + fmt.Print(divider) + fmt.Printf(" Connection Details \n") + fmt.Print(divider + "\n") + fmt.Printf(" Host: 127.0.0.1\n") + fmt.Printf(" Port: %d\n", port) + fmt.Printf(" Protocol: http\n") + for _, row := range [][2]string{{"Database", metadata["database"]}, {"Schema", metadata["schema"]}, {"Warehouse", metadata["warehouse"]}} { + if row[1] != "" { + fmt.Printf(" %-10s %s\n", row[0]+":", row[1]) + } + } + fmt.Print("\n" + divider) + fmt.Printf(" How to Connect \n") + fmt.Print(divider + "\n") + fmt.Printf(" Any username and password work; the gateway authenticates for you.\n\n") + fmt.Printf(" Connection string:\n") + util.PrintfStderr(" jdbc:snowflake://127.0.0.1:%d/?ssl=off&account=%s%s\n\n", port, target, + jdbcExtras(metadata["database"], metadata["schema"], metadata["warehouse"])) + fmt.Printf(" snowsql:\n") + util.PrintfStderr(" $ SNOWSQL_PWD=x snowsql -a %s -u pam -h 127.0.0.1 -p %d\n", target, port) + fmt.Printf("\n Press Ctrl+C to stop the proxy.\n\n") + fmt.Print(rule + "\n") +} + +func jdbcExtras(database, schema, warehouse string) string { + extras := make([]string, 0, 3) + for _, pair := range [][2]string{{"db", database}, {"schema", schema}, {"warehouse", warehouse}} { + if pair[1] != "" { + extras = append(extras, pair[0]+"="+url.QueryEscape(pair[1])) + } + } + if len(extras) == 0 { + return "" + } + return "&" + strings.Join(extras, "&") +} diff --git a/packages/pam/pam-proxy.go b/packages/pam/pam-proxy.go index 2ff2653a..003a532b 100644 --- a/packages/pam/pam-proxy.go +++ b/packages/pam/pam-proxy.go @@ -23,6 +23,7 @@ import ( "github.com/Infisical/infisical-merge/packages/pam/handlers/oracle" "github.com/Infisical/infisical-merge/packages/pam/handlers/rdp" "github.com/Infisical/infisical-merge/packages/pam/handlers/redis" + "github.com/Infisical/infisical-merge/packages/pam/handlers/snowflake" "github.com/Infisical/infisical-merge/packages/pam/handlers/ssh" "github.com/Infisical/infisical-merge/packages/pam/session" "github.com/Infisical/infisical-merge/packages/util" @@ -61,6 +62,7 @@ func GetSupportedResourceTypes() []string { session.ResourceTypeOracledb, session.ResourceTypeGcpServiceAccount, session.ResourceTypeAzureCli, + session.ResourceTypeSnowflake, } // Only advertise RDP when the real bridge is compiled in. A stub // build would otherwise accept RDP session routing and fail every @@ -517,6 +519,39 @@ func HandlePAMProxy(ctx context.Context, conn *tls.Conn, pamConfig *GatewayPAMCo Str("serviceAccountEmail", credentials.ServiceAccountEmail). Msg("Starting GCP Service Account PAM proxy") return proxy.HandleConnection(ctx, handlerConn) + case session.ResourceTypeSnowflake: + var blockedCommands []*regexp.Regexp + if credentials.PolicyRules != nil { + blockedCommands = compilePolicyPatterns(credentials.PolicyRules.CommandBlocking, pamConfig.SessionId, "command-blocking") + } + + proxy := snowflake.NewSnowflakeProxy(snowflake.SnowflakeProxyConfig{ + Account: credentials.Account, + Username: credentials.Username, + AuthMethod: credentials.AuthMethod, + Password: credentials.Password, + Token: credentials.Token, + PrivateKey: credentials.PrivateKey, + PrivateKeyPass: credentials.PrivateKeyPassphrase, + Warehouse: credentials.Warehouse, + Database: credentials.Database, + Schema: credentials.Schema, + Role: credentials.Role, + SessionID: pamConfig.SessionId, + SessionExpiry: pamConfig.ExpiryTime, + SessionLogger: sessionLogger, + BlockedCommands: blockedCommands, + }) + if err := proxy.Connect(ctx); err != nil { + return err + } + defer proxy.Close() + + log.Info(). + Str("sessionId", pamConfig.SessionId). + Str("account", credentials.Account). + Msg("Starting Snowflake PAM proxy") + return proxy.HandleConnection(ctx, handlerConn) case session.ResourceTypeAzureCli: azureConfig := azure.AzureProxyConfig{ Tokens: credentials.Tokens, diff --git a/packages/pam/session/credentials.go b/packages/pam/session/credentials.go index b255d9f2..49544e51 100644 --- a/packages/pam/session/credentials.go +++ b/packages/pam/session/credentials.go @@ -45,6 +45,11 @@ type PAMCredentials struct { Token string Tokens map[string]string ServiceAccountEmail string + Account string + Warehouse string + Schema string + Role string + PrivateKeyPassphrase string PolicyRules *api.PAMPolicyRules } @@ -204,6 +209,11 @@ func (cm *CredentialsManager) GetPAMSessionCredentials(sessionId string, expiryT Token: response.Credentials.Token, Tokens: response.Credentials.Tokens, ServiceAccountEmail: response.Credentials.ServiceAccountEmail, + Account: response.Credentials.Account, + Warehouse: response.Credentials.Warehouse, + Schema: response.Credentials.Schema, + Role: response.Credentials.Role, + PrivateKeyPassphrase: response.Credentials.PrivateKeyPassphrase, PolicyRules: response.PolicyRules, } diff --git a/packages/pam/session/uploader.go b/packages/pam/session/uploader.go index 42f4557c..630feff6 100644 --- a/packages/pam/session/uploader.go +++ b/packages/pam/session/uploader.go @@ -36,8 +36,15 @@ const ( ResourceTypeWindows = "windows" ResourceTypeGcpServiceAccount = "gcp-service-account" ResourceTypeAzureCli = "azure-cli" + ResourceTypeSnowflake = "snowflake" ) +var allResourceTypes = []string{ + ResourceTypeSSH, ResourceTypePostgres, ResourceTypeRedis, ResourceTypeMysql, ResourceTypeMssql, + ResourceTypeKubernetes, ResourceTypeMongodb, ResourceTypeOracledb, ResourceTypeWindows, + ResourceTypeGcpServiceAccount, ResourceTypeAzureCli, ResourceTypeSnowflake, +} + type SessionFileInfo struct { SessionID string ExpiresAt time.Time @@ -82,7 +89,7 @@ func NewSessionUploader(httpClient *resty.Client, credentialsManager *Credential func ParseSessionFilename(filename string) (*SessionFileInfo, error) { // Try new format first: pam_session_{sessionID}_{resourceType}_expires_{timestamp}.enc // Build regex pattern using constants - resourceTypePattern := fmt.Sprintf("(%s|%s|%s|%s|%s|%s|%s|%s|%s|%s|%s)", ResourceTypeSSH, ResourceTypePostgres, ResourceTypeRedis, ResourceTypeMysql, ResourceTypeMssql, ResourceTypeKubernetes, ResourceTypeMongodb, ResourceTypeOracledb, ResourceTypeWindows, ResourceTypeGcpServiceAccount, ResourceTypeAzureCli) + resourceTypePattern := fmt.Sprintf("(%s)", strings.Join(allResourceTypes, "|")) newFormatRegex := regexp.MustCompile(fmt.Sprintf(`^pam_session_(.+)_%s_expires_(\d+)\.enc$`, resourceTypePattern)) matches := newFormatRegex.FindStringSubmatch(filename) diff --git a/packages/pam/session/uploader_test.go b/packages/pam/session/uploader_test.go index 4432fdfa..8d4d25f0 100644 --- a/packages/pam/session/uploader_test.go +++ b/packages/pam/session/uploader_test.go @@ -171,3 +171,21 @@ func TestParseSessionFilename_Invalid(t *testing.T) { }) } } + +// A type missing from the pattern parses as a legacy filename, folding the type into the session ID +func TestParseSessionFilename_CoversEveryResourceType(t *testing.T) { + for _, resourceType := range allResourceTypes { + filename := "pam_session_abc-123_" + resourceType + "_expires_1700000000.enc" + + info, err := ParseSessionFilename(filename) + if err != nil { + t.Fatalf("%s: unexpected error: %v", resourceType, err) + } + if info.SessionID != "abc-123" { + t.Errorf("%s: session ID = %q, want abc-123", resourceType, info.SessionID) + } + if info.ResourceType != resourceType { + t.Errorf("%s: resource type = %q", resourceType, info.ResourceType) + } + } +}