From f996a0065764c92ccf361be002a37fd7b297386a Mon Sep 17 00:00:00 2001 From: BorisTyshkevich <125603338+BorisTyshkevich@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:11:42 +0000 Subject: [PATCH 1/3] Isolate authenticated dynamic catalogs by credential and generation --- CHANGELOG.md | 1 + README.md | 8 +- cmd/altinity-mcp/credential_transport.go | 220 +++++++ cmd/altinity-mcp/credential_transport_test.go | 608 ++++++++++++++++++ cmd/altinity-mcp/main.go | 194 +----- cmd/altinity-mcp/main_test.go | 3 +- docs/jwe_authentication.md | 2 + docs/oauth_authorization.md | 6 + docs/tools.md | 16 +- helm/altinity-mcp/README.md | 2 + pkg/server/multicluster_factory.go | 6 +- pkg/server/server.go | 2 + pkg/server/server_dynamic_tools.go | 43 +- pkg/server/server_dynamic_tools_test.go | 5 +- pkg/server/server_openapi.go | 19 +- pkg/server/server_test.go | 6 + pkg/server/singlecluster_catalog.go | 106 +++ pkg/server/singlecluster_catalog_test.go | 66 ++ 18 files changed, 1084 insertions(+), 229 deletions(-) create mode 100644 cmd/altinity-mcp/credential_transport.go create mode 100644 cmd/altinity-mcp/credential_transport_test.go create mode 100644 pkg/server/singlecluster_catalog.go create mode 100644 pkg/server/singlecluster_catalog_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 32dec10a..37bfc34d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Unreleased BUG FIXES +- security: authenticate the exact `/openapi` schema route and isolate single-cluster HTTP/SSE dynamic tools and schemas by effective request credentials; reset catalog generations on reload and reject SSE messages from other credentials or retired generations (fixes [#182](https://github.com/Altinity/altinity-mcp/issues/182)) - security: OpenAPI `execute_query` always rejects write statements and enforces `max_query_length`; dynamic-tool JSON bodies use the same input-size limit (fixes [#183](https://github.com/Altinity/altinity-mcp/issues/183)) - security: isolate per-request ClickHouse headers, settings, and roles; prevent concurrent OAuth bearer forwarding from crashing the server, confusing caller identities, or leaking a stale bearer into later JWE or static-credential requests (fixes [#181](https://github.com/Altinity/altinity-mcp/issues/181)) diff --git a/README.md b/README.md index 8f9359d7..ebe35b76 100644 --- a/README.md +++ b/README.md @@ -342,10 +342,14 @@ OpenAI MCP/tooling references: ### Authentication - **With JWE**: Add the JWE token to either: - 1. Path parameter: `/{jwe_token}/openapi/...` (now required) + 1. Path parameter: `/{jwe_token}/openapi/...` 2. Authorization header: `Bearer {token}` (alternative) 3. `x-altinity-mcp-key` header (alternative) -- **Without JWE**: Use server-configured credentials (no auth needed in requests) +- **With OAuth**: Send `Authorization: Bearer {token}`. +- **With both**: A self-contained JWE takes priority; otherwise OAuth supplies the credentials. +- **With both disabled**: Use server-configured ClickHouse credentials. + +The exact `/openapi` schema route also requires authentication when either mode is enabled. Dynamic HTTP/SSE tools and OpenAPI schemas use catalogs isolated by the caller's effective credential; server credentials are never a fallback under authentication. Reloading configuration resets catalog generations and requires existing SSE clients to reconnect. See [dynamic discovery](docs/tools.md#dynamic-discovery). ### Available Actions diff --git a/cmd/altinity-mcp/credential_transport.go b/cmd/altinity-mcp/credential_transport.go new file mode 100644 index 00000000..18588d3b --- /dev/null +++ b/cmd/altinity-mcp/credential_transport.go @@ -0,0 +1,220 @@ +package main + +import ( + "context" + "fmt" + "github.com/rs/zerolog/log" + "net/http" + "strings" + "sync" + + "github.com/altinity/altinity-mcp/pkg/config" + altinitymcp "github.com/altinity/altinity-mcp/pkg/server" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// withServerSnapshot captures the immutable parent and its converted tool rules +// together. Authentication, discovery and execution all use this generation. +func (a *application) withServerSnapshot(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + a.configMutex.RLock() + parent := a.mcpServer + a.configMutex.RUnlock() + if parent == nil { + http.Error(w, "Server unavailable", http.StatusServiceUnavailable) + return + } + ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, parent) + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} + +func (a *application) authenticateSnapshot(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + parent := altinitymcp.GetClickHouseJWEServerFromContext(r.Context()) + cfg := parent.Config + // Reuse OAuth challenges and expiry checks with a request-scoped adapter; + // its auth injector cannot accidentally read a later global parent. + if cfg.Server.OAuth.Enabled { + adapter := &application{config: cfg, mcpServer: parent} + adapter.createMCPAuthInjector(cfg)(next).ServeHTTP(w, r) + return + } + if cfg.Server.JWE.Enabled { + token, claims, _, _, err := parent.ValidateAuth(r) + if err != nil || !parent.JWEClaimsHaveCredentials(claims) { + http.Error(w, "Missing or invalid authentication token", http.StatusUnauthorized) + return + } + ctx := context.WithValue(r.Context(), altinitymcp.JWETokenKey, token) + ctx = context.WithValue(ctx, altinitymcp.JWEClaimsKey, claims) + r = r.WithContext(ctx) + } + next.ServeHTTP(w, r) + }) +} + +func (a *application) buildSingleClusterHandler(cfg config.Config, sse bool) http.Handler { + var sdkHandler http.Handler + transport := "" + if sse { + transport = "sse" + a.configMutex.Lock() + pool := &credentialSSEPool{parent: a.mcpServer, buckets: make(map[string]*credentialSSEBucket)} + a.ssePools = append(a.ssePools, pool) + a.configMutex.Unlock() + sdkHandler = pool + } else { + sdkHandler = mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { + return altinitymcp.GetClickHouseJWEServerFromContext(r.Context()).GetServer(r) + }, statelessStreamableOptions()) + } + mux := http.NewServeMux() + openapiHandler := a.withServerSnapshot(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // The reserved prefix can match /{token}/openapi when a tool itself is + // named openapi. It carries no path credential; use the request headers. + if strings.HasPrefix(r.URL.Path, "/openapi/") { + r.SetPathValue("token", "") + } + parent := altinitymcp.GetClickHouseJWEServerFromContext(r.Context()) + parent.OpenAPIHandler(w, r) + })) + transportHandler := a.withServerSnapshot(a.authenticateSnapshot(sdkHandler)) + // /openapi/ and /{token}/openapi/ cannot coexist as Go mux patterns. + // Combined mode reserves the pathless OpenAPI prefix in its fallback + // instead, with the same authentication and body guards as explicit routes. + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if cfg.Server.OpenAPI.Enabled && cfg.Server.JWE.Enabled && cfg.Server.OAuth.Enabled && strings.HasPrefix(r.URL.Path, "/openapi/") { + openapiHandler.ServeHTTP(w, r) + return + } + transportHandler.ServeHTTP(w, r) + }) + for _, pattern := range transportRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled, transport) { + mux.Handle(pattern, handler) + } + if cfg.Server.OpenAPI.Enabled { + + for _, pattern := range openAPIRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled) { + mux.Handle(pattern, openapiHandler) + } + + if sse && cfg.Server.JWE.Enabled && cfg.Server.OAuth.Enabled { + mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasPrefix(r.URL.Path, "/openapi/") { + openapiHandler.ServeHTTP(w, r) + return + } + http.NotFound(w, r) + })) + } + protocol := "http" + if cfg.Server.OpenAPI.TLS { + protocol = "https" + } + path := "/openapi" + if cfg.Server.JWE.Enabled && !cfg.Server.OAuth.Enabled { + path = "/{token}/openapi" + } + log.Info().Str("url", fmt.Sprintf("%s://%s:%d%s", protocol, cfg.Server.Address, cfg.Server.Port, path)).Msg("OpenAPI server listening") + } + mux.HandleFunc("/health", a.healthHandler) + mux.HandleFunc("/livez", a.livenessHandler) + mux.HandleFunc("/jwe-token-generator", a.jweTokenGeneratorHandler) + a.registerOAuthHTTPRoutes(mux) + return finalizeTransportHandler(mux, cfg) +} + +// Each SDK SSE handler owns the sessions of one effective credential. The SDK +// invokes getServer only on GET and retains that GET's context for POST work. +// A POST must find its caller's existing bucket, and then its session inside +// that bucket; misses never allocate or discover anything. +type credentialSSEPool struct { + mu sync.Mutex + parent *altinitymcp.ClickHouseJWEServer + buckets map[string]*credentialSSEBucket + active int + closed bool +} +type credentialSSEBucket struct { + handler *mcp.SSEHandler + refs int + cancels map[*http.Request]context.CancelFunc +} + +func (p *credentialSSEPool) ServeHTTP(w http.ResponseWriter, r *http.Request) { + parent := altinitymcp.GetClickHouseJWEServerFromContext(r.Context()) + key := "static" + if parent.Config.Server.JWE.Enabled || parent.Config.Server.OAuth.Enabled { + var err error + key, _, err = parent.CredentialCatalogKey(r.Context()) + if err != nil { + http.Error(w, "Missing authentication", http.StatusUnauthorized) + return + } + } + p.mu.Lock() + if p.closed || p.parent != parent { + p.mu.Unlock() + http.NotFound(w, r) + return + } + bucket := p.buckets[key] + if r.Method != http.MethodGet { + p.mu.Unlock() + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + if bucket == nil { + http.NotFound(w, r) + return + } + bucket.handler.ServeHTTP(w, r) + return + } + if bucket == nil { + bucket = &credentialSSEBucket{handler: mcp.NewSSEHandler(parent.GetServer, nil), cancels: make(map[*http.Request]context.CancelFunc)} + p.buckets[key] = bucket + } + ctx, cancel := context.WithCancel(r.Context()) + bucket.refs++ + bucket.cancels[r] = cancel + p.active++ + p.mu.Unlock() + defer func() { + cancel() + p.mu.Lock() + delete(bucket.cancels, r) + bucket.refs-- + p.active-- + if bucket.refs == 0 && p.buckets[key] == bucket { + delete(p.buckets, key) + } + p.mu.Unlock() + }() + bucket.handler.ServeHTTP(w, r.WithContext(ctx)) +} + +func (p *credentialSSEPool) Retire(parent *altinitymcp.ClickHouseJWEServer) { + p.mu.Lock() + defer p.mu.Unlock() + for _, bucket := range p.buckets { + for _, cancel := range bucket.cancels { + cancel() + } + } + p.buckets = make(map[string]*credentialSSEBucket) + p.parent = parent +} +func (p *credentialSSEPool) Close() { + p.mu.Lock() + defer p.mu.Unlock() + p.closed = true + for _, bucket := range p.buckets { + for _, cancel := range bucket.cancels { + cancel() + } + } + p.buckets = make(map[string]*credentialSSEBucket) +} diff --git a/cmd/altinity-mcp/credential_transport_test.go b/cmd/altinity-mcp/credential_transport_test.go new file mode 100644 index 00000000..0851f44f --- /dev/null +++ b/cmd/altinity-mcp/credential_transport_test.go @@ -0,0 +1,608 @@ +package main + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + chproto "github.com/ClickHouse/ch-go/proto" + driver "github.com/ClickHouse/clickhouse-go/v2" + "github.com/ClickHouse/clickhouse-go/v2/lib/column" + "github.com/ClickHouse/clickhouse-go/v2/lib/proto" + "github.com/altinity/altinity-mcp/pkg/config" + altinitymcp "github.com/altinity/altinity-mcp/pkg/server" + "github.com/altinity/go-mcp-oauth-sdk/jwe_auth" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/require" +) + +type catalogUpstream struct { + server *httptest.Server + requests atomic.Int64 + marker string + mu sync.Mutex + discovers map[string]int + queries map[string]int +} + +func newCatalogUpstream(t *testing.T) *catalogUpstream { + t.Helper() + u := &catalogUpstream{discovers: make(map[string]int), queries: make(map[string]int)} + u.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + u.requests.Add(1) + query, err := io.ReadAll(r.Body) + if err != nil { + t.Error(err) + return + } + text := string(query) + identity, _, _ := r.BasicAuth() + if strings.HasPrefix(r.Header.Get("Authorization"), "Bearer ") { + token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + switch token { + case "alice": + identity = "alice" + case "bob": + identity = "bob" + default: + parts := strings.Split(token, ".") + if len(parts) == 3 { + // Tests use the forward path; this fake CH is the credential boundary. + if strings.Contains(token, "invalid") { + http.Error(w, "denied", 401) + return + } + identity = "oauth" + } + } + } + if identity == "" { + identity = r.Header.Get("X-ClickHouse-User") + } + if identity == "" { + identity = "static" + } + identity += u.marker + block := proto.NewBlock() + add := func(name, typ string) { + if err := block.AddColumn(name, column.Type(typ)); err != nil { + t.Error(err) + } + } + var row []any + switch { + case strings.Contains(text, "displayName()"): + add("displayName()", "String") + add("version()", "String") + add("revision()", "UInt32") + add("timezone()", "String") + row = []any{"fake", "25.1.1", uint32(driver.ClientTCPProtocolVersion), "UTC"} + case strings.Contains(text, "system.tables"): + u.mu.Lock() + u.discovers[identity]++ + u.mu.Unlock() + add("database", "String") + add("name", "String") + add("create_table_query", "String") + add("comment", "String") + row = []any{"tenant", identity, "CREATE VIEW tenant." + identity + " AS SELECT {" + identity + "_id:UInt64}", "catalog for " + identity} + case strings.Contains(text, "FROM tenant."): + u.mu.Lock() + u.queries[identity]++ + u.mu.Unlock() + if !strings.Contains(text, "tenant."+identity) { + t.Errorf("identity %s executed foreign query %q", identity, text) + http.Error(w, "foreign view", 403) + return + } + add("identity", "String") + row = []any{identity} + default: + add("1", "UInt8") + row = []any{uint8(1)} + } + if err := block.Append(row...); err != nil { + t.Error(err) + return + } + buf := new(chproto.Buffer) + if err := block.Encode(buf, uint64(driver.ClientTCPProtocolVersion)); err != nil { + t.Error(err) + return + } + _, _ = w.Write(buf.Buf) + })) + t.Cleanup(u.server.Close) + return u +} +func (u *catalogUpstream) chConfig(t *testing.T) config.ClickHouseConfig { + t.Helper() + parsed, err := url.Parse(u.server.URL) + require.NoError(t, err) + host, port, err := net.SplitHostPort(parsed.Host) + require.NoError(t, err) + n, err := strconv.Atoi(port) + require.NoError(t, err) + return config.ClickHouseConfig{Host: host, Port: n, Protocol: config.HTTPProtocol, Username: "static", Password: "fake-server-password", ReadOnly: true} +} +func catalogApp(t *testing.T, cfg config.Config) *application { + t.Helper() + parent := altinitymcp.NewClickHouseMCPServer(cfg, "test") + app := &application{config: cfg, mcpServer: parent} + t.Cleanup(app.Close) + return app +} +func catalogConfig(ch config.ClickHouseConfig, jwe, oauth bool) config.Config { + return config.Config{ClickHouse: ch, Logging: config.LoggingConfig{Level: config.InfoLevel}, Server: config.ServerConfig{ + JWE: config.JWEConfig{Enabled: jwe, JWESecretKey: "this-is-a-32-byte-secret-key!!", JWTSecretKey: "fake-jwt-secret"}, + OAuth: config.OAuthConfig{Enabled: oauth}, OpenAPI: config.OpenAPIConfig{Enabled: true}, + Tools: []config.ToolDefinition{{Type: "read", Name: "execute_query"}, {Type: "read", ViewRegexp: "^tenant\\..*", Prefix: "dyn_"}}, + }} +} +func catalogJWE(t *testing.T, cfg config.Config, ch config.ClickHouseConfig, user string) string { + t.Helper() + token, err := jwe_auth.GenerateJWEToken(map[string]interface{}{ + "host": ch.Host, "port": ch.Port, "protocol": "http", "username": user, "password": "fake-password", "exp": time.Now().Add(time.Hour).Unix(), + }, []byte(cfg.Server.JWE.JWESecretKey), []byte(cfg.Server.JWE.JWTSecretKey)) + require.NoError(t, err) + return token +} + +type credentialRoundTripper struct { + token, jwe string + endpointMu sync.Mutex + endpoint string +} + +func (tr *credentialRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { + r = r.Clone(r.Context()) + r.Header = r.Header.Clone() + if tr.token != "" { + r.Header.Set("Authorization", "Bearer "+tr.token) + } + if tr.jwe != "" { + r.Header.Set("X-Altinity-MCP-Key", tr.jwe) + } + if r.URL.Query().Get("sessionid") != "" { + tr.endpointMu.Lock() + tr.endpoint = r.URL.String() + tr.endpointMu.Unlock() + } + return http.DefaultTransport.RoundTrip(r) +} +func catalogSession(t *testing.T, ctx context.Context, endpoint string, tr *credentialRoundTripper, sse bool) *mcp.ClientSession { + t.Helper() + client := mcp.NewClient(&mcp.Implementation{Name: "catalog-test", Version: "1"}, nil) + hc := &http.Client{Transport: tr} + var transport mcp.Transport = &mcp.StreamableClientTransport{Endpoint: endpoint, HTTPClient: hc} + if sse { + transport = &mcp.SSEClientTransport{Endpoint: endpoint, HTTPClient: hc} + } + session, err := client.Connect(ctx, transport, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = session.Close() }) + return session +} +func assertCatalogSession(t *testing.T, ctx context.Context, session *mcp.ClientSession, identity, other string) { + t.Helper() + list, err := session.ListTools(ctx, nil) + require.NoError(t, err) + names := []string{} + for _, tool := range list.Tools { + names = append(names, tool.Name) + } + require.Contains(t, names, "dyn_tenant_"+identity) + require.NotContains(t, names, "dyn_tenant_"+other) + result, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "dyn_tenant_" + identity, Arguments: map[string]any{identity + "_id": 1}}) + require.NoError(t, err) + require.False(t, result.IsError, fmt.Sprint(result.Content)) + data, err := json.Marshal(result) + require.NoError(t, err) + require.Contains(t, string(data), identity) + _, err = session.CallTool(ctx, &mcp.CallToolParams{Name: "dyn_tenant_" + other, Arguments: map[string]any{other + "_id": 1}}) + require.Error(t, err) +} +func catalogSchema(t *testing.T, endpoint string, tr *credentialRoundTripper, identity, other string) { + t.Helper() + hc := &http.Client{Transport: tr} + response, err := hc.Get(endpoint + "/openapi") + require.NoError(t, err) + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + require.Equal(t, 200, response.StatusCode, string(body)) + require.Contains(t, string(body), "dyn_tenant_"+identity) + require.Contains(t, string(body), identity+"_id") + require.NotContains(t, string(body), "dyn_tenant_"+other) +} +func TestCredentialCatalogAnonymousSchema(t *testing.T) { + for _, sse := range []bool{false, true} { + for _, auth := range []string{"jwe", "oauth", "combined"} { + t.Run(fmt.Sprintf("%s/sse=%t", auth, sse), func(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), auth != "oauth", auth != "jwe") + app := catalogApp(t, cfg) + var handler http.Handler = app.buildHTTPHandler(cfg, app.mcpServer.MCPServer) + if sse { + handler = app.buildSSEHandler(cfg, app.mcpServer.MCPServer) + } + web := httptest.NewServer(handler) + defer web.Close() + response, err := http.Get(web.URL + "/openapi") + require.NoError(t, err) + defer response.Body.Close() + require.Equal(t, 401, response.StatusCode) + require.Zero(t, up.requests.Load()) + }) + } + } +} +func TestCredentialCatalogHTTPIsolation(t *testing.T) { + for _, auth := range []string{"jwe", "oauth", "combined"} { + t.Run(auth, func(t *testing.T) { + up := newCatalogUpstream(t) + bobUp := up + if auth == "jwe" { + bobUp = newCatalogUpstream(t) + } + cfg := catalogConfig(up.chConfig(t), auth != "oauth", auth != "jwe") + app := catalogApp(t, cfg) + web := httptest.NewServer(app.buildHTTPHandler(cfg, app.mcpServer.MCPServer)) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + trs := []*credentialRoundTripper{{token: "alice"}, {token: "bob"}} + if auth != "oauth" { + for i, id := range []string{"alice", "bob"} { + ch := cfg.ClickHouse + if i == 1 { + ch = bobUp.chConfig(t) + } + trs[i].jwe = catalogJWE(t, cfg, ch, id) + } + } + if auth == "combined" { + trs[0].token = "bob" + trs[1].token = "alice" + } + endpoints := []string{web.URL, web.URL} + if auth == "jwe" { + endpoints[0] += "/" + trs[0].jwe + endpoints[1] += "/" + trs[1].jwe + trs[0].token = "" + trs[1].token = "" + } + sessions := []*mcp.ClientSession{catalogSession(t, ctx, endpoints[0], trs[0], false), catalogSession(t, ctx, endpoints[1], trs[1], false)} + var wg sync.WaitGroup + for i, id := range []string{"alice", "bob"} { + wg.Add(1) + go func(i int, id string) { + defer wg.Done() + other := []string{"bob", "alice"}[i] + for j := 0; j < 3; j++ { + assertCatalogSession(t, ctx, sessions[i], id, other) + catalogSchema(t, web.URL, trs[i], id, other) + catalogREST(t, web.URL, trs[i], id, other) + } + }(i, id) + } + wg.Wait() + up.mu.Lock() + require.Equal(t, 1, up.discovers["alice"]) + require.Equal(t, 6, up.queries["alice"]) + up.mu.Unlock() + bobUp.mu.Lock() + require.Equal(t, 1, bobUp.discovers["bob"]) + require.Equal(t, 6, bobUp.queries["bob"]) + bobUp.mu.Unlock() + }) + } +} +func TestCredentialCatalogSSEOwnership(t *testing.T) { + for _, auth := range []string{"oauth", "jwe", "combined"} { + t.Run(auth, func(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), auth != "oauth", auth != "jwe") + app := catalogApp(t, cfg) + web := httptest.NewServer(app.buildSSEHandler(cfg, app.mcpServer.MCPServer)) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + aTr := &credentialRoundTripper{token: "alice"} + bTr := &credentialRoundTripper{token: "bob"} + aEndpoint, bEndpoint := web.URL+"/sse", web.URL+"/sse" + if auth != "oauth" { + aTr.jwe = catalogJWE(t, cfg, cfg.ClickHouse, "alice") + bTr.jwe = catalogJWE(t, cfg, cfg.ClickHouse, "bob") + } + if auth == "jwe" { + aEndpoint = web.URL + "/" + aTr.jwe + "/sse" + bEndpoint = web.URL + "/" + bTr.jwe + "/sse" + aTr.token = "" + bTr.token = "" + } + if auth == "combined" { + aTr.token = "bob" + bTr.token = "alice" + } + a := catalogSession(t, ctx, aEndpoint, aTr, true) + a2 := catalogSession(t, ctx, aEndpoint, aTr, true) + b := catalogSession(t, ctx, bEndpoint, bTr, true) + catalogSchema(t, web.URL, aTr, "alice", "bob") + catalogREST(t, web.URL, aTr, "alice", "bob") + assertCatalogSession(t, ctx, a, "alice", "bob") + assertCatalogSession(t, ctx, a2, "alice", "bob") + assertCatalogSession(t, ctx, b, "bob", "alice") + aTr.endpointMu.Lock() + endpoint := aTr.endpoint + aTr.endpointMu.Unlock() + require.NotEmpty(t, endpoint) + if auth == "jwe" { + endpoint = strings.Replace(endpoint, "/"+aTr.jwe+"/sse", "/"+bTr.jwe+"/sse", 1) + } + up.mu.Lock() + before := up.queries["alice"] + up.mu.Unlock() + request, err := http.NewRequest(http.MethodPost, endpoint, strings.NewReader(`{"jsonrpc":"2.0","id":99,"method":"tools/call","params":{"name":"dyn_tenant_alice","arguments":{"alice_id":1}}}`)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + response, err := (&http.Client{Transport: bTr}).Do(request) + require.NoError(t, err) + response.Body.Close() + require.Equal(t, 404, response.StatusCode) + up.mu.Lock() + require.Equal(t, before, up.queries["alice"]) + require.Equal(t, 1, up.discovers["alice"]) + require.Equal(t, 1, up.discovers["bob"]) + up.mu.Unlock() + pool := app.ssePools[0] + pool.mu.Lock() + require.Equal(t, 2, len(pool.buckets)) + require.Equal(t, 3, pool.active) + pool.mu.Unlock() + require.NoError(t, a.Close()) + require.NoError(t, a2.Close()) + require.NoError(t, b.Close()) + require.Eventually(t, func() bool { pool.mu.Lock(); defer pool.mu.Unlock(); return len(pool.buckets) == 0 && pool.active == 0 }, 3*time.Second, 10*time.Millisecond) + if auth == "jwe" { + request.URL, err = url.Parse(strings.Replace(endpoint, "/"+bTr.jwe+"/sse", "/"+aTr.jwe+"/sse", 1)) + require.NoError(t, err) + } + request.Body, err = request.GetBody() + require.NoError(t, err) + response, err = (&http.Client{Transport: aTr}).Do(request) + require.NoError(t, err) + response.Body.Close() + require.Equal(t, 404, response.StatusCode) + pool.mu.Lock() + require.Empty(t, pool.buckets) + pool.mu.Unlock() + }) + } +} +func writeCatalogReload(t *testing.T, app *application, cfg config.Config) { + t.Helper() + if app.configFile == "" { + app.configFile = filepath.Join(t.TempDir(), "config.json") + app.stopConfigReload = make(chan struct{}) + } + data, err := json.Marshal(cfg) + require.NoError(t, err) + require.NoError(t, os.WriteFile(app.configFile, data, 0600)) + require.NoError(t, app.reloadConfig(&mockCommand{flags: map[string]interface{}{}, setFlags: map[string]bool{}, stringMaps: map[string]map[string]string{}})) +} +func TestCredentialCatalogReloadGenerations(t *testing.T) { + for _, sse := range []bool{false, true} { + t.Run(fmt.Sprintf("sse=%t", sse), func(t *testing.T) { + oldUp := newCatalogUpstream(t) + oldUp.marker = "_old" + newUp := newCatalogUpstream(t) + newUp.marker = "_new" + cfg := catalogConfig(oldUp.chConfig(t), false, true) + app := catalogApp(t, cfg) + var handler http.Handler = app.buildHTTPHandler(cfg, app.mcpServer.MCPServer) + endpointSuffix := "" + if sse { + handler = app.buildSSEHandler(cfg, app.mcpServer.MCPServer) + endpointSuffix = "/sse" + } + web := httptest.NewServer(handler) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + tr := &credentialRoundTripper{token: "alice"} + session := catalogSession(t, ctx, web.URL+endpointSuffix, tr, sse) + catalogSchema(t, web.URL, tr, "alice_old", "alice_new") + oldParent := app.mcpServer + tr.endpointMu.Lock() + oldEndpoint := tr.endpoint + tr.endpointMu.Unlock() + next := cfg + next.ClickHouse = newUp.chConfig(t) + next.Server.Tools = []config.ToolDefinition{{Type: "read", Name: "execute_query"}, {Type: "read", ViewRegexp: "^tenant\\..*", Prefix: "next_"}} + writeCatalogReload(t, app, next) + require.NotSame(t, oldParent, app.mcpServer) + if sse { + request, err := http.NewRequest(http.MethodPost, oldEndpoint, strings.NewReader(`{"jsonrpc":"2.0","id":19,"method":"tools/list"}`)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + response, err := (&http.Client{Transport: tr}).Do(request) + require.NoError(t, err) + response.Body.Close() + require.Equal(t, 404, response.StatusCode) + session = catalogSession(t, ctx, web.URL+"/sse", tr, true) + } + list, err := session.ListTools(ctx, nil) + require.NoError(t, err) + names := []string{} + for _, tool := range list.Tools { + names = append(names, tool.Name) + } + require.Contains(t, names, "next_tenant_alice_new") + require.NotContains(t, names, "dyn_tenant_alice_old") + response, err := (&http.Client{Transport: tr}).Get(web.URL + "/openapi") + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + response.Body.Close() + require.NoError(t, err) + require.Equal(t, 200, response.StatusCode) + require.Contains(t, string(body), "next_tenant_alice_new") + require.NotContains(t, string(body), "dyn_tenant_alice_old") + oldUp.mu.Lock() + require.Equal(t, 1, oldUp.discovers["alice_old"]) + oldUp.mu.Unlock() + newUp.mu.Lock() + require.Equal(t, 1, newUp.discovers["alice_new"]) + newUp.mu.Unlock() + }) + } +} +func TestCredentialCatalogExpiredSSEBearer(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), false, true) + app := catalogApp(t, cfg) + web := httptest.NewServer(app.buildSSEHandler(cfg, app.mcpServer.MCPServer)) + defer web.Close() + tr := &credentialRoundTripper{token: makeUnsignedJWT(t, map[string]interface{}{"exp": time.Now().Add(-time.Hour).Unix(), "email": "alice@example.com"})} + request, err := http.NewRequest(http.MethodPost, web.URL+"/sse?sessionid=unowned", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + response, err := (&http.Client{Transport: tr}).Do(request) + require.NoError(t, err) + response.Body.Close() + require.Equal(t, 401, response.StatusCode) + require.Zero(t, up.requests.Load()) + pool := app.ssePools[0] + pool.mu.Lock() + require.Empty(t, pool.buckets) + pool.mu.Unlock() +} + +func catalogREST(t *testing.T, endpoint string, tr *credentialRoundTripper, identity, other string) { + t.Helper() + // The exact schema supports header auth. JWE-only REST routes retain their + // existing tokenized path contract. + prefix := "" + if tr.jwe != "" && tr.token == "" { + prefix = "/" + tr.jwe + } + client := &http.Client{Transport: tr} + response, err := client.Post(endpoint+prefix+"/openapi/dyn_tenant_"+identity, "application/json", strings.NewReader(fmt.Sprintf(`{"%s_id":1}`, identity))) + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + response.Body.Close() + require.NoError(t, err) + require.Equal(t, 200, response.StatusCode, string(body)) + require.Contains(t, string(body), identity) + response, err = client.Post(endpoint+prefix+"/openapi/dyn_tenant_"+other, "application/json", strings.NewReader(fmt.Sprintf(`{"%s_id":1}`, other))) + require.NoError(t, err) + response.Body.Close() + require.Equal(t, 404, response.StatusCode) +} + +// Reload while schema requests are active: an endpoint and its discovery rules +// must always come from the same immutable parent snapshot. +func TestCredentialCatalogConcurrentReloadSnapshots(t *testing.T) { + oldUp := newCatalogUpstream(t) + oldUp.marker = "_old" + newUp := newCatalogUpstream(t) + newUp.marker = "_new" + oldCfg := catalogConfig(oldUp.chConfig(t), false, true) + newCfg := catalogConfig(newUp.chConfig(t), false, true) + newCfg.Server.Tools = []config.ToolDefinition{{Type: "read", ViewRegexp: "^tenant\\..*", Prefix: "next_"}} + app := catalogApp(t, oldCfg) + web := httptest.NewServer(app.buildHTTPHandler(oldCfg, app.mcpServer.MCPServer)) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + var wg sync.WaitGroup + start := make(chan struct{}) + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + for j := 0; j < 15; j++ { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, web.URL+"/openapi", nil) + if err != nil { + t.Error(err) + return + } + req.Header.Set("Authorization", "Bearer alice") + response, err := http.DefaultClient.Do(req) + if err != nil { + t.Error(err) + return + } + data, err := io.ReadAll(response.Body) + response.Body.Close() + if err != nil { + t.Error(err) + return + } + if response.StatusCode != 200 { + t.Errorf("schema status %d", response.StatusCode) + return + } + body := string(data) + old := strings.Contains(body, "dyn_tenant_alice_old") + next := strings.Contains(body, "next_tenant_alice_new") + // A generation can retire before its cache is initialized; static-only + // completion remains safe, but a mixed catalog must never appear. + if strings.Contains(body, "dyn_tenant_alice_new") || strings.Contains(body, "next_tenant_alice_old") || (old && next) { + t.Errorf("mixed catalog generation: %s", body) + return + } + } + }() + } + close(start) + for i := 0; i < 4; i++ { + writeCatalogReload(t, app, newCfg) + writeCatalogReload(t, app, oldCfg) + } + wg.Wait() + catalogSchema(t, web.URL, &credentialRoundTripper{token: "alice"}, "alice_old", "alice_new") +} + +func TestCredentialCatalogReservedOpenAPIToolName(t *testing.T) { + for _, sse := range []bool{false, true} { + t.Run(fmt.Sprintf("sse=%t", sse), func(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), true, true) + cfg.Server.Tools = []config.ToolDefinition{{Type: "read", Name: "execute_query"}, {Type: "read", Name: "openapi", ViewRegexp: "^tenant\\.alice$"}} + app := catalogApp(t, cfg) + var handler http.Handler = app.buildHTTPHandler(cfg, app.mcpServer.MCPServer) + if sse { + handler = app.buildSSEHandler(cfg, app.mcpServer.MCPServer) + } + web := httptest.NewServer(handler) + defer web.Close() + client := &http.Client{Transport: &credentialRoundTripper{token: "alice"}} + response, err := client.Post(web.URL+"/openapi/openapi", "application/json", strings.NewReader(`{"alice_id":1}`)) + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + response.Body.Close() + require.NoError(t, err) + require.Equal(t, 200, response.StatusCode, string(body)) + require.Contains(t, string(body), "alice") + up.mu.Lock() + require.Equal(t, 1, up.discovers["alice"]) + require.Equal(t, 1, up.queries["alice"]) + up.mu.Unlock() + }) + } +} diff --git a/cmd/altinity-mcp/main.go b/cmd/altinity-mcp/main.go index e15901ee..3e316487 100644 --- a/cmd/altinity-mcp/main.go +++ b/cmd/altinity-mcp/main.go @@ -313,11 +313,11 @@ func openAPIRoutePatterns(jweEnabled, oauthEnabled bool) []string { switch { case jweEnabled && oauthEnabled: - // Exact /openapi remains unauthenticated schema discovery in combined mode. + // Exact /openapi uses the same authentication as tool endpoints. // Skip /openapi/ (index 1) — stripTrailingSlash handles it, and it conflicts with /{token}/openapi/. - return append(tokenized, pathless[2:]...) + return append(append(tokenized, pathless[0]), pathless[2:]...) case jweEnabled: - return tokenized + return append(tokenized, pathless[0]) default: return pathless } @@ -503,89 +503,7 @@ func (a *application) startHTTPServer(cfg config.Config, mcpServer *mcp.Server) // Split from startHTTPServer so tests can exercise the exact production // routing without binding a port. func (a *application) buildHTTPHandler(cfg config.Config, mcpServer *mcp.Server) http.Handler { - openAPIProtocol := "http" - if cfg.Server.OpenAPI.TLS { - openAPIProtocol = "https" - } - - authInjector := a.createMCPAuthInjector(cfg) - serverInjector := func(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, a.mcpServer) - next.ServeHTTP(w, r.WithContext(ctx)) - }) - } - serverInjectorOpenAPI := func(w http.ResponseWriter, r *http.Request) { - ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, a.mcpServer) - a.mcpServer.OpenAPIHandler(w, r.WithContext(ctx)) - } - serverInjectorSchema := func(w http.ResponseWriter, r *http.Request) { - ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, a.mcpServer) - a.mcpServer.ServeOpenAPISchema(w, r.WithContext(ctx)) - } - - var httpHandler http.Handler - if cfg.Server.JWE.Enabled { - log.Info().Msg("Using dynamic base path for JWE authentication") - - tokenInjector := a.createTokenInjector() - dtInjector := a.dynamicToolsInjector - httpServer := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { - return mcpServer - }, statelessStreamableOptions()) - - mux := http.NewServeMux() - transportHandler := serverInjector(tokenInjector(dtInjector(httpServer))) - if cfg.Server.OAuth.Enabled { - transportHandler = serverInjector(authInjector(dtInjector(httpServer))) - } - for _, pattern := range transportRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled, "") { - mux.Handle(pattern, transportHandler) - } - if cfg.Server.OpenAPI.Enabled { - mux.HandleFunc("/openapi", serverInjectorSchema) - for _, pattern := range openAPIRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled) { - mux.HandleFunc(pattern, serverInjectorOpenAPI) - } - openAPIPath := "/{token}/openapi" - if cfg.Server.OAuth.Enabled { - openAPIPath = "/openapi" - } - log.Info().Str("url", fmt.Sprintf("%s://%s:%d%s", openAPIProtocol, cfg.Server.Address, cfg.Server.Port, openAPIPath)).Msg("OpenAPI server listening") - } - mux.HandleFunc("/health", a.healthHandler) - mux.HandleFunc("/livez", a.livenessHandler) - mux.HandleFunc("/jwe-token-generator", a.jweTokenGeneratorHandler) - a.registerOAuthHTTPRoutes(mux) - httpHandler = finalizeTransportHandler(mux, cfg) - } else { - // Use standard HTTP server without dynamic paths - httpServer := mcp.NewStreamableHTTPHandler(func(r *http.Request) *mcp.Server { - return mcpServer - }, statelessStreamableOptions()) - dtInjector := a.dynamicToolsInjector - mux := http.NewServeMux() - transportHandler := serverInjector(dtInjector(httpServer)) - if cfg.Server.OAuth.Enabled { - transportHandler = serverInjector(authInjector(dtInjector(httpServer))) - } - for _, pattern := range transportRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled, "") { - mux.Handle(pattern, transportHandler) - } - if cfg.Server.OpenAPI.Enabled { - for _, pattern := range openAPIRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled) { - mux.HandleFunc(pattern, serverInjectorOpenAPI) - } - log.Info().Str("url", fmt.Sprintf("%s://%s:%d/openapi", openAPIProtocol, cfg.Server.Address, cfg.Server.Port)).Msg("OpenAPI server listening") - } - mux.HandleFunc("/health", a.healthHandler) - mux.HandleFunc("/livez", a.livenessHandler) - mux.HandleFunc("/jwe-token-generator", a.jweTokenGeneratorHandler) - a.registerOAuthHTTPRoutes(mux) - httpHandler = finalizeTransportHandler(mux, cfg) - } - - return httpHandler + return a.buildSingleClusterHandler(cfg, false) } // startSSEServer starts the SSE transport server @@ -608,91 +526,7 @@ func (a *application) startSSEServer(cfg config.Config, mcpServer *mcp.Server) e // transport; see buildHTTPHandler for why it is separate from the server // start. func (a *application) buildSSEHandler(cfg config.Config, mcpServer *mcp.Server) http.Handler { - authInjector := a.createMCPAuthInjector(cfg) - serverInjector := func(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, a.mcpServer) - next.ServeHTTP(w, r.WithContext(ctx)) - }) - } - serverInjectorOpenAPI := func(w http.ResponseWriter, r *http.Request) { - ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, a.mcpServer) - a.mcpServer.OpenAPIHandler(w, r.WithContext(ctx)) - } - serverInjectorSchema := func(w http.ResponseWriter, r *http.Request) { - ctx := context.WithValue(r.Context(), altinitymcp.CHJWEServerKey, a.mcpServer) - a.mcpServer.ServeOpenAPISchema(w, r.WithContext(ctx)) - } - - openAPIProtocol := "http" - if cfg.Server.OpenAPI.TLS { - openAPIProtocol = "https" - } - - var sseHandler http.Handler - if cfg.Server.JWE.Enabled { - log.Info().Msg("Using dynamic base path for JWE authentication") - - tokenInjector := a.createTokenInjector() - dtInjector := a.dynamicToolsInjector - - // Use SSEHandler for legacy SSE transport - sseServer := mcp.NewSSEHandler(func(r *http.Request) *mcp.Server { - return mcpServer - }, nil) - - mux := http.NewServeMux() - transportHandler := serverInjector(tokenInjector(dtInjector(sseServer))) - if cfg.Server.OAuth.Enabled { - transportHandler = serverInjector(authInjector(dtInjector(sseServer))) - } - for _, pattern := range transportRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled, "sse") { - mux.Handle(pattern, transportHandler) - } - if cfg.Server.OpenAPI.Enabled { - mux.HandleFunc("/openapi", serverInjectorSchema) - for _, pattern := range openAPIRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled) { - mux.HandleFunc(pattern, serverInjectorOpenAPI) - } - openAPIPath := "/{token}/openapi" - if cfg.Server.OAuth.Enabled { - openAPIPath = "/openapi" - } - log.Info().Str("url", fmt.Sprintf("%s://%s:%d%s", openAPIProtocol, cfg.Server.Address, cfg.Server.Port, openAPIPath)).Msg("OpenAPI server listening") - } - mux.HandleFunc("/health", a.healthHandler) - mux.HandleFunc("/livez", a.livenessHandler) - mux.HandleFunc("/jwe-token-generator", a.jweTokenGeneratorHandler) - a.registerOAuthHTTPRoutes(mux) - sseHandler = finalizeTransportHandler(mux, cfg) - } else { - // Use SSEHandler for legacy SSE transport - sseServer := mcp.NewSSEHandler(func(r *http.Request) *mcp.Server { - return mcpServer - }, nil) - dtInjector := a.dynamicToolsInjector - mux := http.NewServeMux() - transportHandler := serverInjector(dtInjector(sseServer)) - if cfg.Server.OAuth.Enabled { - transportHandler = serverInjector(authInjector(dtInjector(sseServer))) - } - for _, pattern := range transportRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled, "sse") { - mux.Handle(pattern, transportHandler) - } - if cfg.Server.OpenAPI.Enabled { - for _, pattern := range openAPIRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled) { - mux.HandleFunc(pattern, serverInjectorOpenAPI) - } - log.Info().Str("url", fmt.Sprintf("%s://%s:%d/openapi", openAPIProtocol, cfg.Server.Address, cfg.Server.Port)).Msg("OpenAPI server listening") - } - mux.HandleFunc("/health", a.healthHandler) - mux.HandleFunc("/livez", a.livenessHandler) - mux.HandleFunc("/jwe-token-generator", a.jweTokenGeneratorHandler) - a.registerOAuthHTTPRoutes(mux) - sseHandler = finalizeTransportHandler(mux, cfg) - } - - return sseHandler + return a.buildSingleClusterHandler(cfg, true) } // livenessHandler provides a process-level health check endpoint for liveness probes. @@ -1040,6 +874,7 @@ type application struct { // from the running (restart-only) value, so the reload loop warns once // per mismatch instead of on every tick. Guarded by configMutex. metricsReloadWarned bool + ssePools []*credentialSSEPool } // setHTTPServer sets the HTTP server with proper synchronization @@ -1317,6 +1152,14 @@ func (a *application) Close() { close(a.stopConfigReload) } + a.configMutex.Lock() + if a.mcpServer != nil { + a.mcpServer.Close() + } + for _, pool := range a.ssePools { + pool.Close() + } + a.configMutex.Unlock() if a.mcCache != nil { a.mcCache.Close() } @@ -1414,7 +1257,6 @@ func (a *application) reloadConfig(cmd CommandInterface) error { // Update logging level if changed a.configMutex.Lock() oldLogLevel := a.config.Logging.Level - a.config = *newCfg a.configMutex.Unlock() if oldLogLevel != newCfg.Logging.Level { @@ -1433,8 +1275,16 @@ func (a *application) reloadConfig(cmd CommandInterface) error { // Update the server (note: this doesn't restart HTTP servers, only updates the MCP server) a.configMutex.Lock() + oldServer := a.mcpServer + a.config = *newCfg a.mcpServer = newMCPServer + for _, pool := range a.ssePools { + pool.Retire(newMCPServer) + } a.configMutex.Unlock() + if oldServer != nil { + oldServer.Close() + } log.Info().Str("config_file", a.configFile).Msg("Configuration reloaded successfully") return nil diff --git a/cmd/altinity-mcp/main_test.go b/cmd/altinity-mcp/main_test.go index 1243d5bb..4ee5f4ae 100644 --- a/cmd/altinity-mcp/main_test.go +++ b/cmd/altinity-mcp/main_test.go @@ -3615,7 +3615,8 @@ func TestOpenAPIRoutePatterns(t *testing.T) { t.Run("jwe_only", func(t *testing.T) { t.Parallel() patterns := openAPIRoutePatterns(true, false) - require.Equal(t, 5, len(patterns)) + require.Equal(t, 6, len(patterns)) + require.Contains(t, patterns, "/openapi") require.Contains(t, patterns, "/{token}/openapi") }) t.Run("no_jwe", func(t *testing.T) { diff --git a/docs/jwe_authentication.md b/docs/jwe_authentication.md index 9a2224e5..fe024ea2 100644 --- a/docs/jwe_authentication.md +++ b/docs/jwe_authentication.md @@ -194,3 +194,5 @@ http://localhost:8080//sse - Use a strong, random secret key for token signing - Set appropriate token expiration times - To implementing token revocation if needed for additional security, change using jwe-secret-key in `altinity-mcp` configuration + +Dynamic catalogs and OpenAPI schemas are isolated by the effective request credential. The exact `/openapi` route also requires authentication; server credentials cannot supply anonymous discovery. In combined JWE/OAuth mode, a self-contained JWE takes priority and a partial JWE uses OAuth when available. SSE GET and POST must use the same effective credential, and clients must reconnect after configuration reload. See [dynamic discovery](tools.md#dynamic-discovery). diff --git a/docs/oauth_authorization.md b/docs/oauth_authorization.md index 7dd57087..24449bd8 100644 --- a/docs/oauth_authorization.md +++ b/docs/oauth_authorization.md @@ -606,3 +606,9 @@ sidecar binding gotchas) see the [`altinity-oauth-helper`](https://github.com/altinity/altinity-oauth-helper) repo's troubleshooting section. + +## Dynamic catalogs and schema authentication + +With OAuth enabled, `/openapi` and MCP requests require a bearer or a self-contained JWE when JWE is also enabled. Dynamic catalogs are isolated by the effective credential. Discovery and execution never fall back to server credentials. In combined mode, a self-contained JWE wins; a partial JWE falls through to OAuth. + +SSE clients must send the same effective credential on their GET stream and POST messages. Sessions from another credential, or from a retired configuration generation, return 404. Reconnect SSE after configuration reload. See [dynamic discovery](tools.md#dynamic-discovery) for cache lifetimes and STDIO behavior. diff --git a/docs/tools.md b/docs/tools.md index 28648717..9a180e0d 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -3,7 +3,7 @@ Altinity MCP Server exposes ClickHouse functionality to MCP clients through **tools**. There are two categories: - **Static tools** — fixed, registered at startup, always available regardless of ClickHouse state. -- **Dynamic tools** — discovered lazily from ClickHouse (views and tables) on the first authenticated request and surfaced to clients via `notifications/tools/list_changed`. +- **Dynamic tools** — discovered lazily from ClickHouse (views and tables); authenticated HTTP/SSE callers receive their own credential-isolated catalogs. Both kinds are configured under a single unified key: `server.tools`. @@ -262,11 +262,13 @@ The server accepts this form but logs a deprecation warning at startup. Prefer ` ## Dynamic discovery -Discovery is **lazy** — it runs on the first authenticated tool call, not at startup. This matters for OAuth, where ClickHouse credentials can be derived from the user request. +Discovery is **lazy**: the first authenticated MCP request or OpenAPI schema/tool request discovers the caller's catalog. Single-cluster HTTP and SSE catalogs are cached by the effective JWE or OAuth credential, so each caller's `tools/list`, schema, and dynamic-tool lookup use that caller's own metadata. A self-contained JWE takes priority over OAuth; a partial JWE uses OAuth when available. The unified `server.tools` rules and legacy `dynamic_tools` rules both use this isolation. -After a successful discovery pass, the server emits `notifications/tools/list_changed` so compatible MCP clients refresh their tool list. +Successful catalogs expire at the earlier of token expiry and 15 minutes. Authentication failures are cached for 60 seconds; transient discovery failures are retried on the next request. Static tools remain available if discovery fails. With authentication disabled, the server keeps its shared static-credential catalog and emits `notifications/tools/list_changed` after discovery. -If discovery fails (e.g. credential error, network issue), static tools remain available and discovery is retried on the next authenticated call. +Configuration reload starts a new catalog generation. Existing SSE clients must reconnect; old session POST endpoints return 404. SSE GET and POST requests must present the same effective credential, including when multiple sessions share a credential. Closing the final GET stream removes its handler bucket. + +Authenticated STDIO exposes static tools only; it cannot discover or query with server credentials when no request token is available. ### What gets discovered @@ -281,7 +283,7 @@ If discovery fails (e.g. credential error, network issue), static tools remain a | OAuth | The Bearer token from the triggering call. | | Plain / no auth | Static `clickhouse.username` / `clickhouse.password` from config, if set. | -Static credentials are no longer **required** for discovery in JWE or OAuth setups — whichever token arrives with the first authenticated call is used to probe `system.tables`. +Static credentials are used only when both JWE and OAuth are disabled. Enabling either auth mode requires usable request credentials for all database access, including discovery. --- @@ -368,7 +370,7 @@ Dynamic tools have `POST /openapi/{tool_name}` endpoints with JSON request schem `clickhouse.max_query_length` limits both SQL query bytes and the entire dynamic-tool JSON body, including trailing whitespace. The default is 10 MiB; a negative value disables the limit. Oversized inputs return HTTP 413 before any ClickHouse connection; write statements sent to `execute_query` return HTTP 400. -Because dynamic discovery is lazy, dynamic tools only appear in the OpenAPI document **after** the first authenticated call has triggered discovery. +The exact `/openapi` schema route requires authentication whenever JWE or OAuth is enabled and discovers only the caller's catalog. There is no public schema setting. Oversized dynamic-tool bodies are rejected before discovery. --- @@ -376,7 +378,7 @@ Because dynamic discovery is lazy, dynamic tools only appear in the OpenAPI docu ### Dynamic tools don't show up immediately after the server starts -Expected. Discovery runs on the first authenticated tool call, not at startup. Make one call (e.g. `execute_query`) and clients that honor `notifications/tools/list_changed` will refresh to show the new dynamic tools. Clients that don't honor the notification need to re-list tools manually. +Expected before a request is made. With HTTP or SSE authentication enabled, an authenticated `tools/list` request discovers and returns that caller's tools. With authentication disabled, discovery uses the server credentials and clients can refresh after `notifications/tools/list_changed`. ### A view or table isn't being exposed diff --git a/helm/altinity-mcp/README.md b/helm/altinity-mcp/README.md index 37ace0fd..1ac7f926 100644 --- a/helm/altinity-mcp/README.md +++ b/helm/altinity-mcp/README.md @@ -4,6 +4,8 @@ A Helm chart for Altinity MCP Server +When JWE or OAuth is enabled, the exact `/openapi` schema route requires authentication and database access never falls back to the configured static credentials. Single-cluster HTTP/SSE dynamic catalogs are isolated by effective request credentials. Existing SSE clients must reconnect after configuration reload. See [dynamic discovery](../../docs/tools.md#dynamic-discovery). + ## Values | Key | Type | Default | Description | diff --git a/pkg/server/multicluster_factory.go b/pkg/server/multicluster_factory.go index 3880a4be..8fe37c77 100644 --- a/pkg/server/multicluster_factory.go +++ b/pkg/server/multicluster_factory.go @@ -11,9 +11,9 @@ import ( ) // MulticlusterServerFactory mints a per-(bearer, cluster) *mcp.Server on -// each incoming request. Single-cluster mode uses one *mcp.Server for the -// process; multi-cluster mode cannot, because per-tenant dynamic tools -// would poison cross-tenant tools/list. +// each incoming request. Authenticated single-cluster mode also uses fresh +// registries: sharing per-tenant dynamic tools would expose cross-tenant +// metadata through tools/list. // // The factory holds the long-lived (cfg, cache, ClickHouseJWEServer) // triple and exposes a single GetServer(r *http.Request) entry point diff --git a/pkg/server/server.go b/pkg/server/server.go index 2e6ec0b4..78d0fadf 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -40,6 +40,8 @@ type ClickHouseJWEServer struct { // lazily builds it for struct-literal test servers. nil when the feature // is unconfigured. roleFilterRe *regexp.Regexp + catalogOnce sync.Once + catalogCache *CatalogCache } // ToolHandlerFunc is a function type for tool handlers diff --git a/pkg/server/server_dynamic_tools.go b/pkg/server/server_dynamic_tools.go index 0ecb7dd7..c580121f 100644 --- a/pkg/server/server_dynamic_tools.go +++ b/pkg/server/server_dynamic_tools.go @@ -48,22 +48,17 @@ type dynamicToolCommentAnnotations struct { OpenWorldHint *bool `json:"openWorldHint"` } -// EnsureDynamicTools discovers dynamic tools (views for reads, tables for writes) -// from ClickHouse and registers them with the MCP server. It's safe to call on -// every request: the fast path short-circuits once init completes. -// -// Discovery is deferred until the caller has usable credentials. In OAuth -// OAuth the Bearer token only arrives on tools/call, not tools/list — -// so the first tools/list just returns static tools, and the first authenticated -// tools/call triggers discovery. The MCP SDK's AddTool automatically fires -// notifications/tools/list_changed, prompting the client to re-fetch. -// -// Concurrency: discovery does CH round-trips which can be slow. We hold the -// write lock only while discovery is in progress. If another goroutine is -// already discovering we return immediately without blocking — concurrent -// tools/list calls see the current (static-only) tool set and get updated -// when the in-flight discovery notifies. +// EnsureDynamicTools discovers and registers the shared catalog only when +// authentication is disabled. Authenticated HTTP/SSE requests use requestCatalog +// and fresh MCP registries instead; caller metadata must never enter this +// process-wide registry. Successful shared discovery emits tools/list_changed. +// Concurrent discovery leaves other callers with the current registry until +// registration completes. func (s *ClickHouseJWEServer) EnsureDynamicTools(ctx context.Context) error { + // Authenticated catalogs belong to a caller, never the process registry. + if s.Config.Server.JWE.Enabled || s.Config.Server.OAuth.Enabled { + return nil + } // Fast path: already initialized. s.dynamicToolsMu.RLock() if s.dynamicToolsInit { @@ -89,9 +84,8 @@ func (s *ClickHouseJWEServer) EnsureDynamicTools(ctx context.Context) error { return nil } - // In forward-OAuth mode with blank static credentials, the OAuth bearer - // isn't in context on the tools/list handshake. Don't mark dynamicToolsInit - // true here — let the next request retry with a real token. + // Without static credentials, leave discovery pending so a later + // configuration generation with usable credentials can discover. if !s.hasDiscoveryCredentials(ctx) { log.Debug().Msg("dynamic_tools: no credentials available yet; deferring discovery") return nil @@ -114,16 +108,11 @@ func (s *ClickHouseJWEServer) EnsureDynamicTools(ctx context.Context) error { // hasDiscoveryCredentials reports whether the current context has any form // of credentials that can be used to query ClickHouse for tool discovery. func (s *ClickHouseJWEServer) hasDiscoveryCredentials(ctx context.Context) bool { - if s.ExtractTokenFromCtx(ctx) != "" { - return true - } - if s.ExtractOAuthTokenFromCtx(ctx) != "" { - return true - } - if s.Config.ClickHouse.Username != "" { - return true + if s.Config.Server.JWE.Enabled || s.Config.Server.OAuth.Enabled { + _, _, err := s.CredentialCatalogKey(ctx) + return err == nil } - return false + return s.Config.ClickHouse.Username != "" } // getDiscoveryClient returns a ClickHouse client that honors whichever kind diff --git a/pkg/server/server_dynamic_tools_test.go b/pkg/server/server_dynamic_tools_test.go index 302ecc3e..3ef6415d 100644 --- a/pkg/server/server_dynamic_tools_test.go +++ b/pkg/server/server_dynamic_tools_test.go @@ -1115,14 +1115,15 @@ func TestHasDiscoveryCredentials(t *testing.T) { t.Run("jwe_token_present", func(t *testing.T) { t.Parallel() - s := &ClickHouseJWEServer{} + s := &ClickHouseJWEServer{Config: config.Config{Server: config.ServerConfig{JWE: config.JWEConfig{Enabled: true}}}} ctx := context.WithValue(context.Background(), JWETokenKey, "jwe-abc") + ctx = context.WithValue(ctx, JWEClaimsKey, map[string]interface{}{"host": "tenant.example", "username": "alice"}) require.True(t, s.hasDiscoveryCredentials(ctx)) }) t.Run("oauth_token_present", func(t *testing.T) { t.Parallel() - s := &ClickHouseJWEServer{} + s := &ClickHouseJWEServer{Config: config.Config{Server: config.ServerConfig{OAuth: config.OAuthConfig{Enabled: true}}}} ctx := context.WithValue(context.Background(), OAuthTokenKey, "oauth-xyz") require.True(t, s.hasDiscoveryCredentials(ctx)) }) diff --git a/pkg/server/server_openapi.go b/pkg/server/server_openapi.go index 01a59787..9d5fcde9 100644 --- a/pkg/server/server_openapi.go +++ b/pkg/server/server_openapi.go @@ -72,19 +72,14 @@ func (s *ClickHouseJWEServer) OpenAPIHandler(w http.ResponseWriter, r *http.Requ if !ok { return } - // Ensure dynamic tools are loaded - if err := s.EnsureDynamicTools(r.Context()); err != nil { - log.Warn().Err(err).Msg("Failed to ensure dynamic tools in OpenAPI handler") - } + catalog := s.requestCatalog(r.Context()) // dynamic tool endpoint: /openapi/{tool} parts := strings.Split(r.URL.Path, "/openapi/") if len(parts) == 2 { tool := strings.Trim(parts[1], "/") - s.dynamicToolsMu.RLock() - meta, ok := s.dynamicTools[tool] - s.dynamicToolsMu.RUnlock() + meta, ok := catalog[tool] if ok { s.handleDynamicToolOpenAPIWithBody(w, r, meta, body) @@ -99,10 +94,7 @@ func (s *ClickHouseJWEServer) OpenAPIHandler(w http.ResponseWriter, r *http.Requ } func (s *ClickHouseJWEServer) ServeOpenAPISchema(w http.ResponseWriter, r *http.Request) { - // Ensure dynamic tools are loaded - if err := s.EnsureDynamicTools(r.Context()); err != nil { - log.Warn().Err(err).Msg("Failed to ensure dynamic tools in ServeOpenAPISchema") - } + catalog := s.requestCatalog(r.Context()) // Get host URL based on OpenAPI TLS configuration protocol := "http" @@ -210,11 +202,8 @@ func (s *ClickHouseJWEServer) ServeOpenAPISchema(w http.ResponseWriter, r *http. } } - s.dynamicToolsMu.RLock() - defer s.dynamicToolsMu.RUnlock() - for _, prefix := range s.openAPIPathPrefixes() { - for toolName, meta := range s.dynamicTools { + for toolName, meta := range catalog { path := prefix + "/openapi/" + toolName props := map[string]interface{}{} var required []string diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go index 4eb360af..7fe9eaaf 100644 --- a/pkg/server/server_test.go +++ b/pkg/server/server_test.go @@ -432,6 +432,7 @@ func TestOpenAPI_SchemaIncludesCombinedAuthPaths(t *testing.T) { OAuth: config.OAuthConfig{ Enabled: true, }, + DynamicTools: []config.DynamicToolRule{{Regexp: ".*"}}, }, }, dynamicTools: map[string]dynamicToolMeta{ @@ -444,9 +445,14 @@ func TestOpenAPI_SchemaIncludesCombinedAuthPaths(t *testing.T) { }, } + s.catalogOnce.Do(func() { s.catalogCache = NewCatalogCache(config.MulticlusterConfig{}) }) + defer s.Close() + s.catalogCache.insertOK(fullKey(CacheKey("oauth\x00caller"), "singlecluster"), s.dynamicTools, time.Time{}) + s.dynamicTools = nil rr := httptest.NewRecorder() req := httptest.NewRequest("GET", "/openapi", nil) ctx := context.WithValue(req.Context(), CHJWEServerKey, s) + ctx = context.WithValue(ctx, OAuthTokenKey, "caller") req = req.WithContext(ctx) s.ServeOpenAPISchema(rr, req) diff --git a/pkg/server/singlecluster_catalog.go b/pkg/server/singlecluster_catalog.go new file mode 100644 index 00000000..f6e841e5 --- /dev/null +++ b/pkg/server/singlecluster_catalog.go @@ -0,0 +1,106 @@ +package server + +import ( + "context" + "encoding/json" + "fmt" + "maps" + "net/http" + "time" + + "github.com/altinity/altinity-mcp/pkg/clickhouse" + "github.com/altinity/altinity-mcp/pkg/config" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rs/zerolog/log" +) + +// CredentialCatalogKey selects exactly the credential used for database access. +// Token kind is part of the hashed, unambiguous tuple; raw tokens never leave +// the request context or appear in cache diagnostics. +func (s *ClickHouseJWEServer) CredentialCatalogKey(ctx context.Context) (string, time.Time, error) { + if s.Config.Server.JWE.Enabled { + token := s.ExtractTokenFromCtx(ctx) + if token != "" { + claims := s.GetJWEClaimsFromCtx(ctx) + if claims == nil { + var err error + claims, err = s.ParseJWEClaims(token) + if err != nil { + return "", time.Time{}, fmt.Errorf("invalid JWE credential") + } + } + if s.JWEClaimsHaveCredentials(claims) { + exp := time.Time{} + switch v := claims["exp"].(type) { + case float64: + exp = time.Unix(int64(v), 0) + case int64: + exp = time.Unix(v, 0) + case json.Number: + if n, err := v.Int64(); err == nil { + exp = time.Unix(n, 0) + } + } + return CacheKey("jwe\x00" + token), exp, nil + } + } + } + if s.Config.Server.OAuth.Enabled { + if token := s.ExtractOAuthTokenFromCtx(ctx); token != "" { + exp, _ := BearerExp(token) + return CacheKey("oauth\x00" + token), exp, nil + } + } + return "", time.Time{}, fmt.Errorf("missing usable authentication token") +} + +// Close releases this configuration generation's catalog janitor. Requests +// holding the retired immutable generation can complete safely. +func (s *ClickHouseJWEServer) Close() { + s.catalogOnce.Do(func() {}) + if s.catalogCache != nil { + s.catalogCache.Close() + } +} + +func (s *ClickHouseJWEServer) requestCatalog(ctx context.Context) map[string]dynamicToolMeta { + if !s.Config.Server.JWE.Enabled && !s.Config.Server.OAuth.Enabled { + if err := s.EnsureDynamicTools(ctx); err != nil { + log.Warn().Err(err).Msg("dynamic catalog discovery failed") + } + s.dynamicToolsMu.RLock() + defer s.dynamicToolsMu.RUnlock() + return maps.Clone(s.dynamicTools) + } + key, exp, err := s.CredentialCatalogKey(ctx) + if err != nil { + return nil + } + if len(s.Config.Server.DynamicTools) == 0 { + return nil + } + s.catalogOnce.Do(func() { s.catalogCache = NewCatalogCache(config.MulticlusterConfig{}) }) + if s.catalogCache == nil { + return nil + } + tools, err := s.catalogCache.GetOrDiscover(ctx, key, "singlecluster", s.Config.ClickHouse, + func(ctx context.Context, cfg config.ClickHouseConfig) (*clickhouse.Client, error) { + return s.GetClickHouseClientWithOAuthForConfig(ctx, cfg, s.ExtractTokenFromCtx(ctx), s.ExtractOAuthTokenFromCtx(ctx), s.GetOAuthClaimsFromCtx(ctx)) + }, s.Config.Server.DynamicTools, s.Config.ClickHouse.ReadOnly, exp) + if err != nil { + log.Warn().Err(err).Msg("credential catalog discovery failed; static tools remain available") + return nil + } + return tools +} + +// GetServer creates a caller's MCP registry without changing the parent registry. +func (s *ClickHouseJWEServer) GetServer(r *http.Request) *mcp.Server { + if !s.Config.Server.JWE.Enabled && !s.Config.Server.OAuth.Enabled { + _ = s.EnsureDynamicTools(r.Context()) + return s.MCPServer + } + f := &MulticlusterServerFactory{cfg: s.Config, version: s.Version, + instrName: "Altinity ClickHouse MCP Server", instrTitle: "Altinity ClickHouse MCP Server"} + return f.newServer(s.requestCatalog(r.Context())) +} diff --git a/pkg/server/singlecluster_catalog_test.go b/pkg/server/singlecluster_catalog_test.go new file mode 100644 index 00000000..72dfa622 --- /dev/null +++ b/pkg/server/singlecluster_catalog_test.go @@ -0,0 +1,66 @@ +package server + +import ( + "context" + "testing" + "time" + + "github.com/altinity/altinity-mcp/pkg/config" + "github.com/stretchr/testify/require" +) + +func TestCredentialCatalogKeyEffectiveIdentity(t *testing.T) { + s := NewClickHouseMCPServer(config.Config{ClickHouse: config.ClickHouseConfig{Username: "server"}, Server: config.ServerConfig{JWE: config.JWEConfig{Enabled: true}, OAuth: config.OAuthConfig{Enabled: true}}}, "test") + defer s.Close() + ctx := context.WithValue(context.Background(), JWETokenKey, "same-token") + ctx = context.WithValue(ctx, JWEClaimsKey, map[string]interface{}{"host": "tenant.example", "username": "alice", "exp": float64(time.Now().Add(time.Hour).Unix())}) + ctx = context.WithValue(ctx, OAuthTokenKey, "ignored-oauth") + key, exp, err := s.CredentialCatalogKey(ctx) + require.NoError(t, err) + require.Equal(t, CacheKey("jwe\x00same-token"), key) + require.False(t, exp.IsZero()) + ctx = context.WithValue(ctx, OAuthTokenKey, "different-ignored-oauth") + again, _, err := s.CredentialCatalogKey(ctx) + require.NoError(t, err) + require.Equal(t, key, again) + ctx = context.WithValue(ctx, JWEClaimsKey, map[string]interface{}{"host": "tenant.example"}) + key, _, err = s.CredentialCatalogKey(ctx) + require.NoError(t, err) + require.Equal(t, CacheKey("oauth\x00different-ignored-oauth"), key) + ctx = context.WithValue(ctx, OAuthTokenKey, "same-token") + oauthKey, _, err := s.CredentialCatalogKey(ctx) + require.NoError(t, err) + require.NotEqual(t, again, oauthKey) + require.False(t, s.hasDiscoveryCredentials(context.Background()), "static username cannot enable auth discovery") + require.NoError(t, s.EnsureDynamicTools(ctx)) + require.Empty(t, s.dynamicTools) + require.False(t, s.dynamicToolsInit) +} + +func TestCredentialCatalogCacheLifetime(t *testing.T) { + s := NewClickHouseMCPServer(config.Config{Server: config.ServerConfig{OAuth: config.OAuthConfig{Enabled: true}, Tools: []config.ToolDefinition{{Type: "read", ViewRegexp: ".*"}}}}, "test") + s.catalogOnce.Do(func() { s.catalogCache = NewCatalogCache(config.MulticlusterConfig{}) }) + ctx := context.WithValue(context.Background(), OAuthTokenKey, "alice") + key, _, err := s.CredentialCatalogKey(ctx) + require.NoError(t, err) + tools := map[string]dynamicToolMeta{"alice": {ToolName: "alice"}} + s.catalogCache.insertOK(fullKey(key, "singlecluster"), tools, time.Now().Add(time.Minute)) + require.Equal(t, tools, s.requestCatalog(ctx)) + require.Equal(t, uint64(1), s.catalogCache.Metrics.HitsOK.Load()) + s.catalogCache.mu.Lock() + entry := s.catalogCache.entries[fullKey(key, "singlecluster")] + entry.ExpiresAt = time.Now().Add(-time.Second) + s.catalogCache.entries[fullKey(key, "singlecluster")] = entry + s.catalogCache.mu.Unlock() + // An expired entry must be rediscovered; the invalid endpoint cannot return + // the previously cached metadata. + require.Empty(t, s.requestCatalog(ctx)) + require.Equal(t, uint64(1), s.catalogCache.Metrics.Misses.Load()) + s.Close() + s.Close() + require.True(t, s.catalogCache.stopped.Load()) + closed := NewClickHouseMCPServer(s.Config, "test") + closed.Close() + require.Empty(t, closed.requestCatalog(ctx)) + require.Nil(t, closed.catalogCache) +} From 8329e3aeec9ec77f9dfaa2e952bdafd2ee1b68e1 Mon Sep 17 00:00:00 2001 From: BorisTyshkevich <125603338+BorisTyshkevich@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:20:37 +0000 Subject: [PATCH 2/3] Reject authenticated database access without usable request credentials --- CHANGELOG.md | 2 +- cmd/altinity-mcp/credential_transport_test.go | 83 +++++++++++++++++++ pkg/server/server_client.go | 48 +++-------- pkg/server/singlecluster_catalog_test.go | 72 ++++++++++++++++ 4 files changed, 168 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d61a0d66..74107e6c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ # Unreleased BUG FIXES -- security: authenticate the exact `/openapi` schema route and isolate single-cluster HTTP/SSE dynamic tools and schemas by effective request credentials; reset catalog generations on reload and reject SSE messages from other credentials or retired generations (fixes [#182](https://github.com/Altinity/altinity-mcp/issues/182)) +- security: authenticate the exact `/openapi` schema route, reject authenticated database requests without usable tokens, and isolate single-cluster HTTP/SSE dynamic tools and schemas by effective request credentials; reset catalog generations on reload and reject SSE messages from other credentials or retired generations (fixes [#182](https://github.com/Altinity/altinity-mcp/issues/182)) - security: require self-contained JWE ClickHouse connections with host and username; prevent inheritance of operator endpoints, credentials, headers, roles, and TLS material; restrict token TLS file paths to `server.jwe.tls_material_dir` (fixes [#179](https://github.com/Altinity/altinity-mcp/issues/179)). Missing request tokens now fail closed in JWE-only and combined JWE/OAuth client paths, including discovery, instead of using static operator credentials. Regenerate tokens that relied on static connection defaults. - security: OpenAPI `execute_query` always rejects write statements and enforces `max_query_length`; dynamic-tool JSON bodies use the same input-size limit (fixes [#183](https://github.com/Altinity/altinity-mcp/issues/183)) - security: isolate per-request ClickHouse headers, settings, and roles; prevent concurrent OAuth bearer forwarding from crashing the server, confusing caller identities, or leaking a stale bearer into later JWE or static-credential requests (fixes [#181](https://github.com/Altinity/altinity-mcp/issues/181)) diff --git a/cmd/altinity-mcp/credential_transport_test.go b/cmd/altinity-mcp/credential_transport_test.go index 0851f44f..b00f6c93 100644 --- a/cmd/altinity-mcp/credential_transport_test.go +++ b/cmd/altinity-mcp/credential_transport_test.go @@ -606,3 +606,86 @@ func TestCredentialCatalogReservedOpenAPIToolName(t *testing.T) { }) } } + +func TestCredentialCatalogPartialJWEUsesOAuthEndpoint(t *testing.T) { + for _, sse := range []bool{false, true} { + t.Run(fmt.Sprintf("sse=%t", sse), func(t *testing.T) { + operator := newCatalogUpstream(t) + rejected := newCatalogUpstream(t) + cfg := catalogConfig(operator.chConfig(t), true, true) + app := catalogApp(t, cfg) + var handler http.Handler = app.buildHTTPHandler(cfg, app.mcpServer.MCPServer) + suffix := "" + if sse { + handler = app.buildSSEHandler(cfg, app.mcpServer.MCPServer) + suffix = "/sse" + } + web := httptest.NewServer(handler) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + for _, partial := range []string{"host_only", "username_only"} { + t.Run(partial, func(t *testing.T) { + ch := rejected.chConfig(t) + claims := map[string]interface{}{"port": ch.Port, "protocol": "http", "password": "unused-password", "exp": time.Now().Add(time.Hour).Unix()} + if partial == "host_only" { + claims["host"] = ch.Host + } else { + claims["username"] = "unused-user" + } + token, err := jwe_auth.GenerateJWEToken(claims, []byte(cfg.Server.JWE.JWESecretKey), []byte(cfg.Server.JWE.JWTSecretKey)) + require.NoError(t, err) + tr := &credentialRoundTripper{token: "alice", jwe: token} + session := catalogSession(t, ctx, web.URL+suffix, tr, sse) + assertCatalogSession(t, ctx, session, "alice", "bob") + catalogSchema(t, web.URL, tr, "alice", "bob") + catalogREST(t, web.URL, tr, "alice", "bob") + }) + } + require.Zero(t, rejected.requests.Load(), "partial JWE routing and credentials must never be used") + operator.mu.Lock() + require.Equal(t, 1, operator.discovers["alice"], "partial JWEs share the effective OAuth cache key") + require.Equal(t, 4, operator.queries["alice"]) + operator.mu.Unlock() + }) + } +} + +func TestLegacyClientEffectiveCredentials(t *testing.T) { + for _, auth := range []string{"oauth", "jwe", "combined"} { + t.Run(auth, func(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), auth != "oauth", auth != "jwe") + app := catalogApp(t, cfg) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + tokenParam := "" + if auth == "oauth" { + ctx = context.WithValue(ctx, altinitymcp.OAuthTokenKey, "alice") + } else { + tokenParam = catalogJWE(t, cfg, cfg.ClickHouse, "alice") + // The context intentionally contains a different validated JWE. The + // legacy API's explicit argument must retain its historical precedence. + other := catalogJWE(t, cfg, cfg.ClickHouse, "bob") + claims, err := app.mcpServer.ParseJWEClaims(other) + require.NoError(t, err) + ctx = context.WithValue(ctx, altinitymcp.JWETokenKey, other) + ctx = context.WithValue(ctx, altinitymcp.JWEClaimsKey, claims) + if auth == "combined" { + ctx = context.WithValue(ctx, altinitymcp.OAuthTokenKey, "bob") + } + } + client, err := app.mcpServer.GetClickHouseClient(ctx, tokenParam) + require.NoError(t, err) + defer client.Close() + result, err := client.ExecuteQuery(ctx, "SELECT * FROM tenant.alice") + require.NoError(t, err) + require.Equal(t, [][]interface{}{{"alice"}}, result.Rows) + up.mu.Lock() + require.Equal(t, 1, up.queries["alice"]) + require.Zero(t, up.queries["bob"]) + require.Zero(t, up.queries["static"]) + up.mu.Unlock() + }) + } +} diff --git a/pkg/server/server_client.go b/pkg/server/server_client.go index 67e7445c..4e0f8b5a 100644 --- a/pkg/server/server_client.go +++ b/pkg/server/server_client.go @@ -17,43 +17,15 @@ import ( "github.com/rs/zerolog/log" ) -// GetClickHouseClient creates a self-contained JWE connection when JWE is -// enabled, and uses the operator connection only when JWE is disabled. +// GetClickHouseClient uses its explicit JWE token argument and any OAuth +// credential in the request context. All connection paths share the central +// authentication boundary; authenticated modes never use static credentials. func (s *ClickHouseJWEServer) GetClickHouseClient(ctx context.Context, tokenParam string) (*clickhouse.Client, error) { - var chConfig config.ClickHouseConfig - - if !s.Config.Server.JWE.Enabled { - chConfig = s.Config.ClickHouse.Clone() - } else { - if tokenParam == "" { - // JWE auth is enabled but no token provided - return nil, jwe_auth.ErrMissingToken - } - - // Parse and validate JWE token - claims, err := jwe_auth.ParseAndDecryptJWE(tokenParam, []byte(s.Config.Server.JWE.JWESecretKey), []byte(s.Config.Server.JWE.JWTSecretKey)) - if err != nil { - log.Error().Err(err).Msg("failed to parse/decrypt JWE token") - return nil, err - } - - var buildErr error - // Create ClickHouse config from JWE claims - chConfig, buildErr = s.buildConfigFromClaims(claims) - if buildErr != nil { - return nil, buildErr - } - } - - client, err := clickhouse.NewClient(ctx, chConfig) - if err != nil { - if s.Config.Server.JWE.Enabled && (chConfig.TLS.CaCert != "" || chConfig.TLS.ClientCert != "" || chConfig.TLS.ClientKey != "") { - return nil, fmt.Errorf("jwe: failed to create ClickHouse client with TLS material") - } - return nil, fmt.Errorf("failed to create ClickHouse client: %w", err) - } - - return client, nil + // The legacy API has always validated the explicit argument. A parsed JWE + // from another context token must not override it. + ctx = context.WithValue(ctx, JWEClaimsKey, map[string]interface{}(nil)) + return s.GetClickHouseClientWithOAuthForConfig(ctx, CHConfigFromContext(ctx, s.Config.ClickHouse), + tokenParam, s.ExtractOAuthTokenFromCtx(ctx), s.GetOAuthClaimsFromCtx(ctx)) } // buildConfigFromClaims builds a self-contained connection from token claims and @@ -259,6 +231,10 @@ func (s *ClickHouseJWEServer) GetClickHouseClientWithOAuthForConfig(ctx context. return nil, jwe_auth.ErrMissingToken } + if !useJWE && s.Config.Server.OAuth.Enabled && oauthToken == "" { + return nil, ErrMissingOAuthToken + } + // Merge tool-input settings before OAuth so probe configs carry them. if toolSettings := ToolInputSettingsFromContext(ctx); len(toolSettings) > 0 { chConfig = mergeExtraSettings(chConfig, toolSettings) diff --git a/pkg/server/singlecluster_catalog_test.go b/pkg/server/singlecluster_catalog_test.go index 72dfa622..45b9d423 100644 --- a/pkg/server/singlecluster_catalog_test.go +++ b/pkg/server/singlecluster_catalog_test.go @@ -6,6 +6,7 @@ import ( "time" "github.com/altinity/altinity-mcp/pkg/config" + "github.com/altinity/go-mcp-oauth-sdk/jwe_auth" "github.com/stretchr/testify/require" ) @@ -64,3 +65,74 @@ func TestCredentialCatalogCacheLifetime(t *testing.T) { require.Empty(t, closed.requestCatalog(ctx)) require.Nil(t, closed.catalogCache) } + +func TestAuthenticatedClientBoundaryNoStaticFallback(t *testing.T) { + for _, mode := range []string{"jwe", "oauth", "combined"} { + for _, credentials := range []string{"none", "disabled_token"} { + t.Run(mode+"/"+credentials, func(t *testing.T) { + srv, calls := openAPIGuardServer(t, 100) + srv.Config.Server.JWE.Enabled = mode != "oauth" + srv.Config.Server.OAuth.Enabled = mode != "jwe" + ctx := context.Background() + tokenParam := "" + if credentials == "disabled_token" { + if mode == "oauth" { + ctx = context.WithValue(ctx, JWETokenKey, "disabled-jwe-token") + tokenParam = "disabled-jwe-token" + } else if mode == "jwe" { + ctx = context.WithValue(ctx, OAuthTokenKey, "disabled-oauth-token") + } + } + client, err := srv.GetClickHouseClientFromCtx(ctx) + if mode == "oauth" { + require.ErrorIs(t, err, ErrMissingOAuthToken) + } else { + require.ErrorIs(t, err, jwe_auth.ErrMissingToken) + } + require.Error(t, err) + require.Nil(t, client) + require.Zero(t, calls.Load()) + client, err = srv.GetClickHouseClient(ctx, tokenParam) + require.Error(t, err) + require.Nil(t, client) + require.Zero(t, calls.Load()) + client, err = srv.GetClickHouseClientWithOAuthForConfig(ctx, srv.Config.ClickHouse, tokenParam, srv.ExtractOAuthTokenFromCtx(ctx), nil) + require.Error(t, err) + require.Nil(t, client) + require.Zero(t, calls.Load()) + }) + } + } +} + +func TestLegacyClientExplicitJWEArgument(t *testing.T) { + srv, calls := openAPIGuardServer(t, 100) + srv.Config.Server.JWE = config.JWEConfig{Enabled: true, JWESecretKey: "this-is-a-32-byte-secret-key!!", JWTSecretKey: "fake-jwt-secret"} + ctx := context.WithValue(context.Background(), JWETokenKey, "different-token") + ctx = context.WithValue(ctx, JWEClaimsKey, map[string]interface{}{"host": srv.Config.ClickHouse.Host, "port": float64(srv.Config.ClickHouse.Port), "username": "fake-user"}) + client, err := srv.GetClickHouseClient(ctx, "") + require.Error(t, err) + require.Nil(t, client) + require.Zero(t, calls.Load()) + client, err = srv.GetClickHouseClient(ctx, "invalid-explicit-token") + require.Error(t, err) + require.Nil(t, client) + require.Zero(t, calls.Load()) +} + +func TestCombinedClientPartialJWEWithoutOAuthNoIO(t *testing.T) { + srv, calls := openAPIGuardServer(t, 100) + srv.Config.Server.JWE = config.JWEConfig{Enabled: true, JWESecretKey: "this-is-a-32-byte-secret-key!!", JWTSecretKey: "fake-jwt-secret"} + srv.Config.Server.OAuth.Enabled = true + token, err := jwe_auth.GenerateJWEToken(map[string]interface{}{"username": "alice", "exp": time.Now().Add(time.Hour).Unix()}, []byte(srv.Config.Server.JWE.JWESecretKey), []byte(srv.Config.Server.JWE.JWTSecretKey)) + require.NoError(t, err) + ctx := context.WithValue(context.Background(), JWETokenKey, token) + client, err := srv.GetClickHouseClientFromCtx(ctx) + require.ErrorIs(t, err, ErrJWEIncompleteConnection) + require.Nil(t, client) + require.Zero(t, calls.Load()) + client, err = srv.GetClickHouseClient(ctx, token) + require.ErrorIs(t, err, ErrJWEIncompleteConnection) + require.Nil(t, client) + require.Zero(t, calls.Load()) +} From a0fd29af56dbf0308a4e93816571890b5c569603 Mon Sep 17 00:00:00 2001 From: BorisTyshkevich <125603338+BorisTyshkevich@users.noreply.github.com> Date: Wed, 7 Oct 2026 08:39:00 +0000 Subject: [PATCH 3/3] Preserve live catalogs on unchanged reload and fix concurrent logging --- CHANGELOG.md | 2 +- README.md | 2 +- cmd/altinity-mcp/credential_transport.go | 15 +- cmd/altinity-mcp/credential_transport_test.go | 172 +++++++++++++++++- cmd/altinity-mcp/logging_setup_test.go | 48 +++++ cmd/altinity-mcp/main.go | 30 ++- cmd/altinity-mcp/main_test.go | 1 + docs/jwe_authentication.md | 2 +- docs/oauth_authorization.md | 2 +- docs/tools.md | 4 +- helm/altinity-mcp/README.md | 2 +- pkg/server/multicluster_factory.go | 4 +- pkg/server/registry_logs_test.go | 62 +++++++ pkg/server/server.go | 30 ++- pkg/server/server_resources.go | 16 +- pkg/server/singlecluster_catalog.go | 21 ++- 16 files changed, 376 insertions(+), 37 deletions(-) create mode 100644 cmd/altinity-mcp/logging_setup_test.go create mode 100644 pkg/server/registry_logs_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 74107e6c..2ade275f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ # Unreleased BUG FIXES -- security: authenticate the exact `/openapi` schema route, reject authenticated database requests without usable tokens, and isolate single-cluster HTTP/SSE dynamic tools and schemas by effective request credentials; reset catalog generations on reload and reject SSE messages from other credentials or retired generations (fixes [#182](https://github.com/Altinity/altinity-mcp/issues/182)) +- security: authenticate the exact `/openapi` schema route, reject authenticated database requests without usable tokens, and isolate single-cluster HTTP/SSE dynamic tools and schemas by effective request credentials; preserve sessions on unchanged reload polls, safely update logging levels during requests, and reject SSE messages from other credentials or retired generations (fixes [#182](https://github.com/Altinity/altinity-mcp/issues/182)) - security: require self-contained JWE ClickHouse connections with host and username; prevent inheritance of operator endpoints, credentials, headers, roles, and TLS material; restrict token TLS file paths to `server.jwe.tls_material_dir` (fixes [#179](https://github.com/Altinity/altinity-mcp/issues/179)). Missing request tokens now fail closed in JWE-only and combined JWE/OAuth client paths, including discovery, instead of using static operator credentials. Regenerate tokens that relied on static connection defaults. - security: OpenAPI `execute_query` always rejects write statements and enforces `max_query_length`; dynamic-tool JSON bodies use the same input-size limit (fixes [#183](https://github.com/Altinity/altinity-mcp/issues/183)) - security: isolate per-request ClickHouse headers, settings, and roles; prevent concurrent OAuth bearer forwarding from crashing the server, confusing caller identities, or leaking a stale bearer into later JWE or static-credential requests (fixes [#181](https://github.com/Altinity/altinity-mcp/issues/181)) diff --git a/README.md b/README.md index c061a89d..d8764363 100644 --- a/README.md +++ b/README.md @@ -349,7 +349,7 @@ OpenAI MCP/tooling references: - **With both**: A self-contained JWE takes priority; otherwise OAuth supplies the credentials. - **With both disabled**: Use server-configured ClickHouse credentials. -The exact `/openapi` schema route also requires authentication when either mode is enabled. Dynamic HTTP/SSE tools and OpenAPI schemas use catalogs isolated by the caller's effective credential; server credentials are never a fallback under authentication. Reloading configuration resets catalog generations and requires existing SSE clients to reconnect. See [dynamic discovery](docs/tools.md#dynamic-discovery). +The exact `/openapi` schema route also requires authentication when either mode is enabled. Dynamic HTTP/SSE tools and OpenAPI schemas use catalogs isolated by the caller's effective credential; server credentials are never a fallback under authentication. An effective configuration change resets catalog generations and requires existing SSE clients to reconnect. Unchanged reload polls preserve sessions. See [dynamic discovery](docs/tools.md#dynamic-discovery). ### Available Actions diff --git a/cmd/altinity-mcp/credential_transport.go b/cmd/altinity-mcp/credential_transport.go index 18588d3b..720a3cfe 100644 --- a/cmd/altinity-mcp/credential_transport.go +++ b/cmd/altinity-mcp/credential_transport.go @@ -2,8 +2,8 @@ package main import ( "context" + "errors" "fmt" - "github.com/rs/zerolog/log" "net/http" "strings" "sync" @@ -11,6 +11,7 @@ import ( "github.com/altinity/altinity-mcp/pkg/config" altinitymcp "github.com/altinity/altinity-mcp/pkg/server" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rs/zerolog/log" ) // withServerSnapshot captures the immutable parent and its converted tool rules @@ -42,10 +43,18 @@ func (a *application) authenticateSnapshot(next http.Handler) http.Handler { } if cfg.Server.JWE.Enabled { token, claims, _, _, err := parent.ValidateAuth(r) - if err != nil || !parent.JWEClaimsHaveCredentials(claims) { + if errors.Is(err, altinitymcp.ErrJWEIncompleteConnection) { + http.Error(w, altinitymcp.ErrJWEIncompleteConnection.Error(), http.StatusUnauthorized) + return + } + if err != nil { http.Error(w, "Missing or invalid authentication token", http.StatusUnauthorized) return } + if !parent.JWEClaimsHaveCredentials(claims) { + http.Error(w, altinitymcp.ErrJWEIncompleteConnection.Error(), http.StatusUnauthorized) + return + } ctx := context.WithValue(r.Context(), altinitymcp.JWETokenKey, token) ctx = context.WithValue(ctx, altinitymcp.JWEClaimsKey, claims) r = r.WithContext(ctx) @@ -94,11 +103,9 @@ func (a *application) buildSingleClusterHandler(cfg config.Config, sse bool) htt mux.Handle(pattern, handler) } if cfg.Server.OpenAPI.Enabled { - for _, pattern := range openAPIRoutePatterns(cfg.Server.JWE.Enabled, cfg.Server.OAuth.Enabled) { mux.Handle(pattern, openapiHandler) } - if sse && cfg.Server.JWE.Enabled && cfg.Server.OAuth.Enabled { mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if strings.HasPrefix(r.URL.Path, "/openapi/") { diff --git a/cmd/altinity-mcp/credential_transport_test.go b/cmd/altinity-mcp/credential_transport_test.go index b00f6c93..8e9d708b 100644 --- a/cmd/altinity-mcp/credential_transport_test.go +++ b/cmd/altinity-mcp/credential_transport_test.go @@ -522,6 +522,7 @@ func TestCredentialCatalogConcurrentReloadSnapshots(t *testing.T) { newUp.marker = "_new" oldCfg := catalogConfig(oldUp.chConfig(t), false, true) newCfg := catalogConfig(newUp.chConfig(t), false, true) + newCfg.Logging.Level = config.DebugLevel newCfg.Server.Tools = []config.ToolDefinition{{Type: "read", ViewRegexp: "^tenant\\..*", Prefix: "next_"}} app := catalogApp(t, oldCfg) web := httptest.NewServer(app.buildHTTPHandler(oldCfg, app.mcpServer.MCPServer)) @@ -560,9 +561,7 @@ func TestCredentialCatalogConcurrentReloadSnapshots(t *testing.T) { body := string(data) old := strings.Contains(body, "dyn_tenant_alice_old") next := strings.Contains(body, "next_tenant_alice_new") - // A generation can retire before its cache is initialized; static-only - // completion remains safe, but a mixed catalog must never appear. - if strings.Contains(body, "dyn_tenant_alice_new") || strings.Contains(body, "next_tenant_alice_old") || (old && next) { + if strings.Contains(body, "dyn_tenant_alice_new") || strings.Contains(body, "next_tenant_alice_old") || (old == next) { t.Errorf("mixed catalog generation: %s", body) return } @@ -689,3 +688,170 @@ func TestLegacyClientEffectiveCredentials(t *testing.T) { }) } } + +func TestCredentialCatalogUnchangedPeriodicReload(t *testing.T) { + up := newCatalogUpstream(t) + unused := newCatalogUpstream(t) + path := filepath.Join(t.TempDir(), "config.json") + cfg := catalogConfig(up.chConfig(t), false, true) + cfg.ReloadTime = 1 + data, err := json.Marshal(cfg) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, data, 0600)) + cmd := &mockCommand{flags: map[string]interface{}{"clickhouse-port": cfg.ClickHouse.Port}, setFlags: map[string]bool{"clickhouse-port": true}, stringMaps: map[string]map[string]string{}} + initial, err := config.LoadConfigFromFile(path) + require.NoError(t, err) + overrideWithCLIFlags(initial, cmd) + initial.RemovedKeyWarnings = []string{"previous diagnostic only"} + app := catalogApp(t, *initial) + app.configFile = path + app.stopConfigReload = make(chan struct{}) + web := httptest.NewServer(app.buildSSEHandler(*initial, app.mcpServer.MCPServer)) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + tr := &credentialRoundTripper{token: "alice"} + session := catalogSession(t, ctx, web.URL+"/sse", tr, true) + assertCatalogSession(t, ctx, session, "alice", "bob") + initialParent := app.mcpServer + tr.endpointMu.Lock() + initialEndpoint := tr.endpoint + tr.endpointMu.Unlock() + // CLI overrides the on-disk endpoint. Restart-only metrics and multicluster + // settings also must not replace the effective catalog generation. ReloadTime + // is bookkeeping; observing it below proves the periodic tick completed. + next := *initial + next.ReloadTime = 2 + next.ClickHouse.Port = unused.chConfig(t).Port + next.Server.Metrics.Enabled = true + next.Multicluster.CatalogCacheMax = initial.Multicluster.CatalogCacheMax + 1 + data, err = json.Marshal(next) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path, data, 0600)) + loopCtx, stopLoop := context.WithCancel(ctx) + loopDone := make(chan struct{}) + go func() { defer close(loopDone); app.configReloadLoop(loopCtx, cmd) }() + t.Cleanup(func() { stopLoop(); <-loopDone }) + require.Eventually(t, func() bool { return app.GetCurrentConfig().ReloadTime == 2 }, 5*time.Second, 10*time.Millisecond) + app.configMutex.RLock() + parent := app.mcpServer + app.configMutex.RUnlock() + require.Same(t, initialParent, parent) + current := app.GetCurrentConfig() + require.Equal(t, initial.Multicluster, current.Multicluster) + require.False(t, current.Server.Metrics.Enabled) + require.Empty(t, current.RemovedKeyWarnings) + assertCatalogSession(t, ctx, session, "alice", "bob") + catalogSchema(t, web.URL, tr, "alice", "bob") + tr.endpointMu.Lock() + require.Equal(t, initialEndpoint, tr.endpoint) + tr.endpointMu.Unlock() + require.NoError(t, app.reloadConfig(cmd)) + app.configMutex.RLock() + parent = app.mcpServer + app.configMutex.RUnlock() + require.Same(t, initialParent, parent) + up.mu.Lock() + require.Equal(t, 1, up.discovers["alice"], "unchanged polls preserve the credential cache") + up.mu.Unlock() + require.Zero(t, unused.requests.Load()) +} + +func TestCredentialCatalogRetiredSnapshotCompletesDiscovery(t *testing.T) { + oldUp := newCatalogUpstream(t) + oldUp.marker = "_old" + newUp := newCatalogUpstream(t) + newUp.marker = "_new" + cfg := catalogConfig(oldUp.chConfig(t), false, true) + app := catalogApp(t, cfg) + oldParent := app.mcpServer + request := httptest.NewRequest(http.MethodPost, "/openapi/dyn_tenant_alice_old", strings.NewReader(`{"alice_old_id":1}`)) + request.Header.Set("Authorization", "Bearer alice") + response := httptest.NewRecorder() + // Capture through production snapshot middleware, then retire the captured + // generation before it has ever created a catalog cache. + app.withServerSnapshot(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + next := cfg + next.ClickHouse = newUp.chConfig(t) + writeCatalogReload(t, app, next) + captured := altinitymcp.GetClickHouseJWEServerFromContext(r.Context()) + require.Same(t, oldParent, captured) + captured.OpenAPIHandler(w, r) + })).ServeHTTP(response, request) + require.Equal(t, 200, response.Code, response.Body.String()) + require.Contains(t, response.Body.String(), "alice_old") + oldUp.mu.Lock() + require.Equal(t, 1, oldUp.discovers["alice_old"]) + require.Equal(t, 1, oldUp.queries["alice_old"]) + oldUp.mu.Unlock() + require.Zero(t, newUp.requests.Load()) +} + +func TestCredentialCatalogJWEIncompleteTransportMessage(t *testing.T) { + for _, sse := range []bool{false, true} { + for _, partial := range []string{"host_only", "username_only"} { + t.Run(fmt.Sprintf("%s/sse=%t", partial, sse), func(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), true, false) + app := catalogApp(t, cfg) + claims := map[string]interface{}{"exp": time.Now().Add(time.Hour).Unix()} + if partial == "host_only" { + claims["host"] = cfg.ClickHouse.Host + } else { + claims["username"] = "alice" + } + token, err := jwe_auth.GenerateJWEToken(claims, []byte(cfg.Server.JWE.JWESecretKey), []byte(cfg.Server.JWE.JWTSecretKey)) + require.NoError(t, err) + var handler http.Handler = app.buildHTTPHandler(cfg, app.mcpServer.MCPServer) + suffix := "" + method := http.MethodPost + if sse { + handler = app.buildSSEHandler(cfg, app.mcpServer.MCPServer) + suffix = "/sse" + method = http.MethodGet + } + web := httptest.NewServer(handler) + defer web.Close() + request, err := http.NewRequest(method, web.URL+"/"+token+suffix, strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + require.NoError(t, err) + body, err := io.ReadAll(response.Body) + response.Body.Close() + require.NoError(t, err) + require.Equal(t, 401, response.StatusCode) + require.Equal(t, altinitymcp.ErrJWEIncompleteConnection.Error(), strings.TrimSpace(string(body))) + require.Zero(t, up.requests.Load()) + }) + } + } +} + +func TestCredentialCatalogNoAuthUnifiedTransports(t *testing.T) { + for _, sse := range []bool{false, true} { + t.Run(fmt.Sprintf("sse=%t", sse), func(t *testing.T) { + up := newCatalogUpstream(t) + cfg := catalogConfig(up.chConfig(t), false, false) + app := catalogApp(t, cfg) + var handler http.Handler = app.buildHTTPHandler(cfg, app.mcpServer.MCPServer) + suffix := "" + if sse { + handler = app.buildSSEHandler(cfg, app.mcpServer.MCPServer) + suffix = "/sse" + } + web := httptest.NewServer(handler) + defer web.Close() + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + tr := &credentialRoundTripper{} + session := catalogSession(t, ctx, web.URL+suffix, tr, sse) + assertCatalogSession(t, ctx, session, "static", "foreign") + catalogSchema(t, web.URL, tr, "static", "foreign") + up.mu.Lock() + require.Equal(t, 1, up.discovers["static"]) + require.Equal(t, 1, up.queries["static"]) + up.mu.Unlock() + }) + } +} diff --git a/cmd/altinity-mcp/logging_setup_test.go b/cmd/altinity-mcp/logging_setup_test.go new file mode 100644 index 00000000..86e818d1 --- /dev/null +++ b/cmd/altinity-mcp/logging_setup_test.go @@ -0,0 +1,48 @@ +package main + +import ( + "sync" + "testing" + + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" + "github.com/stretchr/testify/require" +) + +func TestSetupLoggingConcurrentEvents(t *testing.T) { + // TestMain initializes formatting before any workers run, as production's + // CLI Before hook does. Concurrent reloads must change only the atomic level. + initialLogger := log.Logger + initialFormat := zerolog.TimeFieldFormat + initialLevel := zerolog.GlobalLevel() + t.Cleanup(func() { zerolog.SetGlobalLevel(initialLevel) }) + start := make(chan struct{}) + var wg sync.WaitGroup + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + for j := 0; j < 50; j++ { + if err := setupLogging([]string{"debug", "info", "warn", "error"}[j%4]); err != nil { + t.Error(err) + return + } + } + }() + } + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + for j := 0; j < 50; j++ { + log.Log().Timestamp().Int("event", j).Msg("concurrent logging setup regression") + } + }() + } + close(start) + wg.Wait() + require.Equal(t, initialLogger, log.Logger, "reloads must preserve the initialized logger and formatter") + require.Equal(t, initialFormat, zerolog.TimeFieldFormat) +} diff --git a/cmd/altinity-mcp/main.go b/cmd/altinity-mcp/main.go index 3e316487..d0c79c5a 100644 --- a/cmd/altinity-mcp/main.go +++ b/cmd/altinity-mcp/main.go @@ -35,8 +35,10 @@ var ( commit = "unknown" date = "unknown" - // loggingMutex protects global zerolog state during setupLogging calls - loggingMutex sync.Mutex + // Initialize the shared formatter before workers start; subsequent level + // changes use zerolog's atomic global level without replacing the logger. + loggingMutex sync.Mutex + loggingFormatOnce sync.Once ) func main() { @@ -117,9 +119,10 @@ func setupLogging(level string) error { loggingMutex.Lock() defer loggingMutex.Unlock() - // Configure zerolog - zerolog.TimeFieldFormat = time.RFC3339 - log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: "15:04:05"}) + loggingFormatOnce.Do(func() { + zerolog.TimeFieldFormat = time.RFC3339 + log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: "15:04:05"}) + }) // Set log level switch strings.ToLower(level) { @@ -1228,6 +1231,7 @@ func (a *application) reloadConfig(cmd CommandInterface) error { a.configMutex.RUnlock() if !reflect.DeepEqual(oldMC, newCfg.Multicluster) { log.Warn().Msg("config reload: multicluster.* fields changed — restart required for these to take effect; routing/cache remain on the previous configuration") + newCfg.Multicluster = oldMC } // server.metrics.* is restart-only for the same reason: the /metrics @@ -1254,9 +1258,21 @@ func (a *application) reloadConfig(cmd CommandInterface) error { newCfg.Server.Metrics = oldMetrics } - // Update logging level if changed + // File polling must not retire catalogs and live SSE streams when the + // effective configuration did not change. Reload bookkeeping is excluded. a.configMutex.Lock() - oldLogLevel := a.config.Logging.Level + current := a.config + oldEffective, newEffective := current, *newCfg + oldEffective.ReloadTime, newEffective.ReloadTime = 0, 0 + oldEffective.RemovedKeyWarnings, newEffective.RemovedKeyWarnings = nil, nil + if a.mcpServer != nil && reflect.DeepEqual(oldEffective, newEffective) { + a.config.ReloadTime = newCfg.ReloadTime + a.config.RemovedKeyWarnings = newCfg.RemovedKeyWarnings + a.configMutex.Unlock() + log.Debug().Msg("Configuration unchanged; preserving catalog generation") + return nil + } + oldLogLevel := current.Logging.Level a.configMutex.Unlock() if oldLogLevel != newCfg.Logging.Level { diff --git a/cmd/altinity-mcp/main_test.go b/cmd/altinity-mcp/main_test.go index 4ee5f4ae..d169a9ba 100644 --- a/cmd/altinity-mcp/main_test.go +++ b/cmd/altinity-mcp/main_test.go @@ -243,6 +243,7 @@ func TestRoutePatterns(t *testing.T) { "/{token}/openapi/list_tables", "/{token}/openapi/describe_table", "/{token}/openapi/execute_query", + "/openapi", "/openapi/list_tables", "/openapi/describe_table", "/openapi/execute_query", diff --git a/docs/jwe_authentication.md b/docs/jwe_authentication.md index 99ba08c3..64a14331 100644 --- a/docs/jwe_authentication.md +++ b/docs/jwe_authentication.md @@ -250,4 +250,4 @@ http://localhost:8080//sse - Set appropriate token expiration times - To implementing token revocation if needed for additional security, change using jwe-secret-key in `altinity-mcp` configuration -Dynamic catalogs and OpenAPI schemas are isolated by the effective request credential. The exact `/openapi` route also requires authentication; server credentials cannot supply anonymous discovery. In combined JWE/OAuth mode, a self-contained JWE takes priority and a partial JWE uses OAuth when available. SSE GET and POST must use the same effective credential, and clients must reconnect after configuration reload. See [dynamic discovery](tools.md#dynamic-discovery). +Dynamic catalogs and OpenAPI schemas are isolated by the effective request credential. The exact `/openapi` route also requires authentication; server credentials cannot supply anonymous discovery. In combined JWE/OAuth mode, a self-contained JWE takes priority and a partial JWE uses OAuth when available. SSE GET and POST must use the same effective credential, and clients must reconnect after an effective configuration change. Unchanged reload polls preserve sessions. See [dynamic discovery](tools.md#dynamic-discovery). diff --git a/docs/oauth_authorization.md b/docs/oauth_authorization.md index 24449bd8..2ca3d01f 100644 --- a/docs/oauth_authorization.md +++ b/docs/oauth_authorization.md @@ -611,4 +611,4 @@ repo's troubleshooting section. With OAuth enabled, `/openapi` and MCP requests require a bearer or a self-contained JWE when JWE is also enabled. Dynamic catalogs are isolated by the effective credential. Discovery and execution never fall back to server credentials. In combined mode, a self-contained JWE wins; a partial JWE falls through to OAuth. -SSE clients must send the same effective credential on their GET stream and POST messages. Sessions from another credential, or from a retired configuration generation, return 404. Reconnect SSE after configuration reload. See [dynamic discovery](tools.md#dynamic-discovery) for cache lifetimes and STDIO behavior. +SSE clients must send the same effective credential on their GET stream and POST messages. Sessions from another credential, or from a retired configuration generation, return 404. Reconnect SSE after an effective configuration change. Unchanged reload polls preserve sessions. See [dynamic discovery](tools.md#dynamic-discovery) for cache lifetimes and STDIO behavior. diff --git a/docs/tools.md b/docs/tools.md index 9a180e0d..9afa6fc4 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -264,9 +264,9 @@ The server accepts this form but logs a deprecation warning at startup. Prefer ` Discovery is **lazy**: the first authenticated MCP request or OpenAPI schema/tool request discovers the caller's catalog. Single-cluster HTTP and SSE catalogs are cached by the effective JWE or OAuth credential, so each caller's `tools/list`, schema, and dynamic-tool lookup use that caller's own metadata. A self-contained JWE takes priority over OAuth; a partial JWE uses OAuth when available. The unified `server.tools` rules and legacy `dynamic_tools` rules both use this isolation. -Successful catalogs expire at the earlier of token expiry and 15 minutes. Authentication failures are cached for 60 seconds; transient discovery failures are retried on the next request. Static tools remain available if discovery fails. With authentication disabled, the server keeps its shared static-credential catalog and emits `notifications/tools/list_changed` after discovery. +Successful catalogs expire at the earlier of token expiry and 15 minutes. Authentication failures are cached for 60 seconds; transient discovery failures are retried on the next HTTP request or SSE reconnect. Static tools remain available if discovery fails. With authentication disabled, the server keeps its shared static-credential catalog and emits `notifications/tools/list_changed` after discovery. -Configuration reload starts a new catalog generation. Existing SSE clients must reconnect; old session POST endpoints return 404. SSE GET and POST requests must present the same effective credential, including when multiple sessions share a credential. Closing the final GET stream removes its handler bucket. +A configuration change on reload starts a new catalog generation. Polling an unchanged effective configuration preserves the catalog and live SSE sessions. After an effective configuration change, existing SSE clients must reconnect; old session POST endpoints return 404. SSE GET and POST requests must present the same effective credential, including when multiple sessions share a credential. Closing the final GET stream removes its handler bucket. Authenticated STDIO exposes static tools only; it cannot discover or query with server credentials when no request token is available. diff --git a/helm/altinity-mcp/README.md b/helm/altinity-mcp/README.md index 5acada71..b7c1365c 100644 --- a/helm/altinity-mcp/README.md +++ b/helm/altinity-mcp/README.md @@ -4,7 +4,7 @@ A Helm chart for Altinity MCP Server -When JWE or OAuth is enabled, the exact `/openapi` schema route requires authentication and database access never falls back to the configured static credentials. Single-cluster HTTP/SSE dynamic catalogs are isolated by effective request credentials. Existing SSE clients must reconnect after configuration reload. See [dynamic discovery](../../docs/tools.md#dynamic-discovery). +When JWE or OAuth is enabled, the exact `/openapi` schema route requires authentication and database access never falls back to the configured static credentials. Single-cluster HTTP/SSE dynamic catalogs are isolated by effective request credentials. Existing SSE clients must reconnect after an effective configuration change. Unchanged reload polls preserve sessions. See [dynamic discovery](../../docs/tools.md#dynamic-discovery). ## Values diff --git a/pkg/server/multicluster_factory.go b/pkg/server/multicluster_factory.go index 8fe37c77..100fdc12 100644 --- a/pkg/server/multicluster_factory.go +++ b/pkg/server/multicluster_factory.go @@ -109,8 +109,8 @@ func (f *MulticlusterServerFactory) newServer(dynamicTools map[string]dynamicToo Version: f.version, }, opts) adapter := NewSDKServerAdapter(srv) - RegisterResources(adapter) - RegisterPrompts(adapter) + registerResourcesOn(adapter, false) + registerPromptsOn(adapter, false) RegisterStaticToolsOn(adapter, &f.cfg) if len(dynamicTools) > 0 { registerDynamicToolsOn(adapter, dynamicTools, f.cfg.Server.ToolInputSettings, nil) diff --git a/pkg/server/registry_logs_test.go b/pkg/server/registry_logs_test.go new file mode 100644 index 00000000..fc1a7fd2 --- /dev/null +++ b/pkg/server/registry_logs_test.go @@ -0,0 +1,62 @@ +package server + +import ( + "context" + "net/http/httptest" + "os" + "os/exec" + "strings" + "testing" + "time" + + "github.com/altinity/altinity-mcp/pkg/config" + "github.com/stretchr/testify/require" +) + +// Isolate logger output in a subprocess instead of replacing the package's +// global logger while unrelated concurrent tests may still be logging. +func TestRequestRegistryKeepsStartupSignalsQuiet(t *testing.T) { + if os.Getenv("ALTINITY_MCP_TEST_REGISTRY_LOGS") == "1" { + for _, legacy := range []bool{false, true} { + cfg := config.Config{Server: config.ServerConfig{OAuth: config.OAuthConfig{Enabled: true}}} + if legacy { + cfg.Server.DynamicTools = []config.DynamicToolRule{{Regexp: ".*"}} + } + parent := NewClickHouseMCPServer(cfg, "test") + if legacy { + parent.catalogOnce.Do(func() { parent.catalogCache = NewCatalogCache(config.MulticlusterConfig{}) }) + parent.catalogCache.insertOK(fullKey(CacheKey("oauth\x00alice"), "singlecluster"), map[string]dynamicToolMeta{}, time.Time{}) + } + _, _ = os.Stderr.WriteString("BEGIN REQUEST REGISTRIES\n") + request := httptest.NewRequest("POST", "/", nil) + request = request.WithContext(context.WithValue(request.Context(), OAuthTokenKey, "alice")) + for i := 0; i < 3; i++ { + parent.GetServer(request) + } + _, _ = os.Stderr.WriteString("END REQUEST REGISTRIES\n") + parent.Close() + } + // Failed no-auth discovery still emits an operator warning. + parent := NewClickHouseMCPServer(config.Config{ClickHouse: config.ClickHouseConfig{Username: "fake-user", Protocol: "invalid"}, Server: config.ServerConfig{DynamicTools: []config.DynamicToolRule{{Regexp: ".*"}}}}, "test") + parent.GetServer(httptest.NewRequest("POST", "/", nil)) + parent.Close() + return + } + child := exec.Command(os.Args[0], "-test.run=^TestRequestRegistryKeepsStartupSignalsQuiet$") + child.Env = append(os.Environ(), "ALTINITY_MCP_TEST_REGISTRY_LOGS=1") + output, err := child.CombinedOutput() + require.NoError(t, err, string(output)) + text := string(output) + require.Contains(t, text, "Static read tool registered") + require.Contains(t, text, "ClickHouse resources registered") + require.Contains(t, text, "dynamic_tools config is deprecated") + parts := strings.Split(text, "BEGIN REQUEST REGISTRIES\n") + require.Len(t, parts, 3) + for _, part := range parts[1:] { + requestLogs := strings.SplitN(part, "END REQUEST REGISTRIES\n", 2) + require.Len(t, requestLogs, 2) + require.Empty(t, requestLogs[0], "request registry creation must preserve startup-only logging") + } + require.Contains(t, text, `"level":"warn"`) + require.Contains(t, text, "Failed to ensure dynamic tools") +} diff --git a/pkg/server/server.go b/pkg/server/server.go index 78d0fadf..7db3dc32 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -181,7 +181,7 @@ func truncateErrForClient(err error) string { // // Returns the number of static tools actually registered. func RegisterStaticToolsOn(srv AltinityMCPServer, cfg *config.Config) int { - toolsToRegister := resolveToolDefinitions(cfg) + toolsToRegister := resolveToolDefinitionsWithLogs(cfg, false) staticToolCount := 0 for _, td := range toolsToRegister { if td.Type != "read" && td.Type != "write" { @@ -194,7 +194,7 @@ func RegisterStaticToolsOn(srv AltinityMCPServer, cfg *config.Config) int { if td.Name == "" { continue } - if registerStaticTool(srv, td, &cfg.Server, cfg.ClickHouse.ReadOnly) { + if registerStaticToolWithLogs(srv, td, &cfg.Server, cfg.ClickHouse.ReadOnly, false) { staticToolCount++ } } @@ -268,11 +268,19 @@ func RegisterTools(srv AltinityMCPServer, cfg *config.Config) { // the new unified Tools array, the legacy DynamicTools slice (with a // deprecation warning), or a sensible default (execute_query + write_query). func resolveToolDefinitions(cfg *config.Config) []config.ToolDefinition { + return resolveToolDefinitionsWithLogs(cfg, true) +} + +// Request registries reuse startup-validated definitions without repeating +// startup deprecation messages. +func resolveToolDefinitionsWithLogs(cfg *config.Config, startupLogs bool) []config.ToolDefinition { if len(cfg.Server.Tools) > 0 { return cfg.Server.Tools } if len(cfg.Server.DynamicTools) > 0 { - log.Warn().Msg("dynamic_tools config is deprecated, use tools instead") + if startupLogs { + log.Warn().Msg("dynamic_tools config is deprecated, use tools instead") + } out := make([]config.ToolDefinition, 0, len(cfg.Server.DynamicTools)) for _, old := range cfg.Server.DynamicTools { td := config.ToolDefinition{ @@ -304,11 +312,17 @@ func resolveToolDefinitions(cfg *config.Config) []config.ToolDefinition { // registerStaticTool registers one of the supported static tools ("execute_query" // or "write_query"). Returns true if the tool was actually added to srv. func registerStaticTool(srv AltinityMCPServer, td config.ToolDefinition, srvCfg *config.ServerConfig, readOnly bool) bool { + return registerStaticToolWithLogs(srv, td, srvCfg, readOnly, true) +} + +func registerStaticToolWithLogs(srv AltinityMCPServer, td config.ToolDefinition, srvCfg *config.ServerConfig, readOnly, startupLogs bool) bool { switch td.Type { case "read": if td.Name == "execute_query" { srv.AddTool(buildExecuteQueryTool(srvCfg), HandleReadOnlyQuery) - log.Info().Str("tool", "execute_query").Msg("Static read tool registered") + if startupLogs { + log.Info().Str("tool", "execute_query").Msg("Static read tool registered") + } return true } log.Warn().Str("tool_name", td.Name).Msg("Unknown static read tool name") @@ -317,11 +331,15 @@ func registerStaticTool(srv AltinityMCPServer, td config.ToolDefinition, srvCfg case "write": if td.Name == "write_query" { if readOnly { - log.Info().Str("tool", "write_query").Msg("Write tool skipped (read-only mode)") + if startupLogs { + log.Info().Str("tool", "write_query").Msg("Write tool skipped (read-only mode)") + } return false } srv.AddTool(buildWriteQueryTool(srvCfg), HandleExecuteQuery) - log.Info().Str("tool", "write_query").Msg("Static write tool registered") + if startupLogs { + log.Info().Str("tool", "write_query").Msg("Static write tool registered") + } return true } log.Warn().Str("tool_name", td.Name).Msg("Unknown static write tool name") diff --git a/pkg/server/server_resources.go b/pkg/server/server_resources.go index d2fb466c..9d4221fa 100644 --- a/pkg/server/server_resources.go +++ b/pkg/server/server_resources.go @@ -12,6 +12,10 @@ import ( // RegisterResources adds ClickHouse resources to the MCP server func RegisterResources(srv AltinityMCPServer) { + registerResourcesOn(srv, true) +} + +func registerResourcesOn(srv AltinityMCPServer, startupLogs bool) { // Database catalog resource schemaResource := &mcp.Resource{ URI: "clickhouse://schema", @@ -32,7 +36,9 @@ func RegisterResources(srv AltinityMCPServer) { srv.AddResourceTemplate(tableTemplate, HandleTableResource) - log.Info().Int("resource_count", 2).Msg("ClickHouse resources registered") + if startupLogs { + log.Info().Int("resource_count", 2).Msg("ClickHouse resources registered") + } } // HandleSchemaResource handles the schema resource @@ -158,6 +164,12 @@ func HandleTableResource(ctx context.Context, req *mcp.ReadResourceRequest) (*mc // RegisterPrompts adds ClickHouse prompts to the MCP server func RegisterPrompts(srv AltinityMCPServer) { + registerPromptsOn(srv, true) +} + +func registerPromptsOn(srv AltinityMCPServer, startupLogs bool) { // No prompts registered - log.Info().Int("prompt_count", 0).Msg("ClickHouse prompts registered") + if startupLogs { + log.Info().Int("prompt_count", 0).Msg("ClickHouse prompts registered") + } } diff --git a/pkg/server/singlecluster_catalog.go b/pkg/server/singlecluster_catalog.go index f6e841e5..22c7e452 100644 --- a/pkg/server/singlecluster_catalog.go +++ b/pkg/server/singlecluster_catalog.go @@ -80,13 +80,20 @@ func (s *ClickHouseJWEServer) requestCatalog(ctx context.Context) map[string]dyn return nil } s.catalogOnce.Do(func() { s.catalogCache = NewCatalogCache(config.MulticlusterConfig{}) }) + factory := func(ctx context.Context, cfg config.ClickHouseConfig) (*clickhouse.Client, error) { + return s.GetClickHouseClientWithOAuthForConfig(ctx, cfg, s.ExtractTokenFromCtx(ctx), s.ExtractOAuthTokenFromCtx(ctx), s.GetOAuthClaimsFromCtx(ctx)) + } + var tools map[string]dynamicToolMeta if s.catalogCache == nil { - return nil + // Close may retire a generation after a request snapshots it but before + // lazy cache initialization. Finish that request with its own generation, + // without creating a janitor for the retired parent. + boundedCtx, cancel := context.WithTimeout(ctx, discoveryTimeout) + defer cancel() + tools, err = DiscoverTools(boundedCtx, s.Config.ClickHouse, factory, s.Config.Server.DynamicTools, s.Config.ClickHouse.ReadOnly) + } else { + tools, err = s.catalogCache.GetOrDiscover(ctx, key, "singlecluster", s.Config.ClickHouse, factory, s.Config.Server.DynamicTools, s.Config.ClickHouse.ReadOnly, exp) } - tools, err := s.catalogCache.GetOrDiscover(ctx, key, "singlecluster", s.Config.ClickHouse, - func(ctx context.Context, cfg config.ClickHouseConfig) (*clickhouse.Client, error) { - return s.GetClickHouseClientWithOAuthForConfig(ctx, cfg, s.ExtractTokenFromCtx(ctx), s.ExtractOAuthTokenFromCtx(ctx), s.GetOAuthClaimsFromCtx(ctx)) - }, s.Config.Server.DynamicTools, s.Config.ClickHouse.ReadOnly, exp) if err != nil { log.Warn().Err(err).Msg("credential catalog discovery failed; static tools remain available") return nil @@ -97,7 +104,9 @@ func (s *ClickHouseJWEServer) requestCatalog(ctx context.Context) map[string]dyn // GetServer creates a caller's MCP registry without changing the parent registry. func (s *ClickHouseJWEServer) GetServer(r *http.Request) *mcp.Server { if !s.Config.Server.JWE.Enabled && !s.Config.Server.OAuth.Enabled { - _ = s.EnsureDynamicTools(r.Context()) + if err := s.EnsureDynamicTools(r.Context()); err != nil { + log.Warn().Err(err).Msg("Failed to ensure dynamic tools") + } return s.MCPServer } f := &MulticlusterServerFactory{cfg: s.Config, version: s.Version,