diff --git a/acceptance/README.md b/acceptance/README.md index 9cd4d2ed0a0..9c5ce55403e 100644 --- a/acceptance/README.md +++ b/acceptance/README.md @@ -15,6 +15,19 @@ The cloud run is keyed off `CLOUD_ENV`; see `getSkipReason` in `acceptance/accep To run tests against a real workspace: `deco env run -i -n aws-prod-ucws -- ` (requires the `deco` tool and access to a test env). +## Fault injection + +Set `Proxy = true` in `test.toml` to use `fault.py` against either the fake server or a real workspace. Add `Cloud = true` for cloud runs. Recording is optional. + +```bash +fault.py "POST /api/2.0/workspace/mkdirs" 403 0 1 +fault.py --after-handler "POST /api/2.2/jobs/run-now" 503 0 1 +``` + +A normal fault prevents the upstream call. `--after-handler` forwards the call first and replaces its response, preserving its effects. Rules match the incoming test token; the proxy uses separate upstream credentials. Offsets and counts include client retries. Use precise patterns or `--body-contains` to select a resource, and clean up real resources after testing a failed response. + +The proxy requires localhost access and cannot be used with `RunsOnDbr`. + ## Authoring To author a test, diff --git a/acceptance/acceptance_test.go b/acceptance/acceptance_test.go index 29f1ccde3a3..a333e17a322 100644 --- a/acceptance/acceptance_test.go +++ b/acceptance/acceptance_test.go @@ -919,6 +919,8 @@ func runTest(t *testing.T, args := []string{"bash", "-euo", "pipefail", EntryPointScript} cmd := exec.CommandContext(ctx, args[0], args[1:]...) + require.False(t, isTruePtr(config.RunsOnDbr) && LogRequests, + "RunsOnDbr cannot be combined with -logrequests: the proxy requires localhost access") cfg, user, workspaceID := internal.PrepareServerAndClient(t, config, LogRequests, outputDir, testEnv) testdiff.PrepareReplacementsUser(t, &repls, user) testdiff.PrepareReplacementsWorkspaceConfig(t, &repls, cfg) diff --git a/acceptance/bin/fault.py b/acceptance/bin/fault.py index 5b75c0a5907..6ea3c748ed0 100755 --- a/acceptance/bin/fault.py +++ b/acceptance/bin/fault.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Set up a fault rule on the testserver for the current test token. +"""Set up a fault rule on the testserver or test proxy for the current test token. Usage: fault.py [--after-handler] [--body-contains SUBSTR] PATTERN STATUS_CODE OFFSET TIMES [ERROR_CODE] @@ -16,6 +16,7 @@ The rule is scoped to the current DATABRICKS_TOKEN so it only affects the test that registers it, even when tests share a server. +Cloud runs require Proxy = true in test.toml. """ import json diff --git a/acceptance/internal/config.go b/acceptance/internal/config.go index 6ea79a9578a..8cbe0f68b42 100644 --- a/acceptance/internal/config.go +++ b/acceptance/internal/config.go @@ -78,8 +78,8 @@ type TestConfig struct { // out.requests.txt RecordRequests *bool - // If true, local runs route the CLI through the recording proxy (libs/testproxy) - // instead of straight to the testserver, matching the cloud topology. + // If true, route the CLI through libs/testproxy for fault injection, + // independently of recording. Supported locally and on cloud. Proxy *bool // List of request headers to include when recording requests. @@ -252,10 +252,10 @@ func LoadConfig(t *testing.T, dir string) (TestConfig, string) { // validateConfig checks for incompatible configuration combinations. func validateConfig(t *testing.T, config TestConfig, configPath string) { - // RunsOnDbr and RecordRequests are incompatible because serverless does not + // RunsOnDbr cannot use the proxy because serverless does not // allow access to localhost ports, which the test proxy server requires. - if isTruePtr(config.RunsOnDbr) && isTruePtr(config.RecordRequests) { - t.Fatalf("Invalid config %s: RunsOnDbr and RecordRequests cannot both be true. "+ + if isTruePtr(config.RunsOnDbr) && (isTruePtr(config.RecordRequests) || isTruePtr(config.Proxy)) { + t.Fatalf("Invalid config %s: RunsOnDbr cannot be combined with RecordRequests or Proxy. "+ "Serverless does not allow access to localhost ports, which the test proxy server requires.", configPath) } diff --git a/acceptance/internal/prepare_server.go b/acceptance/internal/prepare_server.go index 3890aa52d5a..59b564372de 100644 --- a/acceptance/internal/prepare_server.go +++ b/acceptance/internal/prepare_server.go @@ -129,10 +129,10 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o cfg := w.Config // If we are running in a cloud environment AND we need to intercept requests - // (for recording or logging), start a proxy server. - if recordRequests || logRequests { - // An empty config resolves the real workspace from the environment. - host := startProxyServer(t, &sdkconfig.Config{}, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + // (for fault injection, recording or logging), start a proxy server. + if isTruePtr(config.Proxy) || recordRequests || logRequests { + // Reuse the resolved cloud config for upstream auth and workspace routing. + host := startProxyServer(t, cfg, token, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) cfg = &sdkconfig.Config{ Host: host, Token: token, @@ -155,7 +155,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o host := startProxyServer(t, &sdkconfig.Config{ Host: upstream, Token: token, - }, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) + }, token, recordRequests, logRequests, config.IncludeRequestHeaders, outputDir) cfg := &sdkconfig.Config{ Host: host, @@ -284,12 +284,13 @@ func startLocalServer(t *testing.T, func startProxyServer(t *testing.T, upstream *sdkconfig.Config, + token string, recordRequests bool, logRequests bool, includeHeaders []string, outputDir string, ) string { - s := testproxy.New(t, upstream) + s := testproxy.New(t, upstream, token) // Record API requests in out.requests.txt if RecordRequests is true in test.toml. if recordRequests { diff --git a/acceptance/selftest/fault_proxy/out.test.toml b/acceptance/selftest/fault_proxy/out.test.toml new file mode 100644 index 00000000000..c7c6f58ed6e --- /dev/null +++ b/acceptance/selftest/fault_proxy/out.test.toml @@ -0,0 +1 @@ +Cloud = true diff --git a/acceptance/selftest/fault_proxy/output.txt b/acceptance/selftest/fault_proxy/output.txt new file mode 100644 index 00000000000..4ae843483cc --- /dev/null +++ b/acceptance/selftest/fault_proxy/output.txt @@ -0,0 +1,38 @@ + +=== a fault before forwarding prevents the mutation +>>> musterr [CLI] workspace mkdirs /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/blocked +Error: Fault injected by test. + +Host: [DATABRICKS_URL] +Auth type: Personal Access Token (pat) + +Next steps: + - Verify you have the required permissions for this operation + - Check your identity: databricks auth describe + - Consider setting up a profile: databricks auth login --profile +blocked directory does not exist + +=== an after-handler fault preserves the upstream mutation +>>> musterr [CLI] workspace mkdirs /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/committed +Error: Fault injected by test. + +Host: [DATABRICKS_URL] +Auth type: Personal Access Token (pat) + +Next steps: + - Verify you have the required permissions for this operation + - Check your identity: databricks auth describe + - Consider setting up a profile: databricks auth login --profile + +>>> [CLI] workspace get-status /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/committed +{ + "object_type": "DIRECTORY" +} + +=== the CLI recovers from a retryable error after the mutation +>>> [CLI] workspace mkdirs /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/retried + +>>> [CLI] workspace get-status /Users/[USERNAME]/fault-proxy-[UNIQUE_NAME]/retried +{ + "object_type": "DIRECTORY" +} diff --git a/acceptance/selftest/fault_proxy/script b/acceptance/selftest/fault_proxy/script new file mode 100644 index 00000000000..81b0aaef615 --- /dev/null +++ b/acceptance/selftest/fault_proxy/script @@ -0,0 +1,22 @@ +root="/Users/$CURRENT_USER_NAME/fault-proxy-$UNIQUE_NAME" +cleanup() { + $CLI workspace delete "$root" --recursive +} +$CLI workspace mkdirs "$root" +trap cleanup EXIT + +title "a fault before forwarding prevents the mutation" +fault.py --body-contains "$root/blocked" "POST /api/2.0/workspace/mkdirs" 403 0 1 +trace musterr $CLI workspace mkdirs "$root/blocked" +musterr $CLI workspace get-status "$root/blocked" >/dev/null 2>&1 +echo "blocked directory does not exist" + +title "an after-handler fault preserves the upstream mutation" +fault.py --after-handler --body-contains "$root/committed" "POST /api/2.0/workspace/mkdirs" 403 0 1 +trace musterr $CLI workspace mkdirs "$root/committed" +trace $CLI workspace get-status "$root/committed" | jq '{object_type}' + +title "the CLI recovers from a retryable error after the mutation" +fault.py --after-handler --body-contains "$root/retried" "POST /api/2.0/workspace/mkdirs" 503 0 1 +trace $CLI workspace mkdirs "$root/retried" +trace $CLI workspace get-status "$root/retried" | jq '{object_type}' diff --git a/acceptance/selftest/fault_proxy/test.toml b/acceptance/selftest/fault_proxy/test.toml new file mode 100644 index 00000000000..dd645d54e45 --- /dev/null +++ b/acceptance/selftest/fault_proxy/test.toml @@ -0,0 +1,4 @@ +Cloud = true +Proxy = true +RecordRequests = false +Env.MSYS_NO_PATHCONV = '1' diff --git a/libs/testproxy/server.go b/libs/testproxy/server.go index be5bcfd4aa7..6cb547c6f32 100644 --- a/libs/testproxy/server.go +++ b/libs/testproxy/server.go @@ -2,186 +2,172 @@ package testproxy import ( "bytes" - "errors" + "crypto/tls" + "fmt" + "io" + "maps" "net/http" "net/http/httptest" + "net/http/httputil" + "net/url" + "strings" "github.com/databricks/cli/internal/testutil" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/databricks-sdk-go/httpclient" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +const faultEndpoint = "/__testserver/fault" + type ProxyServer struct { *httptest.Server - t testutil.TestingT - - apiClient *httpclient.ApiClient + t testutil.TestingT + upstream *url.URL + transport http.RoundTripper + token string + faults *testserver.FaultRules RequestCallback func(request *testserver.Request) ResponseCallback func(request *testserver.Request, response *testserver.EncodedResponse) } -// New creates a reverse proxy server that sits in front of the upstream -// workspace. This is useful for recording API requests and responses in -// integration tests. -// -// Note: We cannot directly proxy the request from a localhost URL to a real -// workspace as is. This is because auth resolution in the Databricks SDK relies -// what the URL actually looks like to determine the auth method to use. -// For example, in OAuth flows, the SDK can make requests to different Microsoft -// OAuth endpoints based on the nature of the URL. -// For reference, see: +// New creates a reverse proxy with token-scoped fault injection and optional recording. +// The upstream config resolves real workspace credentials; token identifies the test +// locally and is never forwarded. Resolving auth against localhost would break OAuth: // https://github.com/databricks/databricks-sdk-go/blob/79e4b3a6e9b0b7dcb1af9ad4025deb447b01d933/common/environment/environments.go#L57 -// -// An upstream with no fields set resolves a real workspace and its auth from the -// environment; an explicit host and token point at a testserver instead. -func New(t testutil.TestingT, upstream *config.Config) *ProxyServer { - s := &ProxyServer{ - t: t, - } - - clientCfg, err := config.HTTPClientConfigFromConfig(upstream) +func New(t testutil.TestingT, upstream *config.Config, token string) *ProxyServer { + require.NoError(t, upstream.EnsureResolved()) + upstreamURL, err := url.Parse(upstream.Host) require.NoError(t, err) - s.apiClient = httpclient.NewApiClient(clientCfg) - - // Set up the proxy handler as the default handler for all requests. - server := httptest.NewServer(http.HandlerFunc(s.proxyToUpstream)) - t.Cleanup(server.Close) - s.Server = server + transport := upstream.HTTPTransport + if transport == nil { + tr := http.DefaultTransport.(*http.Transport).Clone() + tr.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, InsecureSkipVerify: upstream.InsecureSkipVerify} //nolint:gosec // Honor the upstream SDK config. + t.Cleanup(tr.CloseIdleConnections) + transport = tr + } + s := &ProxyServer{ + t: t, + upstream: upstreamURL, + transport: &authenticatedTransport{upstream: upstream, next: transport}, + token: token, + faults: testserver.NewFaultRules(), + } + s.Server = httptest.NewServer(http.HandlerFunc(s.proxyToUpstream)) + t.Cleanup(s.Close) return s } -func (s *ProxyServer) reqBody(r testserver.Request) any { - // The SDK expects the query parameters to be specified in the "request body" - // argument for GET, DELETE, and HEAD requests in the .Do method. - if r.Method == http.MethodGet || r.Method == http.MethodDelete || r.Method == http.MethodHead { - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] +type authenticatedTransport struct { + upstream *config.Config + next http.RoundTripper +} + +func (tr *authenticatedTransport) RoundTrip(r *http.Request) (*http.Response, error) { + r.Header.Del("Authorization") + if err := tr.upstream.Authenticate(r); err != nil { + return nil, fmt.Errorf("authenticate upstream request: %w", err) + } + for key, value := range auth.WorkspaceIDHeaders(tr.upstream) { + r.Header.Set(key, value) + } + if tr.upstream.Headers != nil { + if err := tr.upstream.Headers(r); err != nil { + return nil, err } - return queryParams } - - // Otherwise, return the raw body. - return r.Body + return tr.next.RoundTrip(r) } func (s *ProxyServer) proxyToUpstream(w http.ResponseWriter, r *http.Request) { request := testserver.NewRequest(s.t, r, nil) - if s.RequestCallback != nil { - s.RequestCallback(&request) - } - - headers := make(map[string]string) - for k, v := range r.Header { - // Authorization headers will be set by the SDK. Do not pass them along here. - if k == "Authorization" { - continue + if strings.HasPrefix(r.URL.Path, "/__testserver/") { + if r.URL.Path != faultEndpoint || r.Method != http.MethodPost { + http.NotFound(w, r) + return } - // The default HTTP client in Go sets the Accept-Encoding header to - // "gzip". Since it is originally meant to be read by the server and - // will be set again when the SDK makes the request to the workspace, do - // not pass it along here. - if k == "Accept-Encoding" { - continue + if request.Token == "" || request.Token != s.token { + http.Error(w, "invalid test token", http.StatusUnauthorized) + return } - headers[k] = v[0] - } - - queryParams := make(map[string]any) - for k, v := range r.URL.Query() { - queryParams[k] = v[0] + if err := s.faults.Register(request.Token, request.Body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + w.WriteHeader(http.StatusOK) + return } - reqBody := s.reqBody(request) - respBody := bytes.Buffer{} - - // List of response headers to pass through to the client. - // We can extend this code to pass all X-Databricks headers if we introduce a - // more powerful visitor like `WithRequestHeaders` to the SDK. That is left - // as a future exercise for now. - includeResponseHeaders := []string{"X-Databricks-Org-Id"} - responseHeaders := make(map[string]*string) - - visitors := []httpclient.DoOption{ - httpclient.WithRequestHeaders(headers), - httpclient.WithQueryParameters(queryParams), - httpclient.WithRequestData(reqBody), - httpclient.WithResponseUnmarshal(&respBody), + if s.RequestCallback != nil { + s.RequestCallback(&request) } - - for _, header := range includeResponseHeaders { - responseHeaders[header] = new(string) - visitors = append(visitors, httpclient.WithResponseHeader(header, responseHeaders[header])) + rule := s.faults.Check(r.Method, r.URL.Path, request.Token, request.Body) + if rule != nil && !rule.AfterHandler { + s.writeResponse(w, &request, injectedResponse(rule)) + return } - err := s.apiClient.Do(s.t.Context(), r.Method, r.URL.Path, - visitors..., - ) - - var encodedResponse *testserver.EncodedResponse - - // API errors from the SDK are of type [apierr.APIError]. Forward the raw - // response bytes verbatim — including any error details — so callers see - // exactly what the workspace returned. Re-marshalling from the parsed - // APIError would drop fields the SDK doesn't surface (e.g. metadata in - // details[]) and silently break callers that inspect them. - if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { - rw := apiErr.ResponseWrapper - if rw == nil { - // The SDK populates ResponseWrapper for every APIError produced - // from a real HTTP response. If this ever fires the SDK changed - // shape and we need to revisit how we forward error bodies. - panic("apierr.APIError has no ResponseWrapper") - } - encodedResponse = &testserver.EncodedResponse{ - StatusCode: rw.Response.StatusCode, - Body: rw.DebugBytes, - } - // Visitors registered via WithResponseHeader are not invoked when - // the SDK returns an error, so populate the include list directly - // from the original response headers. - for _, header := range includeResponseHeaders { - if v := rw.Response.Header.Get(header); v != "" { - *responseHeaders[header] = v + r.Body = io.NopCloser(bytes.NewReader(request.Body)) + proxy := httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.SetURL(s.upstream) + // Let the transport negotiate compression so response logging sees decoded bytes. + pr.Out.Header.Del("Accept-Encoding") + // Preserve the client's query encoding and repeated parameters. + pr.Out.URL.RawQuery = pr.In.URL.RawQuery + }, + Transport: s.transport, + ModifyResponse: func(response *http.Response) error { + body, err := io.ReadAll(response.Body) + response.Body.Close() + if err != nil { + return err } - } - } - - // Something else went wrong. - if encodedResponse == nil && err != nil { - encodedResponse = &testserver.EncodedResponse{ - StatusCode: 500, - Body: []byte(err.Error()), - } - } - - // Successful response - if encodedResponse == nil { - encodedResponse = &testserver.EncodedResponse{ - StatusCode: 200, - Body: respBody.Bytes(), - } + encoded := testserver.EncodedResponse{ + StatusCode: response.StatusCode, + Headers: response.Header, + Body: body, + } + if rule != nil { + encoded = injectedResponse(rule) + } + if s.ResponseCallback != nil { + s.ResponseCallback(&request, &encoded) + } + response.StatusCode = encoded.StatusCode + response.Header = encoded.Headers + response.Body = io.NopCloser(bytes.NewReader(encoded.Body)) + response.ContentLength = int64(len(encoded.Body)) + return nil + }, + ErrorHandler: func(w http.ResponseWriter, _ *http.Request, err error) { + s.writeResponse(w, &request, testserver.EncodedResponse{ + StatusCode: http.StatusBadGateway, + Body: []byte(err.Error()), + }) + }, } + proxy.ServeHTTP(w, r) +} - // Send response to client. - if len(responseHeaders) > 0 { - for k, v := range responseHeaders { - w.Header()[k] = []string{*v} - } +func injectedResponse(rule *testserver.FaultRule) testserver.EncodedResponse { + return testserver.EncodedResponse{ + StatusCode: rule.StatusCode, + Headers: http.Header{"Content-Type": {"application/json"}}, + Body: []byte(rule.Body), } +} - w.WriteHeader(encodedResponse.StatusCode) - - _, err = w.Write(encodedResponse.Body) - assert.NoError(s.t, err) - +func (s *ProxyServer) writeResponse(w http.ResponseWriter, request *testserver.Request, response testserver.EncodedResponse) { if s.ResponseCallback != nil { - s.ResponseCallback(&request, encodedResponse) + s.ResponseCallback(request, &response) } + maps.Copy(w.Header(), response.Headers) + w.WriteHeader(response.StatusCode) + _, err := w.Write(response.Body) + assert.NoError(s.t, err) } diff --git a/libs/testproxy/server_test.go b/libs/testproxy/server_test.go new file mode 100644 index 00000000000..eff5c16b6a3 --- /dev/null +++ b/libs/testproxy/server_test.go @@ -0,0 +1,210 @@ +package testproxy_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "sync/atomic" + "testing" + "time" + + "github.com/databricks/cli/libs/testproxy" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testToken = "test-token" + +func newProxy(t *testing.T, handler http.HandlerFunc) *testproxy.ProxyServer { + t.Helper() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/databricks-config" { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"workspace_id":"123"}`)) + return + } + handler(w, r) + })) + t.Cleanup(upstream.Close) + return testproxy.New(t, &config.Config{ + Host: upstream.URL, Token: "upstream-token", AuthType: "pat", ConfigFile: os.DevNull, + WorkspaceID: "123", + }, testToken) +} + +func send(t *testing.T, proxy *testproxy.ProxyServer, method, path, token string, body []byte) (int, http.Header, []byte) { + t.Helper() + req, err := http.NewRequestWithContext(t.Context(), method, proxy.URL+path, bytes.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Authorization", "Bearer "+token) + response, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer response.Body.Close() + data, err := io.ReadAll(response.Body) + require.NoError(t, err) + return response.StatusCode, response.Header, data +} + +func register(t *testing.T, proxy *testproxy.ProxyServer, after bool) { + t.Helper() + data, err := json.Marshal(map[string]any{ + "pattern": "POST /api/write", "status_code": 503, "body": `{"error_code":"INJECTED"}`, + "offset": 1, "times": 1, "body_contains": "target", "after_handler": after, + }) + require.NoError(t, err) + status, _, _ := send(t, proxy, "POST", "/__testserver/fault", testToken, data) + require.Equal(t, 200, status) +} + +func TestProxyFaults(t *testing.T) { + for _, after := range []bool{false, true} { + t.Run(map[bool]string{false: "before", true: "after"}[after], func(t *testing.T) { + var calls, recorded, responses atomic.Int32 + proxy := newProxy(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/write" { + w.WriteHeader(http.StatusNotFound) + return + } + assert.Equal(t, "Bearer upstream-token", r.Header.Get("Authorization")) + assert.Equal(t, "123", r.Header.Get("X-Databricks-Workspace-Id")) + calls.Add(1) + w.Header().Set("Retry-After", "23") + w.Header().Set("Content-Type", "text/plain") + w.Header().Set("Content-Length", "9") + w.WriteHeader(http.StatusCreated) + _, _ = w.Write([]byte("committed")) + }) + proxy.RequestCallback = func(*testserver.Request) { recorded.Add(1) } + proxy.ResponseCallback = func(*testserver.Request, *testserver.EncodedResponse) { responses.Add(1) } + register(t, proxy, after) + assert.Zero(t, calls.Load()) + assert.Zero(t, recorded.Load()) + + // Neither another token nor another body consumes the matching budget. + for _, request := range []struct{ token, body string }{{"other-token", "target"}, {testToken, "other"}, {testToken, "target"}} { + status, _, body := send(t, proxy, "POST", "/api/write", request.token, []byte(request.body)) + assert.Equal(t, 201, status) + assert.Equal(t, "committed", string(body)) + } + status, headers, body := send(t, proxy, "POST", "/api/write", testToken, []byte("target")) + assert.Equal(t, 503, status) + assert.JSONEq(t, `{"error_code":"INJECTED"}`, string(body)) + assert.Equal(t, "application/json", headers.Get("Content-Type")) + assert.Empty(t, headers.Get("Retry-After")) + if after { + assert.Equal(t, int32(4), calls.Load()) + } else { + assert.Equal(t, int32(3), calls.Load()) + } + status, _, _ = send(t, proxy, "POST", "/api/write", testToken, []byte("target")) + assert.Equal(t, 201, status) + assert.Equal(t, int32(5), recorded.Load()) + assert.Equal(t, int32(5), responses.Load()) + }) + } +} + +func TestProxyControlDoesNotReachUpstream(t *testing.T) { + var calls atomic.Int32 + proxy := newProxy(t, func(w http.ResponseWriter, _ *http.Request) { calls.Add(1); w.WriteHeader(http.StatusNotFound) }) + for _, tc := range []struct { + method, path, token string + status int + }{ + {"POST", "/__testserver/fault", "other", 401}, + {"POST", "/__testserver/fault", "", 401}, + {"POST", "/__testserver/kill", testToken, 404}, + {"GET", "/__testserver/fault", testToken, 404}, + {"POST", "/__testserver/fault", testToken, 400}, + } { + t.Run(tc.method+tc.path+tc.token, func(t *testing.T) { + status, _, _ := send(t, proxy, tc.method, tc.path, tc.token, []byte("invalid JSON")) + assert.Equal(t, tc.status, status) + }) + } + assert.Zero(t, calls.Load()) +} + +func TestProxyPreservesUpstreamResponse(t *testing.T) { + var calls atomic.Int32 + proxy := newProxy(t, func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/write" { + w.WriteHeader(http.StatusNotFound) + return + } + calls.Add(1) + assert.Equal(t, "x=a%2Fb&x=c+d", r.URL.RawQuery) + body, err := io.ReadAll(r.Body) + assert.NoError(t, err) + assert.Equal(t, []byte("multipart\x00bytes\r\n"), body) + w.Header().Set("Retry-After", "7") + w.Header().Add("X-Example", "one") + w.Header().Add("X-Example", "two") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"details":[{"reason":"preserved"}]}`)) + }) + status, headers, body := send(t, proxy, "POST", "/api/write?x=a%2Fb&x=c+d", testToken, []byte("multipart\x00bytes\r\n")) + assert.Equal(t, 429, status) + assert.Equal(t, "7", headers.Get("Retry-After")) + assert.Equal(t, []string{"one", "two"}, headers.Values("X-Example")) + assert.JSONEq(t, `{"details":[{"reason":"preserved"}]}`, string(body)) + assert.Equal(t, int32(1), calls.Load()) +} + +func TestProxyDoesNotMaskTransportFailure(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"workspace_id":"123"}`)) + })) + t.Cleanup(upstream.Close) + proxy := testproxy.New(t, &config.Config{ + Host: upstream.URL, Token: "upstream-token", AuthType: "pat", ConfigFile: os.DevNull, + }, testToken) + upstream.Close() + data := []byte(`{"pattern":"POST /api/write","status_code":403,"body":"injected","after_handler":true,"times":1}`) + status, _, _ := send(t, proxy, "POST", "/__testserver/fault", testToken, data) + require.Equal(t, 200, status) + status, _, body := send(t, proxy, "POST", "/api/write", testToken, []byte("target")) + assert.Equal(t, 502, status) + assert.NotEqual(t, "injected", string(body)) +} + +func TestProxyCancellation(t *testing.T) { + started := make(chan struct{}) + cancelled := make(chan struct{}) + proxy := newProxy(t, func(_ http.ResponseWriter, r *http.Request) { + close(started) + <-r.Context().Done() + close(cancelled) + }) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + req, err := http.NewRequestWithContext(ctx, http.MethodGet, proxy.URL+"/api/wait", nil) + require.NoError(t, err) + done := make(chan error, 1) + go func() { + response, err := http.DefaultClient.Do(req) + if response != nil { + response.Body.Close() + } + done <- err + }() + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("upstream request did not start") + } + cancel() + require.ErrorIs(t, <-done, context.Canceled) + select { + case <-cancelled: + case <-time.After(5 * time.Second): + t.Fatal("upstream request was not cancelled") + } +} diff --git a/libs/testserver/fault.go b/libs/testserver/fault.go index 43d46f3f28c..eea6c117e79 100644 --- a/libs/testserver/fault.go +++ b/libs/testserver/fault.go @@ -109,26 +109,34 @@ func (fr *FaultRules) Check(method, path, token string, body []byte) *FaultRule return nil } +// Register decodes and registers a rule for token using the fault.py wire format. +func (fr *FaultRules) Register(token string, data []byte) error { + var body struct { + Pattern string `json:"pattern"` + StatusCode int `json:"status_code"` + Body string `json:"body"` + Offset int `json:"offset"` + Times int `json:"times"` + AfterHandler bool `json:"after_handler"` + BodyContains string `json:"body_contains"` + } + if err := json.Unmarshal(data, &body); err != nil { + return err + } + if body.AfterHandler { + fr.SetAfterHandler(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) + } else { + fr.Set(token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) + } + return nil +} + // faultEndpointHandler handles POST /__testserver/fault. func faultEndpointHandler(fr *FaultRules) HandlerFunc { return func(req Request) any { - var body struct { - Pattern string `json:"pattern"` - StatusCode int `json:"status_code"` - Body string `json:"body"` - Offset int `json:"offset"` - Times int `json:"times"` - AfterHandler bool `json:"after_handler"` - BodyContains string `json:"body_contains"` - } - if err := json.Unmarshal(req.Body, &body); err != nil { + if err := fr.Register(req.Token, req.Body); err != nil { return Response{StatusCode: 400, Body: map[string]string{"error": err.Error()}} } - if body.AfterHandler { - fr.SetAfterHandler(req.Token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) - } else { - fr.Set(req.Token, body.Pattern, body.StatusCode, body.Body, body.Offset, body.Times, body.BodyContains) - } return Response{StatusCode: 200} } } diff --git a/libs/testserver/server.go b/libs/testserver/server.go index 5170b318818..bb74de25f40 100644 --- a/libs/testserver/server.go +++ b/libs/testserver/server.go @@ -106,6 +106,7 @@ func NewRequest(t testutil.TestingT, r *http.Request, fakeWorkspace *FakeWorkspa return Request{ Method: r.Method, + Token: getToken(r), URL: r.URL, Headers: r.Header, Body: body,