diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index 30380e5d81a..08d3b9ba823 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -2143,6 +2143,27 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { redactedExisting := h.store.RedactMCPClientConfig(existingConfig) headers = mergeMCPHeaders(req.Headers, existingConfig.Headers, redactedExisting.Headers) } + // A real change to the static headers has to be applied to the live + // connection, not just to the stored config: on a sticky client the + // credential is baked onto the transport at dial time, so updating the + // config alone would leave the open connection talking to the upstream + // with the header the admin just replaced. Diffed against the resolved + // (post-redaction-merge) value so a caller that round-trips the map + // unchanged doesn't force a needless reconnect. Only shared auth types + // carry a connection-level header credential — per-user values are + // supplied per caller and never baked into a shared transport. HTTP only, + // matching the transport VerifyHeadersConnection builds for the + // pre-flight below. + // Skipped for a disabled client: there is no live connection to protect, + // and refusing the edit because a deliberately-unreachable upstream can't + // be dialed would make a disabled client's credentials uneditable. It is + // verified on enable instead, which dials for real. + staticHeadersChanged := (existingConfig.AuthType == schemas.MCPAuthTypeHeaders || existingConfig.AuthType == schemas.MCPAuthTypeNone) && + existingConfig.ConnectionType == schemas.MCPConnectionTypeHTTP && + !disabled && + req.Headers != nil && !mcpHeadersEqual(existingConfig.Headers, headers) + headersChanged := staticHeadersChanged && !existingConfig.Disabled + // TLSConfig: if omitted keep existing; if provided, restore raw CACertPEM when the // incoming value is the redacted placeholder returned by the API. tlsConfig := existingConfig.TLSConfig @@ -2398,6 +2419,45 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusBadRequest, "oauth credentials cannot be rotated while disabling a client; send these as two separate requests") return } + // Pre-flight the new static headers against the upstream before anything + // is written. VerifyHeadersConnection dials its own ephemeral transport + // and closes it, touching neither the stored row nor the live connection, + // so a rejected credential leaves the client exactly as it was — still + // serving on the headers that work. Committing first and reconnecting + // afterwards cannot offer that: the failed dial drops the live connection + // (leaving the client unusable with "no active connection") while the bad + // headers are already persisted, so the connection checker has nothing + // good to retry with. Same verify-then-persist order the create path uses + // for per-user headers. + if staticHeadersChanged { + verifyConfig := *existingConfig + verifyConfig.Headers = headers + verifyConfig.TLSConfig = tlsConfig + // Only Authorization goes in this map. VerifyHeadersConnection layers + // StaticConfigHeaders(verifyConfig) in itself, which already carries + // every other header on the config; Authorization is the one entry + // that set deliberately withholds (so plugins never observe the + // bearer), and it gets applied after the plugin gate instead. Passing + // exactly that one header here reproduces the header set a real + // sticky dial builds, since sharedHeadersResolver.ConnectionHeaders + // resolves only Authorization too — including its first-match-wins + // behavior if the map somehow holds two casings of the name. + verifyHeaders := make(map[string]string, 1) + for key, value := range headers { + if strings.EqualFold(key, "Authorization") { + verifyHeaders["Authorization"] = value.GetValue() + break + } + } + bifrostCtx, cancel := lib.ConvertToBifrostContext(ctx, h.store) + _, _, verifyErr := h.mcpManager.VerifyHeadersConnection(bifrostCtx, &verifyConfig, verifyHeaders) + cancel() + if verifyErr != nil { + SendError(ctx, fasthttp.StatusUnprocessableEntity, fmt.Sprintf("The new headers were rejected by the MCP server, so nothing was changed: %v", verifyErr)) + return + } + } + // Rotation is deferred until after the DB and in-memory client updates // below both succeed (see the call site further down): RotateMCPOAuthConfig // opens its own transaction and, once it commits, cascades every existing @@ -2557,6 +2617,29 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { } } + // Swap the verified headers onto the live connection. Only a client that + // was and remains sticky needs this: it holds an open transport with the + // old header baked in, and UpdateMCPClientCredentials is the only thing + // that replaces it. Every other combination is already handled — + // UpdateMCPClient above replaced ExecutionConfig (all a per-call client + // needs, since it reads the current config on every dial) and itself + // re-dialed on a per-call->sticky flip, so re-dialing again here would + // just be a second connect with the same credentials. + // + // The same credentials already completed a full connect during the + // pre-flight above, so a failure here is a transient dial problem rather + // than a bad credential — and UpdateClientCredentials restores the + // previous ExecutionConfig on failure, leaving the connection checker to + // retry. Reported as a partial success for that reason, not a hard + // failure: every other effect of this request already committed. + if headersChanged && sharedConnectErr == nil && !wasPerCallConnection && !isPerCallConnection { + if err := h.updateMCPClientCredentialsWithRetry(ctx, id, schemasConfig); err != nil && + !errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { + logger.Error(fmt.Sprintf("MCP client %s updated, but reconnecting with the new headers failed: %v", id, err)) + sharedConnectErr = err + } + } + // Rotate OAuth credentials only now that both the DB row and the runtime // config update above have succeeded — see the comment where // shouldRotateOAuthConfig was computed for why this can't run earlier. @@ -2942,6 +3025,23 @@ func (h *MCPHandler) updateMCPClientWithRetry(ctx context.Context, id string, co return lastErr } +// mcpHeadersEqual reports whether two resolved static header maps carry the +// same credentials. Compared on the post-merge values (mergeMCPHeaders has +// already restored any unchanged redacted entry to its stored raw value), so +// a caller round-tripping the redacted map it was served reads as unchanged. +func mcpHeadersEqual(a, b map[string]schemas.SecretVar) bool { + if len(a) != len(b) { + return false + } + for key, aVal := range a { + bVal, ok := b[key] + if !ok || !aVal.Equals(&bVal) { + return false + } + } + return true +} + // updateMCPClientCredentialsWithRetry calls mcpManager.UpdateMCPClientCredentials with a short retry loop. func (h *MCPHandler) updateMCPClientCredentialsWithRetry(ctx context.Context, id string, config *schemas.MCPClientConfig) error { const maxAttempts = 3 diff --git a/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go b/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go new file mode 100644 index 00000000000..9e7914074ae --- /dev/null +++ b/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go @@ -0,0 +1,126 @@ +package handlers + +import ( + "context" + "encoding/json" + "errors" + "testing" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore" + configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/transports/bifrost-http/lib" + "github.com/valyala/fasthttp" +) + +// fakeMCPManagerVerifyOnly is a minimal MCPManager test double for this path. +// It records VerifyHeadersConnection and UpdateMCPClient calls so the test can +// assert pre-flight verification ran and no live update happened on failure. +type fakeMCPManagerVerifyOnly struct { + verifyCalls int + updateCalls int +} + +func (m *fakeMCPManagerVerifyOnly) AddMCPClient(_ context.Context, _ *schemas.MCPClientConfig) error { + return nil +} +func (m *fakeMCPManagerVerifyOnly) RemoveMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeMCPManagerVerifyOnly) UpdateMCPClient(_ context.Context, _ string, _ *schemas.MCPClientConfig) error { + m.updateCalls++ + return nil +} +func (m *fakeMCPManagerVerifyOnly) UpdateMCPClientCredentials(_ context.Context, _ string, _ *schemas.MCPClientConfig) error { + return nil +} +func (m *fakeMCPManagerVerifyOnly) ReconnectMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeMCPManagerVerifyOnly) CloseAndMarkNeedsReauth(_ context.Context, _ string) error { + return nil +} +func (m *fakeMCPManagerVerifyOnly) DisableMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeMCPManagerVerifyOnly) EnableMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeMCPManagerVerifyOnly) VerifyPerUserOAuthConnection(_ context.Context, _ *schemas.MCPClientConfig, _ string) (map[string]schemas.ChatTool, map[string]string, error) { + return nil, nil, nil +} +func (m *fakeMCPManagerVerifyOnly) VerifyHeadersConnection(_ context.Context, _ *schemas.MCPClientConfig, _ map[string]string) (map[string]schemas.ChatTool, map[string]string, error) { + m.verifyCalls++ + return nil, nil, errors.New("rejected by upstream") +} +func (m *fakeMCPManagerVerifyOnly) SetClientTools(_ string, _ map[string]schemas.ChatTool, _ map[string]string) { +} +func (m *fakeMCPManagerVerifyOnly) RequiresPerCallConnection(_ *schemas.MCPClientConfig) bool { + return false +} + +// mockUpdateConfigStore embeds the interface so unimplemented methods panic if called. +// The test asserts UpdateMCPClientConfig is NOT called when verification fails. +type mockUpdateConfigStore struct { + configstore.ConfigStore + updates int +} + +func (m *mockUpdateConfigStore) GetMCPClientByID(_ context.Context, id string) (*configstoreTables.TableMCPClient, error) { + return &configstoreTables.TableMCPClient{ClientID: id}, nil +} + +func (m *mockUpdateConfigStore) GetClientConfig(_ context.Context) (*configstore.ClientConfig, error) { + return nil, nil +} + +func (m *mockUpdateConfigStore) UpdateMCPClientConfig(_ context.Context, _ string, _ *configstoreTables.TableMCPClient) error { + m.updates++ + return nil +} + +func TestUpdateMCPClient_DisabledToEnabled_WithInvalidReplacementHeaders_PreflightRejects(t *testing.T) { + SetLogger(&mockLogger{}) + plain := func(v string) schemas.SecretVar { return *schemas.NewSecretVar(v) } + + // Existing client is disabled with some stored headers. + existing := &schemas.MCPClientConfig{ + ID: "client-1", + Name: "Test", + ConnectionType: schemas.MCPConnectionTypeHTTP, + AuthType: schemas.MCPAuthTypeHeaders, + Disabled: true, + Headers: map[string]schemas.SecretVar{"Authorization": plain("Bearer old")}, + PerUserHeaderKeys: nil, + } + store := &lib.Config{MCPConfig: &schemas.MCPConfig{ClientConfigs: []*schemas.MCPClientConfig{existing}}} + store.ClientConfig = &configstore.ClientConfig{} // avoid nil deref in ConvertToBifrostContext + cfgStore := &mockUpdateConfigStore{} + store.ConfigStore = cfgStore + mgr := &fakeMCPManagerVerifyOnly{} + h := &MCPHandler{store: store, mcpManager: mgr} + + // Build request body: enable the client and set replacement headers. + disabled := false + body, err := json.Marshal(MCPClientUpdateRequest{ + Disabled: &disabled, + Headers: map[string]schemas.SecretVar{"Authorization": plain("Bearer new-invalid")}, + }) + if err != nil { + t.Fatalf("marshal: %v", err) + } + ctx := &fasthttp.RequestCtx{} + ctx.SetUserValue("id", existing.ID) + ctx.Request.SetBody(body) + + // Exercise handler. + h.updateMCPClient(ctx) + + if code := ctx.Response.StatusCode(); code != fasthttp.StatusUnprocessableEntity { + t.Fatalf("expected 422, got %d: %s", code, string(ctx.Response.Body())) + } + if mgr.verifyCalls != 1 { + t.Fatalf("expected VerifyHeadersConnection to run once, got %d", mgr.verifyCalls) + } + if mgr.updateCalls != 0 { + t.Fatalf("expected no live UpdateMCPClient on pre-flight failure, got %d", mgr.updateCalls) + } + if cfgStore.updates != 0 { + t.Fatalf("expected no DB UpdateMCPClientConfig on pre-flight failure, got %d", cfgStore.updates) + } + if !existing.Disabled { + t.Fatalf("expected in-memory config to remain disabled after rejection, got Disabled=false") + } +} diff --git a/transports/bifrost-http/handlers/mcp_headers_changed_test.go b/transports/bifrost-http/handlers/mcp_headers_changed_test.go new file mode 100644 index 00000000000..82c5dc215ff --- /dev/null +++ b/transports/bifrost-http/handlers/mcp_headers_changed_test.go @@ -0,0 +1,93 @@ +package handlers + +import ( + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +// TestMCPHeadersEqual pins the diff that decides whether a header edit has to +// re-dial the upstream connection. A sticky client bakes its credential onto +// the transport at connect time, so a false negative here leaves it talking to +// the upstream with the header the admin just replaced; a false positive +// cycles a live connection for an edit that changed nothing. +func TestMCPHeadersEqual(t *testing.T) { + plain := func(v string) schemas.SecretVar { return *schemas.NewSecretVar(v) } + + tests := []struct { + name string + a map[string]schemas.SecretVar + b map[string]schemas.SecretVar + want bool + }{ + { + name: "both empty", + want: true, + }, + { + name: "identical single header", + a: map[string]schemas.SecretVar{"Authorization": plain("Bearer one")}, + b: map[string]schemas.SecretVar{"Authorization": plain("Bearer one")}, + want: true, + }, + { + name: "changed value", + a: map[string]schemas.SecretVar{"Authorization": plain("Bearer one")}, + b: map[string]schemas.SecretVar{"Authorization": plain("Bearer two")}, + want: false, + }, + { + name: "added header", + a: map[string]schemas.SecretVar{"Authorization": plain("Bearer one")}, + b: map[string]schemas.SecretVar{"Authorization": plain("Bearer one"), "X-Tenant": plain("acme")}, + want: false, + }, + { + name: "removed header", + a: map[string]schemas.SecretVar{"Authorization": plain("Bearer one"), "X-Tenant": plain("acme")}, + b: map[string]schemas.SecretVar{"Authorization": plain("Bearer one")}, + want: false, + }, + { + name: "renamed key with same value", + a: map[string]schemas.SecretVar{"X-Api-Key": plain("secret")}, + b: map[string]schemas.SecretVar{"X-Other-Key": plain("secret")}, + want: false, + }, + { + // Header names are case-sensitive in Go maps and the resolvers + // compare them case-insensitively, but the stored map is what gets + // baked onto the transport verbatim — a case change is a real + // change to what is persisted, so re-dialing is the safe answer. + name: "case-differing key", + a: map[string]schemas.SecretVar{"X-Api-Key": plain("secret")}, + b: map[string]schemas.SecretVar{"x-api-key": plain("secret")}, + want: false, + }, + { + name: "empty vs populated", + a: map[string]schemas.SecretVar{}, + b: map[string]schemas.SecretVar{"Authorization": plain("Bearer one")}, + want: false, + }, + { + name: "nil vs empty are both no headers", + a: nil, + b: map[string]schemas.SecretVar{}, + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := mcpHeadersEqual(tt.a, tt.b); got != tt.want { + t.Errorf("mcpHeadersEqual() = %v, want %v", got, tt.want) + } + // The comparison must not depend on argument order: the caller + // passes (stored, resolved) and either side can be the larger map. + if got := mcpHeadersEqual(tt.b, tt.a); got != tt.want { + t.Errorf("mcpHeadersEqual() reversed = %v, want %v", got, tt.want) + } + }) + } +}