diff --git a/agent/internal/serverless/gateway.go b/agent/internal/serverless/gateway.go index 9bc4a7b..cc81b36 100644 --- a/agent/internal/serverless/gateway.go +++ b/agent/internal/serverless/gateway.go @@ -20,8 +20,9 @@ import ( ) const ( - GatewayPort = 18080 - upstreamCacheTTL = 10 * time.Second + GatewayPort = 18080 + upstreamCacheTTL = 10 * time.Second + proxyRetryResolveTimeout = 2 * time.Second ) var ( @@ -183,28 +184,68 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { } upstream := upstreams[g.nextIndex(len(upstreams))] + proxyErr := g.proxyOnce(w, r, upstream) + if proxyErr == nil { + return + } + if requestCancelled(r, proxyErr) { + return + } + + log.Printf("[serverless-gateway] proxy error for host %s to %s: %v", host, upstream.Url, proxyErr) + g.evictUpstreams(host) + if !retryableProxyRequest(r) { + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + + retryCtx, cancel := context.WithTimeout(r.Context(), proxyRetryResolveTimeout) + defer cancel() + retryUpstreams, err := g.getUpstreams(retryCtx, host) + if err != nil || len(retryUpstreams) == 0 { + if r.Context().Err() != nil { + return + } + if err == nil { + err = errors.New("no upstreams after cache eviction") + } + log.Printf("[serverless-gateway] proxy retry resolution failed for host %s: %v", host, err) + http.Error(w, "bad gateway", http.StatusBadGateway) + return + } + + retryUpstream := selectRetryUpstream(retryUpstreams, upstream.Url, g.nextIndex(len(retryUpstreams))) + proxyErr = g.proxyOnce(w, r, retryUpstream) + if proxyErr == nil || requestCancelled(r, proxyErr) { + return + } + log.Printf("[serverless-gateway] proxy retry failed for host %s to %s: %v", host, retryUpstream.Url, proxyErr) + g.evictUpstreams(host) + http.Error(w, "bad gateway", http.StatusBadGateway) +} + +func (g *Gateway) proxyOnce(w http.ResponseWriter, r *http.Request, upstream agenthttp.ServerlessUpstream) error { target, err := url.Parse("http://" + upstream.Url) if err != nil { - log.Printf("[serverless-gateway] invalid upstream %q for host %s: %v", upstream.Url, host, err) - http.Error(w, "bad upstream", http.StatusBadGateway) - return + return fmt.Errorf("invalid upstream %q: %w", upstream.Url, err) } proxy := httputil.NewSingleHostReverseProxy(target) originalDirector := proxy.Director originalHost := r.Host + originalProto := forwardedProto(r) proxy.Director = func(req *http.Request) { originalDirector(req) req.Host = originalHost req.Header.Set("X-Forwarded-Host", originalHost) - req.Header.Set("X-Forwarded-Proto", forwardedProto(r)) + req.Header.Set("X-Forwarded-Proto", originalProto) } - proxy.ErrorHandler = func(w http.ResponseWriter, req *http.Request, err error) { - log.Printf("[serverless-gateway] proxy error for host %s to %s: %v", host, upstream.Url, err) - g.evictUpstreams(host) - http.Error(w, "bad gateway", http.StatusBadGateway) + var proxyErr error + proxy.ErrorHandler = func(_ http.ResponseWriter, _ *http.Request, err error) { + proxyErr = err } proxy.ServeHTTP(w, r) + return proxyErr } func (g *Gateway) getUpstreams(ctx context.Context, host string) ([]agenthttp.ServerlessUpstream, error) { @@ -1307,6 +1348,41 @@ func (g *Gateway) nextIndex(length int) int { return int(atomic.AddUint64(&g.counter, 1)-1) % length } +func retryableProxyRequest(r *http.Request) bool { + if r.Method != http.MethodGet && r.Method != http.MethodHead { + return false + } + if r.ContentLength != 0 || (r.Body != nil && r.Body != http.NoBody) { + return false + } + return !headerContainsToken(r.Header, "Connection", "upgrade") && r.Header.Get("Upgrade") == "" +} + +func requestCancelled(r *http.Request, err error) bool { + return r.Context().Err() != nil || errors.Is(err, context.Canceled) +} + +func selectRetryUpstream(upstreams []agenthttp.ServerlessUpstream, failedURL string, start int) agenthttp.ServerlessUpstream { + for offset := range len(upstreams) { + upstream := upstreams[(start+offset)%len(upstreams)] + if upstream.Url != failedURL { + return upstream + } + } + return upstreams[start%len(upstreams)] +} + +func headerContainsToken(header http.Header, name string, token string) bool { + for _, value := range header.Values(name) { + for part := range strings.SplitSeq(value, ",") { + if strings.EqualFold(strings.TrimSpace(part), token) { + return true + } + } + } + return false +} + func normalizeHost(host string) string { if h, _, err := net.SplitHostPort(host); err == nil { return strings.ToLower(strings.TrimSpace(h)) diff --git a/agent/internal/serverless/gateway_test.go b/agent/internal/serverless/gateway_test.go index 45dc0e1..aa07ec4 100644 --- a/agent/internal/serverless/gateway_test.go +++ b/agent/internal/serverless/gateway_test.go @@ -5,9 +5,14 @@ import ( "context" "errors" "log" + "net" + "net/http" + "net/http/httptest" "slices" + "strconv" "strings" "sync" + "sync/atomic" "testing" "time" @@ -164,6 +169,160 @@ func assertProbeSet(t *testing.T, got []string, want ...string) { } } +func TestServeHTTPRetriesStaleUpstreamForSafeRequest(t *testing.T) { + var requests atomic.Int32 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(backend.Close) + + gateway := testProxyGateway(t, backend.Listener.Addr().String()) + gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}}) + + response := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "http://app.example.com/health", nil) + gateway.ServeHTTP(response, request) + + if response.Code != http.StatusNoContent { + t.Fatalf("status = %d, want %d", response.Code, http.StatusNoContent) + } + if requests.Load() != 1 { + t.Fatalf("backend requests = %d, want 1", requests.Load()) + } +} + +func TestServeHTTPDoesNotRetryUnsafeRequest(t *testing.T) { + var requests atomic.Int32 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(backend.Close) + + gateway := testProxyGateway(t, backend.Listener.Addr().String()) + gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}}) + + response := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "http://app.example.com/items", strings.NewReader("item")) + gateway.ServeHTTP(response, request) + + if response.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want %d", response.Code, http.StatusBadGateway) + } + if requests.Load() != 0 { + t.Fatalf("backend requests = %d, want 0", requests.Load()) + } +} + +func TestServeHTTPDoesNotRetryUpgradeRequest(t *testing.T) { + var requests atomic.Int32 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests.Add(1) + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(backend.Close) + + gateway := testProxyGateway(t, backend.Listener.Addr().String()) + gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}}) + + response := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "http://app.example.com/socket", nil) + request.Header.Set("Connection", "keep-alive, Upgrade") + request.Header.Set("Upgrade", "websocket") + gateway.ServeHTTP(response, request) + + if response.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want %d", response.Code, http.StatusBadGateway) + } + if requests.Load() != 0 { + t.Fatalf("backend requests = %d, want 0", requests.Load()) + } +} + +func TestServeHTTPDoesNotEvictCacheWhenRequestIsCancelled(t *testing.T) { + gateway := testProxyGateway(t, "127.0.0.1:0") + cached := []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}} + gateway.cacheUpstreams("app.example.com", cached) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + request := httptest.NewRequest(http.MethodGet, "http://app.example.com/", nil).WithContext(ctx) + gateway.ServeHTTP(httptest.NewRecorder(), request) + + upstreams, ok := gateway.cachedUpstreams("app.example.com") + if !ok || len(upstreams) != 1 || upstreams[0].Url != cached[0].Url { + t.Fatalf("cached upstreams = %+v, ok=%t; want original cache entry", upstreams, ok) + } +} + +func TestServeHTTPRetriesOnlyOnce(t *testing.T) { + retryAddress, accepts := acceptAndCloseAddress(t) + gateway := testProxyGateway(t, retryAddress) + gateway.cacheUpstreams("app.example.com", []agenthttp.ServerlessUpstream{{Url: "127.0.0.1:0"}}) + + response := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "http://app.example.com/", nil) + gateway.ServeHTTP(response, request) + + if response.Code != http.StatusBadGateway { + t.Fatalf("status = %d, want %d", response.Code, http.StatusBadGateway) + } + if _, ok := gateway.cachedUpstreams("app.example.com"); ok { + t.Fatal("failed retry left upstream cache populated") + } + if accepts.Load() != 1 { + t.Fatalf("retry upstream connections = %d, want 1", accepts.Load()) + } +} + +func testProxyGateway(t *testing.T, upstreamAddress string) *Gateway { + t.Helper() + host, portText, err := net.SplitHostPort(upstreamAddress) + if err != nil { + t.Fatalf("split upstream address: %v", err) + } + port, err := strconv.Atoi(portText) + if err != nil { + t.Fatalf("parse upstream port: %v", err) + } + state := testExpectedState("running") + state.Containers[0].IPAddress = host + state.Serverless.Routes[0].Port = port + state.Serverless.Routes[0].Upstreams = nil + runtime := &fakeRuntime{ + state: state, + containers: []container.Container{{ + ID: "ctr-local", + State: "running", + DeploymentID: "dep_local", + ServiceID: "svc_1", + }}, + } + return NewGateway(runtime) +} + +func acceptAndCloseAddress(t *testing.T) (string, *atomic.Int32) { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen for closing upstream: %v", err) + } + t.Cleanup(func() { listener.Close() }) + var accepts atomic.Int32 + go func() { + for { + conn, err := listener.Accept() + if err != nil { + return + } + accepts.Add(1) + conn.Close() + } + }() + return listener.Addr().String(), &accepts +} + func TestGetUpstreamsReturnsWhenFollowerContextIsCancelled(t *testing.T) { state := testExpectedState("stopped") g := NewGateway(&fakeRuntime{state: state})