Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 129 additions & 4 deletions internal/proxy/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
Expand All @@ -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
}

Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}

Expand Down
147 changes: 147 additions & 0 deletions internal/proxy/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down