diff --git a/browser_routing.go b/browser_routing.go index b5af277..d0f756b 100644 --- a/browser_routing.go +++ b/browser_routing.go @@ -83,10 +83,11 @@ func browserRouteFromRef(ref browserrouting.Ref) (browserrouting.Route, bool) { func browserRoutingSubresourcesFromEnv() []string { raw, ok := os.LookupEnv(browserRoutingSubresourcesEnv) if !ok { - // Path prefixes eligible for direct-to-VM routing. "telemetry/stream" is - // the live SSE endpoint (served by the VM); "telemetry/events" is a - // historical read served by the control plane (S2) and must NOT be here. - return []string{"curl", "telemetry/stream", "computer", "playwright", "process"} + // Path prefixes eligible for direct-to-VM routing. "fs" intentionally + // covers every filesystem endpoint, while "logs/stream" excludes other + // current or future logs endpoints. "telemetry/events" is a historical + // read served by the control plane (S2), unlike the VM's live stream. + return []string{"curl", "telemetry/stream", "computer", "playwright", "process", "fs", "logs/stream"} } if strings.TrimSpace(raw) == "" { return []string{} diff --git a/browser_routing_test.go b/browser_routing_test.go index f39e5f2..3a53308 100644 --- a/browser_routing_test.go +++ b/browser_routing_test.go @@ -1,12 +1,16 @@ package kernel import ( + "bytes" "context" "encoding/json" "errors" + "io" "net/http" "net/http/httptest" "os" + "slices" + "strings" "sync" "testing" "time" @@ -266,8 +270,9 @@ func TestBrowserRoutingSubresourcesFromEnvDefaults(t *testing.T) { } _ = os.Setenv(browserRoutingSubresourcesEnv, original) }) - if got := browserRoutingSubresourcesFromEnv(); len(got) != 5 || got[0] != "curl" || got[1] != "telemetry/stream" || got[2] != "computer" || got[3] != "playwright" || got[4] != "process" { - t.Fatalf("expected default subresources [curl telemetry/stream computer playwright process], got %#v", got) + want := []string{"curl", "telemetry/stream", "computer", "playwright", "process", "fs", "logs/stream"} + if got := browserRoutingSubresourcesFromEnv(); !slices.Equal(got, want) { + t.Fatalf("expected default subresources %v, got %#v", want, got) } t.Setenv(browserRoutingSubresourcesEnv, "") @@ -276,15 +281,9 @@ func TestBrowserRoutingSubresourcesFromEnvDefaults(t *testing.T) { } t.Setenv(browserRoutingSubresourcesEnv, "curl, process") - got := browserRoutingSubresourcesFromEnv() - want := []string{"curl", "process"} - if len(got) != len(want) { - t.Fatalf("expected %v, got %#v", want, got) - } - for i := range want { - if got[i] != want[i] { - t.Fatalf("expected %v, got %#v", want, got) - } + wantConfigured := []string{"curl", "process"} + if got := browserRoutingSubresourcesFromEnv(); !slices.Equal(got, wantConfigured) { + t.Fatalf("expected %v, got %#v", wantConfigured, got) } } @@ -385,46 +384,89 @@ func TestBrowserRoutingDefaultsRouteToVM(t *testing.T) { } } -func TestBrowserRoutingDefaultsKeepFsAndTelemetryEventsOnAPI(t *testing.T) { - original, ok := os.LookupEnv(browserRoutingSubresourcesEnv) - if err := os.Unsetenv(browserRoutingSubresourcesEnv); err != nil { - t.Fatal(err) - } - t.Cleanup(func() { - if !ok { - _ = os.Unsetenv(browserRoutingSubresourcesEnv) - return - } - _ = os.Setenv(browserRoutingSubresourcesEnv, original) - }) +func TestBrowserRoutingDefaultsKeepControlPlaneSubresourcesOnAPI(t *testing.T) { + unsetBrowserRoutingEnv(t) var paths []string var srv *httptest.Server srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { paths = append(paths, r.URL.Path) - switch { - case r.URL.Path == "/browsers": + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/browsers": + _ = json.NewEncoder(w).Encode(map[string]any{ + "session_id": "sess-1", + "base_url": srv.URL + "/browser/kernel", + "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=token-abc", + }) + default: + _ = json.NewEncoder(w).Encode([]any{}) + } + })) + defer srv.Close() + + client := NewClient( + option.WithBaseURL(srv.URL), + option.WithAPIKey("sk_test"), + option.WithHTTPClient(srv.Client()), + ) + + if _, err := client.Browsers.New(context.Background(), BrowserNewParams{}); err != nil { + t.Fatal(err) + } + if _, err := client.Browsers.Telemetry.Events(context.Background(), "sess-1", BrowserTelemetryEventsParams{}); err != nil { + t.Fatal(err) + } + if _, err := client.Browsers.Replays.List(context.Background(), "sess-1"); err != nil { + t.Fatal(err) + } + + want := []string{ + "/browsers", + "/browsers/sess-1/telemetry/events", + "/browsers/sess-1/replays", + } + if !slices.Equal(paths, want) { + t.Fatalf("expected paths %v, got %v", want, paths) + } +} + +func TestBrowserRoutingDefaultsRouteFsAndLogsToVM(t *testing.T) { + unsetBrowserRoutingEnv(t) + + type call struct { + Path string + Auth string + Body string + } + var calls []call + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + calls = append(calls, call{Path: r.URL.RequestURI(), Auth: r.Header.Get("Authorization"), Body: string(body)}) + + switch r.URL.Path { + case "/browsers": w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(map[string]any{ "session_id": "sess-1", "base_url": srv.URL + "/browser/kernel", "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=token-abc", }) - case r.URL.Path == "/browsers/sess-1/fs/read_file": + case "/browser/kernel/fs/read_file": w.Header().Set("Content-Type", "application/octet-stream") w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte("x")) - case r.URL.Path == "/browsers/sess-1/telemetry/events": + _, _ = w.Write([]byte{0x00, 0x01}) + case "/browser/kernel/fs/list_files": w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode([]any{}) + case "/browser/kernel/logs/stream", "/browser/kernel/fs/watch/watch-1/events": + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("data: {\"event\":\"log\",\"message\":\"hello\",\"timestamp\":\"2020-01-01T00:00:00Z\"}\n\n")) default: w.Header().Set("Content-Type", "application/json") - _ = json.NewEncoder(w).Encode(map[string]any{ - "duration_ms": 1, - "exit_code": 0, - "stderr_b64": "", - "stdout_b64": "", - }) + _ = json.NewEncoder(w).Encode(map[string]any{}) } })) defer srv.Close() @@ -438,28 +480,279 @@ func TestBrowserRoutingDefaultsKeepFsAndTelemetryEventsOnAPI(t *testing.T) { if _, err := client.Browsers.New(context.Background(), BrowserNewParams{}); err != nil { t.Fatal(err) } - fsResp, err := client.Browsers.Fs.ReadFile(context.Background(), "sess-1", BrowserFReadFileParams{Path: "/tmp/x"}) + if _, err := client.Browsers.Fs.ListFiles(context.Background(), "sess-1", BrowserFListFilesParams{Path: "/tmp"}); err != nil { + t.Fatal(err) + } + readRes, err := client.Browsers.Fs.ReadFile(context.Background(), "sess-1", BrowserFReadFileParams{Path: "/tmp/x"}) if err != nil { t.Fatal(err) } - if fsResp != nil { - _ = fsResp.Body.Close() + contents, err := io.ReadAll(readRes.Body) + if err != nil { + t.Fatal(err) } - if _, err := client.Browsers.Telemetry.Events(context.Background(), "sess-1", BrowserTelemetryEventsParams{}); err != nil { + _ = readRes.Body.Close() + if !bytes.Equal(contents, []byte{0x00, 0x01}) { + t.Fatalf("expected binary body from VM, got %v", contents) + } + if err := client.Browsers.Fs.WriteFile(context.Background(), "sess-1", strings.NewReader("payload"), BrowserFWriteFileParams{ + Path: "/tmp/x", + Mode: String("600"), + }); err != nil { + t.Fatal(err) + } + if err := client.Browsers.Fs.Upload(context.Background(), "sess-1", BrowserFUploadParams{ + Files: []BrowserFUploadParamsFile{ + {DestPath: "/tmp/one", File: strings.NewReader("one")}, + {DestPath: "/tmp/two", File: strings.NewReader("two")}, + }, + }); err != nil { + t.Fatal(err) + } + watchStream := client.Browsers.Fs.Watch.EventsStreaming(context.Background(), "watch-1", BrowserFWatchEventsParams{IDOrName: "sess-1"}) + for watchStream.Next() { + } + if err := watchStream.Err(); err != nil { + t.Fatal(err) + } + logStream := client.Browsers.Logs.StreamStreaming(context.Background(), "sess-1", BrowserLogStreamParams{ + Source: BrowserLogStreamParamsSourcePath, + Path: String("/var/log/x"), + Follow: Bool(true), + }) + for logStream.Next() { + } + if err := logStream.Err(); err != nil { t.Fatal(err) } - want := []string{ + wantPaths := []string{ "/browsers", - "/browsers/sess-1/fs/read_file", - "/browsers/sess-1/telemetry/events", + "/browser/kernel/fs/list_files?jwt=token-abc&path=%2Ftmp", + "/browser/kernel/fs/read_file?jwt=token-abc&path=%2Ftmp%2Fx", + "/browser/kernel/fs/write_file?jwt=token-abc&mode=600&path=%2Ftmp%2Fx", + "/browser/kernel/fs/upload?jwt=token-abc", + "/browser/kernel/fs/watch/watch-1/events?jwt=token-abc", + "/browser/kernel/logs/stream?follow=true&jwt=token-abc&path=%2Fvar%2Flog%2Fx&source=path", + } + gotPaths := make([]string, 0, len(calls)) + for _, c := range calls { + gotPaths = append(gotPaths, c.Path) + } + if !slices.Equal(gotPaths, wantPaths) { + t.Fatalf("expected paths %v, got %v", wantPaths, gotPaths) + } + for _, c := range calls[1:] { + if c.Auth != "" { + t.Fatalf("expected authorization header removed for %q, got %q", c.Path, c.Auth) + } } - if len(paths) != len(want) { - t.Fatalf("expected paths %v, got %v", want, paths) + if calls[3].Body != "payload" { + t.Fatalf("expected write_file body forwarded to the VM, got %q", calls[3].Body) + } + uploadBody := calls[4].Body + // The VM accepts dot-indexed multipart array names (files..). + for _, want := range []string{`name="files.0.dest_path"`, `name="files.0.file"`, `name="files.1.dest_path"`} { + if !strings.Contains(uploadBody, want) { + t.Fatalf("expected upload body to contain %s, got %q", want, uploadBody) + } + } +} + +func TestBrowserRoutingLogsStreamPreservesContextCancellation(t *testing.T) { + unsetBrowserRoutingEnv(t) + + handlerDone := make(chan struct{}) + requestCanceled := make(chan struct{}) + streamPath := make(chan string, 1) + var requestCanceledOnce sync.Once + var streamPathOnce sync.Once + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/browsers" { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "session_id": "sess-1", + "base_url": srv.URL + "/browser/kernel", + "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=token-abc", + }) + return + } + + streamPathOnce.Do(func() { streamPath <- r.URL.RequestURI() }) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + select { + case <-r.Context().Done(): + requestCanceledOnce.Do(func() { close(requestCanceled) }) + case <-handlerDone: + } + })) + t.Cleanup(func() { + close(handlerDone) + srv.Close() + }) + + client := NewClient( + option.WithBaseURL(srv.URL), + option.WithAPIKey("sk_test"), + option.WithHTTPClient(srv.Client()), + ) + if _, err := client.Browsers.New(context.Background(), BrowserNewParams{}); err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithCancel(context.Background()) + stream := client.Browsers.Logs.StreamStreaming(ctx, "sess-1", BrowserLogStreamParams{ + Source: BrowserLogStreamParamsSourceSupervisor, + SupervisorProcess: String("chromium"), + }) + if got := <-streamPath; got != "/browser/kernel/logs/stream?jwt=token-abc&source=supervisor&supervisor_process=chromium" { + t.Fatalf("expected direct VM logs path, got %q", got) } - for i := range want { - if paths[i] != want[i] { - t.Fatalf("expected paths %v, got %v", want, paths) + next := make(chan bool, 1) + go func() { + next <- stream.Next() + }() + + cancel() + select { + case got := <-next: + if got { + t.Fatal("expected canceled stream to stop") } + if !errors.Is(stream.Err(), context.Canceled) { + t.Fatalf("expected context cancellation, got %v", stream.Err()) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for canceled stream to stop") + } + + select { + case <-requestCanceled: + case <-time.After(time.Second): + t.Fatal("timed out waiting for direct VM request cancellation") + } +} + +func TestBrowserRoutingFsFallsBackToControlPlaneOnStaleJWT(t *testing.T) { + unsetBrowserRoutingEnv(t) + + type call struct { + Path string + Auth string + Body string } + var calls []call + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + calls = append(calls, call{Path: r.URL.RequestURI(), Auth: r.Header.Get("Authorization"), Body: string(body)}) + + switch { + case r.URL.Path == "/browsers": + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "session_id": "sess-1", + "base_url": srv.URL + "/browser/kernel", + "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=token-abc", + }) + case strings.HasPrefix(r.URL.Path, "/browser/kernel/"): + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte("Invalid JWT")) + default: + w.WriteHeader(http.StatusCreated) + } + })) + defer srv.Close() + + client := NewClient( + option.WithBaseURL(srv.URL), + option.WithAPIKey("sk_test"), + option.WithHTTPClient(srv.Client()), + ) + + if _, err := client.Browsers.New(context.Background(), BrowserNewParams{}); err != nil { + t.Fatal(err) + } + if err := client.Browsers.Fs.Upload(context.Background(), "sess-1", BrowserFUploadParams{ + Files: []BrowserFUploadParamsFile{{DestPath: "/tmp/one", File: strings.NewReader("one")}}, + }); err != nil { + t.Fatal(err) + } + + if len(calls) != 3 { + t.Fatalf("expected browser create plus VM and control-plane upload, got %#v", calls) + } + if calls[1].Path != "/browser/kernel/fs/upload?jwt=token-abc" { + t.Fatalf("expected direct VM upload first, got %q", calls[1].Path) + } + if calls[2].Path != "/browsers/sess-1/fs/upload" { + t.Fatalf("expected control-plane fallback, got %q", calls[2].Path) + } + if calls[2].Auth != "Bearer sk_test" { + t.Fatalf("expected API key on the fallback request, got %q", calls[2].Auth) + } + if calls[2].Body != calls[1].Body { + t.Fatalf("expected the fallback to replay the upload body") + } + if _, ok := client.BrowserRouteCache.Load("sess-1"); ok { + t.Fatal("expected the stale route to be evicted") + } +} + +func TestBrowserRoutingEnvOverrideKeepsFsAndLogsOnControlPlane(t *testing.T) { + t.Setenv(browserRoutingSubresourcesEnv, "computer") + + var paths []string + var srv *httptest.Server + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + paths = append(paths, r.URL.Path) + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/browsers": + _ = json.NewEncoder(w).Encode(map[string]any{ + "session_id": "sess-1", + "base_url": srv.URL + "/browser/kernel", + "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=token-abc", + }) + default: + _ = json.NewEncoder(w).Encode([]any{}) + } + })) + defer srv.Close() + + client := NewClient( + option.WithBaseURL(srv.URL), + option.WithAPIKey("sk_test"), + option.WithHTTPClient(srv.Client()), + ) + + if _, err := client.Browsers.New(context.Background(), BrowserNewParams{}); err != nil { + t.Fatal(err) + } + if _, err := client.Browsers.Fs.ListFiles(context.Background(), "sess-1", BrowserFListFilesParams{Path: "/tmp"}); err != nil { + t.Fatal(err) + } + + want := []string{"/browsers", "/browsers/sess-1/fs/list_files"} + if !slices.Equal(paths, want) { + t.Fatalf("expected paths %v, got %v", want, paths) + } +} + +func unsetBrowserRoutingEnv(t *testing.T) { + t.Helper() + original, ok := os.LookupEnv(browserRoutingSubresourcesEnv) + if err := os.Unsetenv(browserRoutingSubresourcesEnv); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if !ok { + _ = os.Unsetenv(browserRoutingSubresourcesEnv) + return + } + _ = os.Setenv(browserRoutingSubresourcesEnv, original) + }) } diff --git a/lib/browserrouting/route_cache.go b/lib/browserrouting/route_cache.go index 2ec3973..1756b1f 100644 --- a/lib/browserrouting/route_cache.go +++ b/lib/browserrouting/route_cache.go @@ -143,13 +143,15 @@ func DirectVMRoutingMiddleware(cache *RouteCache, subresources []string) option. return res, err } if routed && isStaleDirectVMAuthResponse(res, req) { - failedJWT := req.URL.Query().Get("jwt") + if sessionID != "" { + cache.DeleteIfJWT(sessionID, req.URL.Query().Get("jwt")) + } + // Without a replayable body the fallback would send a truncated request, + // so surface the auth failure instead; the route is already evicted, so a + // later call from the caller goes to the control plane. if !prepareControlPlaneFallback(req, origURL, origHost, origAuth) { return res, nil } - if 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 ff5c80e..d46afe1 100644 --- a/lib/browserrouting/route_cache_test.go +++ b/lib/browserrouting/route_cache_test.go @@ -71,7 +71,16 @@ func TestDirectVMRoutingMiddlewareAllowlistMatching(t *testing.T) { {"playwright execute -> VM", "/browsers/sess-1/playwright/execute", true}, {"process exec -> VM", "/browsers/sess-1/process/exec", true}, {"process stdout stream -> VM", "/browsers/sess-1/process/proc-1/stdout/stream", true}, - {"non-allowlisted subresource -> control plane", "/browsers/sess-1/fs/read", false}, + {"fs read_file -> VM", "/browsers/sess-1/fs/read_file", true}, + {"fs watch events -> VM", "/browsers/sess-1/fs/watch/watch-1/events", true}, + {"fs-prefixed segment not matched", "/browsers/sess-1/fsx/read_file", false}, + {"logs stream -> VM", "/browsers/sess-1/logs/stream", true}, + {"logs stream suffix -> VM", "/browsers/sess-1/logs/stream/x", true}, + {"bare logs -> control plane", "/browsers/sess-1/logs", false}, + {"logs history -> control plane", "/browsers/sess-1/logs/history", false}, + {"logs-prefixed segment not matched", "/browsers/sess-1/logstream", false}, + {"extensions -> control plane", "/browsers/sess-1/extensions", false}, + {"replays -> control plane", "/browsers/sess-1/replays/rec-1", false}, } for _, tc := range cases { @@ -82,7 +91,7 @@ func TestDirectVMRoutingMiddlewareAllowlistMatching(t *testing.T) { BaseURL: "https://browser.example/browser/kernel", JWT: "jwt-123", }) - middleware := DirectVMRoutingMiddleware(cache, []string{"curl", "telemetry/stream", "computer", "playwright", "process"}) + middleware := DirectVMRoutingMiddleware(cache, []string{"curl", "telemetry/stream", "computer", "playwright", "process", "fs", "logs/stream"}) reqURL, err := url.Parse("https://api.example" + tc.path) if err != nil { @@ -593,6 +602,9 @@ func TestDirectVMRoutingMiddlewareKeepsAuthResponseWhenBodyCannotRewind(t *testi if calls != 1 { t.Fatalf("expected no control-plane retry without GetBody, got %d calls", calls) } + if _, ok := cache.Load("sess-1"); ok { + t.Fatal("expected the stale route to be evicted even without a retry") + } if res.StatusCode != http.StatusUnauthorized { t.Fatalf("expected original 401, got %d", res.StatusCode) } @@ -644,6 +656,9 @@ func TestDirectVMRoutingMiddlewareKeepsAuthResponseWhenGetBodyFails(t *testing.T if calls != 1 { t.Fatalf("expected no control-plane retry when GetBody fails, got %d calls", calls) } + if _, ok := cache.Load("sess-1"); ok { + t.Fatal("expected the stale route to be evicted even without a retry") + } if res.StatusCode != http.StatusUnauthorized { t.Fatalf("expected original 401, got %d", res.StatusCode) }