From 128110f9b9dc1b18986b059b0befc696bceaee63 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:02:29 +0000 Subject: [PATCH 1/5] Fall back to the control plane on a stale session JWT A direct-to-VM 401/403 with a jwt query param evicts the cached route and retries the original request against the API. --- lib/browserrouting/route_cache.go | 45 ++++++++++++++++++++ lib/browserrouting/route_cache_test.go | 57 ++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index dfd41d7..de6a741 100644 --- a/lib/browserrouting/route_cache.go +++ b/lib/browserrouting/route_cache.go @@ -91,7 +91,11 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. if err != nil { return nil, err } + origURL := cloneURL(req.URL) + origHost := req.Host + origAuth := req.Header.Get("Authorization") sessionID, subresource, suffix, ok := parseDirectVMPath(req.URL.Path) + routed := false if ok { if matchesDirectVMPrefix(subresource+suffix, allowPrefixes) { route, ok := cache.Load(sessionID) @@ -114,6 +118,7 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. req.Host = base.Host req.URL.Path = joinURLPath(base.Path, subresource, suffix) req.URL.RawPath = "" + routed = true } } } @@ -122,6 +127,24 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. if err != nil { return res, err } + if routed && isStaleDirectVMAuthResponse(res, req) { + if sessionID != "" { + cache.Delete(sessionID) + } + req.URL = origURL + req.Host = origHost + if origAuth != "" { + req.Header.Set("Authorization", origAuth) + } + q := req.URL.Query() + q.Del("jwt") + req.URL.RawQuery = q.Encode() + res.Body.Close() + res, err = next(req) + if err != nil { + return res, err + } + } return finalizeResponse(res, cache, lifecycle) } } @@ -333,6 +356,28 @@ func matchesDirectVMPrefix(tail string, prefixes []string) bool { return false } +func isStaleDirectVMAuthResponse(res *http.Response, req *http.Request) bool { + if res == nil || req == nil || req.URL == nil { + return false + } + if res.StatusCode != http.StatusUnauthorized && res.StatusCode != http.StatusForbidden { + return false + } + return req.URL.Query().Get("jwt") != "" +} + +func cloneURL(u *url.URL) *url.URL { + if u == nil { + return nil + } + c := *u + if u.User != nil { + user := *u.User + c.User = &user + } + return &c +} + func joinURLPath(basePath, subresource, suffix string) string { base := "/" + strings.Trim(strings.TrimSpace(basePath), "/") if base == "/" { diff --git a/lib/browserrouting/route_cache_test.go b/lib/browserrouting/route_cache_test.go index 3304bfe..20456ed 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -394,3 +394,60 @@ func TestDirectVMRoutingMiddlewareDeleteWinsOverJSONCacheSniff(t *testing.T) { t.Fatal("expected delete response to leave cached route evicted") } } + +func TestDirectVMRoutingMiddlewareFallsBackOnStaleJWT(t *testing.T) { + cache := NewRouteCache() + cache.Store(Route{ + SessionID: "sess-1", + BaseURL: "https://browser.example/browser/kernel", + JWT: "jwt-123", + }) + + middleware := DirectVMRoutingMiddleware(cache, []string{"computer"}) + reqURL, err := url.Parse("https://api.example/browsers/sess-1/computer/screenshot") + if err != nil { + t.Fatal(err) + } + req := &http.Request{ + Method: http.MethodPost, + URL: reqURL, + Header: http.Header{"Authorization": []string{"Bearer sk_test"}}, + Host: "api.example", + } + + var calls []string + res, err := middleware(req, func(next *http.Request) (*http.Response, error) { + calls = append(calls, next.URL.String()) + if next.URL.Host == "browser.example" { + return &http.Response{ + StatusCode: http.StatusUnauthorized, + Body: io.NopCloser(strings.NewReader("Invalid JWT")), + }, nil + } + if next.Header.Get("Authorization") != "Bearer sk_test" { + t.Fatalf("expected restored authorization, got %q", next.Header.Get("Authorization")) + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("png")), + }, nil + }) + if err != nil { + t.Fatal(err) + } + if res.StatusCode != http.StatusOK { + t.Fatalf("expected 200 after fallback, got %d", res.StatusCode) + } + if len(calls) != 2 { + t.Fatalf("expected vm then control-plane call, got %v", calls) + } + if !strings.Contains(calls[0], "browser.example") || !strings.Contains(calls[0], "jwt=jwt-123") { + t.Fatalf("expected first call on VM with jwt, got %q", calls[0]) + } + if !strings.Contains(calls[1], "api.example/browsers/sess-1/computer/screenshot") { + t.Fatalf("expected second call on control plane, got %q", calls[1]) + } + if _, ok := cache.Load("sess-1"); ok { + t.Fatal("expected stale jwt to evict cached route") + } +} From a956805a830e06ade91c5161e243a7309e50d424 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:20:43 +0000 Subject: [PATCH 2/5] Rewind the request body before control-plane JWT fallback The first metro attempt consumes Body; GetBody restores it so computer/playwright POSTs retry with the original payload. --- lib/browserrouting/route_cache.go | 12 +++++- lib/browserrouting/route_cache_test.go | 55 ++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index de6a741..7ffcdea 100644 --- a/lib/browserrouting/route_cache.go +++ b/lib/browserrouting/route_cache.go @@ -139,7 +139,17 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. q := req.URL.Query() q.Del("jwt") req.URL.RawQuery = q.Encode() - res.Body.Close() + if res.Body != nil { + _ = res.Body.Close() + } + if req.GetBody != nil { + req.Body, err = req.GetBody() + if err != nil { + return nil, err + } + } else if req.Body != nil { + return res, nil + } res, err = next(req) if err != nil { return res, err diff --git a/lib/browserrouting/route_cache_test.go b/lib/browserrouting/route_cache_test.go index 20456ed..846b4e9 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -451,3 +451,58 @@ func TestDirectVMRoutingMiddlewareFallsBackOnStaleJWT(t *testing.T) { t.Fatal("expected stale jwt to evict cached route") } } + +func TestDirectVMRoutingMiddlewareRewindsBodyOnStaleJWTFallback(t *testing.T) { + cache := NewRouteCache() + cache.Store(Route{ + SessionID: "sess-1", + BaseURL: "https://browser.example/browser/kernel", + JWT: "jwt-123", + }) + + body := []byte(`{"code":"return 1"}`) + middleware := DirectVMRoutingMiddleware(cache, []string{"playwright"}) + reqURL, err := url.Parse("https://api.example/browsers/sess-1/playwright/execute") + if err != nil { + t.Fatal(err) + } + req := &http.Request{ + Method: http.MethodPost, + URL: reqURL, + Header: http.Header{"Authorization": []string{"Bearer sk_test"}}, + Host: "api.example", + Body: io.NopCloser(strings.NewReader(string(body))), + GetBody: func() (io.ReadCloser, error) { + return io.NopCloser(strings.NewReader(string(body))), nil + }, + ContentLength: int64(len(body)), + } + + var gotBodies []string + _, err = middleware(req, func(next *http.Request) (*http.Response, error) { + b, readErr := io.ReadAll(next.Body) + if readErr != nil { + return nil, readErr + } + gotBodies = append(gotBodies, string(b)) + if next.URL.Host == "browser.example" { + return &http.Response{ + StatusCode: http.StatusUnauthorized, + Body: io.NopCloser(strings.NewReader("Invalid JWT")), + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{"success":true}`)), + }, nil + }) + if err != nil { + t.Fatal(err) + } + if len(gotBodies) != 2 { + t.Fatalf("expected two bodies, got %v", gotBodies) + } + if gotBodies[0] != string(body) || gotBodies[1] != string(body) { + t.Fatalf("expected rewound body on fallback, got %v", gotBodies) + } +} From 82ad2ea1df38edef18be950131643ab78b42cdb5 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:25:42 +0000 Subject: [PATCH 3/5] Don't close the metro 401 until the body can be rewound If GetBody is missing, return the original auth response still readable. --- lib/browserrouting/route_cache.go | 17 ++++----- lib/browserrouting/route_cache_test.go | 48 ++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 8 deletions(-) diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index 7ffcdea..01aefd2 100644 --- a/lib/browserrouting/route_cache.go +++ b/lib/browserrouting/route_cache.go @@ -128,6 +128,15 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. return res, err } if routed && isStaleDirectVMAuthResponse(res, req) { + if req.GetBody == nil && req.Body != nil { + return res, nil + } + if req.GetBody != nil { + req.Body, err = req.GetBody() + if err != nil { + return res, err + } + } if sessionID != "" { cache.Delete(sessionID) } @@ -142,14 +151,6 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. if res.Body != nil { _ = res.Body.Close() } - if req.GetBody != nil { - req.Body, err = req.GetBody() - if err != nil { - return nil, err - } - } else if req.Body != nil { - return res, nil - } res, err = next(req) if err != nil { return res, err diff --git a/lib/browserrouting/route_cache_test.go b/lib/browserrouting/route_cache_test.go index 846b4e9..6c823f0 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -506,3 +506,51 @@ func TestDirectVMRoutingMiddlewareRewindsBodyOnStaleJWTFallback(t *testing.T) { t.Fatalf("expected rewound body on fallback, got %v", gotBodies) } } + +func TestDirectVMRoutingMiddlewareKeepsAuthResponseWhenBodyCannotRewind(t *testing.T) { + cache := NewRouteCache() + cache.Store(Route{ + SessionID: "sess-1", + BaseURL: "https://browser.example/browser/kernel", + JWT: "jwt-123", + }) + + middleware := DirectVMRoutingMiddleware(cache, []string{"playwright"}) + reqURL, err := url.Parse("https://api.example/browsers/sess-1/playwright/execute") + if err != nil { + t.Fatal(err) + } + req := &http.Request{ + Method: http.MethodPost, + URL: reqURL, + Header: http.Header{"Authorization": []string{"Bearer sk_test"}}, + Host: "api.example", + Body: io.NopCloser(strings.NewReader(`{"code":"return 1"}`)), + } + + var calls int + res, err := middleware(req, func(next *http.Request) (*http.Response, error) { + calls++ + _, _ = io.ReadAll(next.Body) + return &http.Response{ + StatusCode: http.StatusUnauthorized, + Body: io.NopCloser(strings.NewReader("Invalid JWT")), + }, nil + }) + if err != nil { + t.Fatal(err) + } + if calls != 1 { + t.Fatalf("expected no control-plane retry without GetBody, got %d calls", calls) + } + if res.StatusCode != http.StatusUnauthorized { + t.Fatalf("expected original 401, got %d", res.StatusCode) + } + got, err := io.ReadAll(res.Body) + if err != nil { + t.Fatalf("expected readable 401 body, got %v", err) + } + if string(got) != "Invalid JWT" { + t.Fatalf("expected Invalid JWT, got %q", got) + } +} From c342a58c82748d63e7043391323c605c780daebc Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:44:15 +0000 Subject: [PATCH 4/5] Keep the metro 401 when GetBody fails A rewind error must not mask the auth response or leak its body. Only mutate the request and close the 401 after rewind succeeds. --- lib/browserrouting/route_cache.go | 40 ++++++++++++-------- lib/browserrouting/route_cache_test.go | 51 ++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 15 deletions(-) diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index 01aefd2..a608a91 100644 --- a/lib/browserrouting/route_cache.go +++ b/lib/browserrouting/route_cache.go @@ -128,26 +128,12 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. return res, err } if routed && isStaleDirectVMAuthResponse(res, req) { - if req.GetBody == nil && req.Body != nil { + if !prepareControlPlaneFallback(req, origURL, origHost, origAuth) { return res, nil } - if req.GetBody != nil { - req.Body, err = req.GetBody() - if err != nil { - return res, err - } - } if sessionID != "" { cache.Delete(sessionID) } - req.URL = origURL - req.Host = origHost - if origAuth != "" { - req.Header.Set("Authorization", origAuth) - } - q := req.URL.Query() - q.Del("jwt") - req.URL.RawQuery = q.Encode() if res.Body != nil { _ = res.Body.Close() } @@ -367,6 +353,30 @@ func matchesDirectVMPrefix(tail string, prefixes []string) bool { return false } +func prepareControlPlaneFallback(req *http.Request, origURL *url.URL, origHost, origAuth string) bool { + if req.Body != nil && req.GetBody == nil { + return false + } + if req.GetBody != nil { + body, err := req.GetBody() + if err != nil { + return false + } + req.Body = body + } + req.URL = origURL + req.Host = origHost + if origAuth != "" { + req.Header.Set("Authorization", origAuth) + } else { + req.Header.Del("Authorization") + } + q := req.URL.Query() + q.Del("jwt") + req.URL.RawQuery = q.Encode() + return true +} + func isStaleDirectVMAuthResponse(res *http.Response, req *http.Request) bool { if res == nil || req == nil || req.URL == nil { return false diff --git a/lib/browserrouting/route_cache_test.go b/lib/browserrouting/route_cache_test.go index 6c823f0..6e0b403 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -554,3 +554,54 @@ func TestDirectVMRoutingMiddlewareKeepsAuthResponseWhenBodyCannotRewind(t *testi t.Fatalf("expected Invalid JWT, got %q", got) } } + +func TestDirectVMRoutingMiddlewareKeepsAuthResponseWhenGetBodyFails(t *testing.T) { + cache := NewRouteCache() + cache.Store(Route{ + SessionID: "sess-1", + BaseURL: "https://browser.example/browser/kernel", + JWT: "jwt-123", + }) + + middleware := DirectVMRoutingMiddleware(cache, []string{"playwright"}) + reqURL, err := url.Parse("https://api.example/browsers/sess-1/playwright/execute") + if err != nil { + t.Fatal(err) + } + req := &http.Request{ + Method: http.MethodPost, + URL: reqURL, + Header: http.Header{"Authorization": []string{"Bearer sk_test"}}, + Host: "api.example", + Body: io.NopCloser(strings.NewReader(`{"code":"return 1"}`)), + GetBody: func() (io.ReadCloser, error) { + return nil, io.ErrUnexpectedEOF + }, + } + + var calls int + res, err := middleware(req, func(next *http.Request) (*http.Response, error) { + calls++ + _, _ = io.ReadAll(next.Body) + return &http.Response{ + StatusCode: http.StatusUnauthorized, + Body: io.NopCloser(strings.NewReader("Invalid JWT")), + }, nil + }) + if err != nil { + t.Fatalf("expected original auth response, got err %v", err) + } + if calls != 1 { + t.Fatalf("expected no control-plane retry when GetBody fails, got %d calls", calls) + } + if res.StatusCode != http.StatusUnauthorized { + t.Fatalf("expected original 401, got %d", res.StatusCode) + } + got, err := io.ReadAll(res.Body) + if err != nil { + t.Fatalf("expected readable 401 body, got %v", err) + } + if string(got) != "Invalid JWT" { + t.Fatalf("expected Invalid JWT, got %q", got) + } +} From d5e0fc93093b7bcc929ac225808966170900f092 Mon Sep 17 00:00:00 2001 From: tnsardesai <18272584+tnsardesai@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:34:55 +0000 Subject: [PATCH 5/5] Only evict a stale JWT if it is still the cached one A later 401 must not delete a route that was refreshed in flight. --- lib/browserrouting/route_cache.go | 18 +++++++++- lib/browserrouting/route_cache_test.go | 49 ++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index a608a91..2ec3973 100644 --- a/lib/browserrouting/route_cache.go +++ b/lib/browserrouting/route_cache.go @@ -71,6 +71,21 @@ func (c *RouteCache) Delete(sessionID string) { delete(c.routes, sessionID) } +// DeleteIfJWT removes the cached route only when its JWT still matches. +func (c *RouteCache) DeleteIfJWT(sessionID, jwt string) bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + route, ok := c.routes[sessionID] + if !ok || route.JWT != jwt { + return false + } + delete(c.routes, sessionID) + return true +} + // DirectVMRoutingMiddleware rewrites allowlisted browser subresource requests to // the browser VM using cached base_url and jwt data. func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option.Middleware { @@ -128,11 +143,12 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. return res, err } if routed && isStaleDirectVMAuthResponse(res, req) { + failedJWT := req.URL.Query().Get("jwt") if !prepareControlPlaneFallback(req, origURL, origHost, origAuth) { return res, nil } if sessionID != "" { - cache.Delete(sessionID) + cache.DeleteIfJWT(sessionID, failedJWT) } if res.Body != nil { _ = res.Body.Close() diff --git a/lib/browserrouting/route_cache_test.go b/lib/browserrouting/route_cache_test.go index 6e0b403..9d58d86 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -452,6 +452,55 @@ func TestDirectVMRoutingMiddlewareFallsBackOnStaleJWT(t *testing.T) { } } +func TestDirectVMRoutingMiddlewareKeepsRefreshedRouteAfterStaleJWT(t *testing.T) { + cache := NewRouteCache() + cache.Store(Route{ + SessionID: "sess-1", + BaseURL: "https://browser.example/browser/kernel", + JWT: "jwt-123", + }) + + middleware := DirectVMRoutingMiddleware(cache, []string{"computer"}) + reqURL, err := url.Parse("https://api.example/browsers/sess-1/computer/screenshot") + if err != nil { + t.Fatal(err) + } + req := &http.Request{ + Method: http.MethodPost, + URL: reqURL, + Header: http.Header{"Authorization": []string{"Bearer sk_test"}}, + Host: "api.example", + } + + _, err = middleware(req, func(next *http.Request) (*http.Response, error) { + if next.URL.Host == "browser.example" { + cache.Store(Route{ + SessionID: "sess-1", + BaseURL: "https://browser.example/browser/kernel", + JWT: "jwt-FRESH", + }) + return &http.Response{ + StatusCode: http.StatusUnauthorized, + Body: io.NopCloser(strings.NewReader("Invalid JWT")), + }, nil + } + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("png")), + }, nil + }) + if err != nil { + t.Fatal(err) + } + route, ok := cache.Load("sess-1") + if !ok { + t.Fatal("expected refreshed route to survive stale jwt fallback") + } + if route.JWT != "jwt-FRESH" { + t.Fatalf("expected jwt-FRESH, got %q", route.JWT) + } +} + func TestDirectVMRoutingMiddlewareRewindsBodyOnStaleJWTFallback(t *testing.T) { cache := NewRouteCache() cache.Store(Route{