From 3392fe05f3226067317bcfa9e3ff6797f6d1121d Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Fri, 28 Aug 2026 15:07:47 +1000 Subject: [PATCH] Make origin pools rotate again MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `Origin()` expanded ${VAR} against the environment before the pool was consulted, so a variable that exists only as `VAR_POOL` became the empty string and every request against a pooled origin failed — while `pool_test.go`, which exercises `PoolManager` directly, stayed green. The Router now resolves the origin from `RawOrigin()`, with the ${VAR} reference intact, so `ResolvePoolVars` has something left to rotate. `defaultScheme` comes out of `Origin()` so both paths fill in https the same way. `router_pool_test.go` covers this at the Router level, where the bug was; it fails against the old resolution. Co-Authored-By: Claude Opus 5 --- .../snykbroker/acceptfile/accept_file.go | 17 ++- .../acceptfile/origin_contract_test.go | 1 + agent/server/snykbroker/acceptfile/router.go | 34 ++++- .../snykbroker/acceptfile/router_pool_test.go | 120 ++++++++++++++++++ 4 files changed, 165 insertions(+), 7 deletions(-) create mode 100644 agent/server/snykbroker/acceptfile/router_pool_test.go diff --git a/agent/server/snykbroker/acceptfile/accept_file.go b/agent/server/snykbroker/acceptfile/accept_file.go index 0713bf3..10cad1a 100644 --- a/agent/server/snykbroker/acceptfile/accept_file.go +++ b/agent/server/snykbroker/acceptfile/accept_file.go @@ -3,7 +3,6 @@ package acceptfile import ( "encoding/json" "fmt" - "net/url" "os" "sync" @@ -234,19 +233,25 @@ func (r AcceptFileRuleWrapper) Origin() string { return "" } origin := os.ExpandEnv(rawOrigin) - asUrl, err := url.Parse(origin) + withScheme, err := defaultScheme(origin) if err != nil { r.acceptFile.logger.Panic("failed to parse origin URL", zap.String("origin", rawOrigin), zap.Error(err)) } - if asUrl.Scheme == "" { + if withScheme != origin { if _, seen := r.acceptFile.schemeWarned.LoadOrStore(rawOrigin, struct{}{}); !seen { r.acceptFile.logger.Debug("origin URL has no scheme, defaulting to https", zap.String("origin", rawOrigin)) } - asUrl.Scheme = "https" - return asUrl.String() } - return origin + return withScheme +} +// RawOrigin returns the rule's origin exactly as written, with ${VAR} +// references intact. The Router needs this because pool rotation has to happen +// on the reference: expanding ${API} against the environment first turns a +// pool-only variable into the empty string, leaving nothing to rotate. +func (r AcceptFileRuleWrapper) RawOrigin() string { + origin, _ := r.dict["origin"].(string) + return origin } func (r AcceptFileRuleWrapper) Path() string { diff --git a/agent/server/snykbroker/acceptfile/origin_contract_test.go b/agent/server/snykbroker/acceptfile/origin_contract_test.go index ff9a174..b5c86c3 100644 --- a/agent/server/snykbroker/acceptfile/origin_contract_test.go +++ b/agent/server/snykbroker/acceptfile/origin_contract_test.go @@ -116,4 +116,5 @@ func TestOriginOfRuleWithoutOneIsEmpty(t *testing.T) { rules := af.Wrapper().PrivateRules() require.Len(t, rules, 1) require.Equal(t, "", rules[0].Origin()) + require.Equal(t, "", rules[0].RawOrigin()) } diff --git a/agent/server/snykbroker/acceptfile/router.go b/agent/server/snykbroker/acceptfile/router.go index ab3e9a1..9733f46 100644 --- a/agent/server/snykbroker/acceptfile/router.go +++ b/agent/server/snykbroker/acceptfile/router.go @@ -108,7 +108,10 @@ func (rt *Router) Route(method, rawPath string, headers map[string]string) (*Rou } } - origin := rt.pools.ResolvePoolVars(rule.Origin()) + origin, err := rt.resolveOrigin(*rule) + if err != nil { + return nil, err + } targetURL, err := buildTargetURL(origin, pathOnly, decodedPath, query) if err != nil { return nil, fmt.Errorf("failed to build target URL: %w", err) @@ -157,6 +160,35 @@ func (rt *Router) Route(method, rawPath string, headers map[string]string) (*Rou }, nil } +// resolveOrigin turns a rule's raw origin into a concrete URL string, rotating +// any ${VAR} that names a pool. +// +// Resolution has to happen on the reference rather than on an already-expanded +// string: rule.Origin() runs os.ExpandEnv first, so a variable that exists only +// as VAR_POOL becomes the empty string and the PoolManager has nothing left to +// rotate. +func (rt *Router) resolveOrigin(rule AcceptFileRuleWrapper) (string, error) { + raw := rule.RawOrigin() + if raw == "" { + return "", fmt.Errorf("rule has no origin") + } + return defaultScheme(rt.pools.ResolvePoolVars(raw)) +} + +// defaultScheme fills in https for an origin written without one, which the +// accept-file format allows (${GITHUB:github.com}). +func defaultScheme(origin string) (string, error) { + asURL, err := url.Parse(origin) + if err != nil { + return "", fmt.Errorf("invalid origin %q: %w", origin, err) + } + if asURL.Scheme != "" { + return origin, nil + } + asURL.Scheme = "https" + return asURL.String(), nil +} + // buildTargetURL joins the rule's resolved origin with the request path // and query, preserving percent-encoded characters (e.g. %2F in GitLab // project IDs) on the wire. diff --git a/agent/server/snykbroker/acceptfile/router_pool_test.go b/agent/server/snykbroker/acceptfile/router_pool_test.go new file mode 100644 index 0000000..7188125 --- /dev/null +++ b/agent/server/snykbroker/acceptfile/router_pool_test.go @@ -0,0 +1,120 @@ +package acceptfile + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +// Pool rotation has to happen where the origin is resolved for a request. +// Resolving ${VAR} against the environment first turns a pool-only variable +// into the empty string, and the PoolManager then has nothing left to rotate — +// so every request fails while pool_test.go stays green. + +func TestRouterRotatesPoolOrigins(t *testing.T) { + t.Setenv("POOLED_API_POOL", "https://a.example.com,https://b.example.com,https://c.example.com") + + rt := newTestRouter(t, `{"private":[{"method":"any","path":"/*","origin":"${POOLED_API}"}]}`) + + want := []string{ + "https://a.example.com/x", "https://b.example.com/x", "https://c.example.com/x", + "https://a.example.com/x", "https://b.example.com/x", "https://c.example.com/x", + } + for i, w := range want { + routed, err := rt.Route("GET", "/x", nil) + require.NoError(t, err, "request %d", i) + require.Equal(t, w, routed.URL.String(), "request %d", i) + } +} + +func TestRouterPoolTrimsWhitespaceAndEmptyEntries(t *testing.T) { + t.Setenv("POOLED_API_POOL", " https://a.example.com , ,https://b.example.com ") + + rt := newTestRouter(t, `{"private":[{"method":"any","path":"/*","origin":"${POOLED_API}"}]}`) + + for _, w := range []string{"https://a.example.com/x", "https://b.example.com/x", "https://a.example.com/x"} { + routed, err := rt.Route("GET", "/x", nil) + require.NoError(t, err) + require.Equal(t, w, routed.URL.String()) + } +} + +// A pool of one is a plain origin. +func TestRouterPoolOfOne(t *testing.T) { + t.Setenv("POOLED_API_POOL", "https://only.example.com") + + rt := newTestRouter(t, `{"private":[{"method":"any","path":"/*","origin":"${POOLED_API}"}]}`) + for i := 0; i < 3; i++ { + routed, err := rt.Route("GET", "/x", nil) + require.NoError(t, err) + require.Equal(t, "https://only.example.com/x", routed.URL.String()) + } +} + +// With no pool set the plain environment variable still wins. +func TestRouterFallsBackToPlainEnvVar(t *testing.T) { + t.Setenv("SINGLE_API", "https://api.example.com") + + rt := newTestRouter(t, `{"private":[{"method":"any","path":"/*","origin":"${SINGLE_API}"}]}`) + routed, err := rt.Route("GET", "/x", nil) + require.NoError(t, err) + require.Equal(t, "https://api.example.com/x", routed.URL.String()) +} + +// A pool beats the plain variable when both are set, matching the broker. +func TestRouterPoolBeatsPlainEnvVar(t *testing.T) { + t.Setenv("BOTH_API", "https://plain.example.com") + t.Setenv("BOTH_API_POOL", "https://pooled.example.com") + + rt := newTestRouter(t, `{"private":[{"method":"any","path":"/*","origin":"${BOTH_API}"}]}`) + routed, err := rt.Route("GET", "/x", nil) + require.NoError(t, err) + require.Equal(t, "https://pooled.example.com/x", routed.URL.String()) +} + +// Every concurrent request must land on a real pool member; the counter is the +// only shared state on the hot path. +func TestRouterPoolIsSafeUnderConcurrency(t *testing.T) { + t.Setenv("POOLED_API_POOL", "https://a.example.com,https://b.example.com") + + rt := newTestRouter(t, `{"private":[{"method":"any","path":"/*","origin":"${POOLED_API}"}]}`) + + const n = 200 + var wg sync.WaitGroup + var mu sync.Mutex + counts := map[string]int{} + + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + routed, err := rt.Route("GET", "/x", nil) + if err != nil { + return + } + mu.Lock() + counts[routed.URL.Host]++ + mu.Unlock() + }() + } + wg.Wait() + + require.Equal(t, n, counts["a.example.com"]+counts["b.example.com"], + "every request must resolve to a pool member") + require.Equal(t, n/2, counts["a.example.com"], "rotation must stay even") + require.Equal(t, n/2, counts["b.example.com"], "rotation must stay even") +} + +// A pool-only variable satisfies the accept file's required-variable check +// (varIsSet accepts VAR_POOL), so it must also route. +func TestPoolOnlyVariableSatisfiesLoadValidation(t *testing.T) { + t.Setenv("DECLARED_API_POOL", "https://a.example.com") + + rt := newTestRouter(t, `{"$vars":["${DECLARED_API}"], + "private":[{"method":"any","path":"/*","origin":"${DECLARED_API}"}]}`) + + routed, err := rt.Route("GET", "/x", nil) + require.NoError(t, err) + require.Equal(t, "https://a.example.com/x", routed.URL.String()) +}