From c2a14fb67b97e89a769308096506be9abd1d7894 Mon Sep 17 00:00:00 2001 From: Andy Bonventre Date: Tue, 23 Jun 2026 01:49:46 +0200 Subject: [PATCH 1/4] feat(proxy): add HTTP/2 and gRPC support through TLS interception Advertise h2 in ALPN on the client-facing TLS connection so gRPC clients (and any other h2 client) can negotiate HTTP/2 through the CONNECT intercept path. When h2 is negotiated, http2.ConfigureServer wires up correct h2 framing on the inner http.Server and the upstream transport switches to http2.Transport; HTTP/1.1 clients continue to use the original http.Transport unchanged. Credential injection works identically on h2 connections. Also adds black-box test coverage for previously uncovered proxy paths (LLM gateway policy, response transformers, token substitution, tunnel mode) and two new multiHandler slog tests in gatekeeper_test.go. --- CHANGELOG.md | 6 + gatekeeper_test.go | 35 +++ go.mod | 10 +- go.sum | 10 + proxy/coverage_test.go | 603 +++++++++++++++++++++++++++++++++++++++++ proxy/http2_test.go | 212 +++++++++++++++ proxy/proxy.go | 59 ++-- 7 files changed, 913 insertions(+), 22 deletions(-) create mode 100644 proxy/coverage_test.go create mode 100644 proxy/http2_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 65d4225..fe97068 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ Gatekeeper is a standalone credential-injecting TLS-intercepting proxy. It trans Gatekeeper is pre-1.0. The configuration schema and credential source interface may change between minor versions. +## v0.13.0 — 2026-06-23 + +### Added + +- **HTTP/2 and gRPC support through TLS interception** — the CONNECT interception path now negotiates HTTP/2 via ALPN (`h2` advertised first, `http/1.1` as fallback); when a client negotiates h2 (e.g., a gRPC client), the inner `http.Server` handles the connection with `http2.ConfigureServer` for correct h2 framing, and the upstream transport switches to `http2.Transport` so requests are forwarded over h2 end-to-end; credential injection (arbitrary headers such as `x-modal-token-id` / `x-modal-token-secret`) works identically on h2 connections; HTTP/1.1 clients are unaffected — `http.Transport` is used unchanged when h2 is not negotiated ([#34](https://github.com/majorcontext/gatekeeper/pull/34)) + ## v0.12.0 — 2026-06-12 ### Added diff --git a/gatekeeper_test.go b/gatekeeper_test.go index 78b5af5..380b706 100644 --- a/gatekeeper_test.go +++ b/gatekeeper_test.go @@ -14,6 +14,7 @@ import ( "errors" "fmt" "io" + "log/slog" "math/big" "net" "net/http" @@ -2490,3 +2491,37 @@ func TestServerPostgresUnknownResolver(t *testing.T) { t.Errorf("error = %q, want mention of 'unknown resolver'", err) } } + +// ── multiHandler (slog fan-out) ─────────────────────────────────────────────── + +// TestMultiHandler_WithAttrs verifies that WithAttrs propagates to all child +// handlers, exercising the previously uncovered branch. +func TestMultiHandler_WithAttrs(t *testing.T) { + h1 := slog.NewTextHandler(io.Discard, nil) + h2 := slog.NewTextHandler(io.Discard, nil) + mh := newMultiHandler(h1, h2) + + derived := mh.WithAttrs([]slog.Attr{slog.String("k", "v")}) + if derived == nil { + t.Fatal("WithAttrs returned nil") + } + // Derived handler should still respond to Enabled. + if !derived.Enabled(context.Background(), slog.LevelInfo) { + t.Error("derived handler Enabled(Info) = false, want true") + } +} + +// TestMultiHandler_WithGroup verifies that WithGroup propagates to all child +// handlers. +func TestMultiHandler_WithGroup(t *testing.T) { + h1 := slog.NewTextHandler(io.Discard, nil) + mh := newMultiHandler(h1) + + derived := mh.WithGroup("mygroup") + if derived == nil { + t.Fatal("WithGroup returned nil") + } + if !derived.Enabled(context.Background(), slog.LevelInfo) { + t.Error("derived handler Enabled(Info) = false, want true") + } +} diff --git a/go.mod b/go.mod index bedd61b..0dad8d7 100644 --- a/go.mod +++ b/go.mod @@ -23,7 +23,7 @@ require ( go.opentelemetry.io/otel/sdk/log v0.19.0 go.opentelemetry.io/otel/sdk/metric v1.43.0 go.opentelemetry.io/otel/trace v1.43.0 - golang.org/x/sync v0.20.0 + golang.org/x/sync v0.21.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -129,12 +129,12 @@ require ( go.opentelemetry.io/proto/otlp v1.10.0 // indirect go.uber.org/multierr v1.11.0 // indirect go4.org v0.0.0-20230225012048-214862532bf5 // indirect - golang.org/x/crypto v0.49.0 // indirect + golang.org/x/crypto v0.53.0 // indirect golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa // indirect - golang.org/x/net v0.52.0 // indirect + golang.org/x/net v0.56.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/sys v0.42.0 // indirect - golang.org/x/text v0.35.0 // indirect + golang.org/x/sys v0.46.0 // indirect + golang.org/x/text v0.38.0 // indirect golang.org/x/time v0.15.0 // indirect google.golang.org/api v0.274.0 // indirect google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect diff --git a/go.sum b/go.sum index 28266c0..6949cf7 100644 --- a/go.sum +++ b/go.sum @@ -381,6 +381,8 @@ golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8U golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4= golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA= +golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= +golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20190510132918-efd6b22b2522/go.mod h1:ZjyILWgesfNpC6sMxTJOJm9Kp84zZh5NQWvqDGG3Qr8= @@ -427,6 +429,8 @@ golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0= golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= +golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= @@ -443,6 +447,8 @@ golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -466,6 +472,8 @@ golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= @@ -479,6 +487,8 @@ golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8= golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA= +golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= +golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= diff --git a/proxy/coverage_test.go b/proxy/coverage_test.go new file mode 100644 index 0000000..8cd94c8 --- /dev/null +++ b/proxy/coverage_test.go @@ -0,0 +1,603 @@ +package proxy + +// Black-box coverage tests for paths that were previously uncovered. +// Every test drives the proxy through its HTTP interface — no internal +// functions are called directly. + +import ( + "bufio" + "crypto/tls" + "crypto/x509" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + keeplib "github.com/majorcontext/keep" +) + +// ── Tunnel (no-CA) path ────────────────────────────────────────────────────── + +// TestTunnel_ForwardsPlainHTTPS verifies that when the proxy has no CA +// configured, a CONNECT request is forwarded as a raw TCP tunnel without +// TLS interception. Credentials can still be configured but will not be +// injected (the proxy can't see the plaintext). +func TestTunnel_ForwardsPlainHTTPS(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "tunneled") + })) + t.Cleanup(backend.Close) + + // Proxy with no CA — uses handleConnectTunnel. + p := NewProxy() + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + // Client trusts the backend's self-signed cert directly (no MITM). + backendCAs := x509.NewCertPool() + backendCAs.AddCert(backend.Certificate()) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxyServer.URL)), + TLSClientConfig: &tls.Config{RootCAs: backendCAs}, + }, + } + + resp, err := client.Get(backend.URL + "/hello") + if err != nil { + t.Fatalf("GET through tunnel: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if string(body) != "tunneled" { + t.Errorf("body = %q, want %q", string(body), "tunneled") + } + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +// TestTunnel_NetworkPolicyBlocked verifies that the network policy is still +// enforced even when no CA is set (tunnel mode). +func TestTunnel_NetworkPolicyBlocked(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "should not reach") + })) + t.Cleanup(backend.Close) + + p := NewProxy() + p.SetNetworkPolicy("strict", nil, nil) + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + backendCAs := x509.NewCertPool() + backendCAs.AddCert(backend.Certificate()) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxyServer.URL)), + TLSClientConfig: &tls.Config{RootCAs: backendCAs}, + }, + } + + resp, err := client.Get(backend.URL + "/hello") + if err != nil { + // Connection refused or proxy error — also acceptable as a block. + return + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusOK { + t.Errorf("expected blocked response, got 200") + } +} + +// TestTunnel_InvalidHostFormat verifies that a malformed CONNECT target +// (missing port) returns a 400 Bad Request. +func TestTunnel_InvalidHostFormat(t *testing.T) { + p := NewProxy() + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + // Issue a raw CONNECT with a host that has no port — SplitHostPort will fail. + conn, err := net.Dial("tcp", proxyServer.Listener.Addr().String()) + if err != nil { + t.Fatalf("dial proxy: %v", err) + } + defer conn.Close() + + fmt.Fprintf(conn, "CONNECT noporthost HTTP/1.1\r\nHost: noporthost\r\n\r\n") + + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + t.Fatalf("read response: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400", resp.StatusCode) + } +} + +// TestTunnel_PathRulesWarning exercises the code path where TLS interception +// is disabled (no CA) but per-path rules are configured — the proxy should +// log a warning and fall through to the tunnel handler. +func TestTunnel_PathRulesWarning(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "ok") + })) + t.Cleanup(backend.Close) + + p := NewProxy() + // Inject a PathRulesChecker that claims path rules exist for any host. + p.SetContextResolver(func(token string) (*RunContextData, bool) { + return &RunContextData{ + Policy: "permissive", + PathRulesCheck: func(host string, port int) bool { + return true // pretend per-path rules exist + }, + }, true + }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + backendCAs := x509.NewCertPool() + backendCAs.AddCert(backend.Certificate()) + + proxyURL := mustParseURL(proxyServer.URL) + proxyURL.User = url.UserPassword("user", "tok") + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: backendCAs}, + }, + } + + // The proxy has no CA, so it tunnels despite path rules existing. + resp, err := client.Get(backend.URL + "/path") + if err != nil { + t.Fatalf("GET: %v", err) + } + resp.Body.Close() + // Main assertion: we reach the backend (tunnel works) without panicking. + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +// ── LLM gateway policy (evaluateAndReplaceLLMResponse) ────────────────────── + +const llmGatewayDenyEditPolicy = ` +scope: llm-gateway +mode: enforce +rules: + - name: deny-edit + match: + operation: "llm.tool_use" + when: "params.name == 'edit'" + action: deny + message: "Editing blocked." +` + +// newAnthropicInterceptSetup builds an intercept test setup where the backend +// acts as a fake api.anthropic.com. The client sends requests to +// https://api.anthropic.com: and the proxy rewrites the upstream +// dial to 127.0.0.1: via the HostGateway mechanism, so the +// host check (host == "api.anthropic.com") triggers the LLM policy path. +// +// The backend port is added to AllowedHostPorts so the HostGateway policy +// permits the CONNECT (HostGateway traffic requires explicit port allow-listing). +func newAnthropicInterceptSetup(t *testing.T, llmEng *keeplib.Engine, backendHandler http.Handler) (client *http.Client, backendURL string) { + t.Helper() + + ca, err := generateCA() + if err != nil { + t.Fatal(err) + } + + backend := httptest.NewTLSServer(backendHandler) + t.Cleanup(backend.Close) + + upstreamCAs := x509.NewCertPool() + upstreamCAs.AddCert(backend.Certificate()) + + backendAddr := mustParseURL(backend.URL) + backendPort := backendAddr.Port() + backendPortInt := 0 + fmt.Sscanf(backendPort, "%d", &backendPortInt) + + p := NewProxy() + p.SetCA(ca) + p.SetUpstreamCAs(upstreamCAs) + p.SetContextResolver(func(token string) (*RunContextData, bool) { + if token != "llmtok" { + return nil, false + } + rc := &RunContextData{ + Policy: "permissive", + HostGateway: "api.anthropic.com", + HostGatewayIP: "127.0.0.1", + AllowedHostPorts: []int{backendPortInt}, + } + if llmEng != nil { + rc.KeepEngines = map[string]*keeplib.Engine{"llm-gateway": llmEng} + } + return rc, true + }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + clientCAs := x509.NewCertPool() + clientCAs.AppendCertsFromPEM(ca.certPEM) + + proxyURL := mustParseURL(proxyServer.URL) + proxyURL.User = url.UserPassword("user", "llmtok") + + client = &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: clientCAs}, + }, + } + + backendURL = "https://api.anthropic.com:" + backendPort + return client, backendURL +} + +// TestIntercept_LLMPolicy_Deny verifies that the llm-gateway Keep engine +// blocks a tool-use response from api.anthropic.com. +func TestIntercept_LLMPolicy_Deny(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + body := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{"file_path":"/foo"}}],"stop_reason":"tool_use"}` + + client, backendURL := newAnthropicInterceptSetup(t, eng, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, body) + }), + ) + + resp, err := client.Post(backendURL+"/v1/messages", "application/json", + strings.NewReader(`{"model":"claude-opus-4-5"}`)) + if err != nil { + t.Fatalf("POST: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400 (policy denied)", resp.StatusCode) + } + if got := resp.Header.Get("X-Moat-Blocked"); got != "llm-policy" { + t.Errorf("X-Moat-Blocked = %q, want llm-policy", got) + } + respBody, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(respBody), "policy_denied") { + t.Errorf("response body missing policy_denied: %s", respBody) + } +} + +// TestIntercept_LLMPolicy_Allow verifies that a non-matching response passes +// through the llm-gateway engine unchanged. +func TestIntercept_LLMPolicy_Allow(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + allowedBody := `{"content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn"}` + + client, backendURL := newAnthropicInterceptSetup(t, eng, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, allowedBody) + }), + ) + + resp, err := client.Get(backendURL + "/v1/messages") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if string(respBody) != allowedBody { + t.Errorf("body = %q, want %q", string(respBody), allowedBody) + } +} + +// TestIntercept_LLMPolicy_ResponseTooLarge verifies that oversized responses +// from api.anthropic.com are blocked with a size-limit error. +func TestIntercept_LLMPolicy_ResponseTooLarge(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + // Build a body larger than maxLLMResponseSize (10 MiB). + // Use 11 MiB of JSON-ish padding so the size check triggers. + hugeBody := `{"content":[{"type":"text","text":"` + strings.Repeat("x", 11*1024*1024) + `"}]}` + + client, backendURL := newAnthropicInterceptSetup(t, eng, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, hugeBody) + }), + ) + + resp, err := client.Get(backendURL + "/v1/messages") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400 (size-limit)", resp.StatusCode) + } + if got := resp.Header.Get("X-Moat-Blocked"); got != "llm-policy" { + t.Errorf("X-Moat-Blocked = %q, want llm-policy", got) + } + respBody, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(respBody), "size-limit") { + t.Errorf("response body missing size-limit: %s", respBody) + } +} + +// TestIntercept_LLMPolicy_NoEnginePassesThrough verifies that without a +// llm-gateway engine the response is passed through unmodified, even for +// api.anthropic.com. +func TestIntercept_LLMPolicy_NoEnginePassesThrough(t *testing.T) { + rawBody := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{}}],"stop_reason":"tool_use"}` + + // nil engine — no llm-gateway key in KeepEngines. + client, backendURL := newAnthropicInterceptSetup(t, nil, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, rawBody) + }), + ) + + resp, err := client.Get(backendURL + "/v1/messages") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if string(respBody) != rawBody { + t.Errorf("body = %q, want %q", string(respBody), rawBody) + } +} + +// TestIntercept_LLMPolicy_DeniedLogged verifies that a denied LLM response is +// recorded in the canonical request log with Denied=true. +func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + toolBody := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{"file_path":"/f"}}],"stop_reason":"tool_use"}` + + // Build the setup manually so we can attach a log listener to the proxy. + ca, err := generateCA() + if err != nil { + t.Fatal(err) + } + + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, toolBody) + })) + t.Cleanup(backend.Close) + + upstreamCAs := x509.NewCertPool() + upstreamCAs.AddCert(backend.Certificate()) + + backendPort := 0 + fmt.Sscanf(mustParseURL(backend.URL).Port(), "%d", &backendPort) + + p := NewProxy() + p.SetCA(ca) + p.SetUpstreamCAs(upstreamCAs) + p.SetContextResolver(func(token string) (*RunContextData, bool) { + if token != "logtest" { + return nil, false + } + return &RunContextData{ + Policy: "permissive", + HostGateway: "api.anthropic.com", + HostGatewayIP: "127.0.0.1", + AllowedHostPorts: []int{backendPort}, + KeepEngines: map[string]*keeplib.Engine{"llm-gateway": eng}, + }, true + }) + + var logged RequestLogData + p.SetLogger(func(data RequestLogData) { logged = data }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + clientCAs := x509.NewCertPool() + clientCAs.AppendCertsFromPEM(ca.certPEM) + + proxyURL := mustParseURL(proxyServer.URL) + proxyURL.User = url.UserPassword("user", "logtest") + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: clientCAs}, + }, + } + + resp, err := client.Post( + fmt.Sprintf("https://api.anthropic.com:%d/v1/messages", backendPort), + "application/json", + strings.NewReader(`{"model":"claude-opus-4-5"}`), + ) + if err != nil { + t.Fatalf("POST: %v", err) + } + io.ReadAll(resp.Body) + resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400", resp.StatusCode) + } + if !logged.Denied { + t.Errorf("RequestLogData.Denied = false, want true") + } +} + +// ── AddResponseTransformer ─────────────────────────────────────────────────── + +// TestIntercept_ResponseTransformer verifies that a registered response +// transformer runs in the intercept path and can observe the response. +func TestIntercept_ResponseTransformer(t *testing.T) { + var transformerCalled bool + setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"original":true}`) + })) + + setup.Proxy.AddResponseTransformer(setup.BackendHost, func(req, resp any) (any, bool) { + transformerCalled = true + // Return false — observe without modifying (no replacement). + return resp, false + }) + + resp, err := setup.Client.Get(setup.Backend.URL + "/data") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + io.ReadAll(resp.Body) + + if !transformerCalled { + t.Error("response transformer was not called") + } + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +// TestIntercept_ResponseTransformer_NoMatch verifies that a transformer +// registered for a different host does not affect other hosts. +func TestIntercept_ResponseTransformer_NoMatch(t *testing.T) { + originalBody := `{"original":true}` + setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, originalBody) + })) + + // Register transformer for a completely different host. + setup.Proxy.AddResponseTransformer("other.example.com", func(req, resp any) (any, bool) { + // This should never be called for requests to the backend host. + t.Error("transformer called for wrong host") + return resp, false + }) + + resp, err := setup.Client.Get(setup.Backend.URL + "/data") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + if string(body) != originalBody { + t.Errorf("body = %q, want %q (transformer should not apply)", string(body), originalBody) + } +} + +// ── SetTokenSubstitution (proxy-level setter) ──────────────────────────────── + +// TestIntercept_SetTokenSubstitution verifies that the proxy-level +// SetTokenSubstitution setter is wired into the intercept path. +func TestIntercept_SetTokenSubstitution(t *testing.T) { + var receivedPath string + setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath = r.URL.Path + fmt.Fprint(w, "ok") + })) + + // Register token substitution: placeholder in URL → real token. + setup.Proxy.SetTokenSubstitution(setup.BackendHost, "placeholder-token", "real-secret-value") + + resp, err := setup.Client.Get(setup.Backend.URL + "/bot" + "placeholder-token" + "/getUpdates") + if err != nil { + t.Fatalf("GET: %v", err) + } + resp.Body.Close() + + if !strings.Contains(receivedPath, "real-secret-value") { + t.Errorf("path = %q, want real-secret-value substituted", receivedPath) + } + if strings.Contains(receivedPath, "placeholder-token") { + t.Errorf("path = %q, placeholder should have been replaced", receivedPath) + } +} + +// ── Exported API wrappers ──────────────────────────────────────────────────── + +// TestNewTokenSubstitution verifies the exported constructor returns a usable substitution. +func TestNewTokenSubstitution(t *testing.T) { + sub := NewTokenSubstitution("placeholder", "real") + if sub == nil { + t.Fatal("NewTokenSubstitution returned nil") + } +} + +// TestParseHostPattern and TestMatchesHostPattern exercise the exported wrappers +// for parseHostPattern and matchesPattern. +func TestParseAndMatchHostPattern(t *testing.T) { + cases := []struct { + pattern string + host string + port int + want bool + }{ + {"api.example.com", "api.example.com", 443, true}, + {"api.example.com", "other.example.com", 443, false}, + {"*.example.com", "sub.example.com", 443, true}, + {"*.example.com", "example.com", 443, false}, + } + for _, tc := range cases { + p := ParseHostPattern(tc.pattern) + got := MatchesHostPattern(p, tc.host, tc.port) + if got != tc.want { + t.Errorf("MatchesHostPattern(%q, %q, %d) = %v, want %v", tc.pattern, tc.host, tc.port, got, tc.want) + } + } +} + +// TestRegisterGrantHosts verifies that registered grant hosts are retrievable. +func TestRegisterGrantHosts(t *testing.T) { + RegisterGrantHosts("test-grant-coverage", []string{"coverage.example.com"}) + hosts := GetHostsForGrant("test-grant-coverage") + if len(hosts) == 0 { + t.Fatal("GetHostsForGrant returned empty slice after RegisterGrantHosts") + } + if hosts[0] != "coverage.example.com" { + t.Errorf("hosts[0] = %q, want coverage.example.com", hosts[0]) + } +} diff --git a/proxy/http2_test.go b/proxy/http2_test.go new file mode 100644 index 0000000..c67ca43 --- /dev/null +++ b/proxy/http2_test.go @@ -0,0 +1,212 @@ +package proxy + +// Black-box acceptance tests for HTTP/2 (gRPC) support through the +// TLS-intercepting proxy. These tests drive the proxy through its HTTP +// CONNECT interface using a real http2.Transport, verifying that: +// - the proxy negotiates h2 via ALPN on the client-facing TLS connection +// - credential headers are injected into h2 requests +// - the proxy forwards requests upstream over h2 when the backend supports it + +import ( + "bufio" + "context" + "crypto/tls" + "crypto/x509" + "encoding/binary" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + + "golang.org/x/net/http2" +) + +// newGRPCServer returns an httptest.Server that speaks HTTP/2 and handles +// minimal unary gRPC calls on any path. receivedHeaders captures all +// request headers from the first call. +func newGRPCServer(t *testing.T, receivedHeaders *http.Header) *httptest.Server { + t.Helper() + + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + if r.ProtoMajor != 2 { + http.Error(w, "require HTTP/2", http.StatusHTTPVersionNotSupported) + return + } + if *receivedHeaders == nil { + *receivedHeaders = r.Header.Clone() + } + + // Read and discard the gRPC request frame (5-byte length-prefix + body). + frame := make([]byte, 5) + if _, err := io.ReadFull(r.Body, frame); err == nil { + msgLen := binary.BigEndian.Uint32(frame[1:]) + io.CopyN(io.Discard, r.Body, int64(msgLen)) + } + + // gRPC trailers must be declared before WriteHeader. + w.Header().Set("Trailer", "grpc-status") + w.Header().Set("Content-Type", "application/grpc") + w.WriteHeader(http.StatusOK) + // Minimal gRPC response: compressed-flag(0) + message-length(0) = 5 zero bytes. + w.Write([]byte{0, 0, 0, 0, 0}) + w.(http.Flusher).Flush() + // Setting grpc-status after Flush makes it a trailer in HTTP/2. + w.Header().Set("grpc-status", "0") + }) + + srv := httptest.NewUnstartedServer(mux) + srv.EnableHTTP2 = true + srv.StartTLS() + return srv +} + +// newGRPCProxySetup creates a TLS-intercepting proxy configured to inject +// Modal-style credentials for api.modal.com. The fake gRPC backend is +// reachable via the HostGateway mechanism so the credential host check fires. +func newGRPCProxySetup(t *testing.T, receivedHeaders *http.Header) (transport *http2.Transport, backendURL string) { + t.Helper() + + ca, err := generateCA() + if err != nil { + t.Fatal(err) + } + + backend := newGRPCServer(t, receivedHeaders) + t.Cleanup(backend.Close) + + upstreamCAs := x509.NewCertPool() + upstreamCAs.AddCert(backend.Certificate()) + + backendAddr, _ := url.Parse(backend.URL) + backendPort := 0 + fmt.Sscanf(backendAddr.Port(), "%d", &backendPort) + + p := NewProxy() + p.SetCA(ca) + p.SetUpstreamCAs(upstreamCAs) + p.SetContextResolver(func(token string) (*RunContextData, bool) { + if token != "grpctest" { + return nil, false + } + return &RunContextData{ + Policy: "permissive", + HostGateway: "api.modal.com", + HostGatewayIP: "127.0.0.1", + AllowedHostPorts: []int{backendPort}, + Credentials: map[string][]credentialHeader{ + "api.modal.com": { + {Name: "x-modal-token-id", Value: "token-id-test", Grant: "modal"}, + {Name: "x-modal-token-secret", Value: "token-secret-test", Grant: "modal"}, + }, + }, + }, true + }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + clientCAs := x509.NewCertPool() + clientCAs.AppendCertsFromPEM(ca.certPEM) + + proxyAddr := proxyServer.Listener.Addr().String() + authHeader := "Basic " + basicAuth("user", "grpctest") + + // http2.Transport uses DialTLSContext to establish the connection. + // We manually build the CONNECT tunnel first, then negotiate h2 via ALPN. + transport = &http2.Transport{ + DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) { + conn, err := net.Dial("tcp", proxyAddr) + if err != nil { + return nil, fmt.Errorf("dial proxy: %w", err) + } + connectReq := "CONNECT " + addr + " HTTP/1.1\r\n" + + "Host: " + addr + "\r\n" + + "Proxy-Authorization: " + authHeader + "\r\n\r\n" + if _, err := conn.Write([]byte(connectReq)); err != nil { + conn.Close() + return nil, fmt.Errorf("write CONNECT: %w", err) + } + // http.ReadResponse handles partial reads and validates the status line. + cresp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + conn.Close() + return nil, fmt.Errorf("read CONNECT response: %w", err) + } + cresp.Body.Close() + if cresp.StatusCode != http.StatusOK { + conn.Close() + return nil, fmt.Errorf("CONNECT failed: %s", cresp.Status) + } + // Upgrade to TLS, advertising h2 in ALPN. + serverName, _, _ := net.SplitHostPort(addr) + tlsCfg := &tls.Config{ + RootCAs: clientCAs, + ServerName: serverName, + NextProtos: []string{http2.NextProtoTLS}, + } + tlsConn := tls.Client(conn, tlsCfg) + if err := tlsConn.HandshakeContext(ctx); err != nil { + conn.Close() + return nil, fmt.Errorf("TLS handshake: %w", err) + } + return tlsConn, nil + }, + } + + backendURL = fmt.Sprintf("https://api.modal.com:%d", backendPort) + return transport, backendURL +} + +// TestHTTP2_GRPCCredentialInjection verifies that HTTP/2 (gRPC) requests +// through the CONNECT proxy succeed and receive credential injection. +// +// This test is expected to FAIL until HTTP/2 support is implemented in +// handleConnectWithInterception (proxy must advertise h2 in ALPN and use +// http2.ConfigureServer on the inner http.Server). +func TestHTTP2_GRPCCredentialInjection(t *testing.T) { + var receivedHeaders http.Header + transport, backendURL := newGRPCProxySetup(t, &receivedHeaders) + + // Minimal gRPC request frame: compressed=0, message-length=0. + grpcBody := []byte{0, 0, 0, 0, 0} + + req, err := http.NewRequestWithContext(context.Background(), + "POST", backendURL+"/modal.api.v1.AppService/ListApps", + strings.NewReader(string(grpcBody))) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/grpc") + req.Header.Set("te", "trailers") + + resp, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("gRPC request through proxy failed: %v", err) + } + defer resp.Body.Close() + io.Copy(io.Discard, resp.Body) + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + // grpc-status arrives as an HTTP/2 trailer — must read from resp.Trailer + // (only populated after the body is fully consumed). + if got := resp.Trailer.Get("grpc-status"); got != "0" { + t.Errorf("grpc-status trailer = %q, want 0 (header: %q)", got, resp.Header.Get("grpc-status")) + } + + if receivedHeaders == nil { + t.Fatal("backend received no request — proxy may have blocked or h2 ALPN not negotiated") + } + if got := receivedHeaders.Get("x-modal-token-id"); got != "token-id-test" { + t.Errorf("x-modal-token-id = %q, want token-id-test", got) + } + if got := receivedHeaders.Get("x-modal-token-secret"); got != "token-secret-test" { + t.Errorf("x-modal-token-secret = %q, want token-secret-test", got) + } +} diff --git a/proxy/proxy.go b/proxy/proxy.go index e5d17de..57606f9 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -52,6 +52,7 @@ import ( keeplib "github.com/majorcontext/keep" "go.jetify.com/typeid" + "golang.org/x/net/http2" ) // contextKey is the type for request-scoped context values. @@ -2152,9 +2153,14 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req return } + // Advertise h2 first so it is preferred during ALPN negotiation; + // http/1.1 is kept as fallback for non-h2 clients. + // Ordering matters: ConfigureServer only appends missing protos, it + // does not reorder, so h2-preference must be established here. tlsConfig := &tls.Config{ Certificates: []tls.Certificate{*cert}, MinVersion: tls.VersionTLS12, + NextProtos: []string{http2.NextProtoTLS, "http/1.1"}, } tlsClientConn := tls.Server(clientConn, tlsConfig) if err := tlsClientConn.Handshake(); err != nil { @@ -2168,23 +2174,35 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req } }() - transport := &http.Transport{ - Proxy: nil, - DialContext: (&net.Dialer{ - Timeout: 30 * time.Second, - KeepAlive: 30 * time.Second, - }).DialContext, - TLSClientConfig: &tls.Config{ - MinVersion: tls.VersionTLS12, - RootCAs: p.upstreamCAs, // nil means system roots - }, - TLSHandshakeTimeout: 10 * time.Second, - ResponseHeaderTimeout: 5 * time.Minute, - MaxIdleConns: 100, - IdleConnTimeout: 90 * time.Second, - // Note: Do NOT set ForceAttemptHTTP2 here. This transport forwards - // HTTP/1.1 requests read from the intercepted TLS connection. Enabling - // HTTP/2 on the upstream side causes framing mismatches and hangs. + // Shared TLS config for upstream connections (both h2 and h1 paths). + upstreamTLS := &tls.Config{ + MinVersion: tls.VersionTLS12, + RootCAs: p.upstreamCAs, + } + + // Build an upstream transport matching the negotiated protocol. + // When the client negotiated h2 (e.g., gRPC), the request object is an + // h2 request and cannot be round-tripped via an HTTP/1.1 transport + // without framing errors, so we must forward upstream over h2 as well. + var transport http.RoundTripper + if tlsClientConn.ConnectionState().NegotiatedProtocol == http2.NextProtoTLS { + transport = &http2.Transport{TLSClientConfig: upstreamTLS} + } else { + transport = &http.Transport{ + Proxy: nil, + DialContext: (&net.Dialer{ + Timeout: 30 * time.Second, + KeepAlive: 30 * time.Second, + }).DialContext, + TLSClientConfig: upstreamTLS, + TLSHandshakeTimeout: 10 * time.Second, + // Do NOT set ForceAttemptHTTP2: this path handles HTTP/1.1 + // requests. Enabling h2 upstream for h1 clients causes + // framing mismatches. + ResponseHeaderTimeout: 5 * time.Minute, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + } } // Extract port from the CONNECT request for rule checking. @@ -2539,5 +2557,12 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req } }, } + // Enable HTTP/2 on the inner server so h2 clients (e.g., gRPC) get + // proper framing. h1 clients are unaffected — ConfigureServer + // falls back to the normal http.Handler when h2 is not negotiated. + if err := http2.ConfigureServer(srv, nil); err != nil { + slog.Warn("http2.ConfigureServer failed, falling back to HTTP/1.1", + "subsystem", "proxy", "host", host, "error", err) + } _ = srv.Serve(ln) } From c7979be84e4a3c02c7fa2c9b52c5272647fef909 Mon Sep 17 00:00:00 2001 From: Andy Bonventre Date: Tue, 23 Jun 2026 01:57:52 +0200 Subject: [PATCH 2/4] fix(proxy): address HTTP/2 review comments - Add dial timeout, KeepAlive, ReadIdleTimeout, PingTimeout to http2.Transport upstream to match the connection-safety posture of the h1 path - Document the known limitation: http2.Transport never falls back to h1, so h2-negotiating clients hitting h1-only upstreams will fail - Bail out early in handleConnectWithInterception if ConfigureServer fails and the client already negotiated h2, rather than serve h1 to an h2 client - Fix bufferedConn fragility in test: wrap conn with bufferedConn shim so bytes pre-fetched into the bufio.Reader are not lost before tls.Client consumes them --- proxy/http2_test.go | 17 +++++++++++++++-- proxy/proxy.go | 38 ++++++++++++++++++++++++++++++-------- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/proxy/http2_test.go b/proxy/http2_test.go index c67ca43..bf64eb3 100644 --- a/proxy/http2_test.go +++ b/proxy/http2_test.go @@ -25,6 +25,16 @@ import ( "golang.org/x/net/http2" ) +// bufferedConn wraps a net.Conn with a pre-filled bufio.Reader so bytes +// already consumed into the buffer (e.g., from reading the CONNECT response) +// are not lost when the connection is handed to tls.Client. +type bufferedConn struct { + net.Conn + r *bufio.Reader +} + +func (c *bufferedConn) Read(b []byte) (int, error) { return c.r.Read(b) } + // newGRPCServer returns an httptest.Server that speaks HTTP/2 and handles // minimal unary gRPC calls on any path. receivedHeaders captures all // request headers from the first call. @@ -132,7 +142,10 @@ func newGRPCProxySetup(t *testing.T, receivedHeaders *http.Header) (transport *h return nil, fmt.Errorf("write CONNECT: %w", err) } // http.ReadResponse handles partial reads and validates the status line. - cresp, err := http.ReadResponse(bufio.NewReader(conn), nil) + // Wrap conn in a bufferedConn so any bytes pre-fetched by the + // bufio.Reader are not lost before the TLS handshake consumes them. + br := bufio.NewReader(conn) + cresp, err := http.ReadResponse(br, nil) if err != nil { conn.Close() return nil, fmt.Errorf("read CONNECT response: %w", err) @@ -149,7 +162,7 @@ func newGRPCProxySetup(t *testing.T, receivedHeaders *http.Header) (transport *h ServerName: serverName, NextProtos: []string{http2.NextProtoTLS}, } - tlsConn := tls.Client(conn, tlsCfg) + tlsConn := tls.Client(&bufferedConn{Conn: conn, r: br}, tlsCfg) if err := tlsConn.HandshakeContext(ctx); err != nil { conn.Close() return nil, fmt.Errorf("TLS handshake: %w", err) diff --git a/proxy/proxy.go b/proxy/proxy.go index 57606f9..472e7ef 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -2184,21 +2184,38 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req // When the client negotiated h2 (e.g., gRPC), the request object is an // h2 request and cannot be round-tripped via an HTTP/1.1 transport // without framing errors, so we must forward upstream over h2 as well. + // + // Limitation: http2.Transport never falls back to HTTP/1.1, so if the + // upstream only speaks HTTP/1.1 the connection will fail when the client + // has negotiated h2. For gRPC this is always correct (gRPC requires h2); + // for general h2 clients hitting h1-only upstreams it is a known + // limitation of the current implementation. var transport http.RoundTripper + dialer := &net.Dialer{ + Timeout: 30 * time.Second, + KeepAlive: 30 * time.Second, + } if tlsClientConn.ConnectionState().NegotiatedProtocol == http2.NextProtoTLS { - transport = &http2.Transport{TLSClientConfig: upstreamTLS} + transport = &http2.Transport{ + TLSClientConfig: upstreamTLS, + DialTLSContext: func(ctx context.Context, network, addr string, cfg *tls.Config) (net.Conn, error) { + if cfg == nil { + cfg = upstreamTLS + } + return tls.DialWithDialer(dialer, network, addr, cfg) + }, + ReadIdleTimeout: 30 * time.Second, + PingTimeout: 15 * time.Second, + } } else { transport = &http.Transport{ - Proxy: nil, - DialContext: (&net.Dialer{ - Timeout: 30 * time.Second, - KeepAlive: 30 * time.Second, - }).DialContext, - TLSClientConfig: upstreamTLS, - TLSHandshakeTimeout: 10 * time.Second, + Proxy: nil, + DialContext: dialer.DialContext, + TLSClientConfig: upstreamTLS, // Do NOT set ForceAttemptHTTP2: this path handles HTTP/1.1 // requests. Enabling h2 upstream for h1 clients causes // framing mismatches. + TLSHandshakeTimeout: 10 * time.Second, ResponseHeaderTimeout: 5 * time.Minute, MaxIdleConns: 100, IdleConnTimeout: 90 * time.Second, @@ -2563,6 +2580,11 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req if err := http2.ConfigureServer(srv, nil); err != nil { slog.Warn("http2.ConfigureServer failed, falling back to HTTP/1.1", "subsystem", "proxy", "host", host, "error", err) + // If the client already negotiated h2 we cannot serve it correctly + // over h1 — close rather than produce a framing error. + if tlsClientConn.ConnectionState().NegotiatedProtocol == http2.NextProtoTLS { + return + } } _ = srv.Serve(ln) } From 55c385b5a0e26935a91721203bad9dc136c0d0bb Mon Sep 17 00:00:00 2001 From: Andy Bonventre Date: Tue, 23 Jun 2026 02:04:43 +0200 Subject: [PATCH 3/4] fix(proxy): context-aware upstream TLS dial and promote x/net to direct dep Replace tls.DialWithDialer (no context support) with tls.Dialer.DialContext so upstream h2 dials respect cancellation and deadlines from the client request context, matching the h1 path's dialer.DialContext behavior. Run go mod tidy to promote golang.org/x/net to a direct require entry now that proxy/proxy.go imports golang.org/x/net/http2 directly. --- go.mod | 2 +- go.sum | 10 ---------- proxy/proxy.go | 2 +- 3 files changed, 2 insertions(+), 12 deletions(-) diff --git a/go.mod b/go.mod index 0dad8d7..bf3c99d 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,7 @@ require ( go.opentelemetry.io/otel/sdk/log v0.19.0 go.opentelemetry.io/otel/sdk/metric v1.43.0 go.opentelemetry.io/otel/trace v1.43.0 + golang.org/x/net v0.56.0 golang.org/x/sync v0.21.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -131,7 +132,6 @@ require ( go4.org v0.0.0-20230225012048-214862532bf5 // indirect golang.org/x/crypto v0.53.0 // indirect golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa // indirect - golang.org/x/net v0.56.0 // indirect golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sys v0.46.0 // indirect golang.org/x/text v0.38.0 // indirect diff --git a/go.sum b/go.sum index 6949cf7..cff69de 100644 --- a/go.sum +++ b/go.sum @@ -379,8 +379,6 @@ golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8U golang.org/x/crypto v0.0.0-20190605123033-f99c8df09eb5/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= -golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4= -golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA= golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= @@ -427,8 +425,6 @@ golang.org/x/net v0.0.0-20200202094626-16171245cfb2/go.mod h1:z5CRVTTTmAJ677TzLL golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= -golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0= -golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw= golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= @@ -445,8 +441,6 @@ golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -470,8 +464,6 @@ golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= -golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= @@ -485,8 +477,6 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= -golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8= -golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA= golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= diff --git a/proxy/proxy.go b/proxy/proxy.go index 472e7ef..593adb0 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -2202,7 +2202,7 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req if cfg == nil { cfg = upstreamTLS } - return tls.DialWithDialer(dialer, network, addr, cfg) + return (&tls.Dialer{NetDialer: dialer, Config: cfg}).DialContext(ctx, network, addr) }, ReadIdleTimeout: 30 * time.Second, PingTimeout: 15 * time.Second, From d1aff9165f33d7c8f5a4d5c878d5429eeec1b448 Mon Sep 17 00:00:00 2001 From: Andy Bonventre Date: Tue, 23 Jun 2026 02:16:27 +0200 Subject: [PATCH 4/4] refactor(proxy): distribute coverage tests into their logical home files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the omnibus coverage_test.go file and move each test to the file where related tests already live: - Tunnel tests (ForwardsPlainHTTPS, NetworkPolicyBlocked, InvalidHostFormat, NewTokenSubstitution) → proxy_test.go - Intercept path tests (PathRulesWarning, LLMPolicy_DeniedLogged, ResponseTransformer, ResponseTransformer_NoMatch, SetTokenSubstitution) and the newAnthropicInterceptSetup helper → intercept_test.go - LLM gateway integration tests (LLMPolicy_Deny, Allow, ResponseTooLarge, NoEnginePassesThrough) and llmGatewayDenyEditPolicy const → llmpolicy_test.go - Exported host pattern and grant host tests → hosts_test.go --- proxy/coverage_test.go | 603 ---------------------------------------- proxy/hosts_test.go | 35 +++ proxy/intercept_test.go | 269 ++++++++++++++++++ proxy/llmpolicy_test.go | 150 ++++++++++ proxy/proxy_test.go | 103 +++++++ 5 files changed, 557 insertions(+), 603 deletions(-) delete mode 100644 proxy/coverage_test.go diff --git a/proxy/coverage_test.go b/proxy/coverage_test.go deleted file mode 100644 index 8cd94c8..0000000 --- a/proxy/coverage_test.go +++ /dev/null @@ -1,603 +0,0 @@ -package proxy - -// Black-box coverage tests for paths that were previously uncovered. -// Every test drives the proxy through its HTTP interface — no internal -// functions are called directly. - -import ( - "bufio" - "crypto/tls" - "crypto/x509" - "fmt" - "io" - "net" - "net/http" - "net/http/httptest" - "net/url" - "strings" - "testing" - - keeplib "github.com/majorcontext/keep" -) - -// ── Tunnel (no-CA) path ────────────────────────────────────────────────────── - -// TestTunnel_ForwardsPlainHTTPS verifies that when the proxy has no CA -// configured, a CONNECT request is forwarded as a raw TCP tunnel without -// TLS interception. Credentials can still be configured but will not be -// injected (the proxy can't see the plaintext). -func TestTunnel_ForwardsPlainHTTPS(t *testing.T) { - backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fmt.Fprint(w, "tunneled") - })) - t.Cleanup(backend.Close) - - // Proxy with no CA — uses handleConnectTunnel. - p := NewProxy() - proxyServer := httptest.NewServer(p) - t.Cleanup(proxyServer.Close) - - // Client trusts the backend's self-signed cert directly (no MITM). - backendCAs := x509.NewCertPool() - backendCAs.AddCert(backend.Certificate()) - - client := &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyURL(mustParseURL(proxyServer.URL)), - TLSClientConfig: &tls.Config{RootCAs: backendCAs}, - }, - } - - resp, err := client.Get(backend.URL + "/hello") - if err != nil { - t.Fatalf("GET through tunnel: %v", err) - } - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - if string(body) != "tunneled" { - t.Errorf("body = %q, want %q", string(body), "tunneled") - } - if resp.StatusCode != http.StatusOK { - t.Errorf("status = %d, want 200", resp.StatusCode) - } -} - -// TestTunnel_NetworkPolicyBlocked verifies that the network policy is still -// enforced even when no CA is set (tunnel mode). -func TestTunnel_NetworkPolicyBlocked(t *testing.T) { - backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fmt.Fprint(w, "should not reach") - })) - t.Cleanup(backend.Close) - - p := NewProxy() - p.SetNetworkPolicy("strict", nil, nil) - proxyServer := httptest.NewServer(p) - t.Cleanup(proxyServer.Close) - - backendCAs := x509.NewCertPool() - backendCAs.AddCert(backend.Certificate()) - - client := &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyURL(mustParseURL(proxyServer.URL)), - TLSClientConfig: &tls.Config{RootCAs: backendCAs}, - }, - } - - resp, err := client.Get(backend.URL + "/hello") - if err != nil { - // Connection refused or proxy error — also acceptable as a block. - return - } - defer resp.Body.Close() - if resp.StatusCode == http.StatusOK { - t.Errorf("expected blocked response, got 200") - } -} - -// TestTunnel_InvalidHostFormat verifies that a malformed CONNECT target -// (missing port) returns a 400 Bad Request. -func TestTunnel_InvalidHostFormat(t *testing.T) { - p := NewProxy() - proxyServer := httptest.NewServer(p) - t.Cleanup(proxyServer.Close) - - // Issue a raw CONNECT with a host that has no port — SplitHostPort will fail. - conn, err := net.Dial("tcp", proxyServer.Listener.Addr().String()) - if err != nil { - t.Fatalf("dial proxy: %v", err) - } - defer conn.Close() - - fmt.Fprintf(conn, "CONNECT noporthost HTTP/1.1\r\nHost: noporthost\r\n\r\n") - - resp, err := http.ReadResponse(bufio.NewReader(conn), nil) - if err != nil { - t.Fatalf("read response: %v", err) - } - resp.Body.Close() - if resp.StatusCode != http.StatusBadRequest { - t.Errorf("status = %d, want 400", resp.StatusCode) - } -} - -// TestTunnel_PathRulesWarning exercises the code path where TLS interception -// is disabled (no CA) but per-path rules are configured — the proxy should -// log a warning and fall through to the tunnel handler. -func TestTunnel_PathRulesWarning(t *testing.T) { - backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fmt.Fprint(w, "ok") - })) - t.Cleanup(backend.Close) - - p := NewProxy() - // Inject a PathRulesChecker that claims path rules exist for any host. - p.SetContextResolver(func(token string) (*RunContextData, bool) { - return &RunContextData{ - Policy: "permissive", - PathRulesCheck: func(host string, port int) bool { - return true // pretend per-path rules exist - }, - }, true - }) - - proxyServer := httptest.NewServer(p) - t.Cleanup(proxyServer.Close) - - backendCAs := x509.NewCertPool() - backendCAs.AddCert(backend.Certificate()) - - proxyURL := mustParseURL(proxyServer.URL) - proxyURL.User = url.UserPassword("user", "tok") - - client := &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyURL(proxyURL), - TLSClientConfig: &tls.Config{RootCAs: backendCAs}, - }, - } - - // The proxy has no CA, so it tunnels despite path rules existing. - resp, err := client.Get(backend.URL + "/path") - if err != nil { - t.Fatalf("GET: %v", err) - } - resp.Body.Close() - // Main assertion: we reach the backend (tunnel works) without panicking. - if resp.StatusCode != http.StatusOK { - t.Errorf("status = %d, want 200", resp.StatusCode) - } -} - -// ── LLM gateway policy (evaluateAndReplaceLLMResponse) ────────────────────── - -const llmGatewayDenyEditPolicy = ` -scope: llm-gateway -mode: enforce -rules: - - name: deny-edit - match: - operation: "llm.tool_use" - when: "params.name == 'edit'" - action: deny - message: "Editing blocked." -` - -// newAnthropicInterceptSetup builds an intercept test setup where the backend -// acts as a fake api.anthropic.com. The client sends requests to -// https://api.anthropic.com: and the proxy rewrites the upstream -// dial to 127.0.0.1: via the HostGateway mechanism, so the -// host check (host == "api.anthropic.com") triggers the LLM policy path. -// -// The backend port is added to AllowedHostPorts so the HostGateway policy -// permits the CONNECT (HostGateway traffic requires explicit port allow-listing). -func newAnthropicInterceptSetup(t *testing.T, llmEng *keeplib.Engine, backendHandler http.Handler) (client *http.Client, backendURL string) { - t.Helper() - - ca, err := generateCA() - if err != nil { - t.Fatal(err) - } - - backend := httptest.NewTLSServer(backendHandler) - t.Cleanup(backend.Close) - - upstreamCAs := x509.NewCertPool() - upstreamCAs.AddCert(backend.Certificate()) - - backendAddr := mustParseURL(backend.URL) - backendPort := backendAddr.Port() - backendPortInt := 0 - fmt.Sscanf(backendPort, "%d", &backendPortInt) - - p := NewProxy() - p.SetCA(ca) - p.SetUpstreamCAs(upstreamCAs) - p.SetContextResolver(func(token string) (*RunContextData, bool) { - if token != "llmtok" { - return nil, false - } - rc := &RunContextData{ - Policy: "permissive", - HostGateway: "api.anthropic.com", - HostGatewayIP: "127.0.0.1", - AllowedHostPorts: []int{backendPortInt}, - } - if llmEng != nil { - rc.KeepEngines = map[string]*keeplib.Engine{"llm-gateway": llmEng} - } - return rc, true - }) - - proxyServer := httptest.NewServer(p) - t.Cleanup(proxyServer.Close) - - clientCAs := x509.NewCertPool() - clientCAs.AppendCertsFromPEM(ca.certPEM) - - proxyURL := mustParseURL(proxyServer.URL) - proxyURL.User = url.UserPassword("user", "llmtok") - - client = &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyURL(proxyURL), - TLSClientConfig: &tls.Config{RootCAs: clientCAs}, - }, - } - - backendURL = "https://api.anthropic.com:" + backendPort - return client, backendURL -} - -// TestIntercept_LLMPolicy_Deny verifies that the llm-gateway Keep engine -// blocks a tool-use response from api.anthropic.com. -func TestIntercept_LLMPolicy_Deny(t *testing.T) { - eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) - if err != nil { - t.Fatal(err) - } - t.Cleanup(eng.Close) - - body := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{"file_path":"/foo"}}],"stop_reason":"tool_use"}` - - client, backendURL := newAnthropicInterceptSetup(t, eng, - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, body) - }), - ) - - resp, err := client.Post(backendURL+"/v1/messages", "application/json", - strings.NewReader(`{"model":"claude-opus-4-5"}`)) - if err != nil { - t.Fatalf("POST: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusBadRequest { - t.Errorf("status = %d, want 400 (policy denied)", resp.StatusCode) - } - if got := resp.Header.Get("X-Moat-Blocked"); got != "llm-policy" { - t.Errorf("X-Moat-Blocked = %q, want llm-policy", got) - } - respBody, _ := io.ReadAll(resp.Body) - if !strings.Contains(string(respBody), "policy_denied") { - t.Errorf("response body missing policy_denied: %s", respBody) - } -} - -// TestIntercept_LLMPolicy_Allow verifies that a non-matching response passes -// through the llm-gateway engine unchanged. -func TestIntercept_LLMPolicy_Allow(t *testing.T) { - eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) - if err != nil { - t.Fatal(err) - } - t.Cleanup(eng.Close) - - allowedBody := `{"content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn"}` - - client, backendURL := newAnthropicInterceptSetup(t, eng, - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, allowedBody) - }), - ) - - resp, err := client.Get(backendURL + "/v1/messages") - if err != nil { - t.Fatalf("GET: %v", err) - } - defer resp.Body.Close() - respBody, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != http.StatusOK { - t.Errorf("status = %d, want 200", resp.StatusCode) - } - if string(respBody) != allowedBody { - t.Errorf("body = %q, want %q", string(respBody), allowedBody) - } -} - -// TestIntercept_LLMPolicy_ResponseTooLarge verifies that oversized responses -// from api.anthropic.com are blocked with a size-limit error. -func TestIntercept_LLMPolicy_ResponseTooLarge(t *testing.T) { - eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) - if err != nil { - t.Fatal(err) - } - t.Cleanup(eng.Close) - - // Build a body larger than maxLLMResponseSize (10 MiB). - // Use 11 MiB of JSON-ish padding so the size check triggers. - hugeBody := `{"content":[{"type":"text","text":"` + strings.Repeat("x", 11*1024*1024) + `"}]}` - - client, backendURL := newAnthropicInterceptSetup(t, eng, - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, hugeBody) - }), - ) - - resp, err := client.Get(backendURL + "/v1/messages") - if err != nil { - t.Fatalf("GET: %v", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusBadRequest { - t.Errorf("status = %d, want 400 (size-limit)", resp.StatusCode) - } - if got := resp.Header.Get("X-Moat-Blocked"); got != "llm-policy" { - t.Errorf("X-Moat-Blocked = %q, want llm-policy", got) - } - respBody, _ := io.ReadAll(resp.Body) - if !strings.Contains(string(respBody), "size-limit") { - t.Errorf("response body missing size-limit: %s", respBody) - } -} - -// TestIntercept_LLMPolicy_NoEnginePassesThrough verifies that without a -// llm-gateway engine the response is passed through unmodified, even for -// api.anthropic.com. -func TestIntercept_LLMPolicy_NoEnginePassesThrough(t *testing.T) { - rawBody := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{}}],"stop_reason":"tool_use"}` - - // nil engine — no llm-gateway key in KeepEngines. - client, backendURL := newAnthropicInterceptSetup(t, nil, - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, rawBody) - }), - ) - - resp, err := client.Get(backendURL + "/v1/messages") - if err != nil { - t.Fatalf("GET: %v", err) - } - defer resp.Body.Close() - respBody, _ := io.ReadAll(resp.Body) - - if resp.StatusCode != http.StatusOK { - t.Errorf("status = %d, want 200", resp.StatusCode) - } - if string(respBody) != rawBody { - t.Errorf("body = %q, want %q", string(respBody), rawBody) - } -} - -// TestIntercept_LLMPolicy_DeniedLogged verifies that a denied LLM response is -// recorded in the canonical request log with Denied=true. -func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) { - eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) - if err != nil { - t.Fatal(err) - } - t.Cleanup(eng.Close) - - toolBody := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{"file_path":"/f"}}],"stop_reason":"tool_use"}` - - // Build the setup manually so we can attach a log listener to the proxy. - ca, err := generateCA() - if err != nil { - t.Fatal(err) - } - - backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, toolBody) - })) - t.Cleanup(backend.Close) - - upstreamCAs := x509.NewCertPool() - upstreamCAs.AddCert(backend.Certificate()) - - backendPort := 0 - fmt.Sscanf(mustParseURL(backend.URL).Port(), "%d", &backendPort) - - p := NewProxy() - p.SetCA(ca) - p.SetUpstreamCAs(upstreamCAs) - p.SetContextResolver(func(token string) (*RunContextData, bool) { - if token != "logtest" { - return nil, false - } - return &RunContextData{ - Policy: "permissive", - HostGateway: "api.anthropic.com", - HostGatewayIP: "127.0.0.1", - AllowedHostPorts: []int{backendPort}, - KeepEngines: map[string]*keeplib.Engine{"llm-gateway": eng}, - }, true - }) - - var logged RequestLogData - p.SetLogger(func(data RequestLogData) { logged = data }) - - proxyServer := httptest.NewServer(p) - t.Cleanup(proxyServer.Close) - - clientCAs := x509.NewCertPool() - clientCAs.AppendCertsFromPEM(ca.certPEM) - - proxyURL := mustParseURL(proxyServer.URL) - proxyURL.User = url.UserPassword("user", "logtest") - client := &http.Client{ - Transport: &http.Transport{ - Proxy: http.ProxyURL(proxyURL), - TLSClientConfig: &tls.Config{RootCAs: clientCAs}, - }, - } - - resp, err := client.Post( - fmt.Sprintf("https://api.anthropic.com:%d/v1/messages", backendPort), - "application/json", - strings.NewReader(`{"model":"claude-opus-4-5"}`), - ) - if err != nil { - t.Fatalf("POST: %v", err) - } - io.ReadAll(resp.Body) - resp.Body.Close() - - if resp.StatusCode != http.StatusBadRequest { - t.Errorf("status = %d, want 400", resp.StatusCode) - } - if !logged.Denied { - t.Errorf("RequestLogData.Denied = false, want true") - } -} - -// ── AddResponseTransformer ─────────────────────────────────────────────────── - -// TestIntercept_ResponseTransformer verifies that a registered response -// transformer runs in the intercept path and can observe the response. -func TestIntercept_ResponseTransformer(t *testing.T) { - var transformerCalled bool - setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, `{"original":true}`) - })) - - setup.Proxy.AddResponseTransformer(setup.BackendHost, func(req, resp any) (any, bool) { - transformerCalled = true - // Return false — observe without modifying (no replacement). - return resp, false - }) - - resp, err := setup.Client.Get(setup.Backend.URL + "/data") - if err != nil { - t.Fatalf("GET: %v", err) - } - defer resp.Body.Close() - io.ReadAll(resp.Body) - - if !transformerCalled { - t.Error("response transformer was not called") - } - if resp.StatusCode != http.StatusOK { - t.Errorf("status = %d, want 200", resp.StatusCode) - } -} - -// TestIntercept_ResponseTransformer_NoMatch verifies that a transformer -// registered for a different host does not affect other hosts. -func TestIntercept_ResponseTransformer_NoMatch(t *testing.T) { - originalBody := `{"original":true}` - setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - fmt.Fprint(w, originalBody) - })) - - // Register transformer for a completely different host. - setup.Proxy.AddResponseTransformer("other.example.com", func(req, resp any) (any, bool) { - // This should never be called for requests to the backend host. - t.Error("transformer called for wrong host") - return resp, false - }) - - resp, err := setup.Client.Get(setup.Backend.URL + "/data") - if err != nil { - t.Fatalf("GET: %v", err) - } - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - - if string(body) != originalBody { - t.Errorf("body = %q, want %q (transformer should not apply)", string(body), originalBody) - } -} - -// ── SetTokenSubstitution (proxy-level setter) ──────────────────────────────── - -// TestIntercept_SetTokenSubstitution verifies that the proxy-level -// SetTokenSubstitution setter is wired into the intercept path. -func TestIntercept_SetTokenSubstitution(t *testing.T) { - var receivedPath string - setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - receivedPath = r.URL.Path - fmt.Fprint(w, "ok") - })) - - // Register token substitution: placeholder in URL → real token. - setup.Proxy.SetTokenSubstitution(setup.BackendHost, "placeholder-token", "real-secret-value") - - resp, err := setup.Client.Get(setup.Backend.URL + "/bot" + "placeholder-token" + "/getUpdates") - if err != nil { - t.Fatalf("GET: %v", err) - } - resp.Body.Close() - - if !strings.Contains(receivedPath, "real-secret-value") { - t.Errorf("path = %q, want real-secret-value substituted", receivedPath) - } - if strings.Contains(receivedPath, "placeholder-token") { - t.Errorf("path = %q, placeholder should have been replaced", receivedPath) - } -} - -// ── Exported API wrappers ──────────────────────────────────────────────────── - -// TestNewTokenSubstitution verifies the exported constructor returns a usable substitution. -func TestNewTokenSubstitution(t *testing.T) { - sub := NewTokenSubstitution("placeholder", "real") - if sub == nil { - t.Fatal("NewTokenSubstitution returned nil") - } -} - -// TestParseHostPattern and TestMatchesHostPattern exercise the exported wrappers -// for parseHostPattern and matchesPattern. -func TestParseAndMatchHostPattern(t *testing.T) { - cases := []struct { - pattern string - host string - port int - want bool - }{ - {"api.example.com", "api.example.com", 443, true}, - {"api.example.com", "other.example.com", 443, false}, - {"*.example.com", "sub.example.com", 443, true}, - {"*.example.com", "example.com", 443, false}, - } - for _, tc := range cases { - p := ParseHostPattern(tc.pattern) - got := MatchesHostPattern(p, tc.host, tc.port) - if got != tc.want { - t.Errorf("MatchesHostPattern(%q, %q, %d) = %v, want %v", tc.pattern, tc.host, tc.port, got, tc.want) - } - } -} - -// TestRegisterGrantHosts verifies that registered grant hosts are retrievable. -func TestRegisterGrantHosts(t *testing.T) { - RegisterGrantHosts("test-grant-coverage", []string{"coverage.example.com"}) - hosts := GetHostsForGrant("test-grant-coverage") - if len(hosts) == 0 { - t.Fatal("GetHostsForGrant returned empty slice after RegisterGrantHosts") - } - if hosts[0] != "coverage.example.com" { - t.Errorf("hosts[0] = %q, want coverage.example.com", hosts[0]) - } -} diff --git a/proxy/hosts_test.go b/proxy/hosts_test.go index 10346e5..7f0a189 100644 --- a/proxy/hosts_test.go +++ b/proxy/hosts_test.go @@ -538,3 +538,38 @@ func TestMatchesPattern(t *testing.T) { }) } } + +// TestParseAndMatchHostPattern exercises the exported wrappers for +// parseHostPattern and matchesPattern. +func TestParseAndMatchHostPattern(t *testing.T) { + cases := []struct { + pattern string + host string + port int + want bool + }{ + {"api.example.com", "api.example.com", 443, true}, + {"api.example.com", "other.example.com", 443, false}, + {"*.example.com", "sub.example.com", 443, true}, + {"*.example.com", "example.com", 443, false}, + } + for _, tc := range cases { + p := ParseHostPattern(tc.pattern) + got := MatchesHostPattern(p, tc.host, tc.port) + if got != tc.want { + t.Errorf("MatchesHostPattern(%q, %q, %d) = %v, want %v", tc.pattern, tc.host, tc.port, got, tc.want) + } + } +} + +// TestRegisterGrantHosts verifies that registered grant hosts are retrievable. +func TestRegisterGrantHosts(t *testing.T) { + RegisterGrantHosts("test-grant-coverage", []string{"coverage.example.com"}) + hosts := GetHostsForGrant("test-grant-coverage") + if len(hosts) == 0 { + t.Fatal("GetHostsForGrant returned empty slice after RegisterGrantHosts") + } + if hosts[0] != "coverage.example.com" { + t.Errorf("hosts[0] = %q, want coverage.example.com", hosts[0]) + } +} diff --git a/proxy/intercept_test.go b/proxy/intercept_test.go index 39ab0cf..6b7f5ec 100644 --- a/proxy/intercept_test.go +++ b/proxy/intercept_test.go @@ -798,3 +798,272 @@ func TestIntercept_UserID_ContextResolver(t *testing.T) { t.Errorf("UserID = %q, want %q (CONNECT path)", logged.UserID, "alice") } } + +// TestTunnel_PathRulesWarning exercises the code path where TLS interception +// is disabled (no CA) but per-path rules are configured — the proxy should +// fall through to the tunnel handler. +func TestTunnel_PathRulesWarning(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "ok") + })) + t.Cleanup(backend.Close) + + p := NewProxy() + p.SetContextResolver(func(token string) (*RunContextData, bool) { + return &RunContextData{ + Policy: "permissive", + PathRulesCheck: func(host string, port int) bool { + return true + }, + }, true + }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + backendCAs := x509.NewCertPool() + backendCAs.AddCert(backend.Certificate()) + + proxyURL := mustParseURL(proxyServer.URL) + proxyURL.User = url.UserPassword("user", "tok") + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: backendCAs}, + }, + } + + resp, err := client.Get(backend.URL + "/path") + if err != nil { + t.Fatalf("GET: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +// newAnthropicInterceptSetup builds an intercept test setup where the backend +// acts as a fake api.anthropic.com. The client sends requests to +// https://api.anthropic.com: and the proxy rewrites the upstream +// dial to 127.0.0.1: via the HostGateway mechanism, so the +// host check (host == "api.anthropic.com") triggers the LLM policy path. +func newAnthropicInterceptSetup(t *testing.T, llmEng *keeplib.Engine, backendHandler http.Handler) (client *http.Client, backendURL string) { + t.Helper() + + ca, err := generateCA() + if err != nil { + t.Fatal(err) + } + + backend := httptest.NewTLSServer(backendHandler) + t.Cleanup(backend.Close) + + upstreamCAs := x509.NewCertPool() + upstreamCAs.AddCert(backend.Certificate()) + + backendAddr := mustParseURL(backend.URL) + backendPort := backendAddr.Port() + backendPortInt := 0 + fmt.Sscanf(backendPort, "%d", &backendPortInt) + + p := NewProxy() + p.SetCA(ca) + p.SetUpstreamCAs(upstreamCAs) + p.SetContextResolver(func(token string) (*RunContextData, bool) { + if token != "llmtok" { + return nil, false + } + rc := &RunContextData{ + Policy: "permissive", + HostGateway: "api.anthropic.com", + HostGatewayIP: "127.0.0.1", + AllowedHostPorts: []int{backendPortInt}, + } + if llmEng != nil { + rc.KeepEngines = map[string]*keeplib.Engine{"llm-gateway": llmEng} + } + return rc, true + }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + clientCAs := x509.NewCertPool() + clientCAs.AppendCertsFromPEM(ca.certPEM) + + proxyURL := mustParseURL(proxyServer.URL) + proxyURL.User = url.UserPassword("user", "llmtok") + + client = &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: clientCAs}, + }, + } + + backendURL = "https://api.anthropic.com:" + backendPort + return client, backendURL +} + +// TestIntercept_LLMPolicy_DeniedLogged verifies that a denied LLM response is +// recorded in the canonical request log with Denied=true. +func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + toolBody := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{"file_path":"/f"}}],"stop_reason":"tool_use"}` + + ca, err := generateCA() + if err != nil { + t.Fatal(err) + } + + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, toolBody) + })) + t.Cleanup(backend.Close) + + upstreamCAs := x509.NewCertPool() + upstreamCAs.AddCert(backend.Certificate()) + + backendPort := 0 + fmt.Sscanf(mustParseURL(backend.URL).Port(), "%d", &backendPort) + + p := NewProxy() + p.SetCA(ca) + p.SetUpstreamCAs(upstreamCAs) + p.SetContextResolver(func(token string) (*RunContextData, bool) { + if token != "logtest" { + return nil, false + } + return &RunContextData{ + Policy: "permissive", + HostGateway: "api.anthropic.com", + HostGatewayIP: "127.0.0.1", + AllowedHostPorts: []int{backendPort}, + KeepEngines: map[string]*keeplib.Engine{"llm-gateway": eng}, + }, true + }) + + var logged RequestLogData + p.SetLogger(func(data RequestLogData) { logged = data }) + + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + clientCAs := x509.NewCertPool() + clientCAs.AppendCertsFromPEM(ca.certPEM) + + proxyURL := mustParseURL(proxyServer.URL) + proxyURL.User = url.UserPassword("user", "logtest") + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(proxyURL), + TLSClientConfig: &tls.Config{RootCAs: clientCAs}, + }, + } + + resp, err := client.Post( + fmt.Sprintf("https://api.anthropic.com:%d/v1/messages", backendPort), + "application/json", + strings.NewReader(`{"model":"claude-opus-4-5"}`), + ) + if err != nil { + t.Fatalf("POST: %v", err) + } + io.ReadAll(resp.Body) + resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400", resp.StatusCode) + } + if !logged.Denied { + t.Errorf("RequestLogData.Denied = false, want true") + } +} + +// TestIntercept_ResponseTransformer verifies that a registered response +// transformer runs in the intercept path and can observe the response. +func TestIntercept_ResponseTransformer(t *testing.T) { + var transformerCalled bool + setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"original":true}`) + })) + + setup.Proxy.AddResponseTransformer(setup.BackendHost, func(req, resp any) (any, bool) { + transformerCalled = true + return resp, false + }) + + resp, err := setup.Client.Get(setup.Backend.URL + "/data") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + io.ReadAll(resp.Body) + + if !transformerCalled { + t.Error("response transformer was not called") + } + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +// TestIntercept_ResponseTransformer_NoMatch verifies that a transformer +// registered for a different host does not affect other hosts. +func TestIntercept_ResponseTransformer_NoMatch(t *testing.T) { + originalBody := `{"original":true}` + setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, originalBody) + })) + + setup.Proxy.AddResponseTransformer("other.example.com", func(req, resp any) (any, bool) { + t.Error("transformer called for wrong host") + return resp, false + }) + + resp, err := setup.Client.Get(setup.Backend.URL + "/data") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + + if string(body) != originalBody { + t.Errorf("body = %q, want %q (transformer should not apply)", string(body), originalBody) + } +} + +// TestIntercept_SetTokenSubstitution verifies that the proxy-level +// SetTokenSubstitution setter is wired into the intercept path. +func TestIntercept_SetTokenSubstitution(t *testing.T) { + var receivedPath string + setup := newInterceptTestSetup(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath = r.URL.Path + fmt.Fprint(w, "ok") + })) + + setup.Proxy.SetTokenSubstitution(setup.BackendHost, "placeholder-token", "real-secret-value") + + resp, err := setup.Client.Get(setup.Backend.URL + "/bot" + "placeholder-token" + "/getUpdates") + if err != nil { + t.Fatalf("GET: %v", err) + } + resp.Body.Close() + + if !strings.Contains(receivedPath, "real-secret-value") { + t.Errorf("path = %q, want real-secret-value substituted", receivedPath) + } + if strings.Contains(receivedPath, "placeholder-token") { + t.Errorf("path = %q, placeholder should have been replaced", receivedPath) + } +} diff --git a/proxy/llmpolicy_test.go b/proxy/llmpolicy_test.go index 6d5b5c4..d1b2b3c 100644 --- a/proxy/llmpolicy_test.go +++ b/proxy/llmpolicy_test.go @@ -5,7 +5,9 @@ import ( "compress/gzip" "context" "fmt" + "io" "net/http" + "strings" "testing" keeplib "github.com/majorcontext/keep" @@ -272,3 +274,151 @@ rules: result := evaluateLLMResponse(context.Background(), eng, []byte(body), resp) assert.False(t, result.Denied) } + +// llmGatewayDenyEditPolicy is a Keep policy that denies tool_use responses +// where the tool name is "edit". Used by integration tests that drive the +// policy through the proxy's HTTP interface. +const llmGatewayDenyEditPolicy = ` +scope: llm-gateway +mode: enforce +rules: + - name: deny-edit + match: + operation: "llm.tool_use" + when: "params.name == 'edit'" + action: deny + message: "Editing blocked." +` + +// TestIntercept_LLMPolicy_Deny verifies that the llm-gateway Keep engine +// blocks a tool-use response from api.anthropic.com. +func TestIntercept_LLMPolicy_Deny(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + body := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{"file_path":"/foo"}}],"stop_reason":"tool_use"}` + + client, backendURL := newAnthropicInterceptSetup(t, eng, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, body) + }), + ) + + resp, err := client.Post(backendURL+"/v1/messages", "application/json", + strings.NewReader(`{"model":"claude-opus-4-5"}`)) + if err != nil { + t.Fatalf("POST: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400 (policy denied)", resp.StatusCode) + } + if got := resp.Header.Get("X-Moat-Blocked"); got != "llm-policy" { + t.Errorf("X-Moat-Blocked = %q, want llm-policy", got) + } + respBody, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(respBody), "policy_denied") { + t.Errorf("response body missing policy_denied: %s", respBody) + } +} + +// TestIntercept_LLMPolicy_Allow verifies that a non-matching response passes +// through the llm-gateway engine unchanged. +func TestIntercept_LLMPolicy_Allow(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + allowedBody := `{"content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn"}` + + client, backendURL := newAnthropicInterceptSetup(t, eng, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, allowedBody) + }), + ) + + resp, err := client.Get(backendURL + "/v1/messages") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if string(respBody) != allowedBody { + t.Errorf("body = %q, want %q", string(respBody), allowedBody) + } +} + +// TestIntercept_LLMPolicy_ResponseTooLarge verifies that oversized responses +// from api.anthropic.com are blocked with a size-limit error. +func TestIntercept_LLMPolicy_ResponseTooLarge(t *testing.T) { + eng, err := keeplib.LoadFromBytes([]byte(llmGatewayDenyEditPolicy)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(eng.Close) + + hugeBody := `{"content":[{"type":"text","text":"` + strings.Repeat("x", 11*1024*1024) + `"}]}` + + client, backendURL := newAnthropicInterceptSetup(t, eng, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, hugeBody) + }), + ) + + resp, err := client.Get(backendURL + "/v1/messages") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400 (size-limit)", resp.StatusCode) + } + if got := resp.Header.Get("X-Moat-Blocked"); got != "llm-policy" { + t.Errorf("X-Moat-Blocked = %q, want llm-policy", got) + } + respBody, _ := io.ReadAll(resp.Body) + if !strings.Contains(string(respBody), "size-limit") { + t.Errorf("response body missing size-limit: %s", respBody) + } +} + +// TestIntercept_LLMPolicy_NoEnginePassesThrough verifies that without a +// llm-gateway engine the response is passed through unmodified. +func TestIntercept_LLMPolicy_NoEnginePassesThrough(t *testing.T) { + rawBody := `{"content":[{"type":"tool_use","id":"t1","name":"Edit","input":{}}],"stop_reason":"tool_use"}` + + client, backendURL := newAnthropicInterceptSetup(t, nil, + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, rawBody) + }), + ) + + resp, err := client.Get(backendURL + "/v1/messages") + if err != nil { + t.Fatalf("GET: %v", err) + } + defer resp.Body.Close() + respBody, _ := io.ReadAll(resp.Body) + + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } + if string(respBody) != rawBody { + t.Errorf("body = %q, want %q", string(respBody), rawBody) + } +} diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index d7aa5e6..6ddb973 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -3620,3 +3620,106 @@ func TestProxy_CaptureHeaders_AvailableInLogData(t *testing.T) { t.Errorf("RequestHeaders[X-Workspace-Slug] = %q, want sneaky-plum", got) } } + +// TestTunnel_ForwardsPlainHTTPS verifies that when the proxy has no CA +// configured, a CONNECT request is forwarded as a raw TCP tunnel without +// TLS interception. +func TestTunnel_ForwardsPlainHTTPS(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "tunneled") + })) + t.Cleanup(backend.Close) + + p := NewProxy() + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + backendCAs := x509.NewCertPool() + backendCAs.AddCert(backend.Certificate()) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxyServer.URL)), + TLSClientConfig: &tls.Config{RootCAs: backendCAs}, + }, + } + + resp, err := client.Get(backend.URL + "/hello") + if err != nil { + t.Fatalf("GET through tunnel: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if string(body) != "tunneled" { + t.Errorf("body = %q, want %q", string(body), "tunneled") + } + if resp.StatusCode != http.StatusOK { + t.Errorf("status = %d, want 200", resp.StatusCode) + } +} + +// TestTunnel_NetworkPolicyBlocked verifies that the network policy is enforced +// even when no CA is set (tunnel mode). +func TestTunnel_NetworkPolicyBlocked(t *testing.T) { + backend := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, "should not reach") + })) + t.Cleanup(backend.Close) + + p := NewProxy() + p.SetNetworkPolicy("strict", nil, nil) + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + backendCAs := x509.NewCertPool() + backendCAs.AddCert(backend.Certificate()) + + client := &http.Client{ + Transport: &http.Transport{ + Proxy: http.ProxyURL(mustParseURL(proxyServer.URL)), + TLSClientConfig: &tls.Config{RootCAs: backendCAs}, + }, + } + + resp, err := client.Get(backend.URL + "/hello") + if err != nil { + return + } + defer resp.Body.Close() + if resp.StatusCode == http.StatusOK { + t.Errorf("expected blocked response, got 200") + } +} + +// TestTunnel_InvalidHostFormat verifies that a malformed CONNECT target +// (missing port) returns a 400 Bad Request. +func TestTunnel_InvalidHostFormat(t *testing.T) { + p := NewProxy() + proxyServer := httptest.NewServer(p) + t.Cleanup(proxyServer.Close) + + conn, err := net.Dial("tcp", proxyServer.Listener.Addr().String()) + if err != nil { + t.Fatalf("dial proxy: %v", err) + } + defer conn.Close() + + fmt.Fprintf(conn, "CONNECT noporthost HTTP/1.1\r\nHost: noporthost\r\n\r\n") + + resp, err := http.ReadResponse(bufio.NewReader(conn), nil) + if err != nil { + t.Fatalf("read response: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("status = %d, want 400", resp.StatusCode) + } +} + +// TestNewTokenSubstitution verifies the exported constructor returns a usable substitution. +func TestNewTokenSubstitution(t *testing.T) { + sub := NewTokenSubstitution("placeholder", "real") + if sub == nil { + t.Fatal("NewTokenSubstitution returned nil") + } +}