From a2f18622b78cde14a860464947272e9b79a767a2 Mon Sep 17 00:00:00 2001 From: Jacob Clayden Date: Fri, 10 Jul 2026 05:00:43 +0100 Subject: [PATCH] fix: added image proxy handling --- internal/proxy/server.go | 133 +++++++++++++++++++++++++++++- internal/proxy/server_test.go | 147 ++++++++++++++++++++++++++++++++++ 2 files changed, 276 insertions(+), 4 deletions(-) diff --git a/internal/proxy/server.go b/internal/proxy/server.go index 27b7b2c..7c3fb06 100644 --- a/internal/proxy/server.go +++ b/internal/proxy/server.go @@ -30,6 +30,8 @@ const ( codexCompactResponsesPath = "/v1/responses/compact" legacyCodexCompactResponsesPath = "/responses/compact" codexAppServerPath = "/app-server" + codexImagesPathPrefix = "/v1/images/" + legacyCodexImagesPathPrefix = "/images/" ) // RegistryRefresher is the interface for triggering a registry refresh. @@ -121,6 +123,8 @@ func (s *Server) handler() (http.Handler, error) { mux.HandleFunc("POST "+codexCompactResponsesPath, s.handleCodexCompactResponsesRoute) mux.HandleFunc("POST "+legacyCodexCompactResponsesPath, s.handleLegacyCodexCompactResponsesRoute) mux.HandleFunc(codexAppServerPath, s.handleCodexAppServerRoute) + mux.HandleFunc(codexImagesPathPrefix, s.handleCodexImagesRoute) + mux.HandleFunc(legacyCodexImagesPathPrefix, s.handleCodexImagesRoute) mux.HandleFunc("/", s.proxyHandler(upstream)) return mux, nil } @@ -140,7 +144,7 @@ func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) { return } if !s.isValidToken(bearerToken(r)) { - writeError(w, http.StatusForbidden, "authentication_error", "invalid proxy token") + s.rejectInvalidProxyToken(w, r, "models", time.Now(), true) return } @@ -155,6 +159,24 @@ func (s *Server) handleModels(w http.ResponseWriter, r *http.Request) { }) } +func (s *Server) rejectInvalidProxyToken(w http.ResponseWriter, r *http.Request, routeKind string, start time.Time, emitDiag bool) { + fmt.Fprintf(os.Stderr, "cq: reject %s %s: invalid proxy token\n", r.Method, r.URL.Path) + writeError(w, http.StatusForbidden, "authentication_error", "invalid proxy token") + if !emitDiag { + return + } + s.emitDiagnostics(RouteEvent{ + Time: start.UTC(), + Method: r.Method, + Path: r.URL.Path, + Provider: "unknown", + RouteKind: routeKind, + StatusCode: http.StatusForbidden, + LatencyMS: time.Since(start).Milliseconds(), + Error: diagnosticsErrorCode("authentication_error", "invalid proxy token"), + }) +} + // handleCodexNativeModels serves GET /models?client_version=... in the native // Codex ModelsResponse shape. This endpoint is used by Codex CLI to fetch the // available model list. No proxy token is required — the proxy binds to 127.0.0.1 @@ -174,7 +196,7 @@ func (s *Server) handleCodexNativeModels(w http.ResponseWriter, r *http.Request) // snapshot as JSON. Requires a valid local proxy token. func (s *Server) handleRegistry(w http.ResponseWriter, r *http.Request) { if !s.isValidToken(bearerToken(r)) { - writeError(w, http.StatusForbidden, "authentication_error", "invalid proxy token") + s.rejectInvalidProxyToken(w, r, "registry", time.Now(), true) return } if s.Catalog == nil { @@ -192,7 +214,7 @@ func (s *Server) handleRegistry(w http.ResponseWriter, r *http.Request) { // serialisation (Refresher.Refresh acquires its own mutex). func (s *Server) handleRegistryRefresh(w http.ResponseWriter, r *http.Request) { if !s.isValidToken(bearerToken(r)) { - writeError(w, http.StatusForbidden, "authentication_error", "invalid proxy token") + s.rejectInvalidProxyToken(w, r, "registry_refresh", time.Now(), true) return } if s.Refresher == nil { @@ -339,6 +361,109 @@ func (s *Server) handleCodexAppServerRoute(w http.ResponseWriter, r *http.Reques s.proxyCodexAppServer(w, r) } +func (s *Server) handleCodexImagesRoute(w http.ResponseWriter, r *http.Request) { + start := time.Now() + statusCode := 0 + diagError := "" + defer func() { + s.emitDiagnostics(RouteEvent{ + Time: start.UTC(), + Method: r.Method, + Path: r.URL.Path, + Provider: "codex", + RouteKind: "codex_images", + StatusCode: statusCode, + LatencyMS: time.Since(start).Milliseconds(), + Error: diagError, + }) + }() + + if r.Method != http.MethodPost { + statusCode = http.StatusMethodNotAllowed + message := fmt.Sprintf("%s only supports POST", r.URL.Path) + diagError = diagnosticsErrorCode("invalid_request_error", message) + w.Header().Set("Allow", http.MethodPost) + writeError(w, http.StatusMethodNotAllowed, "invalid_request_error", message) + return + } + + if s.CodexTransport == nil { + statusCode = http.StatusServiceUnavailable + diagError = diagnosticsErrorCode("api_error", "no codex accounts configured") + writeError(w, http.StatusServiceUnavailable, "api_error", "no codex accounts configured") + return + } + + var body []byte + if r.Body != nil { + var err error + body, err = io.ReadAll(io.LimitReader(r.Body, maxRequestBody+1)) + r.Body.Close() + if err != nil { + statusCode = http.StatusBadRequest + diagError = diagnosticsErrorCode("invalid_request_error", "failed to read request body") + writeError(w, http.StatusBadRequest, "invalid_request_error", "failed to read request body") + return + } + if len(body) > maxRequestBody { + statusCode = http.StatusRequestEntityTooLarge + diagError = diagnosticsErrorCode("invalid_request_error", "request body exceeds 10 MiB") + writeError(w, http.StatusRequestEntityTooLarge, "invalid_request_error", "request body exceeds 10 MiB") + return + } + } + + upstreamURL := strings.TrimRight(s.Config.CodexUpstream, "/") + codexImagesUpstreamPath(r.URL.EscapedPath()) + if r.URL.RawQuery != "" { + upstreamURL += "?" + r.URL.RawQuery + } + upReq, err := http.NewRequestWithContext(r.Context(), r.Method, upstreamURL, bytes.NewReader(body)) + if err != nil { + statusCode = http.StatusInternalServerError + diagError = diagnosticsErrorCode("api_error", fmt.Sprintf("create upstream request: %v", err)) + writeError(w, http.StatusInternalServerError, "api_error", fmt.Sprintf("create upstream request: %v", err)) + return + } + upReq.ContentLength = int64(len(body)) + upReq.GetBody = func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(body)), nil + } + for key, vals := range r.Header { + for _, v := range vals { + upReq.Header.Add(key, v) + } + } + + fmt.Fprintf(os.Stderr, "cq: route %s %s provider=codex (images)\n", r.Method, r.URL.Path) + resp, err := s.CodexTransport.RoundTrip(upReq) + if err != nil { + statusCode = http.StatusBadGateway + diagError = diagnosticsErrorCode("api_error", fmt.Sprintf("codex upstream error: %v", err)) + writeError(w, http.StatusBadGateway, "api_error", fmt.Sprintf("codex upstream error: %v", err)) + return + } + defer resp.Body.Close() + + fmt.Fprintf(os.Stderr, "cq: proxy %s %s → %d (codex images)\n", r.Method, r.URL.Path, resp.StatusCode) + for key, vals := range resp.Header { + for _, v := range vals { + w.Header().Add(key, v) + } + } + statusCode = resp.StatusCode + w.WriteHeader(resp.StatusCode) + if _, err := io.Copy(w, resp.Body); err != nil { + fmt.Fprintf(os.Stderr, "cq: codex images response copy: %v\n", err) + } +} + +func codexImagesUpstreamPath(path string) string { + if strings.HasPrefix(path, codexImagesPathPrefix) { + return strings.TrimPrefix(path, "/v1") + } + return path +} + func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) { start := time.Now() if wrapped, rec := s.wrapDiagnosticsResponseWriter(w); rec != nil { @@ -973,7 +1098,7 @@ func (s *Server) proxyHandler(upstream *url.URL) http.HandlerFunc { // Auth check: accept local proxy token or a known Claude account token. token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") if !s.isValidToken(token) { - writeError(w, http.StatusForbidden, "authentication_error", "invalid proxy token") + s.rejectInvalidProxyToken(w, r, "proxy_auth", start, diagnosticsAnthropicRouteKind(r.URL.Path) == "") return } diff --git a/internal/proxy/server_test.go b/internal/proxy/server_test.go index c6a0f82..979412d 100644 --- a/internal/proxy/server_test.go +++ b/internal/proxy/server_test.go @@ -1885,6 +1885,153 @@ func TestServer_Handler_CodexResponsesPath_ForwardsWithAuth(t *testing.T) { } } +func TestServer_Handler_CodexImagesPath_ForwardsWithoutProxyToken(t *testing.T) { + tests := []struct { + name string + requestPath string + upstreamPath string + }{ + { + name: "root images path", + requestPath: "/images/generations?api-version=1", + upstreamPath: "/images/generations", + }, + { + name: "openai compatible images path", + requestPath: "/v1/images/generations?api-version=1", + upstreamPath: "/images/generations", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var gotPath, gotQuery, gotAuth, gotAcctID, gotBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotQuery = r.URL.RawQuery + gotAuth = r.Header.Get("Authorization") + gotAcctID = r.Header.Get("ChatGPT-Account-ID") + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("ReadAll() error = %v", err) + } + gotBody = string(body) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"created":123,"data":[{"b64_json":"abc"}]}`)) + })) + defer upstream.Close() + + srv := &Server{ + Config: &Config{ClaudeUpstream: "https://api.anthropic.com", CodexUpstream: upstream.URL, LocalToken: "tok"}, + CodexTransport: &CodexTokenTransport{ + Selector: &fakeCodexSelector{account: &codex.CodexAccount{AccessToken: "codex-tok", AccountID: "acct-1"}}, + Inner: http.DefaultTransport, + }, + } + + handler, err := srv.handler() + if err != nil { + t.Fatalf("handler() error = %v", err) + } + + w := httptest.NewRecorder() + body := `{"model":"gpt-image-1","prompt":"pingy","size":"1024x1024"}` + req := httptest.NewRequest(http.MethodPost, tt.requestPath, strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + handler.ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200, body: %s", w.Code, w.Body.String()) + } + if strings.Contains(w.Body.String(), "invalid proxy token") { + t.Fatalf("body contains invalid proxy token: %s", w.Body.String()) + } + if gotPath != tt.upstreamPath { + t.Errorf("upstream path = %q, want %s", gotPath, tt.upstreamPath) + } + if gotQuery != "api-version=1" { + t.Errorf("upstream query = %q, want api-version=1", gotQuery) + } + if gotAuth != "Bearer codex-tok" { + t.Errorf("upstream auth = %q, want Bearer codex-tok", gotAuth) + } + if gotAcctID != "acct-1" { + t.Errorf("upstream account-id = %q, want acct-1", gotAcctID) + } + if gotBody != body { + t.Errorf("upstream body = %q, want %q", gotBody, body) + } + }) + } +} + +func TestServer_Handler_CodexImagesPath_RequiresPost(t *testing.T) { + srv := &Server{ + Config: &Config{ClaudeUpstream: "https://api.anthropic.com", CodexUpstream: "https://chatgpt.com/backend-api/codex", LocalToken: "tok"}, + CodexTransport: http.DefaultTransport, + } + + handler, err := srv.handler() + if err != nil { + t.Fatalf("handler() error = %v", err) + } + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/images/generations", nil) + handler.ServeHTTP(w, req) + + if w.Code != http.StatusMethodNotAllowed { + t.Fatalf("status = %d, want 405, body: %s", w.Code, w.Body.String()) + } + if strings.Contains(w.Body.String(), "invalid proxy token") { + t.Fatalf("body contains invalid proxy token: %s", w.Body.String()) + } + if allow := w.Header().Get("Allow"); allow != http.MethodPost { + t.Fatalf("Allow = %q, want POST", allow) + } +} + +func TestServer_GenericProxy_InvalidTokenEmitsDiagnosticsForUnknownPath(t *testing.T) { + path := filepath.Join(t.TempDir(), "diag.jsonl") + diag, err := OpenDiagnosticsWriter(path) + if err != nil { + t.Fatalf("OpenDiagnosticsWriter: %v", err) + } + defer diag.Close() + + srv := &Server{ + Config: &Config{ClaudeUpstream: "https://api.anthropic.com", CodexUpstream: "https://chatgpt.com/backend-api/codex", LocalToken: "tok"}, + Diag: diag, + } + + handler, err := srv.handler() + if err != nil { + t.Fatalf("handler() error = %v", err) + } + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/unknown/imagegen/path?secret=nope", nil) + handler.ServeHTTP(w, req) + + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403", w.Code) + } + events := waitForDiagnosticsEvents(t, path, 1) + if len(events) != 1 { + t.Fatalf("events = %d, want 1", len(events)) + } + ev := events[0] + if ev.Path != "/unknown/imagegen/path" { + t.Fatalf("event path = %q, want path without query", ev.Path) + } + if ev.RouteKind != "proxy_auth" { + t.Fatalf("RouteKind = %q, want proxy_auth", ev.RouteKind) + } + if ev.Error != "authentication_error:invalid_proxy_token" { + t.Fatalf("Error = %q, want authentication_error:invalid_proxy_token", ev.Error) + } +} + func TestServer_Handler_LegacyCodexResponsesPost_Compatibility(t *testing.T) { var gotPath string upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {