diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index dfd41d7..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 { @@ -91,7 +106,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 +133,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 +142,22 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. if err != nil { 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.DeleteIfJWT(sessionID, failedJWT) + } + if res.Body != nil { + _ = res.Body.Close() + } + res, err = next(req) + if err != nil { + return res, err + } + } return finalizeResponse(res, cache, lifecycle) } } @@ -333,6 +369,52 @@ 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 + } + 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..9d58d86 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -394,3 +394,263 @@ 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") + } +} + +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{ + 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) + } +} + +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) + } +} + +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) + } +}