diff --git a/internal/credentials/chatgpt.go b/internal/credentials/chatgpt.go index fa964b4..a8265c5 100644 --- a/internal/credentials/chatgpt.go +++ b/internal/credentials/chatgpt.go @@ -119,6 +119,9 @@ func (c *ChatGPTInjector) Inject(req *http.Request) bool { log.Printf("ERROR chatgpt credential initialization failed") return false } + if c.document.tokens.RefreshToken == "" { + return false + } if c.needsRefreshLocked() { if err := c.refreshLocked(); err != nil { log.Printf("ERROR chatgpt credential refresh failed") @@ -147,16 +150,16 @@ func (c *ChatGPTInjector) ensureLoadedLocked() error { } c.loaded = true - if c.config.AuthPath == "" { + if c.config.AuthPath == "" && c.config.StatePath == "" { c.loadErr = errors.New("auth source is not configured") return c.loadErr } - if c.config.StatePath != "" && samePath(c.config.AuthPath, c.config.StatePath) { + if c.config.AuthPath != "" && c.config.StatePath != "" && samePath(c.config.AuthPath, c.config.StatePath) { c.loadErr = errors.New("auth source and state paths must differ") return c.loadErr } - path := c.config.AuthPath + path := "" if c.config.StatePath != "" { if _, err := os.Stat(c.config.StatePath); err == nil { path = c.config.StatePath @@ -165,7 +168,16 @@ func (c *ChatGPTInjector) ensureLoadedLocked() error { return c.loadErr } } + if path == "" { + path = c.config.AuthPath + } + if path == "" { + return nil + } + return c.loadFromFileLocked(path) +} +func (c *ChatGPTInjector) loadFromFileLocked(path string) error { data, err := readPrivateSecretFile(path) if err != nil { c.loadErr = errors.New("auth file is unavailable or has insecure permissions") @@ -288,6 +300,66 @@ func (c *ChatGPTInjector) refreshLocked() error { return nil } +// AcceptLoginTokens processes a successful login exchange response, +// persists the real tokens to StatePath, and updates in-memory state. +func (c *ChatGPTInjector) AcceptLoginTokens(responseBody []byte) error { + c.mu.Lock() + defer c.mu.Unlock() + + var raw struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + IDToken string `json:"id_token"` + ExpiresIn int64 `json:"expires_in"` + } + if err := json.Unmarshal(responseBody, &raw); err != nil || raw.AccessToken == "" || raw.RefreshToken == "" { + return errors.New("login response is malformed or incomplete") + } + + tokens := chatGPTTokens{ + AccessToken: raw.AccessToken, + RefreshToken: raw.RefreshToken, + IDToken: raw.IDToken, + } + doc := chatGPTDocument{ + fields: make(map[string]json.RawMessage), + tokens: tokens, + } + authMode, _ := json.Marshal(chatGPTDefaultAuthMode) + doc.fields["auth_mode"] = authMode + rawTokens, err := json.Marshal(tokens) + if err != nil { + return errors.New("marshal login tokens") + } + doc.fields["tokens"] = rawTokens + lastRefresh, _ := json.Marshal(c.config.Now().UTC().Format(time.RFC3339)) + doc.fields["last_refresh"] = lastRefresh + + if doc.accountID() == "" { + return errors.New("login response missing account id") + } + + if c.config.StatePath == "" { + return errors.New("state path not configured") + } + data, err := json.Marshal(doc.fields) + if err != nil { + return errors.New("marshal login state") + } + if err := atomicWritePrivateSecret(c.config.StatePath, data); err != nil { + return errors.New("persist login state") + } + + c.document = doc + c.expiresAt = jwtExpiry(raw.AccessToken) + if c.expiresAt.IsZero() && raw.ExpiresIn > 0 { + c.expiresAt = c.config.Now().Add(time.Duration(raw.ExpiresIn) * time.Second) + } + c.loaded = true + c.loadErr = nil + return nil +} + func (d chatGPTDocument) accountID() string { if d.tokens.AccountID != "" { return d.tokens.AccountID diff --git a/internal/credentials/chatgpt_test.go b/internal/credentials/chatgpt_test.go index 2284cac..fc1ae93 100644 --- a/internal/credentials/chatgpt_test.go +++ b/internal/credentials/chatgpt_test.go @@ -305,8 +305,12 @@ func TestChatGPTTokenVending(t *testing.T) { if IsChatGPTTokenExchange(&http.Request{Method: http.MethodPost, URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/other"}}) { t.Error("wrong path must not be treated as token exchange") } - vendor := NewChatGPTTokenVendor() - response := vendor.HandleTokenExchange(&http.Request{Method: http.MethodPost, URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}}) + vendor := NewChatGPTTokenVendor(nil) + response := vendor.HandleTokenExchange(&http.Request{ + Method: http.MethodPost, + URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}, + Body: io.NopCloser(strings.NewReader("grant_type=refresh_token")), + }) if response == nil || response.StatusCode != http.StatusOK { t.Fatal("ChatGPT token exchange was not handled") } @@ -319,6 +323,125 @@ func TestChatGPTTokenVending(t *testing.T) { } } +func TestChatGPTStatePathOnly_NoFileYet(t *testing.T) { + dir := t.TempDir() + statePath := filepath.Join(dir, "state", "auth.json") + + injector := NewChatGPTInjector("", statePath) + if !injector.Available() { + t.Fatal("StatePath-only injector should be available even before login completes") + } + + req := &http.Request{Header: make(http.Header)} + req.Header.Set("Authorization", "Bearer agent-dummy") + if injector.Inject(req) { + t.Error("Inject should return false when no tokens are loaded yet") + } + if req.Header.Get("Authorization") != "Bearer agent-dummy" { + t.Error("failed injection should not modify the agent header") + } +} + +func TestChatGPTStatePathOnly_FileExists(t *testing.T) { + dir := t.TempDir() + statePath := filepath.Join(dir, "auth.json") + access := testJWT(map[string]any{"exp": time.Now().Add(time.Hour).Unix()}) + writePrivateAuth(t, statePath, testAuthJSON(access, "refresh", "", "account")) + + injector := NewChatGPTInjector("", statePath) + if !injector.Available() { + t.Fatal("StatePath-only with existing file should be available") + } + + req := &http.Request{Header: make(http.Header)} + if !injector.Inject(req) { + t.Fatal("Inject should succeed when StatePath has valid tokens") + } + if req.Header.Get("Authorization") != "Bearer "+access { + t.Error("access token was not injected") + } + if req.Header.Get(chatGPTAccountHeader) != "account" { + t.Error("account ID was not injected") + } +} + +func TestChatGPTAcceptLoginTokens(t *testing.T) { + dir := t.TempDir() + statePath := filepath.Join(dir, "state", "auth.json") + now := time.Unix(1_700_000_000, 0) + access := testJWT(map[string]any{ + "exp": now.Add(time.Hour).Unix(), + "chatgpt_account_id": "logged-in-account", + }) + + injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{ + StatePath: statePath, + Now: func() time.Time { return now }, + }) + + if !injector.Available() { + t.Fatal("StatePath-only injector should be available") + } + req := &http.Request{Header: make(http.Header)} + if injector.Inject(req) { + t.Error("Inject should return false before login") + } + + loginResp, _ := json.Marshal(map[string]any{ + "access_token": access, + "refresh_token": "real-refresh", + "id_token": testJWT(map[string]any{"chatgpt_account_id": "logged-in-account"}), + "expires_in": 3600, + }) + if err := injector.AcceptLoginTokens(loginResp); err != nil { + t.Fatalf("AcceptLoginTokens failed: %v", err) + } + + req2 := &http.Request{Header: make(http.Header)} + if !injector.Inject(req2) { + t.Fatal("Inject should succeed after AcceptLoginTokens") + } + if req2.Header.Get("Authorization") != "Bearer "+access { + t.Error("access token was not injected after login") + } + if req2.Header.Get(chatGPTAccountHeader) != "logged-in-account" { + t.Error("account ID was not injected after login") + } + + info, err := os.Stat(statePath) + if err != nil { + t.Fatalf("state file was not persisted: %v", err) + } + if info.Mode().Perm() != 0600 { + t.Errorf("state permissions = %o, want 600", info.Mode().Perm()) + } + persisted, _ := os.ReadFile(statePath) + if !bytes.Contains(persisted, []byte("real-refresh")) { + t.Error("refresh token was not persisted") + } +} + +func TestChatGPTAcceptLoginTokens_Malformed(t *testing.T) { + injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{ + StatePath: filepath.Join(t.TempDir(), "auth.json"), + Now: time.Now, + }) + + cases := map[string][]byte{ + "not json": []byte("not-json"), + "missing access": []byte(`{"refresh_token":"r"}`), + "missing refresh": []byte(`{"access_token":"a"}`), + "missing account_id": []byte(`{"access_token":"opaque","refresh_token":"r"}`), + } + for name, body := range cases { + t.Run(name, func(t *testing.T) { + if err := injector.AcceptLoginTokens(body); err == nil { + t.Error("AcceptLoginTokens should reject malformed input") + } + }) + } +} + func TestChatGPTNoCredentialAndNoLogLeakage(t *testing.T) { if NewChatGPTInjector("", "").Available() { t.Error("missing auth path should be unavailable") diff --git a/internal/credentials/config.go b/internal/credentials/config.go index b565995..79a957f 100644 --- a/internal/credentials/config.go +++ b/internal/credentials/config.go @@ -99,15 +99,21 @@ func BuildFromConfig(cfg *CredentialConfig) (*Store, *TokenVendor, map[string][] domainMap := make(map[string][]string) gcpADCJSON := os.Getenv("GCP_ADC_JSON") + chatGPTStatePath := os.Getenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE") for _, entry := range cfg.Credentials { value := os.Getenv(entry.EnvVar) - // For gcloud entries, GCP_ADC_JSON takes precedence over the file path - // and allows processing even if GOOGLE_APPLICATION_CREDENTIALS is unset. - if entry.InjectorType == "gcloud" && gcpADCJSON == "" && value == "" { - continue - } else if value == "" && entry.InjectorType != "gcloud" { + hasSource := value != "" + if !hasSource { + switch entry.InjectorType { + case "gcloud": + hasSource = gcpADCJSON != "" + case "chatgpt": + hasSource = chatGPTStatePath != "" + } + } + if !hasSource { continue } @@ -143,14 +149,15 @@ func BuildFromConfig(cfg *CredentialConfig) (*Store, *TokenVendor, map[string][] tokenVendor.googleEnabled = true log.Println("Token vendor: ENABLED (returns dummy tokens for oauth2.googleapis.com/token)") case "chatgpt": - chatGPTInjector := NewChatGPTInjector(value, os.Getenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE")) + chatGPTInjector := NewChatGPTInjector(value, chatGPTStatePath) if !chatGPTInjector.Available() { - log.Printf("WARN: %s auth file is not loadable", entry.EnvVar) + log.Printf("WARN: chatgpt auth is not loadable") continue } injector = chatGPTInjector tokenVendor = ensureTokenVendor(tokenVendor) tokenVendor.chatGPTEnabled = true + tokenVendor.chatGPTInjector = chatGPTInjector log.Println("Token vendor: ENABLED (returns dummy tokens for auth.openai.com/oauth/token)") } diff --git a/internal/credentials/config_test.go b/internal/credentials/config_test.go index e26c628..ba555a1 100644 --- a/internal/credentials/config_test.go +++ b/internal/credentials/config_test.go @@ -566,4 +566,67 @@ func TestBuildFromConfig_ChatGPT(t *testing.T) { if req.Header.Get("ChatGPT-Account-ID") != "account" { t.Error("ChatGPT route did not inject account ID") } + if tokenVendor.chatGPTInjector == nil { + t.Error("ChatGPT injector should be wired to token vendor") + } +} + +func TestBuildFromConfig_ChatGPT_StateFileOnly(t *testing.T) { + dir := t.TempDir() + statePath := filepath.Join(dir, "state", "auth.json") + t.Setenv("CHATGPT_AUTH_FILE", "") + t.Setenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE", statePath) + + cfg := &CredentialConfig{Credentials: []CredentialEntry{ + { + EnvVar: "CHATGPT_AUTH_FILE", + InjectorType: "chatgpt", + Params: map[string]string{"path_prefix": "/backend-api/codex"}, + Domains: []string{"chatgpt.com"}, + }, + }} + store, tokenVendor, domainMap := BuildFromConfig(cfg) + if tokenVendor == nil { + t.Fatal("ChatGPT token vendor should be enabled with state file only") + } + if !tokenVendor.chatGPTEnabled { + t.Error("chatGPTEnabled should be true") + } + if tokenVendor.chatGPTInjector == nil { + t.Error("chatGPTInjector should be wired to token vendor") + } + if got := domainMap["CHATGPT_AUTH_FILE"]; len(got) < 2 || got[1] != "auth.openai.com" { + t.Errorf("domain map = %v, want API and OAuth domains", got) + } + + req := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Host: "chatgpt.com", Path: "/backend-api/codex/responses"}, + Header: make(http.Header), + } + matched, injected := store.InjectCredentials(req) + if !matched { + t.Fatal("route should match even before login") + } + if injected { + t.Error("should not inject before login (no tokens yet)") + } +} + +func TestBuildFromConfig_ChatGPT_NeitherSet(t *testing.T) { + t.Setenv("CHATGPT_AUTH_FILE", "") + t.Setenv("PAUDE_PROXY_CHATGPT_AUTH_STATE_FILE", "") + + cfg := &CredentialConfig{Credentials: []CredentialEntry{ + { + EnvVar: "CHATGPT_AUTH_FILE", + InjectorType: "chatgpt", + Params: map[string]string{"path_prefix": "/backend-api/codex"}, + Domains: []string{"chatgpt.com"}, + }, + }} + _, tokenVendor, _ := BuildFromConfig(cfg) + if tokenVendor != nil && tokenVendor.chatGPTEnabled { + t.Error("neither env var set should not enable ChatGPT token vending") + } } diff --git a/internal/credentials/token_vending.go b/internal/credentials/token_vending.go index e056f03..f01b7c1 100644 --- a/internal/credentials/token_vending.go +++ b/internal/credentials/token_vending.go @@ -4,9 +4,11 @@ import ( "bytes" "encoding/base64" "encoding/json" + "fmt" "io" "log" "net/http" + "net/url" "strings" ) @@ -34,8 +36,9 @@ func errorResponse(statusCode int, message string) *http.Response { // This means the agent never sees any real credential — not the refresh token, // not the service account key, and not even a short-lived access token. type TokenVendor struct { - googleEnabled bool - chatGPTEnabled bool + googleEnabled bool + chatGPTEnabled bool + chatGPTInjector *ChatGPTInjector } // NewTokenVendor creates a token vendor. @@ -44,9 +47,52 @@ func NewTokenVendor() *TokenVendor { } // NewChatGPTTokenVendor creates a vendor for the Codex ChatGPT OAuth token -// endpoint. The response contains only synthetic values. -func NewChatGPTTokenVendor() *TokenVendor { - return &TokenVendor{chatGPTEnabled: true} +// endpoint. Login-completing exchanges are forwarded to the real endpoint and +// persisted via the injector; refresh requests return synthetic values only. +func NewChatGPTTokenVendor(injector *ChatGPTInjector) *TokenVendor { + return &TokenVendor{chatGPTEnabled: true, chatGPTInjector: injector} +} + +var allowedLoginParams = map[string][]string{ + "authorization_code": { + "grant_type", "code", "redirect_uri", "client_id", "code_verifier", + }, + "urn:ietf:params:oauth:grant-type:device_code": { + "grant_type", "device_code", "client_id", + }, + "urn:ietf:params:oauth:grant-type:token-exchange": { + "grant_type", "client_id", "requested_token", "subject_token", "subject_token_type", + }, +} + +func sanitizeLoginForm(agentValues url.Values, clientID string) (url.Values, error) { + grantType := agentValues.Get("grant_type") + if grantType == "" { + return nil, fmt.Errorf("missing grant_type") + } + allowed, known := allowedLoginParams[grantType] + if !known { + return nil, fmt.Errorf("unsupported grant_type %q", grantType) + } + + allowedSet := make(map[string]bool, len(allowed)) + for _, p := range allowed { + allowedSet[p] = true + } + for key := range agentValues { + if !allowedSet[key] { + log.Printf("LOGIN_SANITIZE stripped disallowed parameter %q from %s grant", key, grantType) + } + } + + sanitized := make(url.Values, len(allowed)) + for _, param := range allowed { + if v := agentValues.Get(param); v != "" { + sanitized.Set(param, v) + } + } + sanitized.Set("client_id", clientID) + return sanitized, nil } // tokenResponse covers the OAuth fields needed by Google Auth and Codex. @@ -107,13 +153,22 @@ func (tv *TokenVendor) HandleTokenExchange(req *http.Request) *http.Response { TokenType: "Bearer", } } else if tv.chatGPTEnabled && IsChatGPTTokenExchange(req) { - resp = &tokenResponse{ - AccessToken: "paude-proxy-managed-access", - RefreshToken: "paude-proxy-managed-refresh", - IDToken: syntheticChatGPTIDToken(), - ExpiresIn: 3600, - TokenType: "Bearer", + bodyBytes, err := io.ReadAll(io.LimitReader(req.Body, 1<<20)) + if err != nil { + return errorResponse(http.StatusBadRequest, "Failed to read request body") + } + + values, err := url.ParseQuery(string(bodyBytes)) + if err != nil { + return errorResponse(http.StatusBadRequest, "Malformed form body") } + grantType := values.Get("grant_type") + + if grantType == "refresh_token" { + log.Printf("TOKEN_VEND host=%s path=%s (returned synthetic token, real injection at request time)", req.URL.Host, req.URL.Path) + return syntheticChatGPTResponse(req) + } + return tv.handleLoginExchange(req, values) } else { return nil } @@ -137,6 +192,84 @@ func (tv *TokenVendor) HandleTokenExchange(req *http.Request) *http.Response { } } +func (tv *TokenVendor) handleLoginExchange(req *http.Request, agentValues url.Values) *http.Response { + if tv.chatGPTInjector == nil { + log.Printf("ERROR login exchange: no ChatGPT injector configured") + return errorResponse(http.StatusInternalServerError, "Internal proxy error") + } + + sanitized, err := sanitizeLoginForm(agentValues, tv.chatGPTInjector.config.ClientID) + if err != nil { + log.Printf("LOGIN_SANITIZE rejected request: %v", err) + return errorResponse(http.StatusBadRequest, "Invalid login exchange request") + } + + forwardReq, err := http.NewRequest(http.MethodPost, chatGPTTokenURL, strings.NewReader(sanitized.Encode())) + if err != nil { + log.Printf("ERROR login exchange: construct forward request") + return errorResponse(http.StatusInternalServerError, "Internal proxy error") + } + forwardReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + + httpClient := tv.chatGPTInjector.config.HTTPClient + upstreamResp, err := httpClient.Do(forwardReq) + if err != nil { + log.Printf("ERROR login exchange: upstream request failed") + return errorResponse(http.StatusBadGateway, "Login exchange failed") + } + defer upstreamResp.Body.Close() + + upstreamBody, err := io.ReadAll(io.LimitReader(upstreamResp.Body, 1<<20)) + if err != nil { + return errorResponse(http.StatusBadGateway, "Login exchange response unreadable") + } + + if upstreamResp.StatusCode < http.StatusOK || upstreamResp.StatusCode >= http.StatusMultipleChoices { + log.Printf("TOKEN_VEND host=%s path=%s (login exchange upstream returned %d, passing through to agent)", req.URL.Host, req.URL.Path, upstreamResp.StatusCode) + ct := upstreamResp.Header.Get("Content-Type") + if ct == "" { + ct = "application/json" + } + return &http.Response{ + StatusCode: upstreamResp.StatusCode, + ProtoMajor: 1, + ProtoMinor: 1, + Header: http.Header{"Content-Type": {ct}}, + Body: io.NopCloser(bytes.NewReader(upstreamBody)), + ContentLength: int64(len(upstreamBody)), + Request: req, + } + } + + if err := tv.chatGPTInjector.AcceptLoginTokens(upstreamBody); err != nil { + log.Printf("ERROR login exchange: token acceptance failed") + return errorResponse(http.StatusInternalServerError, "Login token processing failed") + } + + log.Printf("TOKEN_VEND host=%s path=%s (login exchange completed, real tokens persisted, synthetic response returned)", req.URL.Host, req.URL.Path) + return syntheticChatGPTResponse(req) +} + +func syntheticChatGPTResponse(req *http.Request) *http.Response { + resp := &tokenResponse{ + AccessToken: "paude-proxy-managed-access", + RefreshToken: "paude-proxy-managed-refresh", + IDToken: syntheticChatGPTIDToken(), + ExpiresIn: 3600, + TokenType: "Bearer", + } + body, _ := json.Marshal(resp) + return &http.Response{ + StatusCode: http.StatusOK, + ProtoMajor: 1, + ProtoMinor: 1, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(bytes.NewReader(body)), + ContentLength: int64(len(body)), + Request: req, + } +} + func syntheticChatGPTIDToken() string { header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"none","typ":"JWT"}`)) claims := map[string]any{ diff --git a/internal/credentials/token_vending_test.go b/internal/credentials/token_vending_test.go index 10e8638..cb4cd79 100644 --- a/internal/credentials/token_vending_test.go +++ b/internal/credentials/token_vending_test.go @@ -1,10 +1,17 @@ package credentials import ( + "bytes" + "encoding/json" "io" "net/http" + "net/http/httptest" "net/url" + "os" + "path/filepath" + "strings" "testing" + "time" ) func TestIsTokenExchange_NilRequest(t *testing.T) { @@ -102,3 +109,266 @@ func TestHandleTokenExchange_ValidRequest(t *testing.T) { t.Errorf("response body suspiciously short: %q", bodyStr) } } + +func TestChatGPTTokenVendor_RefreshToken_ReturnsSynthetic(t *testing.T) { + vendor := NewChatGPTTokenVendor(nil) + req := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}, + Body: io.NopCloser(strings.NewReader("grant_type=refresh_token&refresh_token=dummy")), + } + resp := vendor.HandleTokenExchange(req) + if resp == nil || resp.StatusCode != http.StatusOK { + t.Fatal("refresh_token grant should return synthetic response") + } + body, _ := io.ReadAll(resp.Body) + if !bytes.Contains(body, []byte("paude-proxy-managed-access")) { + t.Error("response should contain synthetic access token") + } +} + +func TestChatGPTTokenVendor_LoginExchange_ForwardsAndPersists(t *testing.T) { + dir := t.TempDir() + statePath := filepath.Join(dir, "state", "auth.json") + now := time.Unix(1_700_000_000, 0) + realAccess := testJWT(map[string]any{ + "exp": now.Add(time.Hour).Unix(), + "chatgpt_account_id": "real-account", + }) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + values, _ := url.ParseQuery(string(body)) + if values.Get("grant_type") != "urn:ietf:params:oauth:grant-type:device_code" { + t.Errorf("unexpected grant_type forwarded: %s", values.Get("grant_type")) + } + if values.Get("device_code") != "test-code" { + t.Errorf("device_code = %q, want %q", values.Get("device_code"), "test-code") + } + if values.Get("client_id") != chatGPTClientID { + t.Errorf("client_id = %q, want %q", values.Get("client_id"), chatGPTClientID) + } + if values.Get("audience") != "" { + t.Error("disallowed parameter 'audience' was not stripped") + } + if values.Get("scope") != "" { + t.Error("disallowed parameter 'scope' was not stripped") + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": realAccess, + "refresh_token": "real-refresh", + "id_token": testJWT(map[string]any{"chatgpt_account_id": "real-account"}), + "expires_in": 3600, + }) + })) + defer server.Close() + + injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{ + StatePath: statePath, + HTTPClient: chatGPTTestClient(t, server), + Now: func() time.Time { return now }, + }) + + vendor := NewChatGPTTokenVendor(injector) + req := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}, + Body: io.NopCloser(strings.NewReader("grant_type=urn:ietf:params:oauth:grant-type:device_code&device_code=test-code&audience=evil&scope=admin&client_id=wrong-client")), + } + resp := vendor.HandleTokenExchange(req) + if resp == nil { + t.Fatal("login exchange should return a response") + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("login exchange status = %d, want 200", resp.StatusCode) + } + + body, _ := io.ReadAll(resp.Body) + if !bytes.Contains(body, []byte("paude-proxy-managed-access")) { + t.Error("agent should receive synthetic access token") + } + if bytes.Contains(body, []byte(realAccess)) { + t.Error("agent should NOT receive real access token") + } + if bytes.Contains(body, []byte("real-refresh")) { + t.Error("agent should NOT receive real refresh token") + } + + persisted, err := os.ReadFile(statePath) + if err != nil { + t.Fatalf("state file was not persisted: %v", err) + } + if !bytes.Contains(persisted, []byte("real-refresh")) { + t.Error("real refresh token was not persisted to state file") + } + + injectReq := &http.Request{Header: make(http.Header)} + if !injector.Inject(injectReq) { + t.Fatal("Inject should succeed after login exchange") + } + if injectReq.Header.Get("Authorization") != "Bearer "+realAccess { + t.Error("injected access token does not match login exchange result") + } +} + +func TestChatGPTTokenVendor_LoginExchange_UpstreamError_PassesThrough(t *testing.T) { + dir := t.TempDir() + statePath := filepath.Join(dir, "auth.json") + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + values, _ := url.ParseQuery(string(body)) + if values.Get("audience") != "" { + t.Error("disallowed parameter 'audience' reached upstream") + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{ + "error": "authorization_pending", + "error_description": "The user has not yet completed authorization", + }) + })) + defer server.Close() + + injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{ + StatePath: statePath, + HTTPClient: chatGPTTestClient(t, server), + Now: time.Now, + }) + + vendor := NewChatGPTTokenVendor(injector) + req := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}, + Body: io.NopCloser(strings.NewReader("grant_type=urn:ietf:params:oauth:grant-type:device_code&device_code=test&audience=evil")), + } + resp := vendor.HandleTokenExchange(req) + if resp == nil { + t.Fatal("should return a response") + } + if resp.StatusCode != http.StatusBadRequest { + t.Errorf("upstream error status = %d, want 400", resp.StatusCode) + } + body, _ := io.ReadAll(resp.Body) + if !bytes.Contains(body, []byte("authorization_pending")) { + t.Error("upstream error body should be passed through to agent") + } + + if _, err := os.Stat(statePath); !os.IsNotExist(err) { + t.Error("state file should not be created on upstream error") + } +} + +func TestChatGPTTokenVendor_LoginExchange_RejectsUnknownGrantType(t *testing.T) { + injector := NewChatGPTInjectorWithConfig(ChatGPTOAuthConfig{ + StatePath: filepath.Join(t.TempDir(), "auth.json"), + Now: time.Now, + }) + vendor := NewChatGPTTokenVendor(injector) + req := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Host: "auth.openai.com", Path: "/oauth/token"}, + Body: io.NopCloser(strings.NewReader("grant_type=custom_evil_grant¶m=value")), + } + resp := vendor.HandleTokenExchange(req) + if resp == nil || resp.StatusCode != http.StatusBadRequest { + t.Fatalf("unknown grant_type should be rejected with 400, got %v", resp) + } +} + +func TestSanitizeLoginForm(t *testing.T) { + cases := []struct { + name string + input string + clientID string + wantErr bool + wantKeys []string + checkVals map[string]string + }{ + { + name: "missing grant_type", + input: "device_code=abc", + wantErr: true, + }, + { + name: "unknown grant_type", + input: "grant_type=password&username=admin", + wantErr: true, + }, + { + name: "device_code strips extras", + input: "grant_type=urn:ietf:params:oauth:grant-type:device_code&device_code=dc1&audience=evil&scope=admin", + clientID: "test-client", + wantKeys: []string{"grant_type", "device_code", "client_id"}, + checkVals: map[string]string{ + "grant_type": "urn:ietf:params:oauth:grant-type:device_code", + "device_code": "dc1", + "client_id": "test-client", + }, + }, + { + name: "device_code enforces client_id", + input: "grant_type=urn:ietf:params:oauth:grant-type:device_code&device_code=dc1&client_id=agent-evil", + clientID: "correct-client", + checkVals: map[string]string{ + "client_id": "correct-client", + }, + }, + { + name: "authorization_code allows PKCE params", + input: "grant_type=authorization_code&code=authcode&redirect_uri=http://localhost:1455/callback&code_verifier=verifier&client_id=agent&resource=evil", + clientID: "real-client", + wantKeys: []string{"grant_type", "code", "redirect_uri", "client_id", "code_verifier"}, + checkVals: map[string]string{ + "code": "authcode", + "redirect_uri": "http://localhost:1455/callback", + "code_verifier": "verifier", + "client_id": "real-client", + }, + }, + { + name: "token-exchange allows exchange params", + input: "grant_type=urn:ietf:params:oauth:grant-type:token-exchange&client_id=x&requested_token=openai-api-key&subject_token=tok&subject_token_type=urn:ietf:params:oauth:token-type:id_token&audience=evil", + clientID: "canonical", + wantKeys: []string{"grant_type", "client_id", "requested_token", "subject_token", "subject_token_type"}, + checkVals: map[string]string{ + "requested_token": "openai-api-key", + "subject_token": "tok", + "subject_token_type": "urn:ietf:params:oauth:token-type:id_token", + "client_id": "canonical", + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + vals, _ := url.ParseQuery(tc.input) + result, err := sanitizeLoginForm(vals, tc.clientID) + if tc.wantErr { + if err == nil { + t.Fatal("expected error but got nil") + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if tc.wantKeys != nil { + for _, key := range tc.wantKeys { + if result.Get(key) == "" { + t.Errorf("expected key %q in result", key) + } + } + if len(result) != len(tc.wantKeys) { + t.Errorf("result has %d keys, want %d: %v", len(result), len(tc.wantKeys), result) + } + } + for key, want := range tc.checkVals { + if got := result.Get(key); got != want { + t.Errorf("%s = %q, want %q", key, got, want) + } + } + }) + } +} diff --git a/internal/proxy/integration_test.go b/internal/proxy/integration_test.go index 1f82657..1964482 100644 --- a/internal/proxy/integration_test.go +++ b/internal/proxy/integration_test.go @@ -551,7 +551,7 @@ func TestIntegration_ChatGPTOAuthProxyFlow(t *testing.T) { if err != nil { t.Fatal(err) } - proxyAddr, cleanup := startTestProxy(t, ca, df, store, credentials.NewChatGPTTokenVendor(), pool) + proxyAddr, cleanup := startTestProxy(t, ca, df, store, credentials.NewChatGPTTokenVendor(injector), pool) defer cleanup() client := httpClientViaProxy(t, proxyAddr, ca.Certificate, primary.Certificate()) @@ -577,7 +577,7 @@ func TestIntegration_ChatGPTOAuthProxyFlow(t *testing.T) { t.Error("unrelated request header was not preserved") } - tokenReq, _ := http.NewRequest(http.MethodPost, "http://auth.openai.com/oauth/token", strings.NewReader("refresh_token=agent-dummy")) + tokenReq, _ := http.NewRequest(http.MethodPost, "http://auth.openai.com/oauth/token", strings.NewReader("grant_type=refresh_token&refresh_token=agent-dummy")) tokenResp, err := client.Do(tokenReq) if err != nil { t.Fatalf("dummy token exchange failed: %v", err) @@ -920,3 +920,119 @@ func TestIntegration_ProxyTransport_ResponseHeaderTimeout(t *testing.T) { t.Logf("Request timed out after %v as expected", elapsed) } + +func TestIntegration_ChatGPTLoginFlow(t *testing.T) { + skipIntegration(t) + dir := t.TempDir() + statePath := filepath.Join(dir, "state", "auth.json") + now := time.Unix(1_700_000_000, 0) + realAccess := testProxyJWT(map[string]any{ + "exp": now.Add(time.Hour).Unix(), + "chatgpt_account_id": "login-account", + }) + realID := testProxyJWT(map[string]any{"chatgpt_account_id": "login-account"}) + + loginServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + vals, _ := url.ParseQuery(string(body)) + gt := vals.Get("grant_type") + if gt == "urn:ietf:params:oauth:grant-type:device_code" { + if vals.Get("audience") != "" { + t.Error("disallowed parameter 'audience' was not stripped by proxy") + } + if vals.Get("scope") != "" { + t.Error("disallowed parameter 'scope' was not stripped by proxy") + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": realAccess, + "refresh_token": "login-refresh", + "id_token": realID, + "expires_in": 3600, + }) + } else { + http.Error(w, "unexpected grant_type", http.StatusBadRequest) + } + })) + defer loginServer.Close() + + var primaryHeaders http.Header + primary := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + primaryHeaders = r.Header.Clone() + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"ok":true}`) + })) + defer primary.Close() + primaryURL, _ := url.Parse(primary.URL) + + injector := credentials.NewChatGPTInjectorWithConfig(credentials.ChatGPTOAuthConfig{ + StatePath: statePath, + HTTPClient: chatGPTTestClient(t, loginServer), + Now: func() time.Time { return now }, + }) + store := credentials.NewStore() + store.AddRoute(credentials.Route{ + ExactDomain: primaryURL.Hostname(), + PathPrefix: "/backend-api/codex", + Injector: injector, + }) + df := filter.NewDomainFilter(primaryURL.Hostname() + ",auth.openai.com") + pool := upstreamCertPool(t, primary) + ca, err := GenerateCA() + if err != nil { + t.Fatal(err) + } + proxyAddr, cleanup := startTestProxy(t, ca, df, store, credentials.NewChatGPTTokenVendor(injector), pool) + defer cleanup() + client := httpClientViaProxy(t, proxyAddr, ca.Certificate, primary.Certificate()) + + // Before login: API request should fail with 502 + preLoginReq, _ := http.NewRequest(http.MethodPost, primary.URL+"/backend-api/codex/responses", strings.NewReader("{}")) + preLoginReq.Header.Set("Authorization", "Bearer agent-dummy") + preLoginResp, err := client.Do(preLoginReq) + if err != nil { + t.Fatalf("pre-login request failed: %v", err) + } + _ = preLoginResp.Body.Close() + if preLoginResp.StatusCode != http.StatusBadGateway { + t.Fatalf("pre-login status = %d, want 502", preLoginResp.StatusCode) + } + + // Simulate codex login device-code exchange (with extra params that should be stripped) + loginReq, _ := http.NewRequest(http.MethodPost, "http://auth.openai.com/oauth/token", + strings.NewReader("grant_type=urn:ietf:params:oauth:grant-type:device_code&device_code=test-code&audience=evil&scope=admin")) + loginResp, err := client.Do(loginReq) + if err != nil { + t.Fatalf("login exchange failed: %v", err) + } + loginBody, _ := io.ReadAll(loginResp.Body) + _ = loginResp.Body.Close() + if loginResp.StatusCode != http.StatusOK { + t.Fatalf("login exchange status = %d, want 200", loginResp.StatusCode) + } + if !bytes.Contains(loginBody, []byte("paude-proxy-managed-access")) { + t.Error("agent should receive synthetic access token from login exchange") + } + if bytes.Contains(loginBody, []byte(realAccess)) || bytes.Contains(loginBody, []byte("login-refresh")) { + t.Error("agent should NOT receive real tokens from login exchange") + } + + // After login: API request should succeed with real credentials injected + postLoginReq, _ := http.NewRequest(http.MethodPost, primary.URL+"/backend-api/codex/responses", strings.NewReader("{}")) + postLoginReq.Header.Set("Authorization", "Bearer agent-dummy") + postLoginReq.Header.Set("ChatGPT-Account-ID", "agent-dummy-account") + postLoginResp, err := client.Do(postLoginReq) + if err != nil { + t.Fatalf("post-login request failed: %v", err) + } + _ = postLoginResp.Body.Close() + if postLoginResp.StatusCode != http.StatusOK { + t.Fatalf("post-login status = %d, want 200", postLoginResp.StatusCode) + } + if primaryHeaders.Get("Authorization") != "Bearer "+realAccess { + t.Error("upstream did not receive the login access token") + } + if primaryHeaders.Get("ChatGPT-Account-ID") != "login-account" { + t.Error("upstream did not receive the login account ID") + } +}