From 6865c3e359069d8cf5c3b563bdf3901277ce0e5c Mon Sep 17 00:00:00 2001 From: Andrey Date: Wed, 16 Sep 2026 20:35:20 -0400 Subject: [PATCH 1/7] feat(pam): clickhouse --- .../gateway-v2/test_connection_handler.go | 37 ++ packages/pam/handlers/clickhouse/listener.go | 54 +++ packages/pam/handlers/clickhouse/proxy.go | 439 ++++++++++++++++++ .../pam/handlers/clickhouse/proxy_test.go | 301 ++++++++++++ packages/pam/local/access.go | 27 +- packages/pam/pam-proxy.go | 25 + packages/pam/session/uploader.go | 3 +- 7 files changed, 884 insertions(+), 2 deletions(-) create mode 100644 packages/pam/handlers/clickhouse/listener.go create mode 100644 packages/pam/handlers/clickhouse/proxy.go create mode 100644 packages/pam/handlers/clickhouse/proxy_test.go diff --git a/packages/gateway-v2/test_connection_handler.go b/packages/gateway-v2/test_connection_handler.go index 318d59cc..3395d401 100644 --- a/packages/gateway-v2/test_connection_handler.go +++ b/packages/gateway-v2/test_connection_handler.go @@ -21,6 +21,7 @@ import ( "sync" "time" + clickhousehandler "github.com/Infisical/infisical-merge/packages/pam/handlers/clickhouse" mssqlhandler "github.com/Infisical/infisical-merge/packages/pam/handlers/mssql" oraclehandler "github.com/Infisical/infisical-merge/packages/pam/handlers/oracle" snowflakehandler "github.com/Infisical/infisical-merge/packages/pam/handlers/snowflake" @@ -62,6 +63,7 @@ const ( testConnModeKubernetes = "kubernetes" testConnModeSSH = "ssh" testConnModeSnowflake = "snowflake" + testConnModeClickhouse = "clickhouse" testConnModeTCP = "tcp" ) @@ -118,6 +120,15 @@ type snowflakeTestParams struct { Role string `json:"role"` } +type clickhouseTestParams struct { + Username string `json:"username"` + Password string `json:"password"` + Database string `json:"database"` + SslEnabled bool `json:"sslEnabled"` + SslRejectUnauthorized *bool `json:"sslRejectUnauthorized"` + SslCertificate string `json:"sslCertificate"` +} + type ldapTestParams struct { Username string `json:"username"` Password string `json:"password"` @@ -689,6 +700,32 @@ func handleTestConnection(w http.ResponseWriter, r *http.Request) { defer proxy.Close() return authFailure(proxy.Probe(ctx)) } + case testConnModeClickhouse: + var params clickhouseTestParams + if !decode(¶ms) { + return + } + redactSecrets = append(redactSecrets, params.Password) + op = func() error { + var tlsConfig *tls.Config + if params.SslEnabled { + var err error + if tlsConfig, err = buildTestTLSConfig(target.host, params.SslCertificate, params.SslRejectUnauthorized); err != nil { + return connectFailure(err) + } + } + if err := dialTarget(ctx, target.host, target.port); err != nil { + return connectFailure(err) + } + return authFailure(clickhousehandler.TestConnection(ctx, clickhousehandler.ClickHouseProxyConfig{ + TargetAddr: net.JoinHostPort(target.host, strconv.Itoa(target.port)), + Username: params.Username, + Password: params.Password, + Database: params.Database, + EnableTLS: params.SslEnabled, + TLSConfig: tlsConfig, + })) + } case testConnModeSSH: var params sshTestParams if !decode(¶ms) { diff --git a/packages/pam/handlers/clickhouse/listener.go b/packages/pam/handlers/clickhouse/listener.go new file mode 100644 index 00000000..3aabe6b2 --- /dev/null +++ b/packages/pam/handlers/clickhouse/listener.go @@ -0,0 +1,54 @@ +package clickhouse + +import ( + "errors" + "net" + "net/http" + "sync" +) + +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{})} + 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/clickhouse/proxy.go b/packages/pam/handlers/clickhouse/proxy.go new file mode 100644 index 00000000..8c2987cd --- /dev/null +++ b/packages/pam/handlers/clickhouse/proxy.go @@ -0,0 +1,439 @@ +package clickhouse + +import ( + "bytes" + "compress/flate" + "compress/gzip" + "compress/zlib" + "context" + "crypto/tls" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httputil" + "regexp" + "strconv" + "strings" + "time" + + "github.com/Infisical/infisical-merge/packages/pam/session" + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" +) + +// Brokered over ClickHouse's HTTP interface rather than the native protocol on 9000, where a statement is +// text and can be blocked and recorded without parsing a version-gated binary format. The client's own +// credentials are dropped and the account's injected, so nothing it holds works outside a recorded session. +type ClickHouseProxyConfig struct { + TargetAddr string + Username string + Password string + Database string + EnableTLS bool + TLSConfig *tls.Config + + SessionID string + SessionLogger session.SessionLogger + BlockedCommands []*regexp.Regexp +} + +const ( + // How much of a body is decompressed and matched against the policy. A statement leads the body, so an + // INSERT's rows stream past unbuffered, but padding could bury one: a blocking account refuses the rest. + maxInspectBytes = 1 << 20 + // What of a statement reaches the recording, so one bulk INSERT can't fill a session log + maxLoggedStatementBytes = 8 << 10 + maxLoggedErrorBytes = 8 << 10 + + dialTimeout = 30 * time.Second + testQueryTimeout = 30 * time.Second +) + +// ClickHouse error codes, so a driver reports a gateway refusal the way it reports a server refusal +const ( + codeNotImplemented = 48 + codeNetworkError = 210 + codeAccessDenied = 497 +) + +var errorNames = map[int]string{ + codeNotImplemented: "NOT_IMPLEMENTED", + codeNetworkError: "NETWORK_ERROR", + codeAccessDenied: "ACCESS_DENIED", +} + +// Every way a client can present its own identity +var strippedAuthHeaders = []string{ + "Authorization", + "X-ClickHouse-User", + "X-ClickHouse-Key", + "X-ClickHouse-SSL-Certificate-Auth", +} + +var strippedAuthParams = []string{"user", "password"} + +type ClickHouseProxy struct { + config ClickHouseProxyConfig + reverse *httputil.ReverseProxy +} + +type stateKey struct{} + +type requestState struct { + statement string + started time.Time +} + +func NewClickHouseProxy(config ClickHouseProxyConfig) *ClickHouseProxy { + proxy := &ClickHouseProxy{config: config} + proxy.reverse = &httputil.ReverseProxy{ + Director: proxy.director, + Transport: newTransport(config), + ModifyResponse: proxy.modifyResponse, + ErrorHandler: proxy.handleUpstreamError, + // Progress and long-running results reach the client as they arrive + FlushInterval: -1, + } + return proxy +} + +func newTransport(config ClickHouseProxyConfig) *http.Transport { + return &http.Transport{ + Proxy: http.ProxyFromEnvironment, + DialContext: (&net.Dialer{Timeout: dialTimeout, KeepAlive: 30 * time.Second}).DialContext, + TLSClientConfig: config.TLSConfig, + MaxIdleConns: 10, + MaxIdleConnsPerHost: 10, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } +} + +func (p *ClickHouseProxy) scheme() string { + if p.config.EnableTLS { + return "https" + } + return "http" +} + +func (p *ClickHouseProxy) HandleConnection(ctx context.Context, clientConn net.Conn) error { + defer clientConn.Close() + + l := log.With().Str("sessionId", p.config.SessionID).Str("resourceType", "clickhouse").Logger() + + server := &http.Server{ + Handler: p.handler(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("ClickHouse proxy stopped") + } + return nil +} + +func (p *ClickHouseProxy) handler(l zerolog.Logger) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + statement, body, err := p.inspect(r) + if err != nil { + writeClickHouseError(w, http.StatusBadRequest, codeNotImplemented, err.Error()) + 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") + writeClickHouseError(w, http.StatusForbidden, codeAccessDenied, + "This statement is blocked by the command blocking policy on this account.") + return + } + + r.Body = body + state := &requestState{statement: statement, started: time.Now()} + p.reverse.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), stateKey{}, state))) + }) +} + +type bodyReadCloser struct { + io.Reader + io.Closer +} + +// Returns the statement and a body that still replays in full. ClickHouse concatenates `query` and the body. +func (p *ClickHouseProxy) inspect(r *http.Request) (string, io.ReadCloser, error) { + queryParam := strings.TrimSpace(r.URL.Query().Get("query")) + + if r.Body == nil || r.ContentLength == 0 { + return queryParam, http.NoBody, nil + } + + // ClickHouse's own block compression is opaque to anything but a ClickHouse client + if r.URL.Query().Get("decompress") == "1" { + return "", nil, fmt.Errorf( + "this session cannot read a ClickHouse-compressed request body, so decompress=1 is not supported here. " + + "Send the statement uncompressed or with Content-Encoding: gzip") + } + + encoding := strings.ToLower(strings.TrimSpace(r.Header.Get("Content-Encoding"))) + switch encoding { + case "", "identity", "gzip", "deflate": + default: + return "", nil, fmt.Errorf( + "this session cannot read a %q-encoded request body, so the command blocking policy could not be applied to it. "+ + "Use gzip, deflate, or no compression", encoding) + } + + // One byte past the window is what tells a body that fits from one that was cut short + head := make([]byte, maxInspectBytes+1) + n, err := io.ReadFull(r.Body, head) + if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF { + return "", nil, fmt.Errorf("the gateway could not read the request body: %v", err) + } + head = head[:n] + + if len(head) > maxInspectBytes && len(p.config.BlockedCommands) > 0 { + return "", nil, fmt.Errorf( + "this account blocks commands, so a request body larger than %d MB is refused: the gateway has to read "+ + "the whole statement to apply the policy. Send the data in smaller batches", + maxInspectBytes>>20) + } + + forwarded := bodyReadCloser{Reader: io.MultiReader(bytes.NewReader(head), r.Body), Closer: r.Body} + + decoded, decodeErr := decodeHead(head, encoding) + if decodeErr != nil { + return "", nil, fmt.Errorf( + "the gateway could not decompress the request body to apply the command blocking policy: %v", decodeErr) + } + + return joinStatement(queryParam, string(decoded)), forwarded, nil +} + +func joinStatement(queryParam string, body string) string { + body = strings.TrimSpace(body) + switch { + case queryParam == "": + return body + case body == "": + return queryParam + default: + return queryParam + "\n" + body + } +} + +func decodeHead(head []byte, encoding string) ([]byte, error) { + switch encoding { + case "gzip": + reader, err := gzip.NewReader(bytes.NewReader(head)) + if err != nil { + return nil, err + } + defer reader.Close() + return readTolerant(reader) + case "deflate": + // "deflate" is sent both as zlib and as raw deflate, so both are tried + if reader, err := zlib.NewReader(bytes.NewReader(head)); err == nil { + defer reader.Close() + if decoded, readErr := readTolerant(reader); readErr == nil { + return decoded, nil + } + } + reader := flate.NewReader(bytes.NewReader(head)) + defer reader.Close() + return readTolerant(reader) + default: + return head, nil + } +} + +// The head is a deliberate prefix, so a stream ending mid-frame is expected rather than an error +func readTolerant(r io.Reader) ([]byte, error) { + decoded, err := io.ReadAll(io.LimitReader(r, maxInspectBytes)) + if len(decoded) > 0 || err == nil || err == io.EOF { + return decoded, nil + } + return nil, err +} + +func (p *ClickHouseProxy) director(req *http.Request) { + req.URL.Scheme = p.scheme() + req.URL.Host = p.config.TargetAddr + req.Host = p.config.TargetAddr + + query := req.URL.Query() + for _, param := range strippedAuthParams { + query.Del(param) + } + if p.config.Database != "" && query.Get("database") == "" { + query.Set("database", p.config.Database) + } + req.URL.RawQuery = query.Encode() + + for _, header := range strippedAuthHeaders { + req.Header.Del(header) + } + req.Header.Set("X-ClickHouse-User", p.config.Username) + if p.config.Password != "" { + req.Header.Set("X-ClickHouse-Key", p.config.Password) + } + req.Header["X-Forwarded-For"] = nil +} + +func (p *ClickHouseProxy) modifyResponse(resp *http.Response) error { + state, ok := resp.Request.Context().Value(stateKey{}).(*requestState) + if !ok || state == nil { + return nil + } + elapsed := time.Since(state.started) + + if resp.StatusCode >= http.StatusBadRequest { + original := resp.Body + head, _ := io.ReadAll(io.LimitReader(original, maxLoggedErrorBytes)) + resp.Body = bodyReadCloser{Reader: io.MultiReader(bytes.NewReader(head), original), Closer: original} + p.logStatement(state.statement, fmt.Sprintf("ERROR: %s: %s", resp.Status, strings.TrimSpace(string(head)))) + return nil + } + + p.logStatement(state.statement, summarize(resp, elapsed)) + return nil +} + +func summarize(resp *http.Response, elapsed time.Duration) string { + parts := []string{resp.Status} + + if raw := resp.Header.Get("X-ClickHouse-Summary"); raw != "" { + var summary struct { + ReadRows string `json:"read_rows"` + WrittenRows string `json:"written_rows"` + ResultRows string `json:"result_rows"` + } + if json.Unmarshal([]byte(raw), &summary) == nil { + for _, counter := range [][2]string{ + {summary.ResultRows, "row(s) returned"}, + {summary.WrittenRows, "row(s) written"}, + {summary.ReadRows, "row(s) read"}, + } { + if counter[0] != "" && counter[0] != "0" { + parts = append(parts, counter[0]+" "+counter[1]) + } + } + } + } + + return strings.Join(append(parts, fmt.Sprintf("%dms", elapsed.Milliseconds())), ", ") +} + +func (p *ClickHouseProxy) handleUpstreamError(w http.ResponseWriter, r *http.Request, err error) { + log.Error().Err(err). + Str("sessionId", p.config.SessionID). + Str("path", r.URL.Path). + Msg("Failed to reach ClickHouse") + + if state, ok := r.Context().Value(stateKey{}).(*requestState); ok && state != nil { + p.logStatement(state.statement, fmt.Sprintf("ERROR: %s", err)) + } + + writeClickHouseError(w, http.StatusBadGateway, codeNetworkError, + fmt.Sprintf("The gateway could not reach ClickHouse: %v", err)) +} + +func (p *ClickHouseProxy) blockedBy(statement string) *regexp.Regexp { + if statement == "" { + return nil + } + for _, pattern := range p.config.BlockedCommands { + if pattern.MatchString(statement) { + return pattern + } + } + return nil +} + +func (p *ClickHouseProxy) logStatement(input string, output string) { + if p.config.SessionLogger == nil || strings.TrimSpace(input) == "" { + return + } + if len(input) > maxLoggedStatementBytes { + input = input[:maxLoggedStatementBytes] + "... [truncated]" + } + 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 ClickHouse statement") + } +} + +func writeClickHouseError(w http.ResponseWriter, status int, code int, message string) { + name, ok := errorNames[code] + if !ok { + name = "UNKNOWN" + } + body := fmt.Sprintf("Code: %d. DB::Exception: %s (%s)\n", code, message, name) + + w.Header().Set("Content-Type", "text/plain; charset=UTF-8") + w.Header().Set("X-ClickHouse-Exception-Code", strconv.Itoa(code)) + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) +} + +// Runs one statement as the account, so the test proves the login a session uses, not just an open port. +func TestConnection(ctx context.Context, config ClickHouseProxyConfig) error { + target := (&ClickHouseProxy{config: config}).scheme() + "://" + config.TargetAddr + "/" + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, target, strings.NewReader("SELECT 1")) + if err != nil { + return err + } + + query := req.URL.Query() + if config.Database != "" { + query.Set("database", config.Database) + } + req.URL.RawQuery = query.Encode() + + req.Header.Set("X-ClickHouse-User", config.Username) + if config.Password != "" { + req.Header.Set("X-ClickHouse-Key", config.Password) + } + + transport := newTransport(config) + defer transport.CloseIdleConnections() + + client := &http.Client{Transport: transport, Timeout: testQueryTimeout} + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + body, _ := io.ReadAll(io.LimitReader(resp.Body, maxLoggedErrorBytes)) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("clickhouse rejected the connection: %s", firstLine(string(body), resp.Status)) + } + return nil +} + +func firstLine(body string, fallback string) string { + trimmed := strings.TrimSpace(body) + if trimmed == "" { + return fallback + } + if idx := strings.IndexByte(trimmed, '\n'); idx != -1 { + return strings.TrimSpace(trimmed[:idx]) + } + return trimmed +} diff --git a/packages/pam/handlers/clickhouse/proxy_test.go b/packages/pam/handlers/clickhouse/proxy_test.go new file mode 100644 index 00000000..10754720 --- /dev/null +++ b/packages/pam/handlers/clickhouse/proxy_test.go @@ -0,0 +1,301 @@ +package clickhouse + +import ( + "bytes" + "compress/gzip" + "io" + "net/http" + "net/http/httptest" + "net/url" + "regexp" + "strings" + "testing" + + "github.com/Infisical/infisical-merge/packages/pam/session" + "github.com/rs/zerolog" + "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 } + +type capturedRequest struct { + method string + path string + query url.Values + headers http.Header + body []byte +} + +func newTestProxy(t *testing.T, upstream http.HandlerFunc, blocked ...string) (http.Handler, *recordingLogger, func()) { + t.Helper() + + server := httptest.NewServer(upstream) + + patterns := make([]*regexp.Regexp, 0, len(blocked)) + for _, pattern := range blocked { + patterns = append(patterns, regexp.MustCompile(pattern)) + } + + logger := &recordingLogger{} + proxy := NewClickHouseProxy(ClickHouseProxyConfig{ + TargetAddr: strings.TrimPrefix(server.URL, "http://"), + Username: "pam_svc", + Password: "s3cret", + Database: "analytics", + SessionID: "session-1", + SessionLogger: logger, + BlockedCommands: patterns, + }) + + return proxy.handler(zerolog.Nop()), logger, server.Close +} + +func capturingUpstream(captured *capturedRequest, respond func(http.ResponseWriter)) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + *captured = capturedRequest{ + method: r.Method, + path: r.URL.Path, + query: r.URL.Query(), + headers: r.Header.Clone(), + body: body, + } + if respond != nil { + respond(w) + return + } + w.WriteHeader(http.StatusOK) + } +} + +func gzipped(t *testing.T, payload string) []byte { + t.Helper() + var buffer bytes.Buffer + writer := gzip.NewWriter(&buffer) + _, err := writer.Write([]byte(payload)) + require.NoError(t, err) + require.NoError(t, writer.Close()) + return buffer.Bytes() +} + +func TestReplacesClientCredentialsWithTheAccountsOwn(t *testing.T) { + var captured capturedRequest + handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) + defer closeUpstream() + + req := httptest.NewRequest(http.MethodPost, "/?user=attacker&password=guessed&query=SELECT+1", http.NoBody) + req.Header.Set("Authorization", "Basic YXR0YWNrZXI6Z3Vlc3NlZA==") + req.Header.Set("X-ClickHouse-User", "attacker") + req.Header.Set("X-ClickHouse-Key", "guessed") + handler.ServeHTTP(httptest.NewRecorder(), req) + + require.Equal(t, "pam_svc", captured.headers.Get("X-ClickHouse-User")) + require.Equal(t, "s3cret", captured.headers.Get("X-ClickHouse-Key")) + require.Empty(t, captured.headers.Get("Authorization")) + require.Empty(t, captured.headers.Get("X-Forwarded-For")) + require.Empty(t, captured.query.Get("user")) + require.Empty(t, captured.query.Get("password")) +} + +func TestAppliesTheAccountDatabaseOnlyWhenTheClientNamesNone(t *testing.T) { + var captured capturedRequest + handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) + defer closeUpstream() + + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/?query=SELECT+1", http.NoBody)) + require.Equal(t, "analytics", captured.query.Get("database")) + + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/?database=other&query=SELECT+1", http.NoBody)) + require.Equal(t, "other", captured.query.Get("database")) +} + +func TestBlocksAStatementInTheQueryParameter(t *testing.T) { + reached := false + handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }, `(?i)\bdrop\b`) + defer closeUpstream() + + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/?query=DROP+TABLE+events", http.NoBody)) + + require.False(t, reached) + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Equal(t, "497", recorder.Header().Get("X-ClickHouse-Exception-Code")) + require.Contains(t, recorder.Body.String(), "Code: 497. DB::Exception:") + require.Contains(t, recorder.Body.String(), "(ACCESS_DENIED)") + + require.Len(t, logger.entries, 1) + require.Equal(t, "DROP TABLE events", logger.entries[0].Input) + require.Contains(t, logger.entries[0].Output, "BLOCKED:") +} + +func TestBlocksAStatementInsideAGzippedBody(t *testing.T) { + reached := false + handler, _, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }, `(?i)\btruncate\b`) + defer closeUpstream() + + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(gzipped(t, "TRUNCATE TABLE events"))) + req.Header.Set("Content-Encoding", "gzip") + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, req) + + require.False(t, reached) + require.Equal(t, http.StatusForbidden, recorder.Code) +} + +func TestJoinsTheQueryParameterAndTheBodyTheWayClickHouseDoes(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, capturingUpstream(&capturedRequest{}, nil)) + defer closeUpstream() + + req := httptest.NewRequest(http.MethodPost, "/?query=SELECT+count()+FROM", strings.NewReader("events WHERE id > 5")) + handler.ServeHTTP(httptest.NewRecorder(), req) + + require.Len(t, logger.entries, 1) + require.Equal(t, "SELECT count() FROM\nevents WHERE id > 5", logger.entries[0].Input) +} + +func TestRefusesAPaddedBodyThatWouldPushAStatementPastTheInspectionWindow(t *testing.T) { + reached := false + handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }, `(?i)\bdrop\b`) + defer closeUpstream() + + // A comment long enough to bury the statement behind the window the policy can see + payload := "/*" + strings.Repeat("x", maxInspectBytes) + "*/ DROP TABLE events" + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/", strings.NewReader(payload))) + + require.False(t, reached) + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), "blocks commands") + require.Empty(t, logger.entries) +} + +func TestForwardsABodyLargerThanTheInspectionCapWhenNothingIsBlocked(t *testing.T) { + var captured capturedRequest + handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) + defer closeUpstream() + + payload := "INSERT INTO events FORMAT JSONEachRow\n" + strings.Repeat("x", maxInspectBytes+1024) + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(payload)) + handler.ServeHTTP(httptest.NewRecorder(), req) + + require.Equal(t, len(payload), len(captured.body)) + require.Equal(t, payload, string(captured.body)) +} + +func TestTruncatesWhatALargeStatementWritesToTheRecording(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, capturingUpstream(&capturedRequest{}, nil)) + defer closeUpstream() + + payload := "INSERT INTO events VALUES " + strings.Repeat("a", maxLoggedStatementBytes*2) + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/", strings.NewReader(payload))) + + require.Len(t, logger.entries, 1) + require.True(t, strings.HasSuffix(logger.entries[0].Input, "... [truncated]")) + require.Less(t, len(logger.entries[0].Input), len(payload)) +} + +func TestRejectsABodyEncodingItCannotInspect(t *testing.T) { + reached := false + handler, _, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }) + defer closeUpstream() + + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader("SELECT 1")) + req.Header.Set("Content-Encoding", "zstd") + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, req) + + require.False(t, reached) + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), "command blocking policy") + + req = httptest.NewRequest(http.MethodPost, "/?decompress=1", strings.NewReader("SELECT 1")) + recorder = httptest.NewRecorder() + handler.ServeHTTP(recorder, req) + + require.False(t, reached) + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), "decompress=1") +} + +func TestRecordsTheRowCountsClickHouseReports(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.ReadAll(r.Body) + w.Header().Set("X-ClickHouse-Summary", `{"read_rows":"120","written_rows":"0","result_rows":"7"}`) + w.WriteHeader(http.StatusOK) + }) + defer closeUpstream() + + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/?query=SELECT+1", http.NoBody)) + + require.Len(t, logger.entries, 1) + require.Contains(t, logger.entries[0].Output, "7 row(s) returned") + require.Contains(t, logger.entries[0].Output, "120 row(s) read") + require.NotContains(t, logger.entries[0].Output, "row(s) written") +} + +func TestRecordsAnUpstreamErrorAndStillReturnsItToTheClient(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.ReadAll(r.Body) + w.Header().Set("X-ClickHouse-Exception-Code", "60") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte("Code: 60. DB::Exception: Table analytics.missing does not exist. (UNKNOWN_TABLE)\n")) + }) + defer closeUpstream() + + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/?query=SELECT+1+FROM+missing", http.NoBody)) + + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), "UNKNOWN_TABLE") + + require.Len(t, logger.entries, 1) + require.Contains(t, logger.entries[0].Output, "ERROR:") + require.Contains(t, logger.entries[0].Output, "UNKNOWN_TABLE") +} + +func TestRecordsNothingForARequestThatCarriesNoStatement(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, capturingUpstream(&capturedRequest{}, nil)) + defer closeUpstream() + + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/ping", http.NoBody)) + + require.Empty(t, logger.entries) +} + +func TestReportsAnUnreachableTargetAsAClickHouseError(t *testing.T) { + logger := &recordingLogger{} + proxy := NewClickHouseProxy(ClickHouseProxyConfig{ + // Port 1 on loopback refuses immediately + TargetAddr: "127.0.0.1:1", + Username: "pam_svc", + SessionID: "session-1", + SessionLogger: logger, + }) + + recorder := httptest.NewRecorder() + proxy.handler(zerolog.Nop()).ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/?query=SELECT+1", http.NoBody)) + + require.Equal(t, http.StatusBadGateway, recorder.Code) + require.Equal(t, "210", recorder.Header().Get("X-ClickHouse-Exception-Code")) + require.Contains(t, recorder.Body.String(), "(NETWORK_ERROR)") + require.Len(t, logger.entries, 1) + require.Contains(t, logger.entries[0].Output, "ERROR:") +} diff --git a/packages/pam/local/access.go b/packages/pam/local/access.go index 6b95d671..f795c7b6 100644 --- a/packages/pam/local/access.go +++ b/packages/pam/local/access.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "net/url" "os" "os/signal" "strings" @@ -29,6 +30,7 @@ const ( AccountTypeOracleDB = "oracledb" AccountTypeRedis = "redis" AccountTypeSnowflake = "snowflake" + AccountTypeClickhouse = "clickhouse" AccountTypeKubernetes = "kubernetes" AccountTypeAwsIam = "aws-iam" AccountTypeGcpServiceAccount = "gcp-service-account" @@ -128,7 +130,7 @@ func StartPAMAccess(accessToken string, opts AccessOptions) { // Route based on account type from API response switch pamResponse.AccountType { // Database types - all use the same proxy mechanism with different display configs - case AccountTypePostgres, AccountTypeMySQL, AccountTypeMsSQL, AccountTypeMongoDB, AccountTypeOracleDB: + case AccountTypePostgres, AccountTypeMySQL, AccountTypeMsSQL, AccountTypeMongoDB, AccountTypeOracleDB, AccountTypeClickhouse: startDatabaseProxy(httpClient, &pamResponse, displayPath, durationStr, port) case AccountTypeSSH: @@ -315,6 +317,9 @@ type AccountConnectionDisplay struct { // whose login protocol refuses an empty one. The proxy swaps the real credential in during login. Empty // for every type that accepts no password at all. RequiredPassword string + // Note is guidance shown under "How to Connect", one banner line per entry, for an account type + // whose port does not behave the way a client would assume. Empty for every type that needs none. + Note []string // ConnectionString builds the connection string, and is nil for account types that have none. ConnectionString func(username, database string, port int) string // UsageExamples builds sample CLI commands, and is nil for account types reached without one. @@ -388,6 +393,23 @@ var accountDisplays = map[string]AccountConnectionDisplay{ } }, }, + AccountTypeClickhouse: { + TypeLabel: "ClickHouse", + DefaultPort: 8123, + Note: []string{ + "This is ClickHouse's HTTP interface, so a client that only speaks the", + "native protocol, clickhouse-client included, cannot use this port.", + "JDBC, clickhouse-connect and curl all can.", + }, + ConnectionString: func(username, database string, port int) string { + return fmt.Sprintf("jdbc:clickhouse://127.0.0.1:%d/%s", port, database) + }, + UsageExamples: func(username, database string, port int) []string { + return []string{ + fmt.Sprintf("curl 'http://127.0.0.1:%d/?database=%s&query=SELECT+1'", port, url.QueryEscape(database)), + } + }, + }, AccountTypeRedis: { TypeLabel: "Redis", DefaultPort: 6379, @@ -956,6 +978,9 @@ func printDatabaseSessionInfo(config AccountConnectionDisplay, folder, account s } else { fmt.Printf(" to 127.0.0.1:%d. No password is needed.\n", port) } + for _, line := range config.Note { + fmt.Printf(" %s\n", line) + } fmt.Printf("\n") if examples := config.ConnectionExamples(username, database, port); len(examples) > 0 { fmt.Printf(" Example:\n") diff --git a/packages/pam/pam-proxy.go b/packages/pam/pam-proxy.go index 003a532b..e419c312 100644 --- a/packages/pam/pam-proxy.go +++ b/packages/pam/pam-proxy.go @@ -15,6 +15,7 @@ import ( "github.com/Infisical/infisical-merge/packages/api" "github.com/Infisical/infisical-merge/packages/pam/handlers" "github.com/Infisical/infisical-merge/packages/pam/handlers/azure" + "github.com/Infisical/infisical-merge/packages/pam/handlers/clickhouse" "github.com/Infisical/infisical-merge/packages/pam/handlers/gcp" "github.com/Infisical/infisical-merge/packages/pam/handlers/kubernetes" "github.com/Infisical/infisical-merge/packages/pam/handlers/mongodb" @@ -63,6 +64,7 @@ func GetSupportedResourceTypes() []string { session.ResourceTypeGcpServiceAccount, session.ResourceTypeAzureCli, session.ResourceTypeSnowflake, + session.ResourceTypeClickhouse, } // Only advertise RDP when the real bridge is compiled in. A stub // build would otherwise accept RDP session routing and fail every @@ -552,6 +554,29 @@ func HandlePAMProxy(ctx context.Context, conn *tls.Conn, pamConfig *GatewayPAMCo Str("account", credentials.Account). Msg("Starting Snowflake PAM proxy") return proxy.HandleConnection(ctx, handlerConn) + case session.ResourceTypeClickhouse: + var blockedCommands []*regexp.Regexp + if credentials.PolicyRules != nil { + blockedCommands = compilePolicyPatterns(credentials.PolicyRules.CommandBlocking, pamConfig.SessionId, "command-blocking") + } + + proxy := clickhouse.NewClickHouseProxy(clickhouse.ClickHouseProxyConfig{ + TargetAddr: fmt.Sprintf("%s:%d", credentials.Host, credentials.Port), + Username: credentials.Username, + Password: credentials.Password, + Database: credentials.Database, + EnableTLS: credentials.SSLEnabled, + TLSConfig: tlsConfig, + SessionID: pamConfig.SessionId, + SessionLogger: sessionLogger, + BlockedCommands: blockedCommands, + }) + log.Info(). + Str("sessionId", pamConfig.SessionId). + Str("target", fmt.Sprintf("%s:%d", credentials.Host, credentials.Port)). + Bool("sslEnabled", credentials.SSLEnabled). + Msg("Starting ClickHouse PAM proxy") + return proxy.HandleConnection(ctx, handlerConn) case session.ResourceTypeAzureCli: azureConfig := azure.AzureProxyConfig{ Tokens: credentials.Tokens, diff --git a/packages/pam/session/uploader.go b/packages/pam/session/uploader.go index 630feff6..b1a2752d 100644 --- a/packages/pam/session/uploader.go +++ b/packages/pam/session/uploader.go @@ -37,12 +37,13 @@ const ( ResourceTypeGcpServiceAccount = "gcp-service-account" ResourceTypeAzureCli = "azure-cli" ResourceTypeSnowflake = "snowflake" + ResourceTypeClickhouse = "clickhouse" ) var allResourceTypes = []string{ ResourceTypeSSH, ResourceTypePostgres, ResourceTypeRedis, ResourceTypeMysql, ResourceTypeMssql, ResourceTypeKubernetes, ResourceTypeMongodb, ResourceTypeOracledb, ResourceTypeWindows, - ResourceTypeGcpServiceAccount, ResourceTypeAzureCli, ResourceTypeSnowflake, + ResourceTypeGcpServiceAccount, ResourceTypeAzureCli, ResourceTypeSnowflake, ResourceTypeClickhouse, } type SessionFileInfo struct { From 847eb2fd0c8292ae9ea7bef8807e1dfe11e91bff Mon Sep 17 00:00:00 2001 From: Andrey Date: Wed, 16 Sep 2026 23:58:00 -0400 Subject: [PATCH 2/7] address bot reviews --- packages/pam/handlers/clickhouse/proxy.go | 83 ++++++++++++++----- .../pam/handlers/clickhouse/proxy_test.go | 47 ++++++++++- 2 files changed, 105 insertions(+), 25 deletions(-) diff --git a/packages/pam/handlers/clickhouse/proxy.go b/packages/pam/handlers/clickhouse/proxy.go index 8c2987cd..a91925ee 100644 --- a/packages/pam/handlers/clickhouse/proxy.go +++ b/packages/pam/handlers/clickhouse/proxy.go @@ -8,6 +8,7 @@ import ( "context" "crypto/tls" "encoding/json" + "errors" "fmt" "io" "net" @@ -16,6 +17,7 @@ import ( "regexp" "strconv" "strings" + "sync" "time" "github.com/Infisical/infisical-merge/packages/pam/session" @@ -131,8 +133,13 @@ func (p *ClickHouseProxy) HandleConnection(ctx context.Context, clientConn net.C listener := newSingleConnListener(clientConn) + done := make(chan struct{}) + defer close(done) go func() { - <-ctx.Done() + select { + case <-ctx.Done(): + case <-done: + } listener.Close() server.Close() }() @@ -202,21 +209,21 @@ func (p *ClickHouseProxy) inspect(r *http.Request) (string, io.ReadCloser, error } head = head[:n] - if len(head) > maxInspectBytes && len(p.config.BlockedCommands) > 0 { - return "", nil, fmt.Errorf( - "this account blocks commands, so a request body larger than %d MB is refused: the gateway has to read "+ - "the whole statement to apply the policy. Send the data in smaller batches", - maxInspectBytes>>20) - } - forwarded := bodyReadCloser{Reader: io.MultiReader(bytes.NewReader(head), r.Body), Closer: r.Body} - decoded, decodeErr := decodeHead(head, encoding) + decoded, decodedOverflow, decodeErr := decodeHead(head, encoding) if decodeErr != nil { return "", nil, fmt.Errorf( "the gateway could not decompress the request body to apply the command blocking policy: %v", decodeErr) } + if (len(head) > maxInspectBytes || decodedOverflow) && len(p.config.BlockedCommands) > 0 { + return "", nil, fmt.Errorf( + "this account blocks commands, so a request body larger than %d MB is refused: the gateway has to read "+ + "the whole statement to apply the policy. Send the data in smaller batches", + maxInspectBytes>>20) + } + return joinStatement(queryParam, string(decoded)), forwarded, nil } @@ -232,12 +239,12 @@ func joinStatement(queryParam string, body string) string { } } -func decodeHead(head []byte, encoding string) ([]byte, error) { +func decodeHead(head []byte, encoding string) ([]byte, bool, error) { switch encoding { case "gzip": reader, err := gzip.NewReader(bytes.NewReader(head)) if err != nil { - return nil, err + return nil, false, err } defer reader.Close() return readTolerant(reader) @@ -245,25 +252,29 @@ func decodeHead(head []byte, encoding string) ([]byte, error) { // "deflate" is sent both as zlib and as raw deflate, so both are tried if reader, err := zlib.NewReader(bytes.NewReader(head)); err == nil { defer reader.Close() - if decoded, readErr := readTolerant(reader); readErr == nil { - return decoded, nil + if decoded, overflow, readErr := readTolerant(reader); readErr == nil { + return decoded, overflow, nil } } reader := flate.NewReader(bytes.NewReader(head)) defer reader.Close() return readTolerant(reader) default: - return head, nil + return head, len(head) > maxInspectBytes, nil } } // The head is a deliberate prefix, so a stream ending mid-frame is expected rather than an error -func readTolerant(r io.Reader) ([]byte, error) { - decoded, err := io.ReadAll(io.LimitReader(r, maxInspectBytes)) +func readTolerant(r io.Reader) ([]byte, bool, error) { + decoded, err := io.ReadAll(io.LimitReader(r, maxInspectBytes+1)) + overflow := len(decoded) > maxInspectBytes + if overflow { + decoded = decoded[:maxInspectBytes] + } if len(decoded) > 0 || err == nil || err == io.EOF { - return decoded, nil + return decoded, overflow, nil } - return nil, err + return nil, false, err } func (p *ClickHouseProxy) director(req *http.Request) { @@ -275,7 +286,7 @@ func (p *ClickHouseProxy) director(req *http.Request) { for _, param := range strippedAuthParams { query.Del(param) } - if p.config.Database != "" && query.Get("database") == "" { + if p.config.Database != "" { query.Set("database", p.config.Database) } req.URL.RawQuery = query.Encode() @@ -295,7 +306,6 @@ func (p *ClickHouseProxy) modifyResponse(resp *http.Response) error { if !ok || state == nil { return nil } - elapsed := time.Since(state.started) if resp.StatusCode >= http.StatusBadRequest { original := resp.Body @@ -305,10 +315,41 @@ func (p *ClickHouseProxy) modifyResponse(resp *http.Response) error { return nil } - p.logStatement(state.statement, summarize(resp, elapsed)) + // Recorded once the body has finished, so a transfer that dies mid-stream is not logged as a success + resp.Body = &completionLoggingBody{ReadCloser: resp.Body, onDone: func(err error) { + outcome := summarize(resp, time.Since(state.started)) + if err != nil { + outcome = fmt.Sprintf("INTERRUPTED: %s, %v", outcome, err) + } + p.logStatement(state.statement, outcome) + }} return nil } +type completionLoggingBody struct { + io.ReadCloser + onDone func(error) + once sync.Once +} + +func (b *completionLoggingBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if err != nil { + finished := err + if errors.Is(err, io.EOF) { + finished = nil + } + b.once.Do(func() { b.onDone(finished) }) + } + return n, err +} + +func (b *completionLoggingBody) Close() error { + err := b.ReadCloser.Close() + b.once.Do(func() { b.onDone(errors.New("the response was closed before it finished")) }) + return err +} + func summarize(resp *http.Response, elapsed time.Duration) string { parts := []string{resp.Status} diff --git a/packages/pam/handlers/clickhouse/proxy_test.go b/packages/pam/handlers/clickhouse/proxy_test.go index 10754720..b5f0cfe6 100644 --- a/packages/pam/handlers/clickhouse/proxy_test.go +++ b/packages/pam/handlers/clickhouse/proxy_test.go @@ -50,7 +50,7 @@ func newTestProxy(t *testing.T, upstream http.HandlerFunc, blocked ...string) (h proxy := NewClickHouseProxy(ClickHouseProxyConfig{ TargetAddr: strings.TrimPrefix(server.URL, "http://"), Username: "pam_svc", - Password: "s3cret", + Password: "s3cret", // ggignore Database: "analytics", SessionID: "session-1", SessionLogger: logger, @@ -94,7 +94,7 @@ func TestReplacesClientCredentialsWithTheAccountsOwn(t *testing.T) { defer closeUpstream() req := httptest.NewRequest(http.MethodPost, "/?user=attacker&password=guessed&query=SELECT+1", http.NoBody) - req.Header.Set("Authorization", "Basic YXR0YWNrZXI6Z3Vlc3NlZA==") + req.Header.Set("Authorization", "Basic YXR0YWNrZXI6Z3Vlc3NlZA==") // ggignore req.Header.Set("X-ClickHouse-User", "attacker") req.Header.Set("X-ClickHouse-Key", "guessed") handler.ServeHTTP(httptest.NewRecorder(), req) @@ -107,7 +107,7 @@ func TestReplacesClientCredentialsWithTheAccountsOwn(t *testing.T) { require.Empty(t, captured.query.Get("password")) } -func TestAppliesTheAccountDatabaseOnlyWhenTheClientNamesNone(t *testing.T) { +func TestReplacesWhateverDatabaseTheClientAsksFor(t *testing.T) { var captured capturedRequest handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) defer closeUpstream() @@ -116,7 +116,7 @@ func TestAppliesTheAccountDatabaseOnlyWhenTheClientNamesNone(t *testing.T) { require.Equal(t, "analytics", captured.query.Get("database")) handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/?database=other&query=SELECT+1", http.NoBody)) - require.Equal(t, "other", captured.query.Get("database")) + require.Equal(t, "analytics", captured.query.Get("database")) } func TestBlocksAStatementInTheQueryParameter(t *testing.T) { @@ -210,6 +210,24 @@ func TestTruncatesWhatALargeStatementWritesToTheRecording(t *testing.T) { require.Less(t, len(logger.entries[0].Input), len(payload)) } +func TestRefusesACompressedBodyThatExpandsPastTheInspectionWindow(t *testing.T) { + reached := false + handler, _, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }, `(?i)\bdrop\b`) + defer closeUpstream() + + payload := "/*" + strings.Repeat("x", maxInspectBytes*2) + "*/ DROP TABLE events" + req := httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(gzipped(t, payload))) + req.Header.Set("Content-Encoding", "gzip") + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, req) + + require.False(t, reached) + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.Contains(t, recorder.Body.String(), "blocks commands") +} + func TestRejectsABodyEncodingItCannotInspect(t *testing.T) { reached := false handler, _, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { @@ -251,6 +269,27 @@ func TestRecordsTheRowCountsClickHouseReports(t *testing.T) { require.NotContains(t, logger.entries[0].Output, "row(s) written") } +func TestRecordsAStreamingFailureAsInterruptedRatherThanSuccess(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Length", "64") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("partial")) + + if hijacker, ok := w.(http.Hijacker); ok { + conn, _, err := hijacker.Hijack() + require.NoError(t, err) + _ = conn.Close() + } + }) + defer closeUpstream() + + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, "/?query=SELECT+1", http.NoBody)) + + require.Len(t, logger.entries, 1) + require.Contains(t, logger.entries[0].Output, "INTERRUPTED") +} + func TestRecordsAnUpstreamErrorAndStillReturnsItToTheClient(t *testing.T) { handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { _, _ = io.ReadAll(r.Body) From 365b7f710af6e7c73a642d645350be6efe6e9ae2 Mon Sep 17 00:00:00 2001 From: Andrey Date: Thu, 17 Sep 2026 00:17:59 -0400 Subject: [PATCH 3/7] veria --- packages/pam/handlers/clickhouse/proxy.go | 13 ++++++- .../pam/handlers/clickhouse/proxy_test.go | 35 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/packages/pam/handlers/clickhouse/proxy.go b/packages/pam/handlers/clickhouse/proxy.go index a91925ee..66194c49 100644 --- a/packages/pam/handlers/clickhouse/proxy.go +++ b/packages/pam/handlers/clickhouse/proxy.go @@ -76,6 +76,10 @@ var strippedAuthHeaders = []string{ var strippedAuthParams = []string{"user", "password"} +var strippedExecutionParams = []string{"role"} + +var allowedPaths = map[string]bool{"/": true, "/ping": true} + type ClickHouseProxy struct { config ClickHouseProxyConfig reverse *httputil.ReverseProxy @@ -152,6 +156,13 @@ func (p *ClickHouseProxy) HandleConnection(ctx context.Context, clientConn net.C func (p *ClickHouseProxy) handler(l zerolog.Logger) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !allowedPaths[r.URL.Path] { + l.Info().Str("path", r.URL.Path).Msg("Refused a path outside the query endpoint") + writeClickHouseError(w, http.StatusNotFound, codeNotImplemented, + fmt.Sprintf("This session serves ClickHouse's query endpoint, so %q is not available here.", r.URL.Path)) + return + } + statement, body, err := p.inspect(r) if err != nil { writeClickHouseError(w, http.StatusBadRequest, codeNotImplemented, err.Error()) @@ -283,7 +294,7 @@ func (p *ClickHouseProxy) director(req *http.Request) { req.Host = p.config.TargetAddr query := req.URL.Query() - for _, param := range strippedAuthParams { + for _, param := range append(append([]string{}, strippedAuthParams...), strippedExecutionParams...) { query.Del(param) } if p.config.Database != "" { diff --git a/packages/pam/handlers/clickhouse/proxy_test.go b/packages/pam/handlers/clickhouse/proxy_test.go index b5f0cfe6..3fee4e4f 100644 --- a/packages/pam/handlers/clickhouse/proxy_test.go +++ b/packages/pam/handlers/clickhouse/proxy_test.go @@ -107,6 +107,41 @@ func TestReplacesClientCredentialsWithTheAccountsOwn(t *testing.T) { require.Empty(t, captured.query.Get("password")) } +func TestStripsTheRoleParameterThatWouldActivateARoleWithoutAStatement(t *testing.T) { + var captured capturedRequest + handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) + defer closeUpstream() + + handler.ServeHTTP(httptest.NewRecorder(), + httptest.NewRequest(http.MethodPost, "/?role=privileged&query=SELECT+1", http.NoBody)) + + require.Empty(t, captured.query.Get("role")) + require.Equal(t, "SELECT 1", captured.query.Get("query")) +} + +func TestRefusesAPathOutsideTheQueryEndpoint(t *testing.T) { + reached := false + handler, _, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }) + defer closeUpstream() + + // An SQL-backed custom handler can be mounted at any path, and its statement is never in the request + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/custom_handler", http.NoBody)) + + require.False(t, reached) + require.Equal(t, http.StatusNotFound, recorder.Code) + require.Contains(t, recorder.Body.String(), "not available here") + + // The endpoints a driver actually needs still pass + for _, path := range []string{"/", "/ping"} { + reached = false + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, path, http.NoBody)) + require.True(t, reached, path) + } +} + func TestReplacesWhateverDatabaseTheClientAsksFor(t *testing.T) { var captured capturedRequest handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) From 6a8a56b05356bbf384d645c98ccd33bfb6d46d4d Mon Sep 17 00:00:00 2001 From: Andrey Date: Fri, 18 Sep 2026 05:01:13 -0400 Subject: [PATCH 4/7] address reviews --- packages/pam/handlers/clickhouse/proxy.go | 25 ++++++++++++++-- .../pam/handlers/clickhouse/proxy_test.go | 30 +++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/packages/pam/handlers/clickhouse/proxy.go b/packages/pam/handlers/clickhouse/proxy.go index 66194c49..6a7a19a9 100644 --- a/packages/pam/handlers/clickhouse/proxy.go +++ b/packages/pam/handlers/clickhouse/proxy.go @@ -14,7 +14,9 @@ import ( "net" "net/http" "net/http/httputil" + "net/url" "regexp" + "sort" "strconv" "strings" "sync" @@ -193,7 +195,7 @@ func (p *ClickHouseProxy) inspect(r *http.Request) (string, io.ReadCloser, error queryParam := strings.TrimSpace(r.URL.Query().Get("query")) if r.Body == nil || r.ContentLength == 0 { - return queryParam, http.NoBody, nil + return queryParam + parameterSuffix(r.URL.Query()), http.NoBody, nil } // ClickHouse's own block compression is opaque to anything but a ClickHouse client @@ -235,7 +237,26 @@ func (p *ClickHouseProxy) inspect(r *http.Request) (string, io.ReadCloser, error maxInspectBytes>>20) } - return joinStatement(queryParam, string(decoded)), forwarded, nil + return joinStatement(queryParam, string(decoded)) + parameterSuffix(r.URL.Query()), forwarded, nil +} + +func parameterSuffix(query url.Values) string { + names := make([]string, 0, len(query)) + for name := range query { + if strings.HasPrefix(name, "param_") { + names = append(names, name) + } + } + if len(names) == 0 { + return "" + } + sort.Strings(names) + + pairs := make([]string, 0, len(names)) + for _, name := range names { + pairs = append(pairs, strings.TrimPrefix(name, "param_")+"="+query.Get(name)) + } + return "\n-- parameters: " + strings.Join(pairs, " ") } func joinStatement(queryParam string, body string) string { diff --git a/packages/pam/handlers/clickhouse/proxy_test.go b/packages/pam/handlers/clickhouse/proxy_test.go index 3fee4e4f..0c8a7fcd 100644 --- a/packages/pam/handlers/clickhouse/proxy_test.go +++ b/packages/pam/handlers/clickhouse/proxy_test.go @@ -220,6 +220,36 @@ func TestRefusesAPaddedBodyThatWouldPushAStatementPastTheInspectionWindow(t *tes require.Empty(t, logger.entries) } +func TestMatchesAndRecordsValuesSubstitutedFromQueryParameters(t *testing.T) { + reached := false + handler, logger, closeUpstream := newTestProxy(t, func(w http.ResponseWriter, r *http.Request) { + reached = true + }, `(?i)\busers\b`) + defer closeUpstream() + + // The blocked name is only in param_t, never in the statement text + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/?param_t=users", + strings.NewReader("SELECT count() FROM {t:Identifier}"))) + + require.False(t, reached) + require.Equal(t, http.StatusForbidden, recorder.Code) + require.Len(t, logger.entries, 1) + require.Contains(t, logger.entries[0].Input, "-- parameters: t=users") +} + +func TestRecordsParameterValuesAlongsideTheStatement(t *testing.T) { + handler, logger, closeUpstream := newTestProxy(t, capturingUpstream(&capturedRequest{}, nil)) + defer closeUpstream() + + handler.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodPost, + "/?param_table=events¶m_n=5", strings.NewReader("SELECT {n:UInt8} FROM {table:Identifier}"))) + + require.Len(t, logger.entries, 1) + // Sorted, so a recording reads the same way every time + require.Contains(t, logger.entries[0].Input, "-- parameters: n=5 table=events") +} + func TestForwardsABodyLargerThanTheInspectionCapWhenNothingIsBlocked(t *testing.T) { var captured capturedRequest handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) From b6daded19e89765355068741cf2d793f519ecc2d Mon Sep 17 00:00:00 2001 From: Andrey Date: Fri, 18 Sep 2026 11:54:43 -0400 Subject: [PATCH 5/7] address reviews --- packages/pam/handlers/clickhouse/proxy.go | 6 ++++-- packages/pam/handlers/clickhouse/proxy_test.go | 4 ++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/pam/handlers/clickhouse/proxy.go b/packages/pam/handlers/clickhouse/proxy.go index 6a7a19a9..8efdc982 100644 --- a/packages/pam/handlers/clickhouse/proxy.go +++ b/packages/pam/handlers/clickhouse/proxy.go @@ -69,11 +69,13 @@ var errorNames = map[int]string{ } // Every way a client can present its own identity -var strippedAuthHeaders = []string{ +var strippedClientHeaders = []string{ "Authorization", "X-ClickHouse-User", "X-ClickHouse-Key", "X-ClickHouse-SSL-Certificate-Auth", + "X-ClickHouse-Database", + "X-ClickHouse-Quota-Key", } var strippedAuthParams = []string{"user", "password"} @@ -323,7 +325,7 @@ func (p *ClickHouseProxy) director(req *http.Request) { } req.URL.RawQuery = query.Encode() - for _, header := range strippedAuthHeaders { + for _, header := range strippedClientHeaders { req.Header.Del(header) } req.Header.Set("X-ClickHouse-User", p.config.Username) diff --git a/packages/pam/handlers/clickhouse/proxy_test.go b/packages/pam/handlers/clickhouse/proxy_test.go index 0c8a7fcd..0e12426f 100644 --- a/packages/pam/handlers/clickhouse/proxy_test.go +++ b/packages/pam/handlers/clickhouse/proxy_test.go @@ -97,11 +97,15 @@ func TestReplacesClientCredentialsWithTheAccountsOwn(t *testing.T) { req.Header.Set("Authorization", "Basic YXR0YWNrZXI6Z3Vlc3NlZA==") // ggignore req.Header.Set("X-ClickHouse-User", "attacker") req.Header.Set("X-ClickHouse-Key", "guessed") + req.Header.Set("X-ClickHouse-Database", "other_db") + req.Header.Set("X-ClickHouse-Quota-Key", "someone-elses-quota") handler.ServeHTTP(httptest.NewRecorder(), req) require.Equal(t, "pam_svc", captured.headers.Get("X-ClickHouse-User")) require.Equal(t, "s3cret", captured.headers.Get("X-ClickHouse-Key")) require.Empty(t, captured.headers.Get("Authorization")) + require.Empty(t, captured.headers.Get("X-ClickHouse-Database")) + require.Empty(t, captured.headers.Get("X-ClickHouse-Quota-Key")) require.Empty(t, captured.headers.Get("X-Forwarded-For")) require.Empty(t, captured.query.Get("user")) require.Empty(t, captured.query.Get("password")) From 408563c093135845431297cb3253fa79c141cdd2 Mon Sep 17 00:00:00 2001 From: Andrey Date: Fri, 18 Sep 2026 14:22:24 -0400 Subject: [PATCH 6/7] address veria --- packages/pam/handlers/clickhouse/proxy.go | 8 +++++--- packages/pam/handlers/clickhouse/proxy_test.go | 13 +++++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/packages/pam/handlers/clickhouse/proxy.go b/packages/pam/handlers/clickhouse/proxy.go index 8efdc982..c8348dc7 100644 --- a/packages/pam/handlers/clickhouse/proxy.go +++ b/packages/pam/handlers/clickhouse/proxy.go @@ -99,7 +99,7 @@ type requestState struct { func NewClickHouseProxy(config ClickHouseProxyConfig) *ClickHouseProxy { proxy := &ClickHouseProxy{config: config} proxy.reverse = &httputil.ReverseProxy{ - Director: proxy.director, + Rewrite: proxy.rewrite, Transport: newTransport(config), ModifyResponse: proxy.modifyResponse, ErrorHandler: proxy.handleUpstreamError, @@ -311,7 +311,8 @@ func readTolerant(r io.Reader) ([]byte, bool, error) { return nil, false, err } -func (p *ClickHouseProxy) director(req *http.Request) { +func (p *ClickHouseProxy) rewrite(pr *httputil.ProxyRequest) { + req := pr.Out req.URL.Scheme = p.scheme() req.URL.Host = p.config.TargetAddr req.Host = p.config.TargetAddr @@ -332,7 +333,8 @@ func (p *ClickHouseProxy) director(req *http.Request) { if p.config.Password != "" { req.Header.Set("X-ClickHouse-Key", p.config.Password) } - req.Header["X-Forwarded-For"] = nil + req.Header.Del("X-Forwarded-For") + req.Header.Del("Forwarded") } func (p *ClickHouseProxy) modifyResponse(resp *http.Response) error { diff --git a/packages/pam/handlers/clickhouse/proxy_test.go b/packages/pam/handlers/clickhouse/proxy_test.go index 0e12426f..40a27a5c 100644 --- a/packages/pam/handlers/clickhouse/proxy_test.go +++ b/packages/pam/handlers/clickhouse/proxy_test.go @@ -146,6 +146,19 @@ func TestRefusesAPathOutsideTheQueryEndpoint(t *testing.T) { } } +func TestKeepsInjectedCredentialsWhenTheClientNamesThemHopByHop(t *testing.T) { + var captured capturedRequest + handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) + defer closeUpstream() + + req := httptest.NewRequest(http.MethodPost, "/?query=SELECT+1", http.NoBody) + req.Header.Set("Connection", "X-ClickHouse-User, X-ClickHouse-Key") + handler.ServeHTTP(httptest.NewRecorder(), req) + + require.Equal(t, "pam_svc", captured.headers.Get("X-ClickHouse-User")) + require.Equal(t, "s3cret", captured.headers.Get("X-ClickHouse-Key")) // ggignore +} + func TestReplacesWhateverDatabaseTheClientAsksFor(t *testing.T) { var captured capturedRequest handler, _, closeUpstream := newTestProxy(t, capturingUpstream(&captured, nil)) From 20d8d24a21255e7c651cfd34cb95a857d4c6a36c Mon Sep 17 00:00:00 2001 From: Andrey Date: Sat, 19 Sep 2026 01:17:20 -0400 Subject: [PATCH 7/7] account type for heartbeat --- packages/gateway-v2/capabilities.go | 3 +++ packages/gateway-v2/gateway.go | 1 + 2 files changed, 4 insertions(+) create mode 100644 packages/gateway-v2/capabilities.go diff --git a/packages/gateway-v2/capabilities.go b/packages/gateway-v2/capabilities.go new file mode 100644 index 00000000..b0d6f499 --- /dev/null +++ b/packages/gateway-v2/capabilities.go @@ -0,0 +1,3 @@ +package gatewayv2 + +const CapabilitySupportedAccountTypes = "supported_account_types" diff --git a/packages/gateway-v2/gateway.go b/packages/gateway-v2/gateway.go index b7ca752f..d6afaefa 100644 --- a/packages/gateway-v2/gateway.go +++ b/packages/gateway-v2/gateway.go @@ -460,6 +460,7 @@ func (g *Gateway) registerHeartBeat(ctx context.Context, errCh chan error) { if g.pkcs11Module != nil { capabilities[CapabilityPkcs11] = true } + capabilities[CapabilitySupportedAccountTypes] = pam.GetSupportedResourceTypes() req := api.GatewayHeartbeatRequest{Capabilities: capabilities} if err := api.CallGatewayHeartBeatV2(g.httpClient, req); err != nil { log.Warn().Msgf("Heartbeat failed: %v", err)