Grpc-Tunnel: Fix pool var rotation

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 {

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())
 }

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.

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())
+}