Publish reflector targets by copy-on-write, fix unstable proxy-entry key, gate with -race

.github/workflows/ci.yml

@@ -34,6 +34,9 @@ jobs:
       - name: Run Agent tests
         run: make -C agent test
 
+      - name: Run Agent race detector
+        run: make -C agent test-race
+
       - name: Run Go SDK tests
         run: |
           cd sdks/go

agent/Makefile

@@ -68,6 +68,23 @@ test: mocks
 	go mod tidy
 	go test -v ./...
 
+# Packages that currently pass the race detector. The rest of the tree does
+# not yet: server, server/handler, server/http, server/snykbroker and cmd all
+# have known races (supervisor process lifecycle, handler manager, the fx
+# server stack). Move packages up here as those are fixed.
+RACE_CLEAN_PKGS := ./common/... ./config/... ./scaffold/... ./server/api/... ./server/cron/... ./server/snykbroker/acceptfile/... ./util/...
+
+# The reflector lives in server/snykbroker, whose other tests are not yet
+# race-clean, so it is gated at test granularity instead. Derived from the
+# test files so new reflector tests are covered automatically.
+REFLECTOR_TESTS = $(shell grep -ho '^func Test[A-Za-z0-9_]*' server/snykbroker/reflector*_test.go | sed 's/^func //' | paste -sd'|' -)
+
+test-race: mocks
+	go test -race $(RACE_CLEAN_PKGS)
+	go test -race -run '^($(REFLECTOR_TESTS))$$' ./server/snykbroker/
+
+.PHONY: test-race
+
 validate-filters:
 	test/validate_filters.sh
 

agent/server/http/http_server.go

@@ -269,19 +269,19 @@ func (h *httpServer) Start() (int, error) {
 	if err != nil {
 		panic(err)
 	}
+	// construct the server before spawning the serve goroutine and hand it a
+	// local reference, so Close (which nils h.server) never races with it
+	server := &http.Server{
+		Handler:     h.mux,
+		ReadTimeout: defaultReadTimeout,
+	}
+	h.server = server
 	go func() {
-
-		h.server = &http.Server{
-			Handler:     h.mux,
-			ReadTimeout: defaultReadTimeout,
-		}
-
-		err := h.server.Serve(ln)
+		err := server.Serve(ln)
 		if err != nil && err != http.ErrServerClosed {
 			panic(err)
 		}
 	}()
-	time.Sleep(100 * time.Millisecond)
 	h.listener = ln
 	h.port = ln.Addr().(*net.TCPAddr).Port
 	return h.port, nil

agent/server/snykbroker/reflector.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"net/url"
+	"sort"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -21,10 +22,12 @@ import (
 )
 
 type RegistrationReflector struct {
-	logger          *zap.Logger
-	transport       *http.Transport
-	server          cortexHttp.Server
-	targets         map[string]proxyEntry
+	logger    *zap.Logger
+	transport *http.Transport
+	server    cortexHttp.Server
+	// Never mutated once published; registration swaps in a new map. Readers
+	// hold no lock, so the map a reader is ranging over must stay frozen.
+	targets         atomic.Pointer[map[string]proxyEntry]
 	serverStarted   atomic.Bool
 	mode            config.RelayReflectorMode
 	config          config.AgentConfig
@@ -58,10 +61,10 @@ func NewRegistrationReflector(p RegistrationReflectorParams) *RegistrationReflec
 		transport: p.Transport,
 		server:    server,
 		logger:    httpParams.Logger,
-		targets:   make(map[string]proxyEntry),
 		mode:      p.Config.HttpRelayReflectorMode,
 		config:    p.Config,
 	}
+	rr.targets.Store(&map[string]proxyEntry{})
 
 	// Create WebSocket proxy with callbacks for tunnel lifecycle
 	rr.wsProxy = NewWebSocketProxy(httpParams.Logger, p.Transport)
@@ -152,12 +155,30 @@ func (rr *RegistrationReflector) getProxy(targetURI string, isDefault bool, head
 
 	key := newEntry.key()
 
-	entry, exists := rr.targets[key]
-	if !exists {
-		entry = *newEntry
-		rr.targets[key] = entry
-		newEntry.addResponseHeader("x-axon-relay-instance", rr.config.InstanceId)
+	// Set before the entry is published so the copy in the map carries the
+	// header too, rather than only the reverse proxy's captured newEntry.
+	newEntry.addResponseHeader("x-axon-relay-instance", rr.config.InstanceId)
+
+	// Copy-on-write: publish a new map rather than mutating the live one, so
+	// readers never synchronize. Retried on CAS failure because a concurrent
+	// registration may have published between the load and the swap.
+	for {
+		current := rr.targets.Load()
+		if entry, exists := (*current)[key]; exists {
+			return &entry, nil
+		}
+
+		next := make(map[string]proxyEntry, len(*current)+1)
+		for k, v := range *current {
+			next[k] = v
+		}
+		next[key] = *newEntry
 
+		if !rr.targets.CompareAndSwap(current, &next) {
+			continue
+		}
+
+		entry := next[key]
 		rr.logger.Info("Registered redirector",
 			zap.String("targetURI", entry.TargetURI),
 			zap.String("proxyURI", entry.proxyURI),
@@ -167,7 +188,6 @@ func (rr *RegistrationReflector) getProxy(targetURI string, isDefault bool, head
 		)
 		return &entry, nil
 	}
-	return &entry, nil
 }
 
 func (rr *RegistrationReflector) extractHash(part string) string {
@@ -188,9 +208,12 @@ func (rr *RegistrationReflector) parseTargetUri(proxyPath string) (*proxyEntry,
 		remainder = path[slash:]
 	}
 	hash := rr.extractHash(beforeSlash)
+
+	targets := *rr.targets.Load()
+
 	if hash == "" {
 		// find the default proxy entry
-		if entry, exists := rr.targets["default"]; exists {
+		if entry, exists := targets["default"]; exists {
 			// Found the default proxy entry
 			return &entry, proxyPath, nil
 		} else {
@@ -199,11 +222,9 @@ func (rr *RegistrationReflector) parseTargetUri(proxyPath string) (*proxyEntry,
 		}
 	}
 
-	for _, entry := range rr.targets {
-		if entry.key() == hash {
-			// Found the target URI
-			return &entry, remainder, nil
-		}
+	// the map is keyed by exactly the value key() returns
+	if entry, exists := targets[hash]; exists {
+		return &entry, remainder, nil
 	}
 
 	return nil, "", fmt.Errorf("no proxy entry found for path: %s", proxyPath)
@@ -239,13 +260,15 @@ func WithHeadersResolver(headers acceptfile.ResolverMap) ProxyOption {
 	}
 }
 
+// getUriForTarget scans by target URI rather than by key. Only tests need this
+// direction, so the linear scan is not on any request path.
 func (rr *RegistrationReflector) getUriForTarget(target string) (string, error) {
 
 	if target == "" {
 		return "", fmt.Errorf("target URI cannot be empty")
 	}
 
-	for _, entry := range rr.targets {
+	for _, entry := range *rr.targets.Load() {
 		if entry.TargetURI == target {
 			return entry.proxyURI, nil
 		}
@@ -413,9 +436,15 @@ func (pe *proxyEntry) key() string {
 		key := pe.TargetURI
 
 		if len(pe.headers) > 0 {
-			// Create a unique key that includes headers to allow different header sets for the same URI
-			headerKey := ""
+			// Create a unique key that includes headers to allow different header sets for the same URI.
+			// Sorted so the hash is stable across map iteration order.
+			names := make([]string, 0, len(pe.headers))
 			for k := range pe.headers {
+				names = append(names, k)
+			}
+			sort.Strings(names)
+			headerKey := ""
+			for _, k := range names {
 				headerKey += fmt.Sprintf("|%s=%s", k, pe.headers.ResolverKey(k))
 			}
 			key = key + headerKey

agent/server/snykbroker/reflector_key_test.go

@@ -0,0 +1,101 @@
+package snykbroker
+
+import (
+	"fmt"
+	"net/url"
+	"sync"
+	"testing"
+
+	"github.com/cortexapps/axon/server/snykbroker/acceptfile"
+	"github.com/stretchr/testify/require"
+)
+
+func proxyPath(t *testing.T, proxyURI string) string {
+	u, err := url.Parse(proxyURI)
+	require.NoError(t, err)
+	return u.Path
+}
+
+func TestProxyEntryKeyStableAcrossHeaderOrder(t *testing.T) {
+	headers := map[string]string{
+		"Authorization":        "Bearer ${plugin:gcp-token}",
+		"X-GitHub-Api-Version": "2022-11-28",
+		"X-Third":              "three",
+	}
+	first, err := newProxyEntry("https://example.com", false, 8080, acceptfile.NewResolverMapFromMap(headers), nil)
+	require.NoError(t, err)
+	// map iteration order is randomized per map instance, so repeated
+	// construction flushes out order-dependent hashing
+	for i := 0; i < 20; i++ {
+		next, err := newProxyEntry("https://example.com", false, 8080, acceptfile.NewResolverMapFromMap(headers), nil)
+		require.NoError(t, err)
+		require.Equal(t, first.key(), next.key())
+	}
+}
+
+// TestConcurrentGetProxyAndParseTargetUri reproduces the production race
+// between registration writing rr.targets (the broker start path retries
+// registration on its own goroutine, so it can register entries long after
+// Start returned) and ServeHTTP reading it through parseTargetUri. Mutating
+// the map in place makes this a concurrent map read-and-write, which is a Go
+// runtime throw rather than a tolerable race. Only meaningful under -race.
+func TestConcurrentGetProxyAndParseTargetUri(t *testing.T) {
+	env := newTestReflectorEnv(t)
+
+	// seed entries so readers resolve real hashes while writers add more
+	seedPaths := make([]string, 0, 4)
+	for i := 0; i < 4; i++ {
+		entry, err := env.Reflector.getProxy(fmt.Sprintf("http://seed-%d.example.com", i), false, nil)
+		require.NoError(t, err)
+		seedPaths = append(seedPaths, proxyPath(t, entry.proxyURI))
+	}
+
+	const workers = 16
+	const perWorker = 50
+	var wg sync.WaitGroup
+	errs := make(chan error, 2*workers*perWorker)
+
+	for w := 0; w < workers; w++ {
+		wg.Add(1)
+		go func(w int) {
+			defer wg.Done()
+			for i := 0; i < perWorker; i++ {
+				// a distinct URI per iteration, so every call writes a new entry
+				_, err := env.Reflector.getProxy(fmt.Sprintf("http://w%d-i%d.example.com", w, i), false, nil)
+				if err != nil {
+					errs <- fmt.Errorf("getProxy w=%d i=%d: %w", w, i, err)
+					return
+				}
+			}
+		}(w)
+	}
+
+	for r := 0; r < workers; r++ {
+		wg.Add(1)
+		go func(r int) {
+			defer wg.Done()
+			for i := 0; i < perWorker; i++ {
+				path := seedPaths[(r+i)%len(seedPaths)]
+				entry, _, err := env.Reflector.parseTargetUri(path)
+				if err != nil {
+					errs <- fmt.Errorf("parseTargetUri r=%d i=%d path=%s: %w", r, i, path, err)
+					return
+				}
+				if entry == nil {
+					errs <- fmt.Errorf("parseTargetUri r=%d i=%d path=%s: nil entry", r, i, path)
+					return
+				}
+			}
+		}(r)
+	}
+
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		require.NoError(t, err)
+	}
+
+	// every writer URI registered exactly one entry, plus the seeds. A lost
+	// copy-on-write update would show up here as a short map.
+	require.Len(t, *env.Reflector.targets.Load(), 4+workers*perWorker)
+}

agent/server/snykbroker/reflector_test.go

@@ -32,6 +32,25 @@ func newTestReflectorEnv(t *testing.T) *testReflectorEnv {
 	})
 }
 
+// waitForTunnelDrain blocks until no WebSocket tunnel is still copying.
+// activeConnections is decremented only after both copy goroutines have
+// returned, so reaching zero means none of them can log again.
+func waitForTunnelDrain(t *testing.T, rr *RegistrationReflector) {
+	t.Helper()
+	if rr.wsProxy == nil {
+		return
+	}
+	deadline := time.Now().Add(5 * time.Second)
+	for rr.wsProxy.ActiveConnections() > 0 {
+		if time.Now().After(deadline) {
+			t.Errorf("websocket tunnels did not drain within 5s (%d still active)",
+				rr.wsProxy.ActiveConnections())
+			return
+		}
+		time.Sleep(time.Millisecond)
+	}
+}
+
 func newTestReflectorEnvWithConfig(t *testing.T, cfg config.AgentConfig) *testReflectorEnv {
 	logger := zaptest.NewLogger(t)
 	rr := NewRegistrationReflector(RegistrationReflectorParams{
@@ -41,6 +60,11 @@ func newTestReflectorEnvWithConfig(t *testing.T, cfg config.AgentConfig) *testRe
 	})
 	router := mux.NewRouter()
 	server := httptest.NewServer(router)
+	// Cleanups run LIFO, so this drain runs last - after Stop and server.Close
+	// have torn the connections down. Tunnel goroutines outlive the request
+	// that started them and log through the zaptest logger, which races the
+	// test being marked done if we return while one is still copying.
+	t.Cleanup(func() { waitForTunnelDrain(t, rr) })
 	t.Cleanup(server.Close)
 	t.Cleanup(func() { rr.Stop() })
 	return &testReflectorEnv{