diff --git a/packages/agentvault/cache.go b/packages/agentvault/cache.go index 947c6a0a..233c2e15 100644 --- a/packages/agentvault/cache.go +++ b/packages/agentvault/cache.go @@ -8,6 +8,8 @@ import ( "errors" "golang.org/x/sync/singleflight" "net/http" + "slices" + "strings" "sync" "time" @@ -37,6 +39,8 @@ type credential struct { value []byte username string password []byte + // Only an OAuth token Infisical issued can be refreshed when the upstream rejects it. + refreshable bool } type customHeader struct { @@ -54,8 +58,11 @@ type substitution struct { type resolvedService struct { id string name string + kind string accessBundleName string - hostPatterns []hostPattern + // As Infisical sent it, so a credential refresh can name the destination the rejected token went to. + hostPattern string + hostPatterns []hostPattern // A nil map means every method is allowed; an empty slice of prefixes means every path. allowedMethods map[string]bool allowedPathPrefixes []string @@ -71,6 +78,16 @@ type sessionEntry struct { sessionLog *sessionLogGrant lastSeen time.Time fetchedAt time.Time + // Keyed by service ID. A poll whose resolve began before a swap can return the token that swap replaced, + // so the merge consults this to keep the newer one. + swaps map[string]*credentialSwap +} + +type credentialSwap struct { + installed credential + hostPattern string + replaced map[string]bool + gen uint64 } // The map key is the sha256 of the token, never the token itself, so a heap dump yields no live credential. @@ -83,12 +100,22 @@ type sessionResolver interface { resolve(sessionToken string, held *sessionLogGrant) (*resolveResult, error) } +// Implemented by the Infisical resolver. Kept apart from sessionResolver so resolvers that never broker an +// OAuth credential need not stub it. +type credentialRefresher interface { + refreshCredential(sessionToken, serviceID, hostPattern, rejectedHash string) (credential, string, error) + reportMissingScopes(sessionToken, serviceID, hostPattern, rejectedHash string, requiredScopes []string) error +} + type sessionCache struct { resolver sessionResolver pollInterval func() time.Duration // Concurrent misses for one session share a single resolve rather than each posting to Infisical. inflight singleflight.Group + // Concurrent upstream 401s for one service share a single refresh, so a burst of requests holding the + // same expired token costs Infisical one call. + credentialFlights singleflight.Group // A refresh that outlives the poll interval must not pile a second one on top, nor hold up the heartbeat. refreshing atomic.Bool @@ -98,6 +125,8 @@ type sessionCache struct { // A definitive refusal is remembered briefly, so an agent retrying a dead session costs one resolve per // window rather than one per request. Outages and throttling are never stored here. refused map[string]refusedEntry + // Bumped under mu by every credential swap, so a poll can tell which swaps happened after it began. + swapGen uint64 } type refusedEntry struct { @@ -316,6 +345,7 @@ func (c *sessionCache) refreshOne(key, token string) { if entry, ok := c.entries[key]; ok { held = entry.sessionLog } + startGen := c.swapGen c.mu.Unlock() result, err := c.resolver.resolve(token, held) @@ -329,12 +359,39 @@ func (c *sessionCache) refreshOne(key, token string) { if entry, ok := c.entries[key]; ok { entry.sessionID = result.SessionID entry.expiresAt = result.ExpiresAt - entry.services = result.Services + entry.services = keepNewerSwaps(entry, result.Services, startGen) entry.sessionLog = result.SessionLog entry.fetchedAt = time.Now() } } +// keepNewerSwaps puts back a credential swapped in after the poll began when the poll returned a token that +// swap replaced. Swaps the poll began after are settled by it and forgotten. +func keepNewerSwaps(entry *sessionEntry, polled []*resolvedService, startGen uint64) []*resolvedService { + if len(entry.swaps) == 0 { + return polled + } + services, copied := polled, false + for i, svc := range polled { + swap, ok := entry.swaps[svc.id] + if !ok || swap.gen <= startGen || swap.hostPattern != svc.hostPattern || !svc.credential.refreshable || !swap.replaced[credentialHash(&svc.credential)] { + continue + } + if !copied { + services, copied = append([]*resolvedService(nil), polled...), true + } + updated := *svc + updated.credential = swap.installed + services[i] = &updated + } + for id, swap := range entry.swaps { + if swap.gen <= startGen { + delete(entry.swaps, id) + } + } + return services +} + func (c *sessionCache) handleRefreshFailure(key string, err error) { c.mu.Lock() defer c.mu.Unlock() @@ -382,3 +439,146 @@ func (c *sessionCache) close() { delete(c.tokens, key) } } + +var ( + errCredentialNotRefreshable = errors.New("the refresh did not return a refreshable credential") + errServiceHostChanged = errors.New("the service's host pattern changed while its credential was refreshed") + errServiceRemoved = errors.New("the service was removed from the session while its credential was refreshed") +) + +func credentialHash(cred *credential) string { + sum := sha256.Sum256(cred.value) + return hex.EncodeToString(sum[:]) +} + +// refreshCredential asks Infisical for a replacement of the token the upstream rejected and swaps it into +// the session's entry, so later requests inject it. hostPattern is that of the cached service whose +// credential the rejected request carried: a token minted for wherever the service points now must never be +// sent to where it used to point. +func (c *sessionCache) refreshCredential(sessionToken, serviceID, hostPattern string, rejected *credential) (credential, error) { + refresher, ok := c.resolver.(credentialRefresher) + if !ok { + return credential{}, errCredentialNotRefreshable + } + key := sessionKey(sessionToken) + rejectedHash := credentialHash(rejected) + + result, err, _ := c.credentialFlights.Do(credentialFlightKey(key, serviceID, hostPattern, rejectedHash, nil), func() (any, error) { + // Checked inside the flight so a caller that arrives after another refresh finished also reuses + // its replacement instead of asking Infisical again. + if current, ok := c.currentCredential(key, serviceID, hostPattern); ok && current.refreshable && credentialHash(¤t) != rejectedHash { + return current, nil + } + fresh, status, err := refresher.refreshCredential(sessionToken, serviceID, hostPattern, rejectedHash) + if err != nil { + return nil, err + } + if status != "" && status != oauthStatusConnected { + log.Warn().Str("serviceId", serviceID).Str("status", status). + Msg("agent-vault: Infisical reports this service's OAuth connection needs attention") + } + if fresh.kind != credentialBearer || !fresh.refreshable { + return nil, errCredentialNotRefreshable + } + return c.swapCredential(key, serviceID, hostPattern, rejectedHash, fresh) + }) + if err != nil { + return credential{}, err + } + return result.(credential), nil +} + +// reportMissingScopes tells Infisical which scopes the upstream's 403 insufficient_scope challenge asked for, +// so the service shows what reauthorizing has to request. Nothing is refreshed, so the cache is left alone. +func (c *sessionCache) reportMissingScopes(sessionToken, serviceID, hostPattern string, rejected *credential, requiredScopes []string) error { + refresher, ok := c.resolver.(credentialRefresher) + if !ok { + return errCredentialNotRefreshable + } + rejectedHash := credentialHash(rejected) + _, err, _ := c.credentialFlights.Do(credentialFlightKey(sessionKey(sessionToken), serviceID, hostPattern, rejectedHash, requiredScopes), func() (any, error) { + return nil, refresher.reportMissingScopes(sessionToken, serviceID, hostPattern, rejectedHash, requiredScopes) + }) + return err +} + +// Scope reports are bound to the rejected token, so neither different tokens nor different scope sets may +// share a call. The set is canonical so another order or repeats still share one call. +func credentialFlightKey(key, serviceID, hostPattern, rejectedHash string, requiredScopes []string) string { + base := key + "\x00" + serviceID + "\x00" + hostPattern + "\x00" + rejectedHash + if len(requiredScopes) == 0 { + return base + "\x00refresh" + } + scopes := append([]string(nil), requiredScopes...) + slices.Sort(scopes) + scopes = slices.Compact(scopes) + return base + "\x00scopes\x00" + strings.Join(scopes, " ") +} + +func (c *sessionCache) currentCredential(key, serviceID, hostPattern string) (credential, bool) { + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[key] + if !ok { + return credential{}, false + } + for _, svc := range entry.services { + if svc.id == serviceID && svc.hostPattern == hostPattern { + return svc.credential, true + } + } + return credential{}, false +} + +// Copy on write: requests in flight hold pointers into the old slice and read it without the lock. A newer +// cached credential wins over a delayed refresh or scope report, including the credential returned for retry. +// A poll that removed the session or the service during the refresh revoked that access, so the fresh token +// is refused rather than handed to a retry. +func (c *sessionCache) swapCredential(key, serviceID, hostPattern, replacedHash string, cred credential) (credential, error) { + c.mu.Lock() + defer c.mu.Unlock() + entry, ok := c.entries[key] + if !ok { + return credential{}, errSessionGone + } + for i, svc := range entry.services { + if svc.id != serviceID { + continue + } + if svc.hostPattern != hostPattern { + return credential{}, errServiceHostChanged + } + if !svc.credential.refreshable { + return credential{}, errCredentialNotRefreshable + } + if credentialHash(&svc.credential) != replacedHash { + return svc.credential, nil + } + updated := *svc + updated.credential = cred + services := append([]*resolvedService(nil), entry.services...) + services[i] = &updated + entry.services = services + + c.swapGen++ + if entry.swaps == nil { + entry.swaps = make(map[string]*credentialSwap) + } + swap, ok := entry.swaps[serviceID] + if !ok { + swap = &credentialSwap{replaced: make(map[string]bool)} + entry.swaps[serviceID] = swap + } + swap.installed = cred + if swap.hostPattern != hostPattern { + swap.hostPattern = hostPattern + clear(swap.replaced) + } + swap.replaced[replacedHash] = true + swap.gen = c.swapGen + // A scope report swaps in the token it reported, which is current rather than stale. + delete(swap.replaced, credentialHash(&cred)) + return cred, nil + } + return credential{}, errServiceRemoved +} diff --git a/packages/agentvault/mcp_inspect.go b/packages/agentvault/mcp_inspect.go new file mode 100644 index 00000000..ba8b39de --- /dev/null +++ b/packages/agentvault/mcp_inspect.go @@ -0,0 +1,343 @@ +package agentvault + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "mime" + "net/http" + "regexp" + "strings" +) + +const ( + serviceKindHTTP = "http" + serviceKindMCP = "mcp" + + maxMCPInspectBody = 1 << 20 + // The ring holds thousands of records per session, so one record must stay small. A longer batch is + // still read and validated in full, but only its first calls are kept and the record says so. + maxMCPCallsPerRecord = 32 + + mcpMethodToolsCall = "tools/call" +) + +var mcpNamePattern = regexp.MustCompile(`^[A-Za-z0-9_.:/-]{1,128}$`) + +type mcpCall struct { + Method string `json:"method"` + Tool string `json:"tool,omitempty"` +} + +type mcpDetail struct { + Calls []mcpCall `json:"calls"` + Parsed bool `json:"parsed"` + Truncated bool `json:"truncated,omitempty"` +} + +func (s *resolvedService) serviceType() string { + if s.kind == serviceKindMCP { + return serviceKindMCP + } + return serviceKindHTTP +} + +func unparsedMCP() *mcpDetail { + return &mcpDetail{Calls: []mcpCall{}} +} + +// A body already read into memory, so applyBodySubstitutions takes the bytes rather than reading them again. +type bufferedBody struct { + *bytes.Reader + data []byte +} + +func (*bufferedBody) Close() error { return nil } + +type rechainedBody struct { + io.Reader + io.Closer +} + +type erroringReader struct{ err error } + +func (r erroringReader) Read([]byte) (int, error) { return 0, r.err } + +// Returns nil when the request is not one the proxy parses, so the record carries no mcp detail at all. +// Never refuses or alters the request: whatever is read is put back so the upstream gets the same bytes. +func inspectMCPRequest(req *http.Request) *mcpDetail { + body, applicable, whole := bufferMCPRequest(req) + if !applicable { + return nil + } + if !whole { + return unparsedMCP() + } + return parseMCPBody(body) +} + +// Reads a JSON-RPC POST body into memory, up to maxMCPInspectBody, so it can be parsed for the log and sent +// again on an OAuth retry. applicable is false when the request is not one the proxy buffers, and whole is +// false when the body could not be held in full. Whatever is read is put back so the upstream gets the same +// bytes. +func bufferMCPRequest(req *http.Request) (body []byte, applicable, whole bool) { + if !strings.EqualFold(req.Method, http.MethodPost) { + return nil, false, false + } + mediaType, _, err := mime.ParseMediaType(req.Header.Get("Content-Type")) + if err != nil || mediaType != "application/json" { + return nil, false, false + } + if req.Header.Get("Content-Encoding") != "" || req.ContentLength > maxMCPInspectBody { + return nil, true, false + } + if req.Body == nil || req.Body == http.NoBody { + return nil, true, true + } + + original := req.Body + body, err = io.ReadAll(io.LimitReader(original, maxMCPInspectBody+1)) + if err != nil { + // Put back as it failed, so whatever reads the body next meets the same error and a truncated body + // is refused or fails exactly as it would have without the inspection. + req.Body = rechainedBody{Reader: io.MultiReader(bytes.NewReader(body), erroringReader{err}), Closer: original} + return nil, true, false + } + if len(body) > maxMCPInspectBody { + req.Body = rechainedBody{Reader: io.MultiReader(bytes.NewReader(body), original), Closer: original} + return nil, true, false + } + _ = original.Close() + if len(body) == 0 { + req.Body = http.NoBody + } else { + req.Body = &bufferedBody{Reader: bytes.NewReader(body), data: body} + } + return body, true, true +} + +func parseMCPBody(body []byte) *mcpDetail { + calls, truncated, ok := parseJSONRPCCalls(body) + if !ok { + return unparsedMCP() + } + return &mcpDetail{Calls: calls, Parsed: true, Truncated: truncated} +} + +var errJSONRPCUnreadable = errors.New("not a JSON-RPC message the log can name") + +// Returns the calls to log, whether the batch held more than were kept, and whether the body was readable. +func parseJSONRPCCalls(body []byte) ([]mcpCall, bool, bool) { + dec := json.NewDecoder(bytes.NewReader(body)) + // Numbers are only skipped, and as float64 one past its range would fail an otherwise valid body. + dec.UseNumber() + + tok, err := dec.Token() + if err != nil { + return nil, false, false + } + calls := []mcpCall{} + truncated := false + switch tok { + case json.Delim('{'): + call, ok, err := readJSONRPCMessage(dec) + if err != nil { + return nil, false, false + } + if ok { + calls = append(calls, call) + } + case json.Delim('['): + for dec.More() { + elem, err := dec.Token() + if err != nil { + return nil, false, false + } + if elem != json.Delim('{') { + if err := skipJSONValue(dec, elem); err != nil { + return nil, false, false + } + continue + } + call, ok, err := readJSONRPCMessage(dec) + if err != nil { + return nil, false, false + } + if !ok { + continue + } + if len(calls) == maxMCPCallsPerRecord { + truncated = true + } else { + calls = append(calls, call) + } + } + if _, err := dec.Token(); err != nil { + return nil, false, false + } + default: + return nil, false, false + } + + // A server parsing the whole body rejects trailing content, so a second message hidden after the first + // must not leave the log naming only the first. + if _, err := dec.Token(); err != io.EOF { + return nil, false, false + } + return calls, truncated, true +} + +// Called with the opening brace consumed. A key repeated, or written in another case, makes the message +// unreadable: servers disagree on which copy wins, and Go's own decoder matches field names ignoring case, +// so the logged name could otherwise differ from the one the server runs. +func readJSONRPCMessage(dec *json.Decoder) (mcpCall, bool, error) { + var method, tool string + var hasMethod, hasParams, hasTool, toolNotString bool + + for dec.More() { + key, err := objectKey(dec) + if err != nil { + return mcpCall{}, false, err + } + switch { + case strings.EqualFold(key, "method"): + if key != "method" || hasMethod { + return mcpCall{}, false, errJSONRPCUnreadable + } + hasMethod = true + if method, err = stringValue(dec); err != nil { + return mcpCall{}, false, err + } + case strings.EqualFold(key, "params"): + if key != "params" || hasParams { + return mcpCall{}, false, errJSONRPCUnreadable + } + hasParams = true + if tool, hasTool, toolNotString, err = readParamsName(dec); err != nil { + return mcpCall{}, false, err + } + default: + if err := skipNextJSONValue(dec); err != nil { + return mcpCall{}, false, err + } + } + } + if _, err := dec.Token(); err != nil { + return mcpCall{}, false, err + } + + if !hasMethod { + return mcpCall{}, false, nil + } + if !mcpNamePattern.MatchString(method) { + return mcpCall{}, false, errJSONRPCUnreadable + } + call := mcpCall{Method: method} + if method == mcpMethodToolsCall && hasTool { + if toolNotString || !mcpNamePattern.MatchString(tool) { + return mcpCall{}, false, errJSONRPCUnreadable + } + call.Tool = tool + } + return call, true, nil +} + +// Only name is kept. Every other member, arguments included, is stepped over without being stored. A name +// that is not a string is reported rather than refused here, since it only matters when the method is +// tools/call. +func readParamsName(dec *json.Decoder) (name string, hasName, nameNotString bool, err error) { + tok, err := dec.Token() + if err != nil { + return "", false, false, err + } + if tok != json.Delim('{') { + return "", false, false, skipJSONValue(dec, tok) + } + + for dec.More() { + key, err := objectKey(dec) + if err != nil { + return "", false, false, err + } + if !strings.EqualFold(key, "name") { + if err := skipNextJSONValue(dec); err != nil { + return "", false, false, err + } + continue + } + if key != "name" || hasName { + return "", false, false, errJSONRPCUnreadable + } + hasName = true + value, err := dec.Token() + if err != nil { + return "", false, false, err + } + if str, ok := value.(string); ok { + name = str + } else { + nameNotString = true + if err := skipJSONValue(dec, value); err != nil { + return "", false, false, err + } + } + } + if _, err := dec.Token(); err != nil { + return "", false, false, err + } + return name, hasName, nameNotString, nil +} + +func objectKey(dec *json.Decoder) (string, error) { + tok, err := dec.Token() + if err != nil { + return "", err + } + key, ok := tok.(string) + if !ok { + return "", errJSONRPCUnreadable + } + return key, nil +} + +func stringValue(dec *json.Decoder) (string, error) { + tok, err := dec.Token() + if err != nil { + return "", err + } + s, ok := tok.(string) + if !ok { + return "", errJSONRPCUnreadable + } + return s, nil +} + +func skipNextJSONValue(dec *json.Decoder) error { + tok, err := dec.Token() + if err != nil { + return err + } + return skipJSONValue(dec, tok) +} + +// Called with the value's first token already read; consumes the rest of it when it is an object or array. +func skipJSONValue(dec *json.Decoder, first json.Token) error { + if first != json.Delim('{') && first != json.Delim('[') { + return nil + } + depth := 1 + for depth > 0 { + tok, err := dec.Token() + if err != nil { + return err + } + switch tok { + case json.Delim('{'), json.Delim('['): + depth++ + case json.Delim('}'), json.Delim(']'): + depth-- + } + } + return nil +} diff --git a/packages/agentvault/mcp_inspect_test.go b/packages/agentvault/mcp_inspect_test.go new file mode 100644 index 00000000..fb0e71b6 --- /dev/null +++ b/packages/agentvault/mcp_inspect_test.go @@ -0,0 +1,435 @@ +package agentvault + +import ( + "bytes" + "crypto/sha256" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "reflect" + "strings" + "testing" + + "github.com/Infisical/infisical-merge/packages/config" +) + +func TestParseJSONRPCCalls(t *testing.T) { + cases := []struct { + name string + body string + parsed bool + calls []mcpCall + }{ + {"single request", `{"jsonrpc":"2.0","id":1,"method":"tools/list"}`, true, []mcpCall{{Method: "tools/list"}}}, + {"tool call names the tool", `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"create_issue","arguments":{"title":"x"}}}`, true, + []mcpCall{{Method: "tools/call", Tool: "create_issue"}}}, + {"params before method", `{"params":{"arguments":{},"name":"search"},"method":"tools/call","id":3}`, true, + []mcpCall{{Method: "tools/call", Tool: "search"}}}, + {"notification", `{"jsonrpc":"2.0","method":"notifications/initialized"}`, true, []mcpCall{{Method: "notifications/initialized"}}}, + {"name ignored outside tools/call", `{"method":"prompts/get","params":{"name":"greet"}}`, true, []mcpCall{{Method: "prompts/get"}}}, + {"tools/call without params", `{"method":"tools/call","id":1}`, true, []mcpCall{{Method: "tools/call"}}}, + {"batch", `[{"method":"tools/call","params":{"name":"a"}},{"method":"notifications/cancelled"},{"jsonrpc":"2.0","id":9,"result":{}}]`, true, + []mcpCall{{Method: "tools/call", Tool: "a"}, {Method: "notifications/cancelled"}}}, + {"response only", `{"jsonrpc":"2.0","id":1,"result":{"method":"tools/call"}}`, true, []mcpCall{}}, + {"empty batch", `[]`, true, []mcpCall{}}, + {"duplicate keys inside arguments are not ours to judge", `{"method":"tools/call","params":{"name":"a","arguments":{"k":1,"k":2}}}`, true, + []mcpCall{{Method: "tools/call", Tool: "a"}}}, + {"huge number skipped", `{"method":"ping","id":1e999}`, true, []mcpCall{{Method: "ping"}}}, + + {"duplicate method", `{"method":"tools/list","method":"tools/call","params":{"name":"a"}}`, false, nil}, + {"duplicate params", `{"method":"tools/call","params":{"name":"a"},"params":{"name":"b"}}`, false, nil}, + {"duplicate name", `{"method":"tools/call","params":{"name":"a","name":"b"}}`, false, nil}, + {"case variant of method", `{"method":"tools/list","Method":"tools/call"}`, false, nil}, + {"case variant of name", `{"method":"tools/call","params":{"name":"a","NAME":"b"}}`, false, nil}, + {"duplicate in one batch element", `[{"method":"ping"},{"method":"a","method":"b"}]`, false, nil}, + {"method with a space", `{"method":"tools call"}`, false, nil}, + {"tool with a quote", `{"method":"tools/call","params":{"name":"a\"b"}}`, false, nil}, + {"tool too long", `{"method":"tools/call","params":{"name":"` + strings.Repeat("a", 129) + `"}}`, false, nil}, + {"tool not a string", `{"method":"tools/call","params":{"name":{"x":1}}}`, false, nil}, + {"method not a string", `{"method":7}`, false, nil}, + {"malformed", `{"method":"tools/call",`, false, nil}, + {"trailing message", `{"method":"ping"} {"method":"tools/call","params":{"name":"x"}}`, false, nil}, + {"top level scalar", `"tools/call"`, false, nil}, + {"empty", ``, false, nil}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := parseMCPBody([]byte(tc.body)) + if got.Parsed != tc.parsed { + t.Fatalf("parsed = %v, want %v (calls %+v)", got.Parsed, tc.parsed, got.Calls) + } + want := tc.calls + if want == nil { + want = []mcpCall{} + } + if !reflect.DeepEqual(got.Calls, want) { + t.Fatalf("calls = %+v, want %+v", got.Calls, want) + } + }) + } +} + +func pingBatch(n int, tail ...string) string { + elems := make([]string, 0, n+len(tail)) + for range n { + elems = append(elems, `{"method":"ping"}`) + } + return "[" + strings.Join(append(elems, tail...), ",") + "]" +} + +func TestParseJSONRPCCallsTruncatesLongBatches(t *testing.T) { + kept := make([]mcpCall, maxMCPCallsPerRecord) + for i := range kept { + kept[i] = mcpCall{Method: "ping"} + } + cases := []struct { + name string + body string + want *mcpDetail + }{ + {"exactly the limit", pingBatch(maxMCPCallsPerRecord), &mcpDetail{Calls: kept, Parsed: true}}, + {"one past the limit", pingBatch(maxMCPCallsPerRecord + 1), &mcpDetail{Calls: kept, Parsed: true, Truncated: true}}, + {"far past the limit", pingBatch(100), &mcpDetail{Calls: kept, Parsed: true, Truncated: true}}, + {"responses past the limit are not calls", pingBatch(maxMCPCallsPerRecord, `{"id":1,"result":{}}`), &mcpDetail{Calls: kept, Parsed: true}}, + {"tool call past the limit is still validated", pingBatch(maxMCPCallsPerRecord+5, `{"method":"tools/call","params":{"name":{"x":1}}}`), unparsedMCP()}, + {"duplicate key past the limit", pingBatch(maxMCPCallsPerRecord+1, `{"method":"a","method":"b"}`), unparsedMCP()}, + {"malformed past the limit", pingBatch(maxMCPCallsPerRecord+1, `{"method":`), unparsedMCP()}, + {"trailing content after a long batch", pingBatch(maxMCPCallsPerRecord+1) + ` {"method":"ping"}`, unparsedMCP()}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := parseMCPBody([]byte(tc.body)) + if !reflect.DeepEqual(got, tc.want) { + t.Fatalf("got %+v, want %+v", got, tc.want) + } + }) + } +} + +func TestMCPDetailOmitsTruncatedWhenFalse(t *testing.T) { + out, err := json.Marshal(&mcpDetail{Calls: []mcpCall{{Method: "ping"}}, Parsed: true}) + if err != nil { + t.Fatal(err) + } + if string(out) != `{"calls":[{"method":"ping"}],"parsed":true}` { + t.Fatalf("got %s", out) + } + out, err = json.Marshal(parseMCPBody([]byte(pingBatch(maxMCPCallsPerRecord + 1)))) + if err != nil { + t.Fatal(err) + } + if !bytes.HasSuffix(out, []byte(`"parsed":true,"truncated":true}`)) { + t.Fatalf("got %s", out) + } +} + +type mcpUpstream struct { + server *httptest.Server + bodies chan []byte + auth chan string +} + +func newMCPUpstream(t *testing.T) *mcpUpstream { + t.Helper() + u := &mcpUpstream{bodies: make(chan []byte, 1), auth: make(chan string, 1)} + u.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + u.bodies <- body + u.auth <- r.Header.Get("Authorization") + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(u.server.Close) + return u +} + +func (u *mcpUpstream) host() string { return strings.TrimPrefix(u.server.URL, "http://") } + +func mcpService(host string, subs ...substitution) *resolvedService { + svc := policyService(host, nil, nil, nil, subs) + svc.name = "linear" + svc.kind = serviceKindMCP + svc.credential = credential{kind: credentialBearer, headerPrefix: "Bearer", value: []byte("real-oauth-token")} + return svc +} + +func proxiedDo(t *testing.T, front *httptest.Server, method, target, contentType string, body io.Reader, header http.Header) *http.Response { + t.Helper() + proxyURL, err := url.Parse(front.URL) + if err != nil { + t.Fatal(err) + } + proxyURL.User = url.UserPassword("infisical", "agv_test-token") + client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyURL)}} + + req, err := http.NewRequest(method, target, body) + if err != nil { + t.Fatal(err) + } + for name, values := range header { + req.Header[name] = values + } + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } + res, err := client.Do(req) + if err != nil { + t.Fatalf("request through the proxy failed: %v", err) + } + _, _ = io.Copy(io.Discard, res.Body) + _ = res.Body.Close() + return res +} + +func recordJSON(t *testing.T, rec sessionLogRecord) string { + t.Helper() + b, err := json.Marshal(rec) + if err != nil { + t.Fatal(err) + } + return string(b) +} + +func TestAnMCPToolCallIsRecordedAndForwardedUnchanged(t *testing.T) { + up := newMCPUpstream(t) + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, []*resolvedService{mcpService(up.host())}) + + body := `{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"create_issue","arguments":{"title":"Bug","Authorization":"x"}}}` + proxiedDo(t, front, http.MethodPost, up.server.URL+"/mcp", "application/json; charset=utf-8", strings.NewReader(body), nil) + + if got := <-up.bodies; string(got) != body { + t.Fatalf("upstream body changed:\n got %s\nwant %s", got, body) + } + if got := <-up.auth; got != "Bearer real-oauth-token" { + t.Fatalf("credential was not injected, upstream saw %q", got) + } + + rec := drainOneRecord(t, log) + if rec.ServiceType == nil || *rec.ServiceType != serviceKindMCP { + t.Fatalf("serviceType = %v", rec.ServiceType) + } + want := &mcpDetail{Calls: []mcpCall{{Method: "tools/call", Tool: "create_issue"}}, Parsed: true} + if !reflect.DeepEqual(rec.MCP, want) { + t.Fatalf("mcp = %+v, want %+v", rec.MCP, want) + } + encoded := recordJSON(t, rec) + if !strings.Contains(encoded, `"serviceType":"mcp"`) || + !strings.Contains(encoded, `"mcp":{"calls":[{"method":"tools/call","tool":"create_issue"}],"parsed":true}`) { + t.Fatalf("record JSON is %s", encoded) + } + for _, leaked := range []string{"real-oauth-token", "Bug", "arguments"} { + if strings.Contains(encoded, leaked) { + t.Fatalf("record carries %q: %s", leaked, encoded) + } + } +} + +func TestAnMCPBatchAndNotificationAreRecorded(t *testing.T) { + up := newMCPUpstream(t) + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, []*resolvedService{mcpService(up.host())}) + + body := `[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"search"}}]` + proxiedDo(t, front, http.MethodPost, up.server.URL+"/mcp", "application/json", strings.NewReader(body), nil) + <-up.bodies + + rec := drainOneRecord(t, log) + want := &mcpDetail{Calls: []mcpCall{{Method: "notifications/initialized"}, {Method: "tools/call", Tool: "search"}}, Parsed: true} + if !reflect.DeepEqual(rec.MCP, want) { + t.Fatalf("mcp = %+v, want %+v", rec.MCP, want) + } +} + +func TestAnMCPBodyTheProxyCannotTrustIsRecordedUnparsedAndForwarded(t *testing.T) { + oversized := `{"method":"tools/call","params":{"name":"big","arguments":{"blob":"` + strings.Repeat("x", maxMCPInspectBody) + `"}}}` + cases := []struct { + name string + body string + header http.Header + chunked bool + }{ + {"duplicate keys", `{"method":"tools/call","params":{"name":"safe","name":"rm_rf"}}`, nil, false}, + {"malformed", `{"method":"tools/call","params":`, nil, false}, + {"oversized", oversized, nil, false}, + {"oversized and chunked", oversized, nil, true}, + {"encoded", `{"method":"tools/call","params":{"name":"x"}}`, http.Header{"Content-Encoding": {"gzip"}}, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + up := newMCPUpstream(t) + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, []*resolvedService{mcpService(up.host())}) + + var reader io.Reader = strings.NewReader(tc.body) + if tc.chunked { + reader = io.MultiReader(reader) + } + res := proxiedDo(t, front, http.MethodPost, up.server.URL+"/mcp", "application/json", reader, tc.header) + if res.StatusCode != http.StatusOK { + t.Fatalf("status %d; parsing must never block a request", res.StatusCode) + } + if got := <-up.bodies; sha256.Sum256(got) != sha256.Sum256([]byte(tc.body)) { + t.Fatalf("upstream body changed (%d bytes, sent %d)", len(got), len(tc.body)) + } + + rec := drainOneRecord(t, log) + if !reflect.DeepEqual(rec.MCP, &mcpDetail{Calls: []mcpCall{}, Parsed: false}) { + t.Fatalf("mcp = %+v", rec.MCP) + } + if encoded := recordJSON(t, rec); !strings.Contains(encoded, `"mcp":{"calls":[],"parsed":false}`) { + t.Fatalf("record JSON is %s", encoded) + } + }) + } +} + +func TestMCPDetailIsOmittedWhereNoParseIsAttempted(t *testing.T) { + cases := []struct { + name string + method string + contentType string + body string + kind string + serviceType string + }{ + {"GET on an mcp service", http.MethodGet, "", "", serviceKindMCP, serviceKindMCP}, + {"DELETE on an mcp service", http.MethodDelete, "", "", serviceKindMCP, serviceKindMCP}, + {"non-JSON content type", http.MethodPost, "text/plain", `{"method":"tools/call","params":{"name":"x"}}`, serviceKindMCP, serviceKindMCP}, + {"http service", http.MethodPost, "application/json", `{"method":"tools/call","params":{"name":"x"}}`, serviceKindHTTP, serviceKindHTTP}, + {"service with no kind", http.MethodPost, "application/json", `{"method":"tools/call","params":{"name":"x"}}`, "", serviceKindHTTP}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + up := newMCPUpstream(t) + svc := mcpService(up.host()) + svc.kind = tc.kind + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, []*resolvedService{svc}) + + var body io.Reader + if tc.body != "" { + body = strings.NewReader(tc.body) + } + proxiedDo(t, front, tc.method, up.server.URL+"/mcp", tc.contentType, body, nil) + if got := <-up.bodies; string(got) != tc.body { + t.Fatalf("upstream body changed: %q", got) + } + + rec := drainOneRecord(t, log) + if rec.MCP != nil { + t.Fatalf("mcp = %+v, expected none", rec.MCP) + } + if rec.ServiceType == nil || *rec.ServiceType != tc.serviceType { + t.Fatalf("serviceType = %v, want %q", rec.ServiceType, tc.serviceType) + } + if encoded := recordJSON(t, rec); strings.Contains(encoded, `"mcp":`) { + t.Fatalf("record JSON is %s", encoded) + } + }) + } +} + +func TestAnUnmatchedHostRecordsANullServiceType(t *testing.T) { + up := newMCPUpstream(t) + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, nil) + + body := `{"method":"tools/call","params":{"name":"x"}}` + proxiedDo(t, front, http.MethodPost, up.server.URL+"/mcp", "application/json", strings.NewReader(body), nil) + <-up.bodies + + rec := drainOneRecord(t, log) + if rec.ServiceType != nil || rec.MCP != nil { + t.Fatalf("passthrough record has serviceType %v and mcp %+v", rec.ServiceType, rec.MCP) + } + encoded := recordJSON(t, rec) + if !strings.Contains(encoded, `"serviceType":null`) || strings.Contains(encoded, `"mcp":`) { + t.Fatalf("record JSON is %s", encoded) + } +} + +func TestMCPParsingSeesTheBodyBeforeABodySubstitution(t *testing.T) { + up := newMCPUpstream(t) + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, []*resolvedService{ + mcpService(up.host(), subOn("__TOOL__", "secret_tool", surfaceBody)), + }) + + body := `{"method":"tools/call","params":{"name":"lookup","arguments":{"key":"__TOOL__"}}}` + proxiedDo(t, front, http.MethodPost, up.server.URL+"/mcp", "application/json", strings.NewReader(body), nil) + + want := strings.Replace(body, "__TOOL__", "secret_tool", 1) + if got := <-up.bodies; string(got) != want { + t.Fatalf("upstream body is %s, want %s", got, want) + } + + rec := drainOneRecord(t, log) + if !reflect.DeepEqual(rec.MCP, &mcpDetail{Calls: []mcpCall{{Method: "tools/call", Tool: "lookup"}}, Parsed: true}) { + t.Fatalf("mcp = %+v", rec.MCP) + } + if encoded := recordJSON(t, rec); strings.Contains(encoded, "secret_tool") { + t.Fatalf("a substituted value reached the record: %s", encoded) + } +} + +func TestASubstitutedToolNameIsLoggedAsTheAgentSentIt(t *testing.T) { + up := newMCPUpstream(t) + front, log, _ := newRecordingProxy(t, TrafficPolicyAnyHost, []*resolvedService{ + mcpService(up.host(), subOn("__TOOL__", "secret_tool", surfaceBody)), + }) + + proxiedDo(t, front, http.MethodPost, up.server.URL+"/mcp", "application/json", + strings.NewReader(`{"method":"tools/call","params":{"name":"__TOOL__"}}`), nil) + if got := <-up.bodies; !bytes.Contains(got, []byte("secret_tool")) { + t.Fatalf("the substitution never fired, so this proves nothing: %s", got) + } + + rec := drainOneRecord(t, log) + if len(rec.MCP.Calls) != 1 || rec.MCP.Calls[0].Tool != "__TOOL__" { + t.Fatalf("mcp = %+v", rec.MCP) + } +} + +func TestAnUnreadableMCPBodyIsStillRefusedByABodySubstitution(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "http://example.test/mcp", io.MultiReader(strings.NewReader(`{"method"`), erroringReader{io.ErrUnexpectedEOF})) + req.Header.Set("Content-Type", "application/json") + req.ContentLength = -1 + + if got := inspectMCPRequest(req); got == nil || got.Parsed { + t.Fatalf("mcp = %+v", got) + } + if _, err := applySubstitutions(req, "linear", []substitution{subOn("__X__", "y", surfaceBody)}); err != errBodyUnreadable { + t.Fatalf("err = %v, want errBodyUnreadable", err) + } +} + +func TestResolveMapsTheServiceType(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + services := []map[string]any{ + {"id": "1", "name": "linear", "type": "mcp", "hostPattern": "mcp.linear.app", "credential": map[string]any{"type": "passthrough"}}, + {"id": "2", "name": "github", "type": "http", "hostPattern": "api.github.com", "credential": map[string]any{"type": "passthrough"}}, + {"id": "3", "name": "future", "type": "grpc", "hostPattern": "a.example", "credential": map[string]any{"type": "passthrough"}}, + {"id": "4", "name": "older", "hostPattern": "b.example", "credential": map[string]any{"type": "passthrough"}}, + } + _ = json.NewEncoder(w).Encode(map[string]any{"sessionId": "s1", "services": services}) + })) + defer srv.Close() + old := config.INFISICAL_URL + config.INFISICAL_URL = srv.URL + "/api" + defer func() { config.INFISICAL_URL = old }() + + resolver, err := newInfisicalResolver(func() string { return "ptok" }) + if err != nil { + t.Fatal(err) + } + res, err := resolver.resolve("tok", nil) + if err != nil { + t.Fatal(err) + } + var got []string + for _, svc := range res.Services { + got = append(got, fmt.Sprintf("%s=%s/%s", svc.name, svc.kind, svc.serviceType())) + } + want := []string{"linear=mcp/mcp", "github=http/http", "future=http/http", "older=http/http"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("kinds = %v, want %v", got, want) + } +} diff --git a/packages/agentvault/proxy.go b/packages/agentvault/proxy.go index c3300721..c31c809b 100644 --- a/packages/agentvault/proxy.go +++ b/packages/agentvault/proxy.go @@ -421,9 +421,10 @@ func (ps *proxyServer) forwardHTTP(w http.ResponseWriter, r *http.Request, schem // reqPath was taken before forward, so a credential substituted into the path never reaches the record. if outcome.sessionLog != nil { - var service, bundle *string + var service, bundle, serviceType *string if matched != nil { - service, bundle = &matched.name, &matched.accessBundleName + kind := matched.serviceType() + service, bundle, serviceType = &matched.name, &matched.accessBundleName, &kind } ps.sessionLogs.record(outcome.sessionLog, sessionLogRecord{ Method: reqMethod, @@ -434,6 +435,8 @@ func (ps *proxyServer) forwardHTTP(w http.ResponseWriter, r *http.Request, schem Decision: decision, Service: service, AccessBundle: bundle, + ServiceType: serviceType, + MCP: outcome.mcp, }) } @@ -477,6 +480,7 @@ type forwardOutcome struct { brokered bool substituted []string sessionLog *sessionLogGrant + mcp *mcpDetail } func (ps *proxyServer) forward(req *http.Request, scheme, hostname, port, sessionToken string) (*http.Response, *resolvedService, forwardOutcome, error) { @@ -503,6 +507,9 @@ func (ps *proxyServer) forward(req *http.Request, scheme, hostname, port, sessio } matched := bestMatch(services, hostname, port) + // Set only when the proxy injected a credential Infisical can refresh, which is what decides whether an + // upstream 401 or 403 is handled here or passed to the agent. + var sent *credential if ps.blocksOffBundle(matched, hostname, port) { return nil, nil, outcome, fmt.Errorf("no service covers host %q: %w", hostname, errHostBlocked) @@ -533,6 +540,17 @@ func (ps *proxyServer) forward(req *http.Request, scheme, hostname, port, sessio } if matched != nil { + // Before anything is injected, so the names it records can only have come from the agent. Parsed + // only for a session log; with logs off the body is still buffered for an OAuth credential, so a + // rejected token's request can be sent again. + if matched.kind == serviceKindMCP { + if grant != nil { + outcome.mcp = inspectMCPRequest(req) + } else if matched.credential.refreshable { + bufferMCPRequest(req) + } + } + // Substitutions first, so an injected real value can never itself be rewritten. The credential last, // so a custom header naming the credential's own header loses rather than replacing the token. substituted, err := applySubstitutions(req, matched.name, matched.substitutions) @@ -547,6 +565,9 @@ func (ps *proxyServer) forward(req *http.Request, scheme, hostname, port, sessio } if injectCredential(req, &matched.credential) { outcome.brokered = true + if matched.credential.refreshable { + sent = &matched.credential + } } // After the brokered headers rather than before them, so a custom header cannot reintroduce an @@ -569,10 +590,23 @@ func (ps *proxyServer) forward(req *http.Request, scheme, hostname, port, sessio } } + // Captured before the round trip, which consumes the body. + var replay []byte + replayable := false + if sent != nil { + replay, replayable = replayableBody(req) + } + resp, err := ps.transport.RoundTrip(req) if err != nil { return nil, matched, outcome, err } + if sent != nil { + resp, err = ps.handleUpstreamAuth(req, resp, matched, sessionToken, sent, replay, replayable) + if err != nil { + return nil, matched, outcome, err + } + } return resp, matched, outcome, nil } diff --git a/packages/agentvault/resolve.go b/packages/agentvault/resolve.go index 3545d789..0b756160 100644 --- a/packages/agentvault/resolve.go +++ b/packages/agentvault/resolve.go @@ -64,7 +64,9 @@ func (r *infisicalResolver) resolve(sessionToken string, held *sessionLogGrant) services = append(services, &resolvedService{ id: wire.ID, name: wire.Name, + kind: toServiceKind(wire.Type), accessBundleName: wire.AccessBundleName, + hostPattern: wire.HostPattern, hostPatterns: parseHostPatterns(wire.HostPattern), allowedMethods: toMethodSet(wire.AllowedMethods), allowedPathPrefixes: toPathPrefixes(wire.AllowedPathPrefixes), @@ -82,6 +84,27 @@ func (r *infisicalResolver) resolve(sessionToken string, held *sessionLogGrant) }, nil } +// Bounded by the resolver client's timeout, so a stalled control plane fails the refresh rather than +// holding the agent's request open. +func (r *infisicalResolver) refreshCredential(sessionToken, serviceID, hostPattern, rejectedHash string) (credential, string, error) { + res, err := api.CallRefreshAgentVaultServiceCredential(r.client, sessionToken, serviceID, api.RefreshAgentVaultServiceCredentialRequest{ + HostPattern: hostPattern, + RejectedTokenHash: rejectedHash, + }) + if err != nil { + return credential{}, "", err + } + return toCredential(res.Credential), res.Status, nil +} + +func (r *infisicalResolver) reportMissingScopes(sessionToken, serviceID, hostPattern, rejectedHash string, requiredScopes []string) error { + return api.CallReportAgentVaultServiceMissingScopes(r.client, sessionToken, serviceID, api.ReportAgentVaultServiceMissingScopesRequest{ + HostPattern: hostPattern, + RejectedTokenHash: rejectedHash, + RequiredScopes: requiredScopes, + }) +} + func toSessionLogGrant(sessionID string, wire api.AgentVaultSessionLogGrant, held *sessionLogGrant) *sessionLogGrant { if !wire.Enabled { return nil @@ -103,15 +126,27 @@ func toSessionLogGrant(sessionID string, wire api.AgentVaultSessionLogGrant, hel return newSessionLogGrant(sessionID, key) } +// Anything unrecognised is http, so a newer backend's type is brokered the way this proxy already knows how. +func toServiceKind(wire string) string { + if wire == serviceKindMCP { + return serviceKindMCP + } + return serviceKindHTTP +} + func toCredential(wire api.AgentVaultCredential) credential { switch wire.Type { case credentialBearer: - return credential{ + cred := credential{ kind: credentialBearer, headerName: wire.HeaderName, headerPrefix: wire.HeaderPrefix, value: []byte(wire.Value), } + if wire.Source == credentialSourceOAuth { + cred.refreshable = true + } + return cred case credentialBasic: return credential{ kind: credentialBasic, diff --git a/packages/agentvault/rewrite.go b/packages/agentvault/rewrite.go index 3ffbdb50..e55798c2 100644 --- a/packages/agentvault/rewrite.go +++ b/packages/agentvault/rewrite.go @@ -240,7 +240,13 @@ func applyBodySubstitutions(req *http.Request, serviceName string, subs []substi return false, nil } - body, err := io.ReadAll(io.LimitReader(req.Body, maxBodyRewriteSize+1)) + var body []byte + var err error + if buffered, ok := req.Body.(*bufferedBody); ok { + body = buffered.data + } else { + body, err = io.ReadAll(io.LimitReader(req.Body, maxBodyRewriteSize+1)) + } if err != nil { // Refused outright rather than forwarded short. This used to leave ContentLength disagreeing with the // bytes so http.Transport would refuse, but a chunked upload declares -1 and there is nothing to leave @@ -286,7 +292,7 @@ func applyBodySubstitutions(req *http.Request, serviceName string, subs []substi // POST into a chunked request. Signing schemes and some gateways reject that. req.Body = http.NoBody } else { - req.Body = io.NopCloser(bytes.NewReader(rewritten)) + req.Body = &bufferedBody{Reader: bytes.NewReader(rewritten), data: rewritten} } req.ContentLength = int64(len(rewritten)) req.Header.Set("Content-Length", fmt.Sprintf("%d", len(rewritten))) diff --git a/packages/agentvault/session_log_crypto_test.go b/packages/agentvault/session_log_crypto_test.go index 81780465..9462c93f 100644 --- a/packages/agentvault/session_log_crypto_test.go +++ b/packages/agentvault/session_log_crypto_test.go @@ -20,7 +20,7 @@ const ( vectorIVHex = "aabbccddeeff001122334455" vectorAADHex = "38d8f3b86fbbd1061f9d2a8bd91e6c49c51b6231936a7ecfc21c56e129f37a1c" vectorIVBase64 = "qrvM3e7/ABEiM0RV" - vectorCiphertext = "PLRwxBbgu+W68Br1N9gY1oUy8wjJxQClAtBh0NfJS1UcWOCPn3laS615sIqwFONhPIPNWRI3CA+a5tUJ7aoim0sQkE4d9gzou2mc/AWiCdToVBJPtdumA9jIzh3yAI81YPwcoDXEVnq2+7ooNNJShGdLX95itbrna/t4nFKRKSSgNzbH23eMtSMcSo72puk/2iwh4sVbTKzC2kwvbf1U6Mgd21zkIq2jDKKwhcT6mTfjPivW4FzmmkspQVMoWwANRX+QVyXzrMipZfoq5N/UcUI6rCvRUkqg+3ST5GVMelW0mjOO" + vectorCiphertext = "PLRwxBbgu+W68Br1N9gY1oUy8wjJxQClAtBh0NfJS1UcWOCPn3laS615sIqwFONhPIPNWRI3CA+a5tUJ7aoim0sQkE4d9gzou2mc/AWiCdToVBJPtdumA9jIzh3yAI81YPwcoDXEVnq2+7ooNNJShGdLX95itbrna/t4nFKRKSSgNzbH23eMtSMcSo72puk/2iwh4sVbTKzC2kwvbf1U6Mgd21zkIq2jDKKwhcT6mTfjPivW4FzmmkspQVMoWwANRX+QVyXzrMipZfoq5N/UcUI6/VTBSsIZvF08UTnPeIW+yM1I5hHCoi0WQT6QRe0bdx5+BvGuvJQB" ) // The uploaded object is IV ‖ ciphertext ‖ tag. The 12-byte IV encodes without padding, so the object's base64 @@ -33,7 +33,7 @@ var vectorContext = struct{ sessionID, chunkID string }{ } func vectorRecords() []sessionLogRecord { - service, bundle := "github", "code-review" + service, bundle, serviceType := "github", "code-review", serviceKindHTTP return []sessionLogRecord{{ Ts: "2026-09-16T10:31:04.221Z", Seq: 1, @@ -46,6 +46,7 @@ func vectorRecords() []sessionLogRecord { Decision: decisionBrokered, Service: &service, AccessBundle: &bundle, + ServiceType: &serviceType, }} } diff --git a/packages/agentvault/session_log_spool.go b/packages/agentvault/session_log_spool.go index e5042911..d59c3291 100644 --- a/packages/agentvault/session_log_spool.go +++ b/packages/agentvault/session_log_spool.go @@ -7,7 +7,11 @@ import ( "github.com/Infisical/infisical-merge/packages/api" ) -// Never add headers, bodies or the query string: they can carry the injected credential. +// Never add headers, bodies or the query string: they can carry the injected credential. The one exception +// is mcp, which holds only JSON-RPC method and tool names from the agent's own body. They are read before +// any substitution or credential reaches the request, so nothing injected can be among them, and only names +// matching mcpNamePattern are kept, so a value that happens to hold a secret cannot be logged as one. Tool +// arguments are never read. type sessionLogRecord struct { Ts string `json:"ts"` Seq uint64 `json:"seq"` @@ -20,6 +24,9 @@ type sessionLogRecord struct { Decision string `json:"decision"` Service *string `json:"service"` AccessBundle *string `json:"accessBundle"` + // Null when no service matched, so a reader can tell that apart from an http service. + ServiceType *string `json:"serviceType"` + MCP *mcpDetail `json:"mcp,omitempty"` at time.Time } diff --git a/packages/agentvault/upstream_auth.go b/packages/agentvault/upstream_auth.go new file mode 100644 index 00000000..4c7784d8 --- /dev/null +++ b/packages/agentvault/upstream_auth.go @@ -0,0 +1,208 @@ +package agentvault + +import ( + "bytes" + "io" + "net/http" + "strconv" + "strings" + + "github.com/rs/zerolog/log" +) + +const ( + credentialSourceOAuth = "oauth" + oauthStatusConnected = "connected" + + // The backend's limits on requiredScopes; anything past them would fail the whole report. + maxReportedScopes = 50 + maxReportedScopeLen = 256 +) + +// No WWW-Authenticate goes with it: passing the upstream's challenge on would start the agent's MCP client +// on an OAuth flow of its own, when the credential belongs to Infisical. +const upstreamAuthFailedBody = `{"error":"upstream_auth_failed","message":"The MCP server rejected the credential Infisical provided. Reauthorize this service in Infisical Agent Vault."}` + +// The request body as it will go upstream, when it can be sent a second time. Only a body already held in +// memory, or none at all, can be: a streamed one is consumed by the first attempt. +func replayableBody(req *http.Request) ([]byte, bool) { + if req.Body == nil || req.Body == http.NoBody { + return nil, true + } + if buffered, ok := req.Body.(*bufferedBody); ok { + return buffered.data, true + } + return nil, false +} + +// handleUpstreamAuth runs only for a refreshable credential the proxy injected, so every other credential +// sees the upstream's response untouched. +func (ps *proxyServer) handleUpstreamAuth(req *http.Request, resp *http.Response, svc *resolvedService, sessionToken string, sent *credential, body []byte, replayable bool) (*http.Response, error) { + switch resp.StatusCode { + case http.StatusUnauthorized: + // An auth-error body can stay open indefinitely; draining it would stall the refresh. + _ = resp.Body.Close() + fresh, err := ps.cache.refreshCredential(sessionToken, svc.id, svc.hostPattern, sent) + if err != nil { + log.Warn().Err(err).Str("service", svc.name).Str("serviceId", svc.id). + Msg("agent-vault: the upstream rejected the OAuth token and Infisical could not refresh it") + return upstreamAuthFailed(req), nil + } + if !replayable || bytes.Equal(fresh.value, sent.value) { + log.Warn().Str("service", svc.name).Str("serviceId", svc.id).Bool("replayable", replayable). + Msg("agent-vault: the upstream rejected the OAuth token, not retrying") + return upstreamAuthFailed(req), nil + } + + retry := req.Clone(req.Context()) + if len(body) == 0 { + retry.Body = http.NoBody + } else { + retry.Body = &bufferedBody{Reader: bytes.NewReader(body), data: body} + } + injectCredential(retry, &fresh) + retryResp, err := ps.transport.RoundTrip(retry) + if err != nil { + return nil, err + } + // Only a 401 or an insufficient_scope challenge is about the token. Any other 403 is the upstream's + // answer to the request, which the refreshed token may only now have reached. + if retryResp.StatusCode != http.StatusUnauthorized && !ps.reportInsufficientScope(retryResp, sessionToken, svc, &fresh) { + return retryResp, nil + } + _ = retryResp.Body.Close() + log.Warn().Str("service", svc.name).Str("serviceId", svc.id).Int("retryStatus", retryResp.StatusCode). + Msg("agent-vault: the upstream rejected the refreshed OAuth token too") + return upstreamAuthFailed(req), nil + + case http.StatusForbidden: + if !ps.reportInsufficientScope(resp, sessionToken, svc, sent) { + return resp, nil + } + _ = resp.Body.Close() + log.Warn().Str("service", svc.name).Str("serviceId", svc.id). + Msg("agent-vault: the upstream needs OAuth scopes this service was not granted") + return upstreamAuthFailed(req), nil + } + return resp, nil +} + +// Reports whether the response was an insufficient_scope challenge. The scopes go to Infisical so the +// service shows what reauthorizing has to request; a challenge naming none has nothing to record. +func (ps *proxyServer) reportInsufficientScope(resp *http.Response, sessionToken string, svc *resolvedService, sent *credential) bool { + if resp.StatusCode != http.StatusForbidden { + return false + } + params := bearerChallengeParams(resp.Header.Values("WWW-Authenticate")) + if params["error"] != "insufficient_scope" { + return false + } + scopes := challengedScopes(params["scope"]) + if len(scopes) == 0 { + return true + } + if err := ps.cache.reportMissingScopes(sessionToken, svc.id, svc.hostPattern, sent, scopes); err != nil { + log.Warn().Err(err).Str("service", svc.name).Str("serviceId", svc.id). + Msg("agent-vault: could not report the upstream's required scopes to Infisical") + } + return true +} + +func challengedScopes(raw string) []string { + var scopes []string + for _, scope := range strings.Fields(raw) { + if len(scope) > maxReportedScopeLen { + continue + } + scopes = append(scopes, scope) + if len(scopes) == maxReportedScopes { + break + } + } + return scopes +} + +func upstreamAuthFailed(req *http.Request) *http.Response { + header := make(http.Header) + header.Set("Content-Type", "application/json") + header.Set("Content-Length", strconv.Itoa(len(upstreamAuthFailedBody))) + return &http.Response{ + Status: "502 Bad Gateway", + StatusCode: http.StatusBadGateway, + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: header, + Body: io.NopCloser(strings.NewReader(upstreamAuthFailedBody)), + ContentLength: int64(len(upstreamAuthFailedBody)), + Request: req, + } +} + +// The auth-params of the Bearer challenges in WWW-Authenticate (RFC 9110 section 11.6.1), names +// lower-cased, first occurrence kept. Params of other schemes are skipped. +func bearerChallengeParams(values []string) map[string]string { + params := map[string]string{} + for _, value := range values { + inBearer := false + i := 0 + for i < len(value) { + for i < len(value) && (value[i] == ' ' || value[i] == '\t' || value[i] == ',') { + i++ + } + start := i + for i < len(value) && isTokenChar(value[i]) { + i++ + } + if i == start { + i++ + continue + } + name := value[start:i] + j := i + for j < len(value) && (value[j] == ' ' || value[j] == '\t') { + j++ + } + if j >= len(value) || value[j] != '=' { + inBearer = strings.EqualFold(name, "bearer") + continue + } + i = j + 1 + for i < len(value) && (value[i] == ' ' || value[i] == '\t') { + i++ + } + var paramValue string + if i < len(value) && value[i] == '"' { + var b strings.Builder + i++ + for i < len(value) && value[i] != '"' { + if value[i] == '\\' && i+1 < len(value) { + i++ + } + b.WriteByte(value[i]) + i++ + } + i++ + paramValue = b.String() + } else { + start := i + for i < len(value) && isTokenChar(value[i]) { + i++ + } + paramValue = value[start:i] + } + key := strings.ToLower(name) + if _, seen := params[key]; inBearer && !seen { + params[key] = paramValue + } + } + } + return params +} + +func isTokenChar(c byte) bool { + if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' { + return true + } + return strings.IndexByte("!#$%&'*+-.^_`|~", c) >= 0 +} diff --git a/packages/agentvault/upstream_auth_test.go b/packages/agentvault/upstream_auth_test.go new file mode 100644 index 00000000..7eb6fc19 --- /dev/null +++ b/packages/agentvault/upstream_auth_test.go @@ -0,0 +1,1174 @@ +package agentvault + +import ( + "bufio" + "crypto/sha256" + "crypto/tls" + "crypto/x509" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Infisical/infisical-merge/packages/api" + "github.com/Infisical/infisical-merge/packages/config" +) + +type refreshCall struct { + sessionToken string + serviceID string + hostPattern string + rejectedHash string + requiredScopes []string +} + +type oauthResolver struct { + services []*resolvedService + + mu sync.Mutex + calls []refreshCall + issue credential + status string + err error + // Closed by the test to release refreshes held open, so concurrent 401s can be made to overlap one. + gate chan struct{} +} + +func (r *oauthResolver) resolve(string, *sessionLogGrant) (*resolveResult, error) { + r.mu.Lock() + defer r.mu.Unlock() + return &resolveResult{SessionID: "s1", Services: r.services}, nil +} + +func (r *oauthResolver) setServices(services ...*resolvedService) { + r.mu.Lock() + defer r.mu.Unlock() + r.services = services +} + +func (r *oauthResolver) refreshCredential(sessionToken, serviceID, hostPattern, rejectedHash string) (credential, string, error) { + if r.gate != nil { + <-r.gate + } + r.mu.Lock() + defer r.mu.Unlock() + r.calls = append(r.calls, refreshCall{sessionToken, serviceID, hostPattern, rejectedHash, nil}) + if r.err != nil { + return credential{}, "", r.err + } + return r.issue, r.status, nil +} + +func (r *oauthResolver) reportMissingScopes(sessionToken, serviceID, hostPattern, rejectedHash string, requiredScopes []string) error { + if r.gate != nil { + <-r.gate + } + r.mu.Lock() + defer r.mu.Unlock() + r.calls = append(r.calls, refreshCall{sessionToken, serviceID, hostPattern, rejectedHash, requiredScopes}) + return r.err +} + +func (r *oauthResolver) refreshCalls() []refreshCall { + r.mu.Lock() + defer r.mu.Unlock() + return append([]refreshCall(nil), r.calls...) +} + +func oauthToken(value string) credential { + return credential{kind: credentialBearer, headerPrefix: "Bearer", value: []byte(value), refreshable: true} +} + +func tokenHash(value string) string { + sum := sha256.Sum256([]byte(value)) + return hex.EncodeToString(sum[:]) +} + +type upstreamHit struct { + auth string + body string +} + +type authUpstream struct { + mu sync.Mutex + hits []upstreamHit +} + +func (u *authUpstream) record(r *http.Request) { + body, _ := io.ReadAll(r.Body) + u.mu.Lock() + defer u.mu.Unlock() + u.hits = append(u.hits, upstreamHit{auth: r.Header.Get("Authorization"), body: string(body)}) +} + +func (u *authUpstream) recorded() []upstreamHit { + u.mu.Lock() + defer u.mu.Unlock() + return append([]upstreamHit(nil), u.hits...) +} + +// newPolicyFixture with an upstream the test answers for and a resolver that can refresh. +func newOAuthFixture(t *testing.T, resolver *oauthResolver, build func(host string) *resolvedService, handler func(w http.ResponseWriter, r *http.Request)) (*http.Client, string, *authUpstream) { + t.Helper() + + up := &authUpstream{} + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + up.record(r) + handler(w, r) + })) + t.Cleanup(upstream.Close) + + upstreamURL, err := url.Parse(upstream.URL) + if err != nil { + t.Fatal(err) + } + host := "127.0.0.1:" + upstreamURL.Port() + + upstreamPool := x509.NewCertPool() + upstreamPool.AddCert(upstream.Certificate()) + + key, cert, err := generateRootCa() + if err != nil { + t.Fatal(err) + } + + transport := newUpstreamTransport() + transport.TLSClientConfig = &tls.Config{RootCAs: upstreamPool} + + resolver.services = []*resolvedService{build(host)} + ps := &proxyServer{transport: transport, ca: newCaManager(key, cert)} + ps.setConfig(ProxyConfig{TrafficPolicy: TrafficPolicyAnyHost}) + ps.cache = newSessionCache(resolver, ps.pollInterval) + + front := httptest.NewServer(http.HandlerFunc(ps.dispatch)) + t.Cleanup(front.Close) + + clientPool := x509.NewCertPool() + clientPool.AddCert(cert) + proxyURL, _ := url.Parse(front.URL) + proxyURL.User = url.UserPassword(ProxyAuthUsername, "agv_tok") + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: clientPool}, + }, + } + return client, host, up +} + +func oauthMCPService(host string, token string) *resolvedService { + svc := mcpService(host) + svc.id = "svc-1" + svc.hostPattern = host + svc.credential = oauthToken(token) + return svc +} + +// Accepts only the token named, and answers anything else with the challenge an MCP server sends. +func acceptOnly(token string) func(w http.ResponseWriter, r *http.Request) { + return func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer "+token { + w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token", resource_metadata="https://mcp.example/.well-known/oauth-protected-resource"`) + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"invalid_token"}`)) + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + } +} + +func doWith(t *testing.T, c *http.Client, method, target, contentType string, body io.Reader) (*http.Response, string) { + t.Helper() + req, err := http.NewRequest(method, target, body) + if err != nil { + t.Fatal(err) + } + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } + resp, err := c.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + payload, _ := io.ReadAll(resp.Body) + return resp, string(payload) +} + +func assertUpstreamAuthFailed(t *testing.T, resp *http.Response, body string) { + t.Helper() + if resp.StatusCode != http.StatusBadGateway { + t.Fatalf("status = %d, want 502: %s", resp.StatusCode, body) + } + if got := resp.Header.Values("WWW-Authenticate"); len(got) != 0 { + t.Fatalf("the upstream's challenge reached the agent: %v", got) + } + if body != upstreamAuthFailedBody { + t.Fatalf("body = %s", body) + } + if ct := resp.Header.Get("Content-Type"); ct != "application/json" { + t.Fatalf("content type = %q", ct) + } +} + +const mcpCallBody = `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"create_issue","arguments":{}}}` + +func TestAn401RefreshesTheTokenAndRetriesTheSameRequestOnce(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token"), status: oauthStatusConnected} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, acceptOnly("new-token")) + + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/mcp", host), "application/json", strings.NewReader(mcpCallBody)) + if resp.StatusCode != http.StatusOK || body != "ok" { + t.Fatalf("the retried request should have succeeded, got %d: %s", resp.StatusCode, body) + } + + hits := up.recorded() + if len(hits) != 2 { + t.Fatalf("upstream saw %d requests, want the original and one retry", len(hits)) + } + if hits[0].auth != "Bearer old-token" || hits[1].auth != "Bearer new-token" { + t.Fatalf("tokens sent = %q, %q", hits[0].auth, hits[1].auth) + } + if hits[1].body != mcpCallBody { + t.Fatalf("the retry carried a different body: %q", hits[1].body) + } + + calls := resolver.refreshCalls() + if len(calls) != 1 { + t.Fatalf("refresh called %d times, want 1", len(calls)) + } + if calls[0].serviceID != "svc-1" || calls[0].hostPattern != host || calls[0].sessionToken != "agv_tok" || calls[0].rejectedHash != tokenHash("old-token") || calls[0].requiredScopes != nil { + t.Fatalf("refresh call = %+v", calls[0]) + } + + resp, _ = doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + if resp.StatusCode != http.StatusOK { + t.Fatalf("a later request should inject the refreshed token, got %d", resp.StatusCode) + } + if hits := up.recorded(); hits[len(hits)-1].auth != "Bearer new-token" || len(hits) != 3 { + t.Fatalf("the refreshed token was not swapped into the cache: %+v", hits) + } + if got := len(resolver.refreshCalls()); got != 1 { + t.Fatalf("refresh called %d times after the swap, want 1", got) + } +} + +func TestAn401TwiceGivesA502WithoutTheChallenge(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token")} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, acceptOnly("nothing")) + + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/mcp", host), "application/json", strings.NewReader(mcpCallBody)) + assertUpstreamAuthFailed(t, resp, body) + if got := len(up.recorded()); got != 2 { + t.Fatalf("upstream saw %d requests, want exactly one retry", got) + } + if got := len(resolver.refreshCalls()); got != 1 { + t.Fatalf("refresh called %d times, want 1", got) + } +} + +func TestAnUnchangedTokenIsNotRetried(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("old-token"), status: "needs-reauthorization"} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, acceptOnly("nothing")) + + resp, body := doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + assertUpstreamAuthFailed(t, resp, body) + if got := len(up.recorded()); got != 1 { + t.Fatalf("upstream saw %d requests; the same token must not be sent again", got) + } +} + +func TestANonReplayable401StillRefreshesButIsNotRetried(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token")} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, acceptOnly("new-token")) + + // Not JSON, so the MCP inspection never buffers it, and with no length so it streams. + stream := io.MultiReader(strings.NewReader("chunk-1"), strings.NewReader("chunk-2")) + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/upload", host), "application/octet-stream", stream) + assertUpstreamAuthFailed(t, resp, body) + if got := len(up.recorded()); got != 1 { + t.Fatalf("upstream saw %d requests; a streamed body cannot be replayed", got) + } + if got := len(resolver.refreshCalls()); got != 1 { + t.Fatalf("refresh called %d times, want 1", got) + } + + resp, _ = doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + if resp.StatusCode != http.StatusOK { + t.Fatalf("the refreshed token should still be swapped in, got %d", resp.StatusCode) + } +} + +func TestARefreshFailureFallsThroughToA502(t *testing.T) { + for name, err := range map[string]error{ + "not found": &api.APIError{StatusCode: http.StatusNotFound, Name: "NotFound"}, + "host moved": &api.APIError{StatusCode: http.StatusConflict, Name: "Conflict"}, + "server": &api.APIError{StatusCode: http.StatusInternalServerError}, + "unreachable": errors.New("dial tcp: connection refused"), + } { + t.Run(name, func(t *testing.T) { + resolver := &oauthResolver{err: err} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, acceptOnly("new-token")) + + resp, body := doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + assertUpstreamAuthFailed(t, resp, body) + if got := len(up.recorded()); got != 1 { + t.Fatalf("upstream saw %d requests, want no retry", got) + } + if got := len(resolver.refreshCalls()); got != 1 { + t.Fatalf("refresh called %d times, want 1", got) + } + + doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + if hits := up.recorded(); hits[len(hits)-1].auth != "Bearer old-token" { + t.Fatalf("a failed refresh changed the cached token: %+v", hits) + } + }) + } +} + +func TestInsufficientScopeIsReportedAndNotRetried(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("old-token"), status: "needs-additional-scopes"} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, + func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("WWW-Authenticate", `Bearer error="insufficient_scope", scope="issues:write comments:write", error_description="needs \"write\""`) + w.WriteHeader(http.StatusForbidden) + }) + + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/mcp", host), "application/json", strings.NewReader(mcpCallBody)) + assertUpstreamAuthFailed(t, resp, body) + if got := len(up.recorded()); got != 1 { + t.Fatalf("upstream saw %d requests, want no retry", got) + } + calls := resolver.refreshCalls() + if len(calls) != 1 { + t.Fatalf("refresh called %d times, want 1", len(calls)) + } + if !reflect.DeepEqual(calls[0].requiredScopes, []string{"issues:write", "comments:write"}) || calls[0].rejectedHash != tokenHash("old-token") || calls[0].hostPattern != host { + t.Fatalf("scope report = %+v", calls[0]) + } +} + +func TestAPlain403PassesThrough(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token")} + client, host, _ := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, + func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte("not your project")) + }) + + resp, body := doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + if resp.StatusCode != http.StatusForbidden || body != "not your project" { + t.Fatalf("got %d: %s", resp.StatusCode, body) + } + if got := len(resolver.refreshCalls()); got != 0 { + t.Fatalf("refresh called %d times for a 403 that is not about scopes", got) + } +} + +func TestANonOAuthCredentialKeepsTheUpstreams401(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token")} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { + svc := oauthMCPService(h, "static-token") + svc.credential.refreshable = false + return svc + }, acceptOnly("new-token")) + + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/mcp", host), "application/json", strings.NewReader(mcpCallBody)) + if resp.StatusCode != http.StatusUnauthorized || body != `{"error":"invalid_token"}` { + t.Fatalf("got %d: %s", resp.StatusCode, body) + } + if resp.Header.Get("WWW-Authenticate") == "" { + t.Fatal("a static credential's 401 lost its challenge") + } + if got := len(resolver.refreshCalls()); got != 0 || len(up.recorded()) != 1 { + t.Fatalf("a static credential was refreshed (%d) or retried (%d)", got, len(up.recorded())) + } +} + +func TestConcurrent401sShareOneRefresh(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token"), gate: make(chan struct{})} + var rejected atomic.Int32 + const agents = 8 + allRejected := make(chan struct{}) + client, host, _ := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, + func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") == "Bearer new-token" { + w.WriteHeader(http.StatusOK) + return + } + if rejected.Add(1) == agents { + close(allRejected) + } + w.WriteHeader(http.StatusUnauthorized) + }) + + go func() { + <-allRejected + close(resolver.gate) + }() + + var wg sync.WaitGroup + statuses := make(chan int, agents) + for i := 0; i < agents; i++ { + wg.Add(1) + go func() { + defer wg.Done() + req, _ := http.NewRequest(http.MethodGet, fmt.Sprintf("https://%s/mcp", host), nil) + resp, err := client.Do(req) + if err != nil { + t.Error(err) + return + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + statuses <- resp.StatusCode + }() + } + wg.Wait() + close(statuses) + for status := range statuses { + if status != http.StatusOK { + t.Fatalf("an agent got %d, want every retry to succeed", status) + } + } + if got := len(resolver.refreshCalls()); got != 1 { + t.Fatalf("%d concurrent 401s made %d refresh calls, want 1", agents, got) + } +} + +func TestBearerChallengeParams(t *testing.T) { + cases := []struct { + name string + values []string + want map[string]string + }{ + {"insufficient scope", []string{`Bearer error="insufficient_scope", scope="a b"`}, map[string]string{"error": "insufficient_scope", "scope": "a b"}}, + {"token form and case", []string{`bearer Error=insufficient_scope,Scope=read`}, map[string]string{"error": "insufficient_scope", "scope": "read"}}, + {"escaped quote", []string{`Bearer error_description="say \"hi\", ok", error="invalid_token"`}, map[string]string{"error_description": `say "hi", ok`, "error": "invalid_token"}}, + {"other scheme first", []string{`Basic realm="x", Bearer error="insufficient_scope"`}, map[string]string{"error": "insufficient_scope"}}, + {"other scheme only", []string{`Basic realm="x" error="insufficient_scope"`}, map[string]string{}}, + {"second header", []string{`Basic realm="x"`, `Bearer scope="s"`}, map[string]string{"scope": "s"}}, + {"unterminated", []string{`Bearer error="insuff`}, map[string]string{"error": "insuff"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := bearerChallengeParams(tc.values); !reflect.DeepEqual(got, tc.want) { + t.Fatalf("got %v, want %v", got, tc.want) + } + }) + } +} + +func TestChallengedScopesStayWithinTheBackendsLimits(t *testing.T) { + raw := strings.Repeat("s ", maxReportedScopes+5) + strings.Repeat("x", maxReportedScopeLen+1) + if got := challengedScopes(raw); len(got) != maxReportedScopes { + t.Fatalf("got %d scopes, want %d", len(got), maxReportedScopes) + } + if got := challengedScopes("ok " + strings.Repeat("x", maxReportedScopeLen+1)); !reflect.DeepEqual(got, []string{"ok"}) { + t.Fatalf("got %v", got) + } +} + +func TestResolveTreatsAnUnknownTypeAsHTTPAndMarksOAuthRefreshable(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "sessionId": "s1", + "services": []any{ + map[string]any{"id": "a", "name": "future", "type": "websocket", "hostPattern": "a.example", + "credential": map[string]any{"type": "bearer", "value": "t"}}, + map[string]any{"id": "b", "name": "linear", "type": "mcp", "hostPattern": "b.example", + "credential": map[string]any{"type": "bearer", "headerName": "Authorization", "headerPrefix": "Bearer", + "value": "t", "source": "oauth", "expiresAt": "2026-10-15T12:00:00Z"}}, + map[string]any{"id": "c", "name": "other", "type": "mcp", "hostPattern": "c.example", + "credential": map[string]any{"type": "bearer", "value": "t", "source": "something-else"}}, + }, + "sessionLogs": map[string]any{"enabled": false}, + }) + })) + defer srv.Close() + old := config.INFISICAL_URL + config.INFISICAL_URL = srv.URL + "/api" + defer func() { config.INFISICAL_URL = old }() + + resolver, err := newInfisicalResolver(func() string { return "ptok" }) + if err != nil { + t.Fatal(err) + } + result, err := resolver.resolve("tok", nil) + if err != nil { + t.Fatal(err) + } + if len(result.Services) != 3 { + t.Fatalf("got %d services", len(result.Services)) + } + future, linear, other := result.Services[0], result.Services[1], result.Services[2] + if future.kind != serviceKindHTTP || future.credential.kind != credentialBearer || future.credential.refreshable { + t.Fatalf("unknown type = %q, credential %+v", future.kind, future.credential) + } + if linear.kind != serviceKindMCP || !linear.credential.refreshable { + t.Fatalf("oauth credential = %+v", linear.credential) + } + if other.credential.refreshable { + t.Fatal("only source oauth is refreshable") + } +} + +func TestRefreshCredentialCallsTheContractEndpoint(t *testing.T) { + var gotPath, gotSession, gotAuth string + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath, gotSession, gotAuth = r.URL.Path, r.Header.Get(api.AgentVaultSessionHeader), r.Header.Get("Authorization") + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"credential":{"type":"bearer","headerName":"Authorization","headerPrefix":"Bearer","value":"fresh","source":"oauth","expiresAt":null},"status":"connected"}`)) + })) + defer srv.Close() + old := config.INFISICAL_URL + config.INFISICAL_URL = srv.URL + "/api" + defer func() { config.INFISICAL_URL = old }() + + resolver, err := newInfisicalResolver(func() string { return "ptok" }) + if err != nil { + t.Fatal(err) + } + cred, status, err := resolver.refreshCredential("tok", "svc-1", "mcp.example", tokenHash("stale")) + if err != nil { + t.Fatal(err) + } + if gotPath != "/api/v1/agent-vault/proxy/services/svc-1/credential-refresh" || gotSession != "tok" || gotAuth != "Bearer ptok" { + t.Fatalf("path %q session %q auth %q", gotPath, gotSession, gotAuth) + } + want := map[string]any{"hostPattern": "mcp.example", "rejectedTokenHash": tokenHash("stale")} + if !reflect.DeepEqual(gotBody, want) { + t.Fatalf("body = %v", gotBody) + } + if string(cred.value) != "fresh" || !cred.refreshable || status != "connected" { + t.Fatalf("credential %+v status %q", cred, status) + } +} + +func TestReportMissingScopesCallsTheContractEndpoint(t *testing.T) { + var gotPath, gotSession, gotAuth string + var gotBody map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath, gotSession, gotAuth = r.URL.Path, r.Header.Get(api.AgentVaultSessionHeader), r.Header.Get("Authorization") + _ = json.NewDecoder(r.Body).Decode(&gotBody) + w.WriteHeader(http.StatusNoContent) + })) + defer srv.Close() + old := config.INFISICAL_URL + config.INFISICAL_URL = srv.URL + "/api" + defer func() { config.INFISICAL_URL = old }() + + resolver, err := newInfisicalResolver(func() string { return "ptok" }) + if err != nil { + t.Fatal(err) + } + if err := resolver.reportMissingScopes("tok", "svc-1", "mcp.example", tokenHash("stale"), []string{"read"}); err != nil { + t.Fatal(err) + } + if gotPath != "/api/v1/agent-vault/proxy/services/svc-1/missing-scopes" || gotSession != "tok" || gotAuth != "Bearer ptok" { + t.Fatalf("path %q session %q auth %q", gotPath, gotSession, gotAuth) + } + want := map[string]any{"hostPattern": "mcp.example", "rejectedTokenHash": tokenHash("stale"), "requiredScopes": []any{"read"}} + if !reflect.DeepEqual(gotBody, want) { + t.Fatalf("body = %v", gotBody) + } +} + +// Serves resolves from a token the test sets, and can hold one resolve open so a swap lands while it runs. +type pollRaceResolver struct { + mu sync.Mutex + polled string + holdOne chan struct{} + entered chan struct{} + issue credential +} + +func (r *pollRaceResolver) resolve(string, *sessionLogGrant) (*resolveResult, error) { + r.mu.Lock() + value, hold := r.polled, r.holdOne + r.holdOne = nil + r.mu.Unlock() + if hold != nil { + close(r.entered) + <-hold + } + svc := &resolvedService{id: "svc-1", name: "linear", kind: serviceKindMCP, hostPattern: "mcp.example", credential: oauthToken(value)} + return &resolveResult{SessionID: "s1", Services: []*resolvedService{svc}}, nil +} + +func (r *pollRaceResolver) refreshCredential(string, string, string, string) (credential, string, error) { + return r.issue, oauthStatusConnected, nil +} + +func (r *pollRaceResolver) reportMissingScopes(string, string, string, string, []string) error { + return nil +} + +func (r *pollRaceResolver) setPolled(value string) { + r.mu.Lock() + defer r.mu.Unlock() + r.polled = value +} + +func injectedToken(t *testing.T, cache *sessionCache) string { + t.Helper() + services, err := cache.get("tok") + if err != nil { + t.Fatal(err) + } + return string(services[0].credential.value) +} + +func TestAPollThatStartedBeforeASwapDoesNotRestoreTheReplacedToken(t *testing.T) { + resolver := &pollRaceResolver{polled: "old-token", issue: oauthToken("new-token")} + cache := newTestCache(resolver) + if got := injectedToken(t, cache); got != "old-token" { + t.Fatalf("seeded %q", got) + } + + release := make(chan struct{}) + resolver.mu.Lock() + resolver.holdOne, resolver.entered = release, make(chan struct{}) + resolver.mu.Unlock() + pollDone := make(chan struct{}) + go func() { + defer close(pollDone) + cache.refresh() + }() + <-resolver.entered + + old := oauthToken("old-token") + if _, err := cache.refreshCredential("tok", "svc-1", "mcp.example", &old); err != nil { + t.Fatal(err) + } + close(release) + <-pollDone + + if got := injectedToken(t, cache); got != "new-token" { + t.Fatalf("after a stale poll the proxy injects %q, want new-token", got) + } + + resolver.setPolled("newer-token") + cache.refresh() + if got := injectedToken(t, cache); got != "newer-token" { + t.Fatalf("a later poll's token was not taken, got %q", got) + } +} + +type scopeReportRefresher struct { + oauthResolver + entered chan []string +} + +func (r *scopeReportRefresher) refreshCredential(sessionToken, serviceID, hostPattern, rejectedHash string) (credential, string, error) { + r.entered <- nil + return r.oauthResolver.refreshCredential(sessionToken, serviceID, hostPattern, rejectedHash) +} + +func (r *scopeReportRefresher) reportMissingScopes(sessionToken, serviceID, hostPattern, rejectedHash string, requiredScopes []string) error { + r.entered <- requiredScopes + return r.oauthResolver.reportMissingScopes(sessionToken, serviceID, hostPattern, rejectedHash, requiredScopes) +} + +func TestConcurrentScopeReportsEachReachInfisicalWhenTheirTokenOrScopesDiffer(t *testing.T) { + cases := []struct { + name string + token string + scopes []string + }{ + {"different scopes", "tok-value", []string{"write"}}, + {"different tokens", "new-token", []string{"read"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resolver := &scopeReportRefresher{ + oauthResolver: oauthResolver{issue: oauthToken("tok-value"), gate: make(chan struct{})}, + entered: make(chan []string, 2), + } + resolver.setServices(oauthMCPService("mcp.example", "tok-value")) + cache := newTestCache(resolver) + if _, err := cache.get("tok"); err != nil { + t.Fatal(err) + } + sent := oauthToken("tok-value") + + var wg sync.WaitGroup + report := func(token string, scopes []string) { + wg.Add(1) + go func() { + defer wg.Done() + sent := oauthToken(token) + if err := cache.reportMissingScopes("tok", "svc-1", "mcp.example", &sent, scopes); err != nil { + t.Error(err) + } + }() + } + report(string(sent.value), []string{"read"}) + <-resolver.entered + report(tc.token, tc.scopes) + select { + case <-resolver.entered: + case <-time.After(2 * time.Second): + close(resolver.gate) + wg.Wait() + t.Fatal("the second report was merged into the first one's call") + } + close(resolver.gate) + wg.Wait() + + reported := map[string]bool{} + for _, call := range resolver.refreshCalls() { + reported[call.rejectedHash+" "+strings.Join(call.requiredScopes, " ")] = true + } + if !reported[tokenHash("tok-value")+" read"] || !reported[tokenHash(tc.token)+" "+strings.Join(tc.scopes, " ")] || len(reported) != 2 { + t.Fatalf("Infisical received scope reports %v, want both token/scope pairs", reported) + } + }) + } +} + +func TestCredentialFlightKeyIsCanonicalPerScopeSet(t *testing.T) { + a := credentialFlightKey("k", "svc", "a.example", "hash", []string{"write", "read", "read"}) + b := credentialFlightKey("k", "svc", "a.example", "hash", []string{"read", "write"}) + if a != b { + t.Fatalf("same scope set gave different keys: %q, %q", a, b) + } + if a == credentialFlightKey("k", "svc", "a.example", "hash", []string{"read"}) { + t.Fatal("different scope sets share a key") + } + if credentialFlightKey("k", "svc", "a.example", "hash", nil) == credentialFlightKey("k", "svc", "a.example", "hash", []string{""}) { + t.Fatal("a 401 refresh shares a key with a scope report") + } + if credentialFlightKey("k", "svc", "a.example", "hash", nil) == credentialFlightKey("k", "svc", "b.example", "hash", nil) { + t.Fatal("refreshes for different destinations share a key") + } +} + +func TestAPlain403AfterARefreshPassesThrough(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token")} + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { return oauthMCPService(h, "old-token") }, + func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Authorization") != "Bearer new-token" { + acceptOnly("new-token")(w, r) + return + } + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte("not your project")) + }) + + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/mcp", host), "application/json", strings.NewReader(mcpCallBody)) + if resp.StatusCode != http.StatusForbidden || body != "not your project" { + t.Fatalf("got %d: %s", resp.StatusCode, body) + } + if got := len(up.recorded()); got != 2 { + t.Fatalf("upstream saw %d requests, want the original and one retry", got) + } + if calls := resolver.refreshCalls(); len(calls) != 1 || calls[0].requiredScopes != nil { + t.Fatalf("refresh calls = %+v, want only the 401's refresh", calls) + } +} + +func TestARefreshedTokenIsRefusedWhenAPollRemovedTheAccessMeanwhile(t *testing.T) { + cases := []struct { + name string + remove func(cache *sessionCache, resolver *oauthResolver) + want error + }{ + {"session removed", func(cache *sessionCache, _ *oauthResolver) { + cache.mu.Lock() + delete(cache.entries, sessionKey("tok")) + cache.mu.Unlock() + }, errSessionGone}, + {"service removed", func(cache *sessionCache, resolver *oauthResolver) { + resolver.setServices() + cache.refresh() + }, errServiceRemoved}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token"), gate: make(chan struct{})} + resolver.setServices(oauthMCPService("mcp.example", "old-token")) + cache := newTestCache(resolver) + injectedToken(t, cache) + + sent := oauthToken("old-token") + done := make(chan error, 1) + go func() { + _, err := cache.refreshCredential("tok", "svc-1", "mcp.example", &sent) + done <- err + }() + + tc.remove(cache, resolver) + close(resolver.gate) + + if err := <-done; !errors.Is(err, tc.want) { + t.Fatalf("err = %v, want %v", err, tc.want) + } + }) + } +} + +func TestARefreshedTokenIsDroppedWhenAPollMovedTheServiceMeanwhile(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token"), gate: make(chan struct{})} + resolver.setServices(&resolvedService{id: "svc-1", kind: serviceKindMCP, hostPattern: "a.example", credential: oauthToken("a-token")}) + cache := newTestCache(resolver) + if got := injectedToken(t, cache); got != "a-token" { + t.Fatalf("seeded %q", got) + } + + sent := oauthToken("a-token") + done := make(chan error, 1) + go func() { + _, err := cache.refreshCredential("tok", "svc-1", "a.example", &sent) + done <- err + }() + + resolver.setServices(&resolvedService{id: "svc-1", kind: serviceKindMCP, hostPattern: "b.example", credential: oauthToken("b-token")}) + cache.refresh() + close(resolver.gate) + + if err := <-done; !errors.Is(err, errServiceHostChanged) { + t.Fatalf("err = %v, want errServiceHostChanged", err) + } + if got := injectedToken(t, cache); got != "b-token" { + t.Fatalf("the proxy injects %q, want the polled b-token", got) + } + if calls := resolver.refreshCalls(); len(calls) != 1 || calls[0].hostPattern != "a.example" { + t.Fatalf("refresh calls = %+v", calls) + } +} + +func TestA409FromCredentialRefreshLeavesTheCacheAlone(t *testing.T) { + var refreshes atomic.Int32 + var gotHost atomic.Value + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if strings.HasSuffix(r.URL.Path, "/credential-refresh") { + refreshes.Add(1) + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + gotHost.Store(body["hostPattern"]) + w.WriteHeader(http.StatusConflict) + _, _ = w.Write([]byte(`{"statusCode":409,"error":"Conflict","message":"The service changed; wait for the next poll."}`)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "sessionId": "s1", + "services": []any{map[string]any{"id": "svc-1", "name": "linear", "type": "mcp", "hostPattern": "a.example", + "credential": map[string]any{"type": "bearer", "headerName": "Authorization", "headerPrefix": "Bearer", "value": "a-token", "source": "oauth"}}}, + "sessionLogs": map[string]any{"enabled": false}, + }) + })) + defer srv.Close() + old := config.INFISICAL_URL + config.INFISICAL_URL = srv.URL + "/api" + defer func() { config.INFISICAL_URL = old }() + + resolver, err := newInfisicalResolver(func() string { return "ptok" }) + if err != nil { + t.Fatal(err) + } + cache := newTestCache(resolver) + if got := injectedToken(t, cache); got != "a-token" { + t.Fatalf("seeded %q", got) + } + + sent := oauthToken("a-token") + _, err = cache.refreshCredential("tok", "svc-1", "a.example", &sent) + var apiErr *api.APIError + if !errors.As(err, &apiErr) || apiErr.StatusCode != http.StatusConflict { + t.Fatalf("err = %v, want the 409", err) + } + if got := refreshes.Load(); got != 1 { + t.Fatalf("refresh called %d times, want 1", got) + } + if got := gotHost.Load(); got != "a.example" { + t.Fatalf("hostPattern sent = %v", got) + } + if got := injectedToken(t, cache); got != "a-token" { + t.Fatalf("a 409 changed the cached token to %q", got) + } +} + +func TestALateCredentialResponseKeepsThePolledCredential(t *testing.T) { + cases := []struct { + name string + issue credential + polled credential + err error + }{ + {"reauthorized", oauthToken("refreshed-old-token"), oauthToken("reauthorized-token"), nil}, + {"static bearer", oauthToken("refreshed-old-token"), credential{kind: credentialBearer, headerName: "X-Api-Key", value: []byte("old-token")}, errCredentialNotRefreshable}, + {"basic", oauthToken("refreshed-old-token"), credential{kind: credentialBasic, username: "user", password: []byte("password")}, errCredentialNotRefreshable}, + {"passthrough", oauthToken("refreshed-old-token"), credential{kind: credentialPassthrough}, errCredentialNotRefreshable}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resolver := &scopeReportRefresher{ + oauthResolver: oauthResolver{issue: tc.issue, gate: make(chan struct{})}, + entered: make(chan []string, 1), + } + defer func() { + select { + case <-resolver.gate: + default: + close(resolver.gate) + } + }() + resolver.setServices(oauthMCPService("mcp.example", "old-token")) + cache := newTestCache(resolver) + injectedToken(t, cache) + + type result struct { + cred credential + err error + } + done := make(chan result, 1) + sent := oauthToken("old-token") + go func() { + cred, err := cache.refreshCredential("tok", "svc-1", "mcp.example", &sent) + done <- result{cred, err} + }() + <-resolver.entered + + polled := oauthMCPService("mcp.example", "unused") + polled.credential = tc.polled + resolver.setServices(polled) + cache.refresh() + close(resolver.gate) + got := <-done + if !errors.Is(got.err, tc.err) { + t.Fatalf("refresh error = %v, want %v", got.err, tc.err) + } + if got.err == nil && !reflect.DeepEqual(got.cred, tc.polled) { + t.Fatalf("retry credential = %+v, want the newer credential %+v", got.cred, tc.polled) + } + services, err := cache.get("tok") + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(services[0].credential, tc.polled) { + t.Fatalf("late response replaced the polled credential: %+v", services[0].credential) + } + }) + } +} + +func TestALateScopeReportKeepsAConcurrentRefresh(t *testing.T) { + resolver := &scopeReportRefresher{ + oauthResolver: oauthResolver{issue: oauthToken("old-token"), gate: make(chan struct{})}, + entered: make(chan []string, 1), + } + resolver.setServices(oauthMCPService("mcp.example", "old-token")) + cache := newTestCache(resolver) + injectedToken(t, cache) + sent := oauthToken("old-token") + done := make(chan error, 1) + go func() { + err := cache.reportMissingScopes("tok", "svc-1", "mcp.example", &sent, []string{"write"}) + done <- err + }() + <-resolver.entered + _, swapErr := cache.swapCredential(sessionKey("tok"), "svc-1", "mcp.example", credentialHash(&sent), oauthToken("new-token")) + close(resolver.gate) + if err := <-done; err != nil { + t.Fatal(err) + } + if swapErr != nil { + t.Fatal(swapErr) + } + if got := injectedToken(t, cache); got != "new-token" { + t.Fatalf("late scope report restored %q, want new-token", got) + } +} + +func TestAPollReplacingOAuthWithStaticAuthKeepsItsCredential(t *testing.T) { + static := &resolvedService{ + id: "svc-1", + hostPattern: "mcp.example", + credential: credential{ + kind: credentialBearer, + headerName: "X-Api-Key", + value: []byte("old-token"), + }, + } + rejected := oauthToken("old-token") + entry := &sessionEntry{swaps: map[string]*credentialSwap{ + "svc-1": { + installed: oauthToken("new-token"), + hostPattern: "mcp.example", + replaced: map[string]bool{credentialHash(&rejected): true}, + gen: 1, + }, + }} + services := keepNewerSwaps(entry, []*resolvedService{static}, 0) + if services[0] != static { + t.Fatalf("a poll changing authentication was overwritten with %+v", services[0].credential) + } +} + +func TestOpenAuthErrorBodiesDoNotBlockRecovery(t *testing.T) { + cases := []struct { + name string + first int + retry int + want int + refreshes int + }{ + {"401 then success", http.StatusUnauthorized, http.StatusOK, http.StatusOK, 1}, + {"401 twice", http.StatusUnauthorized, http.StatusUnauthorized, http.StatusBadGateway, 1}, + {"insufficient scope", http.StatusForbidden, 0, http.StatusBadGateway, 1}, + {"401 then insufficient scope", http.StatusUnauthorized, http.StatusForbidden, http.StatusBadGateway, 2}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("new-token")} + var attempts atomic.Int32 + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { + return oauthMCPService(h, "old-token") + }, func(w http.ResponseWriter, r *http.Request) { + status := tc.first + if attempts.Add(1) > 1 { + status = tc.retry + } + if status == http.StatusOK { + _, _ = io.WriteString(w, "ok") + return + } + if status == http.StatusForbidden { + w.Header().Set("WWW-Authenticate", `Bearer error="insufficient_scope", scope="write"`) + } else { + w.Header().Set("WWW-Authenticate", `Bearer error="invalid_token"`) + } + w.WriteHeader(status) + w.(http.Flusher).Flush() + <-r.Context().Done() + }) + client.Timeout = 2 * time.Second + resp, body := doWith(t, client, http.MethodGet, fmt.Sprintf("https://%s/mcp", host), "", nil) + if resp.StatusCode != tc.want { + t.Fatalf("status = %d, want %d: %s", resp.StatusCode, tc.want, body) + } + if tc.want == http.StatusBadGateway { + assertUpstreamAuthFailed(t, resp, body) + } + if got := len(resolver.refreshCalls()); got != tc.refreshes { + t.Fatalf("refresh calls = %d, want %d", got, tc.refreshes) + } + wantAttempts := 2 + if tc.first == http.StatusForbidden { + wantAttempts = 1 + } + if got := len(up.recorded()); got != wantAttempts { + t.Fatalf("upstream attempts = %d, want %d", got, wantAttempts) + } + }) + } +} + +func TestHTTPAuthenticationKeepsUpstreamAuthResponses(t *testing.T) { + cases := []struct { + name string + cred credential + auth string + }{ + {"bearer", credential{kind: credentialBearer, headerPrefix: "Bearer", value: []byte("static-token")}, "Bearer static-token"}, + {"custom bearer header", credential{kind: credentialBearer, headerName: "X-Api-Key", value: []byte("static-token")}, ""}, + {"basic", credential{kind: credentialBasic, username: "user", password: []byte("pass")}, "Basic dXNlcjpwYXNz"}, + {"passthrough", credential{kind: credentialPassthrough}, ""}, + } + for _, tc := range cases { + for _, status := range []int{http.StatusUnauthorized, http.StatusForbidden} { + t.Run(fmt.Sprintf("%s/%d", tc.name, status), func(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("unexpected-refresh")} + const challenge = `Bearer error="insufficient_scope", scope="write"` + client, host, up := newOAuthFixture(t, resolver, func(h string) *resolvedService { + svc := oauthMCPService(h, "unused") + svc.kind = serviceKindHTTP + svc.credential = tc.cred + return svc + }, func(w http.ResponseWriter, r *http.Request) { + if tc.cred.headerName == "X-Api-Key" && r.Header.Get("X-Api-Key") != "static-token" { + t.Error("custom bearer header was not injected") + } + w.Header().Set("WWW-Authenticate", challenge) + w.WriteHeader(status) + _, _ = io.WriteString(w, "permission denied") + }) + resp, body := doWith(t, client, http.MethodPost, fmt.Sprintf("https://%s/api", host), "application/json", strings.NewReader(mcpCallBody)) + if resp.StatusCode != status || body != "permission denied" || resp.Header.Get("WWW-Authenticate") != challenge { + t.Fatalf("upstream response changed: status %d, challenge %q, body %q", resp.StatusCode, resp.Header.Get("WWW-Authenticate"), body) + } + if got := len(resolver.refreshCalls()); got != 0 { + t.Fatalf("static authentication triggered %d refreshes", got) + } + if hits := up.recorded(); len(hits) != 1 || hits[0].auth != tc.auth || hits[0].body != mcpCallBody { + t.Fatalf("upstream requests changed: %+v", hits) + } + }) + } + } +} + +func TestAnMCPEventStreamIsForwardedBeforeItCompletes(t *testing.T) { + resolver := &oauthResolver{issue: oauthToken("unexpected-refresh")} + release := make(chan struct{}) + defer func() { + select { + case <-release: + default: + close(release) + } + }() + client, host, _ := newOAuthFixture(t, resolver, func(h string) *resolvedService { + return oauthMCPService(h, "valid-token") + }, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Mcp-Session-Id", "upstream-session") + _, _ = io.WriteString(w, "event: message\ndata: first\n\n") + w.(http.Flusher).Flush() + select { + case <-release: + _, _ = io.WriteString(w, "data: last\n\n") + case <-r.Context().Done(): + } + }) + client.Timeout = 2 * time.Second + resp, err := client.Get(fmt.Sprintf("https://%s/mcp", host)) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK || resp.Header.Get("Content-Type") != "text/event-stream" || resp.Header.Get("Mcp-Session-Id") != "upstream-session" { + t.Fatalf("event-stream response changed: %d %v", resp.StatusCode, resp.Header) + } + reader := bufio.NewReader(resp.Body) + for _, want := range []string{"event: message\n", "data: first\n", "\n"} { + if got, err := reader.ReadString('\n'); err != nil || got != want { + t.Fatalf("streamed line = %q, err %v, want %q before the upstream finishes", got, err, want) + } + } + close(release) + if tail, err := io.ReadAll(reader); err != nil || string(tail) != "data: last\n\n" { + t.Fatalf("stream tail = %q, err %v", tail, err) + } + if got := len(resolver.refreshCalls()); got != 0 { + t.Fatalf("a successful event stream triggered %d refreshes", got) + } +} diff --git a/packages/api/agent_vault.go b/packages/api/agent_vault.go index f649d7ea..acf360d2 100644 --- a/packages/api/agent_vault.go +++ b/packages/api/agent_vault.go @@ -3,6 +3,7 @@ package api import ( "context" "fmt" + "net/url" "github.com/Infisical/infisical-merge/packages/config" "github.com/go-resty/resty/v2" @@ -76,6 +77,8 @@ type AgentVaultCredential struct { Value string `json:"value,omitempty"` Username string `json:"username,omitempty"` Password string `json:"password,omitempty"` + // "oauth" marks a token Infisical can refresh; every other credential leaves it empty. + Source string `json:"source,omitempty"` } type AgentVaultCustomHeader struct { @@ -93,6 +96,7 @@ type AgentVaultSubstitution struct { type AgentVaultService struct { ID string `json:"id"` Name string `json:"name"` + Type string `json:"type"` AccessBundleName string `json:"accessBundleName"` HostPattern string `json:"hostPattern"` // A nil slice means unrestricted, which is what JSON null decodes to. @@ -138,6 +142,58 @@ func CallResolveAgentVaultSession(httpClient *resty.Client, sessionToken string, return res, nil } +type RefreshAgentVaultServiceCredentialRequest struct { + HostPattern string `json:"hostPattern"` + RejectedTokenHash string `json:"rejectedTokenHash"` +} + +type RefreshAgentVaultServiceCredentialResponse struct { + Credential AgentVaultCredential `json:"credential"` + Status string `json:"status"` +} + +func CallRefreshAgentVaultServiceCredential(httpClient *resty.Client, sessionToken, serviceID string, request RefreshAgentVaultServiceCredentialRequest) (RefreshAgentVaultServiceCredentialResponse, error) { + var res RefreshAgentVaultServiceCredentialResponse + response, err := httpClient. + R(). + SetResult(&res). + SetHeader("User-Agent", USER_AGENT). + SetHeader(AgentVaultSessionHeader, sessionToken). + SetBody(request). + Post(fmt.Sprintf("%v/v1/agent-vault/proxy/services/%s/credential-refresh", config.INFISICAL_URL, url.PathEscape(serviceID))) + + if err != nil { + return RefreshAgentVaultServiceCredentialResponse{}, NewGenericRequestError("CallRefreshAgentVaultServiceCredential", err) + } + if response.IsError() { + return RefreshAgentVaultServiceCredentialResponse{}, NewAPIErrorWithResponse("CallRefreshAgentVaultServiceCredential", response, nil) + } + return res, nil +} + +type ReportAgentVaultServiceMissingScopesRequest struct { + HostPattern string `json:"hostPattern"` + RejectedTokenHash string `json:"rejectedTokenHash"` + RequiredScopes []string `json:"requiredScopes"` +} + +func CallReportAgentVaultServiceMissingScopes(httpClient *resty.Client, sessionToken, serviceID string, request ReportAgentVaultServiceMissingScopesRequest) error { + response, err := httpClient. + R(). + SetHeader("User-Agent", USER_AGENT). + SetHeader(AgentVaultSessionHeader, sessionToken). + SetBody(request). + Post(fmt.Sprintf("%v/v1/agent-vault/proxy/services/%s/missing-scopes", config.INFISICAL_URL, url.PathEscape(serviceID))) + + if err != nil { + return NewGenericRequestError("CallReportAgentVaultServiceMissingScopes", err) + } + if response.IsError() { + return NewAPIErrorWithResponse("CallReportAgentVaultServiceMissingScopes", response, nil) + } + return nil +} + type CreateAgentVaultSessionLogChunkRequest struct { ChunkID string `json:"chunkId"` EndedAt string `json:"endedAt"`