diff --git a/internal/runtime/executor/antigravity_executor.go b/internal/runtime/executor/antigravity_executor.go index cb7b000f..945005c1 100644 --- a/internal/runtime/executor/antigravity_executor.go +++ b/internal/runtime/executor/antigravity_executor.go @@ -6,12 +6,13 @@ package executor import ( "bytes" "context" + "crypto/sha256" "crypto/tls" + "encoding/hex" "encoding/json" "fmt" "net/http" "strings" - "sync" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/cache" @@ -21,6 +22,7 @@ import ( antigravityclaude "github.com/router-for-me/CLIProxyAPI/v7/internal/translator/antigravity/claude" "github.com/router-for-me/CLIProxyAPI/v7/internal/util" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" + "github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" log "github.com/sirupsen/logrus" "github.com/tidwall/gjson" @@ -72,13 +74,58 @@ func (e *AntigravityExecutor) obfuscateSensitiveWords(payload []byte) []byte { // Each Antigravity credential gets its own HTTP/1.1 connection pool. Sessions routed // to the same auth reuse that pool, while different OAuth identities never share a // TCP/TLS connection, matching the native client's one-credential process model. +// The cache is bounded so pools cannot accumulate when keys churn. var ( antigravityBaseTransport = defaultAntigravityBaseTransport() - antigravityTransports sync.Map // antigravityTransportKey -> *http.Transport + antigravityTransports = helps.NewTransportCache[antigravityTransportKey](antigravityTransportCacheCapacity) ) +const ( + // antigravityTransportCacheCapacity caps how many Antigravity connection pools stay + // alive. The bound exists only to stop entries from accumulating when keys churn, for + // example when a credential's proxy is rotated through the management API or when an + // SDK embedder supplies a freshly built base transport per request. + // + // It is sized for large deployments on purpose. An unused cache entry costs under 1 KB + // and no goroutines, so capacity is close to free, whereas evicting a pool that is + // still in active use forces the next request on that credential to redo the TCP + TLS + // handshake and defeats the point of caching. Credential counts in the low thousands + // are expected once Home-managed pools are included. + // + // Capacity is therefore NOT the lever for bounding memory: an idle pooled connection + // costs roughly 38 KB plus three goroutines, and that total is driven by live traffic + // and reclaimed by IdleConnTimeout. Shrinking this number does not save that memory, + // it only causes pool thrashing. + antigravityTransportCacheCapacity = 8192 + + // antigravityMaxIdleConnsPerHost mirrors the value that + // cloud.google.com/go/auth/httptransport and google.golang.org/api/transport/http + // set on their base transport, which is the stack the native Antigravity client + // uses. Both raise Go's DefaultMaxIdleConnsPerHost of 2 to 100 because the low + // default forces concurrent requests to re-handshake instead of reusing pooled + // connections. + antigravityMaxIdleConnsPerHost = 100 + + // antigravityIdleConnTimeout keeps pooled connections usable far longer than Go's + // 90s default. Captured native traffic reuses a connection after idle gaps with a + // p90 of roughly six minutes, and a 90s timeout would discard about an eighth of + // the reuses the native client actually performs. + antigravityIdleConnTimeout = 10 * time.Minute + + // antigravityAnonymousTransportScope is the pool scope for auth objects that carry + // no identity at all. Reaching it means the auth has no ID, no source path and no + // token of any kind, so there is no credential to keep isolated and a single shared + // pool is safe. Allocating a private pool per request instead would leak a + // connection pool, and the goroutines managing it, on every call. + antigravityAnonymousTransportScope = "anonymous" +) + +// antigravityTransportKey identifies one connection pool. At most one of proxy and +// base is set: proxy for a credential-scoped proxy pool, base for a transport handed +// in through the request context, and neither for a direct pool. type antigravityTransportKey struct { credential string + proxy string base *http.Transport } @@ -106,54 +153,134 @@ func cloneTransportWithHTTP11(base *http.Transport) *http.Transport { // Native Antigravity sends no ALPN extension. With HTTP/2 disabled above, // an empty NextProtos keeps the wire shape aligned while using HTTP/1.1. clone.TLSClientConfig.NextProtos = nil + applyAntigravityPoolLimits(clone) return clone } -// antigravityHTTP11Transport returns the HTTP/1.1 pool for one credential and -// base transport. The base is either the process default, a credential-scoped -// proxy transport, or a context-provided transport. +// applyAntigravityPoolLimits widens the connection pool so keep-alive actually +// survives concurrency and idle periods. Limits are only ever raised, so an +// operator-supplied base transport with a larger pool keeps its own settings. +func applyAntigravityPoolLimits(transport *http.Transport) { + if transport == nil { + return + } + // Go treats 0 as DefaultMaxIdleConnsPerHost (2) and a negative value as "never pool + // an idle connection". Raise the default and smaller positive values, but leave a + // negative value alone so an operator can still disable pooling outright. + if transport.MaxIdleConnsPerHost >= 0 && transport.MaxIdleConnsPerHost < antigravityMaxIdleConnsPerHost { + transport.MaxIdleConnsPerHost = antigravityMaxIdleConnsPerHost + } + // MaxIdleConns caps the pool across all hosts. Leaving it below the per-host limit + // would silently throttle Antigravity, which talks to a single host at a time. + // Zero means unlimited, so it must not be lowered. + if transport.MaxIdleConns > 0 && transport.MaxIdleConns < transport.MaxIdleConnsPerHost { + transport.MaxIdleConns = transport.MaxIdleConnsPerHost + } + // Zero already means "never expire idle connections", which is strictly longer. + if transport.IdleConnTimeout > 0 && transport.IdleConnTimeout < antigravityIdleConnTimeout { + transport.IdleConnTimeout = antigravityIdleConnTimeout + } +} + +// antigravityHTTP11Transport returns the HTTP/1.1 pool shared by every request that +// uses the same credential and the same base transport. The base is either the +// process default or a transport provided through the request context. func antigravityHTTP11Transport(auth *cliproxyauth.Auth, base *http.Transport) *http.Transport { if base == nil { return nil } - credential, shareable := antigravityTransportScope(auth) - if !shareable { - // Without a stable credential identity there is no safe cache key: pointer - // identity is reused once the old auth is collected, which would silently - // merge two unrelated OAuth identities onto the same TCP/TLS connections. - // Fall back to a private pool instead of risking cross-credential sharing. - return cloneTransportWithHTTP11(base) - } key := antigravityTransportKey{ - credential: credential, + credential: antigravityTransportScope(auth), base: base, } - if cached, ok := antigravityTransports.Load(key); ok { - return cached.(*http.Transport) + transport, errGet := antigravityTransports.Get(key, func() (*http.Transport, error) { + return cloneTransportWithHTTP11(base), nil + }) + if errGet != nil { + // Defensive only: the builder above cannot fail. Never return nil here, because a + // nil Transport makes http.Client fall back to http.DefaultTransport, which + // advertises h2 over ALPN and would break the Antigravity wire fingerprint. + log.Debugf("antigravity executor: cache HTTP/1.1 transport failed: %v", errGet) + return cloneTransportWithHTTP11(base) + } + return transport +} + +// antigravityProxiedHTTP11Transport returns the credential-scoped HTTP/1.1 pool for +// one proxy setting, or nil when the proxy setting cannot be turned into a +// transport. Keying on the normalized proxy string rather than on a prebuilt +// transport keeps one pool per credential and proxy instead of one per request. +func antigravityProxiedHTTP11Transport(auth *cliproxyauth.Auth, proxyURL string) *http.Transport { + proxyURL = strings.TrimSpace(proxyURL) + if proxyURL == "" { + return nil + } + key := antigravityTransportKey{ + credential: antigravityTransportScope(auth), + proxy: proxyURL, } - clone := cloneTransportWithHTTP11(base) - actual, _ := antigravityTransports.LoadOrStore(key, clone) - stored := actual.(*http.Transport) - if stored != clone { - // Another goroutine won the race; drop the redundant pool. - clone.CloseIdleConnections() + transport, errGet := antigravityTransports.Get(key, func() (*http.Transport, error) { + base, _, errBuild := proxyutil.BuildHTTPTransport(proxyURL) + if errBuild != nil { + return nil, errBuild + } + if base == nil { + return nil, fmt.Errorf("antigravity executor: proxy setting produced no transport") + } + return cloneTransportWithHTTP11(base), nil + }) + if errGet != nil { + // The caller falls back to NewProxyAwareHTTPClient, which reports the failure + // and applies the context transport fallback. + return nil } - return stored + return transport } -// antigravityTransportScope returns the connection-pool scope for one credential -// and reports whether that scope is stable enough to share a pool across requests. -// Runtime auths always carry an ID; incomplete test or plugin auth objects do not -// and must never be grouped together. -func antigravityTransportScope(auth *cliproxyauth.Auth) (string, bool) { +// antigravityTransportScope returns the connection-pool scope for one credential. +// Runtime auths always carry an ID. Incomplete auth objects, such as those built by +// tests, plugins or SDK embedders, fall back to another stable credential marker so +// they neither share a pool with an unrelated OAuth identity nor allocate a fresh +// pool, and with it a fresh set of pool goroutines, on every single request. +func antigravityTransportScope(auth *cliproxyauth.Auth) string { if auth == nil { - return "", false + return antigravityAnonymousTransportScope } - id := strings.TrimSpace(auth.ID) - if id == "" { - return "", false + if id := strings.TrimSpace(auth.ID); id != "" { + return "id:" + id } - return "id:" + id, true + if auth.Attributes != nil { + if path := strings.TrimSpace(auth.Attributes[cliproxyauth.AttributePath]); path != "" { + return "path:" + path + } + if source := strings.TrimSpace(auth.Attributes[cliproxyauth.AttributeSource]); source != "" { + return "source:" + source + } + } + // Fall back to the credential material itself. Auth.Label is deliberately not used: + // it is documented as an optional human readable label for logging and carries no + // uniqueness guarantee, so two different OAuth identities sharing one label would + // wrongly share a TCP/TLS pool. + // + // The refresh token is preferred over the access token because it stays stable + // across token rotation. Keying on the access token would move a credential to a new + // pool on every refresh, and would also strand refresh requests themselves, which + // run before any access token exists. + if refresh := strings.TrimSpace(metaStringValue(auth.Metadata, "refresh_token")); refresh != "" { + return antigravityCredentialScope("refresh:", refresh) + } + if access := strings.TrimSpace(metaStringValue(auth.Metadata, "access_token")); access != "" { + return antigravityCredentialScope("token:", access) + } + return antigravityAnonymousTransportScope +} + +// antigravityCredentialScope derives a pool scope from secret credential material. +// Only a short digest is retained, and it is never logged, so a pool key cannot be +// used to recover the credential it came from. +func antigravityCredentialScope(prefix, secret string) string { + digest := sha256.Sum256([]byte(secret)) + return prefix + hex.EncodeToString(digest[:8]) } // newAntigravityHTTPClient creates an HTTP client specifically for Antigravity, @@ -162,15 +289,15 @@ func antigravityTransportScope(auth *cliproxyauth.Auth) (string, bool) { // The underlying Transport is always shared so keep-alive connections survive across // requests instead of forcing a fresh TCP + TLS handshake every time. func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client { - credential, _ := antigravityTransportScope(auth) - // Native Antigravity reuses one transport across requests. Opt into a // credential-scoped proxy transport only here so other providers keep their // existing lifecycle and different OAuth identities remain isolated. if proxyURL := antigravityProxyURL(cfg, auth); proxyURL != "" { - if transport, _, errProxy := helps.SharedProxyTransport(credential, proxyURL); errProxy == nil && transport != nil { - return &http.Client{Transport: antigravityHTTP11Transport(auth, transport), Timeout: timeout} + if transport := antigravityProxiedHTTP11Transport(auth, proxyURL); transport != nil { + return &http.Client{Transport: transport, Timeout: timeout} } + // Fall through so NewProxyAwareHTTPClient reports the failure and applies the + // context transport fallback, preserving the previous behavior. } client := helps.NewProxyAwareHTTPClient(ctx, cfg, auth, timeout) @@ -182,9 +309,19 @@ func newAntigravityHTTPClient(ctx context.Context, cfg *config.Config, auth *cli // Preserve a context-provided transport while forcing HTTP/1.1. The cache key // includes credential identity, so sharing the base does not share TLS pools. - if transport, ok := client.Transport.(*http.Transport); ok { - client.Transport = antigravityHTTP11Transport(auth, transport) + transport, ok := client.Transport.(*http.Transport) + if !ok { + // A RoundTripper that is not an *http.Transport owns its own protocol behavior. + return client + } + if transport == nil { + // A typed-nil *http.Transport still satisfies the interface nil check in + // NewProxyAwareHTTPClient. Leaving it in place would make http.Client fall back + // to http.DefaultTransport, which advertises h2 over ALPN and breaks the + // Antigravity fingerprint, so substitute the process base transport. + transport = antigravityBaseTransport } + client.Transport = antigravityHTTP11Transport(auth, transport) return client } diff --git a/internal/runtime/executor/antigravity_executor_transport_test.go b/internal/runtime/executor/antigravity_executor_transport_test.go index b8f6b74e..378f02f1 100644 --- a/internal/runtime/executor/antigravity_executor_transport_test.go +++ b/internal/runtime/executor/antigravity_executor_transport_test.go @@ -2,14 +2,21 @@ package executor import ( "context" + "crypto/sha256" "crypto/tls" + "encoding/hex" + "errors" + "fmt" + "io" "net/http" "net/http/httptest" + "strings" "sync" "testing" "time" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" ) @@ -73,16 +80,157 @@ func TestNewAntigravityHTTPClientSharesTransport(t *testing.T) { if len(transport.TLSClientConfig.NextProtos) != 0 { t.Fatalf("Antigravity must omit ALPN like the native client, got %v", transport.TLSClientConfig.NextProtos) } - defaultTransport := http.DefaultTransport.(*http.Transport) - if transport.MaxIdleConns != defaultTransport.MaxIdleConns || - transport.MaxIdleConnsPerHost != defaultTransport.MaxIdleConnsPerHost || - transport.IdleConnTimeout != defaultTransport.IdleConnTimeout { - t.Fatal("Antigravity transport must preserve the standard Go pool settings") + // Go's DefaultMaxIdleConnsPerHost of 2 would force concurrent sessions on one + // credential to re-handshake. The native Antigravity stack raises it to 100. + if transport.MaxIdleConnsPerHost < antigravityMaxIdleConnsPerHost { + t.Fatalf("MaxIdleConnsPerHost = %d, want >= %d", transport.MaxIdleConnsPerHost, antigravityMaxIdleConnsPerHost) + } + if transport.MaxIdleConns > 0 && transport.MaxIdleConns < transport.MaxIdleConnsPerHost { + t.Fatalf("MaxIdleConns = %d must not throttle MaxIdleConnsPerHost = %d", transport.MaxIdleConns, transport.MaxIdleConnsPerHost) + } + if transport.IdleConnTimeout > 0 && transport.IdleConnTimeout < antigravityIdleConnTimeout { + t.Fatalf("IdleConnTimeout = %v, want >= %v", transport.IdleConnTimeout, antigravityIdleConnTimeout) } }) } } +// TestAntigravityPoolLimitsOnlyWiden guards that an operator-supplied base transport +// with a larger pool keeps its own settings, and that "unlimited" sentinels are not +// narrowed into finite limits. +func TestAntigravityPoolLimitsOnlyWiden(t *testing.T) { + wide := &http.Transport{ + MaxIdleConns: 512, + MaxIdleConnsPerHost: 256, + IdleConnTimeout: time.Hour, + } + applyAntigravityPoolLimits(wide) + if wide.MaxIdleConnsPerHost != 256 || wide.MaxIdleConns != 512 || wide.IdleConnTimeout != time.Hour { + t.Fatalf("wider pool settings must be preserved, got perHost=%d total=%d idle=%v", + wide.MaxIdleConnsPerHost, wide.MaxIdleConns, wide.IdleConnTimeout) + } + + // Zero means unlimited for both MaxIdleConns and IdleConnTimeout. + unlimited := &http.Transport{MaxIdleConns: 0, IdleConnTimeout: 0} + applyAntigravityPoolLimits(unlimited) + if unlimited.MaxIdleConns != 0 { + t.Fatalf("MaxIdleConns = %d, want 0 (unlimited) to stay unlimited", unlimited.MaxIdleConns) + } + if unlimited.IdleConnTimeout != 0 { + t.Fatalf("IdleConnTimeout = %v, want 0 (never expire) to stay unlimited", unlimited.IdleConnTimeout) + } + + // A negative MaxIdleConnsPerHost is how an operator disables idle pooling; Go never + // pools a connection in that case, so the intent must survive. + disabled := &http.Transport{MaxIdleConnsPerHost: -1} + applyAntigravityPoolLimits(disabled) + if disabled.MaxIdleConnsPerHost != -1 { + t.Fatalf("MaxIdleConnsPerHost = %d, want -1 (pooling disabled) to be preserved", disabled.MaxIdleConnsPerHost) + } + + // Go's zero value means DefaultMaxIdleConnsPerHost (2), which must be raised. + defaulted := &http.Transport{} + applyAntigravityPoolLimits(defaulted) + if defaulted.MaxIdleConnsPerHost != antigravityMaxIdleConnsPerHost { + t.Fatalf("MaxIdleConnsPerHost = %d, want %d", defaulted.MaxIdleConnsPerHost, antigravityMaxIdleConnsPerHost) + } + + applyAntigravityPoolLimits(nil) // must not panic +} + +// TestNewAntigravityHTTPClientRejectsTypedNilContextTransport guards the fingerprint: +// a typed-nil *http.Transport satisfies the interface nil check in +// NewProxyAwareHTTPClient, and leaving it in place would make http.Client fall back to +// http.DefaultTransport, which advertises h2 over ALPN. +func TestNewAntigravityHTTPClientRejectsTypedNilContextTransport(t *testing.T) { + var typedNil *http.Transport + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(typedNil)) + client := newAntigravityHTTPClient(ctx, &config.Config{}, antigravityAuthWithIDAndProxy("typed-nil", ""), 0) + + transport, ok := client.Transport.(*http.Transport) + if !ok || transport == nil { + t.Fatalf("expected a usable *http.Transport, got %#v", client.Transport) + } + if transport.ForceAttemptHTTP2 { + t.Fatal("fallback transport must not attempt HTTP/2") + } + if len(transport.TLSClientConfig.NextProtos) != 0 { + t.Fatalf("fallback transport must omit ALPN, got %v", transport.TLSClientConfig.NextProtos) + } +} + +// TestNewAntigravityHTTPClientKeepsForeignRoundTripper verifies a RoundTripper that is +// not an *http.Transport is left untouched instead of being replaced. +func TestNewAntigravityHTTPClientKeepsForeignRoundTripper(t *testing.T) { + foreign := roundTripperFunc(func(*http.Request) (*http.Response, error) { + return nil, errors.New("unused") + }) + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(foreign)) + client := newAntigravityHTTPClient(ctx, &config.Config{}, antigravityAuthWithIDAndProxy("foreign-rt", ""), 0) + if _, isTransport := client.Transport.(*http.Transport); isTransport { + t.Fatal("a non-*http.Transport RoundTripper must be preserved as-is") + } +} + +// TestAntigravityConcurrentRequestsReusePooledConnections is the regression test for +// the pool limit: with Go's default of 2 idle connections per host, repeated waves of +// concurrent requests on one credential keep re-handshaking. +func TestAntigravityConcurrentRequestsReusePooledConnections(t *testing.T) { + var mu sync.Mutex + remotes := map[string]struct{}{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + remotes[r.RemoteAddr] = struct{}{} + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + auth := antigravityAuthWithIDAndProxy("concurrent-reuse", "") + client := &http.Client{Transport: antigravityHTTP11Transport(auth, http.DefaultTransport.(*http.Transport))} + + const ( + waves = 3 + perWave = 8 + totalConns = waves * perWave + ) + for wave := 0; wave < waves; wave++ { + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(perWave) + for i := 0; i < perWave; i++ { + go func() { + defer wg.Done() + <-start + resp, errDo := client.Get(srv.URL) + if errDo != nil { + t.Error(errDo) + return + } + if _, errDrain := io.Copy(io.Discard, resp.Body); errDrain != nil { + t.Error(errDrain) + } + if errClose := resp.Body.Close(); errClose != nil { + t.Error(errClose) + } + }() + } + close(start) + wg.Wait() + } + + mu.Lock() + distinct := len(remotes) + mu.Unlock() + // The first wave legitimately opens perWave connections. Later waves must reuse + // them; with MaxIdleConnsPerHost=2 only two survive each wave and distinct grows + // towards totalConns instead. + if distinct > perWave { + t.Fatalf("%d waves of %d concurrent requests opened %d connections, want at most %d (unpooled worst case is %d)", + waves, perWave, distinct, perWave, totalConns) + } +} + func TestNewAntigravityHTTPClientDistinctProxiesUseDistinctPools(t *testing.T) { cfg := &config.Config{} a := newAntigravityHTTPClient(context.Background(), cfg, antigravityAuthWithProxy("http://127.0.0.1:18090"), 0) @@ -113,11 +261,11 @@ func TestNewAntigravityHTTPClientScopesPoolsByAuthIdentity(t *testing.T) { } } -// TestAntigravityHTTP11TransportNeverSharesWithoutStableIdentity guards the pool -// cache against auths that carry no ID. Keying such auths by pointer identity is -// unsafe: the address is reused once the previous auth is collected, which would -// silently place two unrelated OAuth identities on the same TCP/TLS connections. -func TestAntigravityHTTP11TransportNeverSharesWithoutStableIdentity(t *testing.T) { +// TestAntigravityHTTP11TransportReusesPoolWithoutAuthID guards the pool cache +// against auths that carry no ID. Allocating a private pool per call would leak a +// connection pool, and the goroutines managing it, on every request, which is the +// pattern the original singleton transport was introduced to remove. +func TestAntigravityHTTP11TransportReusesPoolWithoutAuthID(t *testing.T) { base := http.DefaultTransport.(*http.Transport) anonymous := &cliproxyauth.Auth{} @@ -126,16 +274,25 @@ func TestAntigravityHTTP11TransportNeverSharesWithoutStableIdentity(t *testing.T if first == nil || second == nil { t.Fatal("expected a transport for an auth without an ID") } - if first == second { - t.Fatal("an auth without a stable ID must not be cached, otherwise a reused address grants another credential its pool") + if first != second { + t.Fatal("an auth without any identity must reuse one shared pool instead of leaking a new pool per request") + } + if nilAuth := antigravityHTTP11Transport(nil, base); nilAuth != first { + t.Fatal("a nil auth carries no credential to isolate and must share the same pool") } - blankID := antigravityHTTP11Transport(&cliproxyauth.Auth{ID: " "}, base) - if blankID == first || blankID == second { - t.Fatal("a blank auth ID must not resolve to an existing pool") + // An auth without an ID but with credential material stays isolated from both the + // anonymous pool and from a different credential. + tokenA := antigravityHTTP11Transport(&cliproxyauth.Auth{Metadata: map[string]any{"access_token": "token-a"}}, base) + tokenB := antigravityHTTP11Transport(&cliproxyauth.Auth{Metadata: map[string]any{"access_token": "token-b"}}, base) + if tokenA == first || tokenB == first { + t.Fatal("a credential with an access token must not fall back to the anonymous pool") } - if nilAuth := antigravityHTTP11Transport(nil, base); nilAuth == nil || nilAuth == first || nilAuth == blankID { - t.Fatal("a nil auth must receive its own private pool") + if tokenA == tokenB { + t.Fatal("different access tokens must not share a connection pool") + } + if again := antigravityHTTP11Transport(&cliproxyauth.Auth{Metadata: map[string]any{"access_token": "token-a"}}, base); again != tokenA { + t.Fatal("the same access token must resolve to the same pool across requests") } // Identified auths keep sharing their pool. @@ -145,28 +302,175 @@ func TestAntigravityHTTP11TransportNeverSharesWithoutStableIdentity(t *testing.T } } -func TestAntigravityTransportScopeRequiresStableID(t *testing.T) { +func TestAntigravityTransportScopeFallsBackToStableMarkers(t *testing.T) { + digest := func(prefix, secret string) string { + sum := sha256.Sum256([]byte(secret)) + return prefix + hex.EncodeToString(sum[:8]) + } cases := []struct { - name string - auth *cliproxyauth.Auth - wantScope string - wantOK bool + name string + auth *cliproxyauth.Auth + want string }{ - {"nil auth", nil, "", false}, - {"missing id", &cliproxyauth.Auth{}, "", false}, - {"blank id", &cliproxyauth.Auth{ID: " \t "}, "", false}, - {"stable id", &cliproxyauth.Auth{ID: " auth-1 "}, "id:auth-1", true}, + {"nil auth", nil, antigravityAnonymousTransportScope}, + {"empty auth", &cliproxyauth.Auth{}, antigravityAnonymousTransportScope}, + {"blank id", &cliproxyauth.Auth{ID: " \t "}, antigravityAnonymousTransportScope}, + {"stable id", &cliproxyauth.Auth{ID: " auth-1 "}, "id:auth-1"}, + { + "id wins over path", + &cliproxyauth.Auth{ID: "auth-1", Attributes: map[string]string{cliproxyauth.AttributePath: "/a.json"}}, + "id:auth-1", + }, + { + "path fallback", + &cliproxyauth.Auth{Attributes: map[string]string{cliproxyauth.AttributePath: " /auths/a.json "}}, + "path:/auths/a.json", + }, + { + "source fallback", + &cliproxyauth.Auth{Attributes: map[string]string{cliproxyauth.AttributeSource: "/auths/b.json"}}, + "source:/auths/b.json", + }, + { + // Auth.Label is a logging label with no uniqueness guarantee, so it must never + // become a pool scope on its own. + "label alone is not an identity", + &cliproxyauth.Auth{Label: "account-c"}, + antigravityAnonymousTransportScope, + }, + { + "refresh token preferred over access token", + &cliproxyauth.Auth{Metadata: map[string]any{"refresh_token": "r-1", "access_token": "a-1"}}, + digest("refresh:", "r-1"), + }, + { + "access token fallback", + &cliproxyauth.Auth{Metadata: map[string]any{"access_token": "secret-token"}}, + digest("token:", "secret-token"), + }, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - scope, ok := antigravityTransportScope(tc.auth) - if scope != tc.wantScope || ok != tc.wantOK { - t.Fatalf("antigravityTransportScope() = (%q, %v), want (%q, %v)", scope, ok, tc.wantScope, tc.wantOK) + if got := antigravityTransportScope(tc.auth); got != tc.want { + t.Fatalf("antigravityTransportScope() = %q, want %q", got, tc.want) + } + }) + } +} + +// TestAntigravityTransportScopeIgnoresNonUniqueLabel is the regression test for using +// Auth.Label as an identity: two different credentials that happen to share a label +// must not end up on the same TCP/TLS pool. +func TestAntigravityTransportScopeIgnoresNonUniqueLabel(t *testing.T) { + first := &cliproxyauth.Auth{Label: "shared-label", Metadata: map[string]any{"refresh_token": "refresh-a"}} + second := &cliproxyauth.Auth{Label: "shared-label", Metadata: map[string]any{"refresh_token": "refresh-b"}} + if antigravityTransportScope(first) == antigravityTransportScope(second) { + t.Fatal("credentials sharing only a label must not share a pool scope") + } + + base := http.DefaultTransport.(*http.Transport) + if antigravityHTTP11Transport(first, base) == antigravityHTTP11Transport(second, base) { + t.Fatal("credentials sharing only a label must not share a connection pool") + } +} + +// TestAntigravityTransportScopeSurvivesAccessTokenRotation covers the refresh flow: +// refreshing an access token must not move a credential onto a new pool, and a refresh +// request that runs before any access token exists must resolve to the same scope. +func TestAntigravityTransportScopeSurvivesAccessTokenRotation(t *testing.T) { + refreshOnly := &cliproxyauth.Auth{Metadata: map[string]any{"refresh_token": "stable-refresh"}} + beforeRotation := &cliproxyauth.Auth{Metadata: map[string]any{"refresh_token": "stable-refresh", "access_token": "access-1"}} + afterRotation := &cliproxyauth.Auth{Metadata: map[string]any{"refresh_token": "stable-refresh", "access_token": "access-2"}} + + want := antigravityTransportScope(refreshOnly) + if got := antigravityTransportScope(beforeRotation); got != want { + t.Fatalf("scope before rotation = %q, want %q", got, want) + } + if got := antigravityTransportScope(afterRotation); got != want { + t.Fatalf("scope after rotation = %q, want %q (access token rotation must not churn pools)", got, want) + } +} + +// TestAntigravityTransportScopeNeverLeaksToken ensures the credential-derived scope +// only carries a short digest, so a pool key can never reveal the credential. +func TestAntigravityTransportScopeNeverLeaksToken(t *testing.T) { + const ( + accessToken = "ya29.super-secret-access-token" + refreshToken = "1//super-secret-refresh-token" + ) + for _, tc := range []struct { + name string + auth *cliproxyauth.Auth + secret string + prefix string + }{ + {"access token", &cliproxyauth.Auth{Metadata: map[string]any{"access_token": accessToken}}, accessToken, "token:"}, + {"refresh token", &cliproxyauth.Auth{Metadata: map[string]any{"refresh_token": refreshToken}}, refreshToken, "refresh:"}, + } { + t.Run(tc.name, func(t *testing.T) { + scope := antigravityTransportScope(tc.auth) + if strings.Contains(scope, tc.secret) { + t.Fatalf("scope %q must not embed the credential", scope) + } + if !strings.HasPrefix(scope, tc.prefix) || len(scope) != len(tc.prefix)+16 { + t.Fatalf("scope = %q, want a short %s digest", scope, tc.prefix) } }) } } +// TestAntigravityTransportCacheEvictsStalePools covers the bounded cache: rotating a +// credential's proxy must not accumulate pools forever. +func TestAntigravityTransportCacheEvictsStalePools(t *testing.T) { + original := antigravityTransports + antigravityTransports = helps.NewTransportCache[antigravityTransportKey](4) + t.Cleanup(func() { + antigravityTransports.Purge() + antigravityTransports = original + }) + + cfg := &config.Config{} + for i := 0; i < 40; i++ { + auth := antigravityAuthWithIDAndProxy("rotating-auth", fmt.Sprintf("http://127.0.0.1:%d", 19000+i)) + if client := newAntigravityHTTPClient(context.Background(), cfg, auth, 0); client.Transport == nil { + t.Fatalf("request %d: expected a transport", i) + } + } + if got := antigravityTransports.Len(); got > 4 { + t.Fatalf("cache holds %d pools, want at most the capacity of 4", got) + } + + // A per-request base transport from the request context must not grow the cache + // without bound either. + auth := antigravityAuthWithIDAndProxy("ctx-auth", "") + for i := 0; i < 40; i++ { + fresh := http.DefaultTransport.(*http.Transport).Clone() + ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", http.RoundTripper(fresh)) + if client := newAntigravityHTTPClient(ctx, cfg, auth, 0); client.Transport == nil { + t.Fatalf("ctx request %d: expected a transport", i) + } + } + if got := antigravityTransports.Len(); got > 4 { + t.Fatalf("cache holds %d pools after context transports, want at most 4", got) + } +} + +// TestAntigravityProxiedHTTP11TransportRejectsInvalidProxy verifies the caller can +// fall back instead of caching a broken pool. +func TestAntigravityProxiedHTTP11TransportRejectsInvalidProxy(t *testing.T) { + auth := antigravityAuthWithIDAndProxy("invalid-proxy", "ftp://127.0.0.1:1") + if transport := antigravityProxiedHTTP11Transport(auth, "ftp://127.0.0.1:1"); transport != nil { + t.Fatal("an unsupported proxy scheme must not produce a transport") + } + if transport := antigravityProxiedHTTP11Transport(auth, " "); transport != nil { + t.Fatal("a blank proxy must not produce a transport") + } + // A failed build must not occupy a cache slot, so a later valid setting still works. + if transport := antigravityProxiedHTTP11Transport(auth, "http://127.0.0.1:18099"); transport == nil { + t.Fatal("a valid proxy must produce a transport") + } +} + func TestAntigravityTransportMatchesNativeTLSProfile(t *testing.T) { var clientHelloProtos []string var requestProto string diff --git a/internal/runtime/executor/antigravity_refresh_test.go b/internal/runtime/executor/antigravity_refresh_test.go index d393d341..647b6996 100644 --- a/internal/runtime/executor/antigravity_refresh_test.go +++ b/internal/runtime/executor/antigravity_refresh_test.go @@ -35,10 +35,10 @@ func useAntigravityRefreshTestTransport(t *testing.T, targetHost string) { } originalBase := antigravityBaseTransport antigravityBaseTransport = transport - antigravityTransports = sync.Map{} + antigravityTransports.Purge() t.Cleanup(func() { antigravityBaseTransport = originalBase - antigravityTransports = sync.Map{} + antigravityTransports.Purge() }) } diff --git a/internal/runtime/executor/helps/transport_cache.go b/internal/runtime/executor/helps/transport_cache.go index 96edc966..9450482a 100644 --- a/internal/runtime/executor/helps/transport_cache.go +++ b/internal/runtime/executor/helps/transport_cache.go @@ -1,49 +1,125 @@ package helps import ( + "container/list" + "errors" "net/http" - "strings" "sync" - - "github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil" ) -// sharedProxyTransports memoizes one transport per credential scope and normalized -// proxy setting. Antigravity uses auth.ID as the scope so different OAuth identities -// never share a TCP/TLS pool, matching the native client's one-credential process model. -var sharedProxyTransports sync.Map // sharedProxyTransportKey -> *sharedProxyTransportEntry +// DefaultTransportCacheCapacity bounds how many transports a TransportCache keeps +// alive at once. Every cached transport owns an independent connection pool, so an +// unbounded cache would let idle sockets and the goroutines managing them grow +// without limit whenever keys churn, for example when a credential's proxy is +// rotated through the management API or when an SDK embedder supplies a freshly +// built base transport per request. +const DefaultTransportCacheCapacity = 64 -type sharedProxyTransportKey struct { - scope string - proxy string +// TransportCache memoizes HTTP transports under a comparable key using a bounded +// LRU. Evicting an entry closes its idle connections so neither the pool nor its +// background goroutines outlive the cache entry. +// +// The key type is generic so callers can mix value identity (a normalized proxy +// URL) with pointer identity (a base transport supplied by the caller) without the +// cache retaining either beyond the LRU window. +type TransportCache[K comparable] struct { + mu sync.Mutex + capacity int + // order keeps the most recently used entry at the front. + order *list.List + items map[K]*list.Element } -type sharedProxyTransportEntry struct { - once sync.Once +type transportCacheEntry[K comparable] struct { + key K transport *http.Transport - mode proxyutil.Mode - err error } -// SharedProxyTransport returns a stable transport for one credential scope and proxy -// setting. It preserves Go's standard connection-pool settings; callers must treat -// the result as read-only and clone it before changing protocol behavior. -func SharedProxyTransport(scope, raw string) (*http.Transport, proxyutil.Mode, error) { - key := sharedProxyTransportKey{ - scope: strings.TrimSpace(scope), - proxy: strings.TrimSpace(raw), - } - value, _ := sharedProxyTransports.LoadOrStore(key, &sharedProxyTransportEntry{}) - entry := value.(*sharedProxyTransportEntry) - entry.once.Do(func() { - entry.transport, entry.mode, entry.err = proxyutil.BuildHTTPTransport(key.proxy) - }) - return entry.transport, entry.mode, entry.err +// NewTransportCache returns a cache holding at most capacity transports. A +// non-positive capacity falls back to DefaultTransportCacheCapacity. +func NewTransportCache[K comparable](capacity int) *TransportCache[K] { + if capacity <= 0 { + capacity = DefaultTransportCacheCapacity + } + return &TransportCache[K]{ + capacity: capacity, + order: list.New(), + items: make(map[K]*list.Element, capacity), + } +} + +// Get returns the transport cached under key, calling build on the first use of +// that key. Concurrent callers observe the same instance. +// +// A build error is propagated without being cached, so a later call can retry and +// a failed lookup never occupies a cache slot. build must not call back into the +// same cache. +func (c *TransportCache[K]) Get(key K, build func() (*http.Transport, error)) (*http.Transport, error) { + if c == nil { + return nil, errors.New("transport cache: nil cache") + } + if build == nil { + return nil, errors.New("transport cache: nil build function") + } + + c.mu.Lock() + defer c.mu.Unlock() + + if element, ok := c.items[key]; ok { + c.order.MoveToFront(element) + return element.Value.(*transportCacheEntry[K]).transport, nil + } + + transport, errBuild := build() + if errBuild != nil { + return nil, errBuild + } + if transport == nil { + return nil, errors.New("transport cache: build returned no transport") + } + + c.items[key] = c.order.PushFront(&transportCacheEntry[K]{key: key, transport: transport}) + c.evictLocked() + return transport, nil +} + +// evictLocked drops least recently used entries until the cache fits its capacity. +// Closing idle connections is what actually releases the evicted pool; in-flight +// requests still holding the transport are unaffected because CloseIdleConnections +// only reaps connections that are currently idle. +func (c *TransportCache[K]) evictLocked() { + for c.order.Len() > c.capacity { + oldest := c.order.Back() + if oldest == nil { + return + } + c.order.Remove(oldest) + entry := oldest.Value.(*transportCacheEntry[K]) + delete(c.items, entry.key) + entry.transport.CloseIdleConnections() + } } -func resetSharedProxyTransportsForTest() { - sharedProxyTransports.Range(func(key, _ any) bool { - sharedProxyTransports.Delete(key) - return true - }) +// Len reports how many transports the cache currently holds. +func (c *TransportCache[K]) Len() int { + if c == nil { + return 0 + } + c.mu.Lock() + defer c.mu.Unlock() + return c.order.Len() +} + +// Purge drops every entry and closes the idle connections it was holding. +func (c *TransportCache[K]) Purge() { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + for element := c.order.Front(); element != nil; element = element.Next() { + element.Value.(*transportCacheEntry[K]).transport.CloseIdleConnections() + } + c.order.Init() + c.items = make(map[K]*list.Element, c.capacity) } diff --git a/internal/runtime/executor/helps/transport_cache_test.go b/internal/runtime/executor/helps/transport_cache_test.go index 64eb38e8..d0050194 100644 --- a/internal/runtime/executor/helps/transport_cache_test.go +++ b/internal/runtime/executor/helps/transport_cache_test.go @@ -1,59 +1,116 @@ package helps import ( + "errors" "net/http" "sync" "testing" - - "github.com/router-for-me/CLIProxyAPI/v7/sdk/proxyutil" ) -func TestSharedProxyTransportCachesNormalizedProxy(t *testing.T) { - resetSharedProxyTransportsForTest() - t.Cleanup(resetSharedProxyTransportsForTest) +type cacheKey struct { + scope string + proxy string +} + +func TestTransportCacheReusesEntriesPerKey(t *testing.T) { + cache := NewTransportCache[cacheKey](8) - first, mode, errFirst := SharedProxyTransport("auth-a", " http://127.0.0.1:3128 ") + builds := 0 + build := func() (*http.Transport, error) { + builds++ + return &http.Transport{}, nil + } + + first, errFirst := cache.Get(cacheKey{"auth-a", "p1"}, build) if errFirst != nil { - t.Fatalf("SharedProxyTransport() error = %v", errFirst) + t.Fatalf("Get() error = %v", errFirst) } - second, _, errSecond := SharedProxyTransport("auth-a", "http://127.0.0.1:3128") + second, errSecond := cache.Get(cacheKey{"auth-a", "p1"}, build) if errSecond != nil { - t.Fatalf("SharedProxyTransport() second error = %v", errSecond) - } - if mode != proxyutil.ModeProxy { - t.Fatalf("mode = %v, want %v", mode, proxyutil.ModeProxy) + t.Fatalf("Get() second error = %v", errSecond) } if first == nil || first != second { - t.Fatalf("expected the same cached transport, got %p and %p", first, second) + t.Fatalf("expected one cached transport, got %p and %p", first, second) } - - other, _, errOther := SharedProxyTransport("auth-a", "http://127.0.0.1:3129") - if errOther != nil { - t.Fatalf("SharedProxyTransport() other proxy error = %v", errOther) + if builds != 1 { + t.Fatalf("build called %d times, want 1", builds) } - if other == first { + + otherProxy, _ := cache.Get(cacheKey{"auth-a", "p2"}, build) + if otherProxy == first { t.Fatal("distinct proxies must not share a transport") } + otherScope, _ := cache.Get(cacheKey{"auth-b", "p1"}, build) + if otherScope == first { + t.Fatal("distinct credential scopes must not share a transport") + } + if got := cache.Len(); got != 3 { + t.Fatalf("cache Len() = %d, want 3", got) + } +} + +// TestTransportCacheBoundsEntries is the regression test for unbounded pool growth: +// every cached transport owns a connection pool, so churning keys must evict. +func TestTransportCacheBoundsEntries(t *testing.T) { + const capacity = 4 + cache := NewTransportCache[cacheKey](capacity) - otherAuth, _, errOtherAuth := SharedProxyTransport("auth-b", "http://127.0.0.1:3128") - if errOtherAuth != nil { - t.Fatalf("SharedProxyTransport() other auth error = %v", errOtherAuth) + for i := 0; i < 100; i++ { + key := cacheKey{"auth", string(rune('a' + i%97))} + if _, err := cache.Get(key, func() (*http.Transport, error) { return &http.Transport{}, nil }); err != nil { + t.Fatalf("Get() error = %v", err) + } + if got := cache.Len(); got > capacity { + t.Fatalf("cache grew to %d entries, want at most %d", got, capacity) + } } - if otherAuth == first { - t.Fatal("distinct credential scopes must not share a transport") +} + +// TestTransportCacheEvictsLeastRecentlyUsed proves recency is honoured, so a hot +// credential is not evicted by a burst of one-off keys. +func TestTransportCacheEvictsLeastRecentlyUsed(t *testing.T) { + cache := NewTransportCache[cacheKey](2) + build := func() (*http.Transport, error) { return &http.Transport{}, nil } + + hot, _ := cache.Get(cacheKey{"hot", ""}, build) + cache.Get(cacheKey{"cold", ""}, build) + // Touch hot so cold becomes the least recently used entry. + if again, _ := cache.Get(cacheKey{"hot", ""}, build); again != hot { + t.Fatal("expected the hot entry to still be cached") } + cache.Get(cacheKey{"new", ""}, build) - defaultTransport := http.DefaultTransport.(*http.Transport) - if first.MaxIdleConns != defaultTransport.MaxIdleConns || - first.MaxIdleConnsPerHost != defaultTransport.MaxIdleConnsPerHost || - first.IdleConnTimeout != defaultTransport.IdleConnTimeout { - t.Fatal("shared transport must preserve the standard Go connection pool settings") + if again, _ := cache.Get(cacheKey{"hot", ""}, build); again != hot { + t.Fatal("the most recently used entry must survive eviction") } } -func TestSharedProxyTransportConcurrentCallersShareOneInstance(t *testing.T) { - resetSharedProxyTransportsForTest() - t.Cleanup(resetSharedProxyTransportsForTest) +// TestTransportCacheDoesNotCacheBuildFailures ensures a transient failure neither +// occupies a cache slot nor becomes permanent. +func TestTransportCacheDoesNotCacheBuildFailures(t *testing.T) { + cache := NewTransportCache[cacheKey](4) + key := cacheKey{"auth", "broken"} + + if _, err := cache.Get(key, func() (*http.Transport, error) { return nil, errors.New("boom") }); err == nil { + t.Fatal("expected the build error to be propagated") + } + if got := cache.Len(); got != 0 { + t.Fatalf("a failed build must not occupy a cache slot, Len() = %d", got) + } + // A build returning (nil, nil) must be reported rather than cached as usable. + if _, err := cache.Get(key, func() (*http.Transport, error) { return nil, nil }); err == nil { + t.Fatal("expected an error when build returns no transport") + } + + transport, err := cache.Get(key, func() (*http.Transport, error) { return &http.Transport{}, nil }) + if err != nil || transport == nil { + t.Fatalf("retry after failure must succeed, got (%p, %v)", transport, err) + } +} + +func TestTransportCacheConcurrentCallersShareOneInstance(t *testing.T) { + cache := NewTransportCache[cacheKey](8) + key := cacheKey{"auth-concurrent", "socks5://127.0.0.1:1080"} const callers = 32 results := make([]*http.Transport, callers) @@ -62,7 +119,7 @@ func TestSharedProxyTransportConcurrentCallersShareOneInstance(t *testing.T) { for i := 0; i < callers; i++ { go func(index int) { defer wg.Done() - results[index], _, _ = SharedProxyTransport("auth-concurrent", "socks5://127.0.0.1:1080") + results[index], _ = cache.Get(key, func() (*http.Transport, error) { return &http.Transport{}, nil }) }(i) } wg.Wait() @@ -73,3 +130,43 @@ func TestSharedProxyTransportConcurrentCallersShareOneInstance(t *testing.T) { } } } + +func TestTransportCachePurgeAndNilSafety(t *testing.T) { + cache := NewTransportCache[cacheKey](4) + build := func() (*http.Transport, error) { return &http.Transport{}, nil } + cache.Get(cacheKey{"a", ""}, build) + cache.Get(cacheKey{"b", ""}, build) + if got := cache.Len(); got != 2 { + t.Fatalf("Len() = %d, want 2", got) + } + cache.Purge() + if got := cache.Len(); got != 0 { + t.Fatalf("Len() after Purge() = %d, want 0", got) + } + // The cache stays usable after a purge. + if transport, err := cache.Get(cacheKey{"a", ""}, build); err != nil || transport == nil { + t.Fatalf("Get() after Purge() = (%p, %v)", transport, err) + } + + var nilCache *TransportCache[cacheKey] + if _, err := nilCache.Get(cacheKey{}, build); err == nil { + t.Fatal("expected an error from a nil cache") + } + if got := nilCache.Len(); got != 0 { + t.Fatalf("nil cache Len() = %d, want 0", got) + } + nilCache.Purge() // must not panic + + if _, err := cache.Get(cacheKey{"nil-build", ""}, nil); err == nil { + t.Fatal("expected an error for a nil build function") + } +} + +func TestNewTransportCacheDefaultsCapacity(t *testing.T) { + for _, capacity := range []int{0, -1} { + cache := NewTransportCache[cacheKey](capacity) + if cache.capacity != DefaultTransportCacheCapacity { + t.Fatalf("NewTransportCache(%d).capacity = %d, want %d", capacity, cache.capacity, DefaultTransportCacheCapacity) + } + } +}