diff --git a/internal/auth/codex/openai_auth.go b/internal/auth/codex/openai_auth.go index 2c1eac086f..fad2ce577e 100644 --- a/internal/auth/codex/openai_auth.go +++ b/internal/auth/codex/openai_auth.go @@ -61,6 +61,15 @@ func NewCodexAuthWithProxyURL(cfg *config.Config, proxyURL string) *CodexAuth { } } +// SetHTTPClient replaces the client used for token exchange and refresh. +// A nil client is ignored. +func (o *CodexAuth) SetHTTPClient(client *http.Client) { + if o == nil || client == nil { + return + } + o.httpClient = client +} + // GenerateAuthURL creates the OAuth authorization URL with PKCE (Proof Key for Code Exchange). // It constructs the URL with the necessary parameters, including the client ID, // response type, redirect URI, scopes, and PKCE challenge. diff --git a/internal/auth/meta/meta.go b/internal/auth/meta/meta.go index 3eb3554192..495286dd06 100644 --- a/internal/auth/meta/meta.go +++ b/internal/auth/meta/meta.go @@ -250,6 +250,15 @@ func NewMetaAuthWithProxyURL(cfg *config.Config, proxyURL string) *MetaAuth { } } +// SetHTTPClient replaces the client used for device flow, token poll, and key minting. +// A nil client is ignored. +func (a *MetaAuth) SetHTTPClient(client *http.Client) { + if a == nil || client == nil { + return + } + a.httpClient = client +} + // StartDeviceFlow initiates the device authorization flow with Meta. func (a *MetaAuth) StartDeviceFlow(ctx context.Context) (*DeviceCodeResponse, error) { return a.StartDeviceFlowWithEndpoint(ctx, DeviceAuthorizationEndpoint) @@ -315,11 +324,6 @@ func (a *MetaAuth) WaitForAuthorization(ctx context.Context, dcr *DeviceCodeResp return nil, fmt.Errorf("meta auth: missing device code response") } - tokenEndpoint := dcr.TokenEndpoint - if tokenEndpoint == "" { - tokenEndpoint = TokenEndpoint - } - interval := time.Duration(dcr.Interval) * time.Second if interval <= 0 { interval = DefaultPollInterval @@ -344,81 +348,112 @@ func (a *MetaAuth) WaitForAuthorization(ctx context.Context, dcr *DeviceCodeResp case <-ctx.Done(): return nil, fmt.Errorf("meta auth: authorization timed out or canceled: %w", ctx.Err()) case <-ticker.C: - form := url.Values{ - "grant_type": {DeviceCodeGrantType}, - "device_code": {dcr.DeviceCode}, - "client_id": {ClientID}, - } - req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(form.Encode())) + bundle, next, err := a.PollDeviceCodeOnce(ctx, dcr, interval) if err != nil { - return nil, fmt.Errorf("meta auth: create token request: %w", err) + return nil, err } - req.Header.Set("Content-Type", "application/x-www-form-urlencoded") - req.Header.Set("Accept", "application/json") - req.Header.Set("User-Agent", "muse-code/1.0.2") - - resp, err := a.httpClient.Do(req) - if err != nil { - log.Warnf("meta auth: poll request error: %v (retrying)", err) - continue + if bundle != nil { + return bundle, nil } - - body, errRead := io.ReadAll(resp.Body) - _ = resp.Body.Close() - if errRead != nil { - log.Warnf("meta auth: read token response error: %v (retrying)", errRead) - continue + if next > interval { + interval = next + ticker.Reset(interval) } + } + } +} - if resp.StatusCode == http.StatusOK { - var tokenData TokenData - if err := json.Unmarshal(body, &tokenData); err != nil { - return nil, fmt.Errorf("meta auth: parse token response: %w", err) - } - if tokenData.AccessToken == "" { - return nil, fmt.Errorf("meta auth: response missing access_token") - } - if tokenData.ExpiresIn > 0 { - tokenData.ExpiresAt = time.Now().Add(time.Duration(tokenData.ExpiresIn) * time.Second).Unix() - } - bundle := &MetaAuthBundle{ - TokenData: &tokenData, - } - minted, errMint := a.MintAPIKey(ctx, tokenData.AccessToken) - if errMint != nil { - log.Warnf("meta auth: could not mint api_key from dca_token: %v", errMint) - } else if minted != nil { - bundle.MintedKey = minted - if minted.UserEmail != "" { - bundle.Email = minted.UserEmail - } - if minted.UserFullName != "" { - bundle.Name = minted.UserFullName - } - } - return bundle, nil - } +// PollDeviceCodeOnce performs one device-code token request and, on approval, mints an API key. +// A nil bundle and a nil error means the user has not authorized yet. +// nextInterval is the delay before another attempt; slow_down increases it. +func (a *MetaAuth) PollDeviceCodeOnce(ctx context.Context, dcr *DeviceCodeResponse, interval time.Duration) (*MetaAuthBundle, time.Duration, error) { + if dcr == nil || dcr.DeviceCode == "" { + return nil, interval, fmt.Errorf("meta auth: missing device code response") + } + if ctx == nil { + ctx = context.Background() + } + if interval <= 0 { + interval = DefaultPollInterval + } - var errResp TokenData - _ = json.Unmarshal(body, &errResp) - switch errResp.Error { - case "authorization_pending": - continue - case "slow_down": - interval += 5 * time.Second - ticker.Reset(interval) - continue - case "access_denied": - return nil, fmt.Errorf("meta auth: access was denied by user") - case "expired_token": - return nil, fmt.Errorf("meta auth: device code has expired") - default: - if errResp.Error != "" { - return nil, fmt.Errorf("meta auth: error from authorization server: %s: %s", errResp.Error, errResp.ErrorDescription) - } - log.Warnf("meta auth: unexpected response %d: %s", resp.StatusCode, string(body)) + tokenEndpoint := dcr.TokenEndpoint + if tokenEndpoint == "" { + tokenEndpoint = TokenEndpoint + } + + form := url.Values{ + "grant_type": {DeviceCodeGrantType}, + "device_code": {dcr.DeviceCode}, + "client_id": {ClientID}, + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, strings.NewReader(form.Encode())) + if err != nil { + return nil, interval, fmt.Errorf("meta auth: create token request: %w", err) + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "muse-code/1.0.2") + + resp, err := a.httpClient.Do(req) + if err != nil { + log.Warnf("meta auth: poll request error: %v (retrying)", err) + return nil, interval, nil + } + + body, errRead := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if errRead != nil { + log.Warnf("meta auth: read token response error: %v (retrying)", errRead) + return nil, interval, nil + } + + if resp.StatusCode == http.StatusOK { + var tokenData TokenData + if err := json.Unmarshal(body, &tokenData); err != nil { + return nil, interval, fmt.Errorf("meta auth: parse token response: %w", err) + } + if tokenData.AccessToken == "" { + return nil, interval, fmt.Errorf("meta auth: response missing access_token") + } + if tokenData.ExpiresIn > 0 { + tokenData.ExpiresAt = time.Now().Add(time.Duration(tokenData.ExpiresIn) * time.Second).Unix() + } + bundle := &MetaAuthBundle{ + TokenData: &tokenData, + } + minted, errMint := a.MintAPIKey(ctx, tokenData.AccessToken) + if errMint != nil { + log.Warnf("meta auth: could not mint api_key from dca_token: %v", errMint) + } else if minted != nil { + bundle.MintedKey = minted + if minted.UserEmail != "" { + bundle.Email = minted.UserEmail + } + if minted.UserFullName != "" { + bundle.Name = minted.UserFullName } } + return bundle, interval, nil + } + + var errResp TokenData + _ = json.Unmarshal(body, &errResp) + switch errResp.Error { + case "authorization_pending": + return nil, interval, nil + case "slow_down": + return nil, interval + 5*time.Second, nil + case "access_denied": + return nil, interval, fmt.Errorf("meta auth: access was denied by user") + case "expired_token": + return nil, interval, fmt.Errorf("meta auth: device code has expired") + default: + if errResp.Error != "" { + return nil, interval, fmt.Errorf("meta auth: error from authorization server: %s: %s", errResp.Error, errResp.ErrorDescription) + } + log.Warnf("meta auth: unexpected response %d: %s", resp.StatusCode, string(body)) + return nil, interval, nil } } diff --git a/internal/auth/xai/xai.go b/internal/auth/xai/xai.go index a0b76ea838..86fbdaf72c 100644 --- a/internal/auth/xai/xai.go +++ b/internal/auth/xai/xai.go @@ -44,6 +44,15 @@ func NewXAIAuthWithProxyURL(cfg *config.Config, proxyURL string) *XAIAuth { return &XAIAuth{httpClient: util.SetProxy(&sdkCfg, &http.Client{Timeout: httpClientTimeout})} } +// SetHTTPClient replaces the client used for discovery, device flow, and refresh. +// A nil client is ignored. +func (a *XAIAuth) SetHTTPClient(client *http.Client) { + if a == nil || client == nil { + return + } + a.httpClient = client +} + // ValidateOAuthEndpoint validates an endpoint returned by xAI discovery. func ValidateOAuthEndpoint(rawURL string, field string) (string, error) { rawURL = strings.TrimSpace(rawURL) @@ -194,22 +203,53 @@ func (a *XAIAuth) WaitForAuthorization(ctx context.Context, deviceCode *DeviceCo }, nil } -// PollForToken polls the token endpoint until the user authorizes or the device code expires. -func (a *XAIAuth) PollForToken(ctx context.Context, deviceCode *DeviceCodeResponse) (*TokenData, error) { +// PollDeviceCodeOnce performs one device-code token request. +// A nil token and a nil error means the user has not authorized yet. +// nextInterval is the delay before another attempt; slow_down increases it. +func (a *XAIAuth) PollDeviceCodeOnce(ctx context.Context, deviceCode *DeviceCodeResponse, interval time.Duration) (*TokenData, time.Duration, error) { if deviceCode == nil { - return nil, fmt.Errorf("xai device code: response is nil") + return nil, interval, fmt.Errorf("xai device code: response is nil") } if ctx == nil { ctx = context.Background() } + if interval <= 0 { + interval = defaultPollInterval + if a != nil && a.minPollInterval > 0 { + interval = a.minPollInterval + } + } tokenEndpoint := strings.TrimSpace(deviceCode.TokenEndpoint) if tokenEndpoint == "" { discovery, errDiscover := a.Discover(ctx) if errDiscover != nil { - return nil, errDiscover + return nil, interval, errDiscover } tokenEndpoint = discovery.TokenEndpoint + deviceCode.TokenEndpoint = tokenEndpoint + } + + token, pollErr, nextInterval, shouldContinue := a.exchangeDeviceCode(ctx, tokenEndpoint, deviceCode.DeviceCode, interval) + if token != nil { + return token, nextInterval, nil + } + if shouldContinue { + return nil, nextInterval, nil + } + if pollErr == nil { + pollErr = fmt.Errorf("xai device token exchange failed") + } + return nil, nextInterval, pollErr +} + +// PollForToken polls the token endpoint until the user authorizes or the device code expires. +func (a *XAIAuth) PollForToken(ctx context.Context, deviceCode *DeviceCodeResponse) (*TokenData, error) { + if deviceCode == nil { + return nil, fmt.Errorf("xai device code: response is nil") + } + if ctx == nil { + ctx = context.Background() } minInterval := defaultPollInterval @@ -247,11 +287,11 @@ func (a *XAIAuth) PollForToken(ctx context.Context, deviceCode *DeviceCodeRespon } firstAttempt = false - token, pollErr, nextInterval, shouldContinue := a.exchangeDeviceCode(ctx, tokenEndpoint, deviceCode.DeviceCode, interval) + token, nextInterval, pollErr := a.PollDeviceCodeOnce(ctx, deviceCode, interval) if token != nil { return token, nil } - if !shouldContinue { + if pollErr != nil { return nil, pollErr } interval = nextInterval diff --git a/sdk/auth/codex.go b/sdk/auth/codex.go index 79448991b8..a436364f7d 100644 --- a/sdk/auth/codex.go +++ b/sdk/auth/codex.go @@ -20,6 +20,10 @@ import ( // CodexAuthenticator implements the OAuth login flow for Codex accounts. type CodexAuthenticator struct { CallbackPort int + httpClient *http.Client + // pollIntervalOverride, when set, replaces the wait between device-flow polls. + // Production leaves it zero so the provider interval is used. + pollIntervalOverride time.Duration } // NewCodexAuthenticator constructs a Codex authenticator with default settings. diff --git a/sdk/auth/codex_device.go b/sdk/auth/codex_device.go index d7ea4e1fe9..b412c4169b 100644 --- a/sdk/auth/codex_device.go +++ b/sdk/auth/codex_device.go @@ -61,46 +61,151 @@ func shouldUseCodexDeviceFlow(opts *LoginOptions) bool { return strings.EqualFold(strings.TrimSpace(opts.Metadata[codexLoginModeMetadataKey]), codexLoginModeDevice) } +func (a *CodexAuthenticator) deviceHTTPClient(cfg *config.Config) *http.Client { + if a != nil && a.httpClient != nil { + return a.httpClient + } + var sdkCfg config.SDKConfig + if cfg != nil { + sdkCfg = cfg.SDKConfig + } + return util.SetProxy(&sdkCfg, &http.Client{}) +} + +func (a *CodexAuthenticator) effectiveInterval(d time.Duration) time.Duration { + if a != nil && a.pollIntervalOverride > 0 { + return a.pollIntervalOverride + } + if d <= 0 { + return time.Duration(codexDeviceDefaultPollIntervalSeconds) * time.Second + } + return d +} + func (a *CodexAuthenticator) loginWithDeviceFlow(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) { if ctx == nil { ctx = context.Background() } + if opts == nil { + opts = &LoginOptions{} + } - httpClient := util.SetProxy(&cfg.SDKConfig, &http.Client{}) - - userCodeResp, err := requestCodexDeviceUserCode(ctx, httpClient) + flow, err := a.StartDeviceFlow(ctx, cfg) if err != nil { return nil, err } - deviceCode := strings.TrimSpace(userCodeResp.UserCode) - if deviceCode == "" { - deviceCode = strings.TrimSpace(userCodeResp.UserCodeAlt) - } - deviceAuthID := strings.TrimSpace(userCodeResp.DeviceAuthID) - if deviceCode == "" || deviceAuthID == "" { - return nil, fmt.Errorf("codex device flow did not return required fields") - } - - pollInterval := parseCodexDevicePollInterval(userCodeResp.Interval) - fmt.Println("Starting Codex device authentication...") - fmt.Printf("Codex device URL: %s\n", codexDeviceVerificationURL) - fmt.Printf("Codex device code: %s\n", deviceCode) + fmt.Printf("Codex device URL: %s\n", flow.VerificationURI) + fmt.Printf("Codex device code: %s\n", flow.UserCode) if !opts.NoBrowser { if !browser.IsAvailable() { log.Warn("No browser available; please open the device URL manually") - } else if errOpen := browser.OpenURL(codexDeviceVerificationURL); errOpen != nil { + } else if errOpen := browser.OpenURL(flow.VerificationURI); errOpen != nil { log.Warnf("Failed to open browser automatically: %v", errOpen) } } - tokenResp, err := pollCodexDeviceToken(ctx, httpClient, deviceAuthID, deviceCode, pollInterval) + return a.CompleteDeviceFlow(ctx, cfg, flow) +} + +// StartDeviceFlow requests a Codex device code and returns it without waiting +// for the user to authorize. +func (a *CodexAuthenticator) StartDeviceFlow(ctx context.Context, cfg *config.Config) (*DeviceFlow, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if ctx == nil { + ctx = context.Background() + } + + userCodeResp, err := requestCodexDeviceUserCode(ctx, a.deviceHTTPClient(cfg)) if err != nil { return nil, err } + userCode := strings.TrimSpace(userCodeResp.UserCode) + if userCode == "" { + userCode = strings.TrimSpace(userCodeResp.UserCodeAlt) + } + deviceAuthID := strings.TrimSpace(userCodeResp.DeviceAuthID) + if userCode == "" || deviceAuthID == "" { + return nil, fmt.Errorf("codex device flow did not return required fields") + } + + return &DeviceFlow{ + Provider: a.Provider(), + UserCode: userCode, + DeviceCode: deviceAuthID, + DeviceAuthID: deviceAuthID, + VerificationURI: codexDeviceVerificationURL, + Interval: parseCodexDevicePollInterval(userCodeResp.Interval), + }, nil +} + +// PollDeviceFlow exchanges a Codex device code once. +// A pending result means the user has not approved yet. Approval exchanges +// the returned authorization code for a Codex credential. +func (a *CodexAuthenticator) PollDeviceFlow(ctx context.Context, cfg *config.Config, flow *DeviceFlow) (*DeviceFlowPoll, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if flow == nil { + return nil, fmt.Errorf("codex device flow is required") + } + if ctx == nil { + ctx = context.Background() + } + deviceAuthID := strings.TrimSpace(flow.DeviceAuthID) + if deviceAuthID == "" { + deviceAuthID = strings.TrimSpace(flow.DeviceCode) + } + userCode := strings.TrimSpace(flow.UserCode) + if deviceAuthID == "" || userCode == "" { + return nil, fmt.Errorf("codex device flow is missing device fields") + } + + tokenResp, pending, err := pollCodexDeviceTokenOnce(ctx, a.deviceHTTPClient(cfg), deviceAuthID, userCode) + if err != nil { + return nil, err + } + interval := a.effectiveInterval(flow.Interval) + if pending { + return &DeviceFlowPoll{Pending: true, Interval: interval}, nil + } + record, errExchange := a.exchangeCodexDeviceAuthorization(ctx, cfg, tokenResp) + if errExchange != nil { + return nil, errExchange + } + return &DeviceFlowPoll{Interval: interval, Auth: record}, nil +} + +// CompleteDeviceFlow polls until the Codex device code is authorized or the flow fails. +func (a *CodexAuthenticator) CompleteDeviceFlow(ctx context.Context, cfg *config.Config, flow *DeviceFlow) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if flow == nil { + return nil, fmt.Errorf("codex device flow is required") + } + interval := a.effectiveInterval(flow.Interval) + deadline := time.Now().Add(codexDeviceTimeout) + return pollUntilDeviceAuthorized(ctx, interval, deadline, fmt.Errorf("codex device authentication timed out after 15 minutes"), func() (*DeviceFlowPoll, error) { + polled, err := a.PollDeviceFlow(ctx, cfg, flow) + if err != nil || polled == nil || polled.Auth != nil { + return polled, err + } + polled.Interval = a.effectiveInterval(polled.Interval) + flow.Interval = polled.Interval + return polled, nil + }) +} + +func (a *CodexAuthenticator) exchangeCodexDeviceAuthorization(ctx context.Context, cfg *config.Config, tokenResp *codexDeviceTokenResponse) (*coreauth.Auth, error) { + if tokenResp == nil { + return nil, fmt.Errorf("codex device flow token response missing required fields") + } authCode := strings.TrimSpace(tokenResp.AuthorizationCode) codeVerifier := strings.TrimSpace(tokenResp.CodeVerifier) codeChallenge := strings.TrimSpace(tokenResp.CodeChallenge) @@ -109,6 +214,9 @@ func (a *CodexAuthenticator) loginWithDeviceFlow(ctx context.Context, cfg *confi } authSvc := codex.NewCodexAuth(cfg) + if a != nil && a.httpClient != nil { + authSvc.SetHTTPClient(a.httpClient) + } authBundle, err := authSvc.ExchangeCodeForTokensWithRedirect( ctx, authCode, @@ -121,7 +229,6 @@ func (a *CodexAuthenticator) loginWithDeviceFlow(ctx context.Context, cfg *confi if err != nil { return nil, codex.NewAuthenticationError(codex.ErrCodeExchangeFailed, err) } - return a.buildAuthRecord(authSvc, authBundle) } @@ -168,61 +275,48 @@ func requestCodexDeviceUserCode(ctx context.Context, client *http.Client) (*code return &parsed, nil } -func pollCodexDeviceToken(ctx context.Context, client *http.Client, deviceAuthID, userCode string, interval time.Duration) (*codexDeviceTokenResponse, error) { - deadline := time.Now().Add(codexDeviceTimeout) - - for { - if time.Now().After(deadline) { - return nil, fmt.Errorf("codex device authentication timed out after 15 minutes") - } +func pollCodexDeviceTokenOnce(ctx context.Context, client *http.Client, deviceAuthID, userCode string) (*codexDeviceTokenResponse, bool, error) { + body, err := json.Marshal(codexDeviceTokenRequest{ + DeviceAuthID: deviceAuthID, + UserCode: userCode, + }) + if err != nil { + return nil, false, fmt.Errorf("failed to encode codex device poll request: %w", err) + } - body, err := json.Marshal(codexDeviceTokenRequest{ - DeviceAuthID: deviceAuthID, - UserCode: userCode, - }) - if err != nil { - return nil, fmt.Errorf("failed to encode codex device poll request: %w", err) - } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, codexDeviceTokenURL, bytes.NewReader(body)) + if err != nil { + return nil, false, fmt.Errorf("failed to create codex device poll request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") - req, err := http.NewRequestWithContext(ctx, http.MethodPost, codexDeviceTokenURL, bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("failed to create codex device poll request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Accept", "application/json") + resp, err := client.Do(req) + if err != nil { + return nil, false, fmt.Errorf("failed to poll codex device token: %w", err) + } - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to poll codex device token: %w", err) - } + respBody, readErr := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if readErr != nil { + return nil, false, fmt.Errorf("failed to read codex device poll response: %w", readErr) + } - respBody, readErr := io.ReadAll(resp.Body) - _ = resp.Body.Close() - if readErr != nil { - return nil, fmt.Errorf("failed to read codex device poll response: %w", readErr) + switch { + case codexDeviceIsSuccessStatus(resp.StatusCode): + var parsed codexDeviceTokenResponse + if err := json.Unmarshal(respBody, &parsed); err != nil { + return nil, false, fmt.Errorf("failed to decode codex device token response: %w", err) } - - switch { - case codexDeviceIsSuccessStatus(resp.StatusCode): - var parsed codexDeviceTokenResponse - if err := json.Unmarshal(respBody, &parsed); err != nil { - return nil, fmt.Errorf("failed to decode codex device token response: %w", err) - } - return &parsed, nil - case resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusNotFound: - select { - case <-ctx.Done(): - return nil, ctx.Err() - case <-time.After(interval): - continue - } - default: - trimmed := strings.TrimSpace(string(respBody)) - if trimmed == "" { - trimmed = "empty response body" - } - return nil, fmt.Errorf("codex device token polling failed with status %d: %s", resp.StatusCode, trimmed) + return &parsed, false, nil + case resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusNotFound: + return nil, true, nil + default: + trimmed := strings.TrimSpace(string(respBody)) + if trimmed == "" { + trimmed = "empty response body" } + return nil, false, fmt.Errorf("codex device token polling failed with status %d: %s", resp.StatusCode, trimmed) } } diff --git a/sdk/auth/device_flow.go b/sdk/auth/device_flow.go new file mode 100644 index 0000000000..83140ab1aa --- /dev/null +++ b/sdk/auth/device_flow.go @@ -0,0 +1,77 @@ +package auth + +import ( + "context" + "time" + + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +// DeviceFlow is a device authorization that has been started and not yet exchanged. +// StartDeviceFlow returns it without waiting for the user to approve. +type DeviceFlow struct { + Provider string + UserCode string + DeviceCode string + DeviceAuthID string + VerificationURI string + VerificationURIComplete string + ExpiresIn time.Duration + Interval time.Duration + TokenEndpoint string +} + +// DeviceFlowPoll is one device-code exchange attempt. +// Pending is set when the user has not approved yet. Auth is set when the +// exchange finished and produced a credential. +type DeviceFlowPoll struct { + Pending bool + Interval time.Duration + Auth *coreauth.Auth +} + +// pollUntilDeviceAuthorized calls poll until it returns an auth or an error. +// The first attempt runs immediately. Later attempts wait the poll interval +// and stop once deadline has passed. +func pollUntilDeviceAuthorized(ctx context.Context, interval time.Duration, deadline time.Time, timeoutErr error, poll func() (*DeviceFlowPoll, error)) (*coreauth.Auth, error) { + if ctx == nil { + ctx = context.Background() + } + if interval <= 0 { + interval = 5 * time.Second + } + first := true + for { + if err := ctx.Err(); err != nil { + return nil, err + } + if !first && !deadline.IsZero() && time.Now().After(deadline) { + if timeoutErr != nil { + return nil, timeoutErr + } + return nil, context.DeadlineExceeded + } + first = false + + result, err := poll() + if err != nil { + return nil, err + } + if result != nil && result.Auth != nil { + return result.Auth, nil + } + next := interval + if result != nil && result.Interval > 0 { + next = result.Interval + } + interval = next + + timer := time.NewTimer(next) + select { + case <-ctx.Done(): + timer.Stop() + return nil, ctx.Err() + case <-timer.C: + } + } +} diff --git a/sdk/auth/device_flow_test.go b/sdk/auth/device_flow_test.go new file mode 100644 index 0000000000..489e5f1030 --- /dev/null +++ b/sdk/auth/device_flow_test.go @@ -0,0 +1,451 @@ +package auth + +import ( + "context" + "encoding/base64" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + codexauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/codex" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" +) + +func TestDeviceFlowStartDoesNotExchange(t *testing.T) { + t.Run("xai", func(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/oauth2/token" { + polls.Add(1) + } + writeXAIDeviceServer(w, r, 0) + })) + defer server.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + flow, err := (XAIAuthenticator{httpClient: rewriteClient(server)}).StartDeviceFlow(ctx, &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + if flow.UserCode != "ABCD-1234" || flow.DeviceCode != "device-abc" || flow.VerificationURIComplete == "" { + t.Fatalf("flow = %+v", flow) + } + if polls.Load() != 0 { + t.Fatalf("start exchanged the device code (%d token calls)", polls.Load()) + } + }) + + t.Run("meta", func(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.Contains(r.URL.Path, "/token") || strings.Contains(r.URL.Path, "/muse-code/") { + polls.Add(1) + } + writeMetaDeviceServer(w, r, 0) + })) + defer server.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + flow, err := (MetaAuthenticator{httpClient: rewriteClient(server)}).StartDeviceFlow(ctx, &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + if flow.UserCode != "TEST-1234" || flow.DeviceCode != "meta-device" { + t.Fatalf("flow = %+v", flow) + } + if polls.Load() != 0 { + t.Fatalf("start exchanged the device code (%d calls)", polls.Load()) + } + }) + + t.Run("codex", func(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/accounts/deviceauth/usercode" { + polls.Add(1) + } + writeCodexDeviceServer(w, r, 0) + })) + defer server.Close() + + authn := NewCodexAuthenticator() + authn.httpClient = rewriteClient(server) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + flow, err := authn.StartDeviceFlow(ctx, &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + if flow.UserCode != "CODE-1234" || flow.DeviceAuthID != "auth-id-1" || flow.VerificationURI == "" { + t.Fatalf("flow = %+v", flow) + } + if polls.Load() != 0 { + t.Fatalf("start exchanged the device code (%d calls)", polls.Load()) + } + }) +} + +func TestDeviceFlowPollPendingDoesNotBlock(t *testing.T) { + t.Run("xai", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writeXAIDeviceServer(w, r, 1) + })) + defer server.Close() + authn := XAIAuthenticator{httpClient: rewriteClient(server)} + flow, err := authn.StartDeviceFlow(context.Background(), &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + started := time.Now() + polled, err := authn.PollDeviceFlow(context.Background(), &config.Config{}, flow) + if err != nil { + t.Fatalf("PollDeviceFlow: %v", err) + } + if polled == nil || !polled.Pending || polled.Auth != nil { + t.Fatalf("poll = %+v, want pending", polled) + } + if time.Since(started) > 2*time.Second { + t.Fatalf("pending poll blocked for %s", time.Since(started)) + } + }) + + t.Run("meta", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writeMetaDeviceServer(w, r, 1) + })) + defer server.Close() + authn := MetaAuthenticator{httpClient: rewriteClient(server)} + flow, err := authn.StartDeviceFlow(context.Background(), &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + started := time.Now() + polled, err := authn.PollDeviceFlow(context.Background(), &config.Config{}, flow) + if err != nil { + t.Fatalf("PollDeviceFlow: %v", err) + } + if polled == nil || !polled.Pending || polled.Auth != nil { + t.Fatalf("poll = %+v, want pending", polled) + } + if time.Since(started) > 2*time.Second { + t.Fatalf("pending poll blocked for %s", time.Since(started)) + } + }) + + t.Run("codex", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writeCodexDeviceServer(w, r, 1) + })) + defer server.Close() + authn := NewCodexAuthenticator() + authn.httpClient = rewriteClient(server) + flow, err := authn.StartDeviceFlow(context.Background(), &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + started := time.Now() + polled, err := authn.PollDeviceFlow(context.Background(), &config.Config{}, flow) + if err != nil { + t.Fatalf("PollDeviceFlow: %v", err) + } + if polled == nil || !polled.Pending || polled.Auth != nil { + t.Fatalf("poll = %+v, want pending", polled) + } + if time.Since(started) > 2*time.Second { + t.Fatalf("pending poll blocked for %s", time.Since(started)) + } + }) +} + +func TestDeviceFlowLoginUsesStartAndComplete(t *testing.T) { + t.Run("xai", func(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pendingPolls := 0 + if r.URL.Path == "/oauth2/token" { + pendingPolls = int(polls.Add(1)) + } + writeXAIDeviceServer(w, r, pendingPolls) + })) + defer server.Close() + + authn := XAIAuthenticator{httpClient: rewriteClient(server), pollIntervalOverride: time.Millisecond} + record, err := authn.Login(context.Background(), &config.Config{}, &LoginOptions{NoBrowser: true}) + if err != nil { + t.Fatalf("Login: %v", err) + } + if record.Metadata["access_token"] != "access-1" || record.Metadata["refresh_token"] != "refresh-1" { + t.Fatalf("metadata = %#v", record.Metadata) + } + if record.Provider != "xai" { + t.Fatalf("provider = %q", record.Provider) + } + if polls.Load() < 2 { + t.Fatalf("token polls = %d, want pending then exchange", polls.Load()) + } + }) + + t.Run("meta", func(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pendingPolls := 0 + if strings.Contains(r.URL.Path, "/oidc/device/token") { + pendingPolls = int(polls.Add(1)) + } + writeMetaDeviceServer(w, r, pendingPolls) + })) + defer server.Close() + + authn := MetaAuthenticator{httpClient: rewriteClient(server), pollIntervalOverride: time.Millisecond} + record, err := authn.Login(context.Background(), &config.Config{}, &LoginOptions{NoBrowser: true}) + if err != nil { + t.Fatalf("Login: %v", err) + } + if record.Metadata["api_key"] != "LLM|minted" || record.Metadata["dca_token"] != "dca-access" { + t.Fatalf("metadata = %#v", record.Metadata) + } + if record.Metadata["email"] != "user@meta.com" { + t.Fatalf("email = %#v", record.Metadata["email"]) + } + if polls.Load() < 2 { + t.Fatalf("token polls = %d, want pending then exchange", polls.Load()) + } + }) + + t.Run("codex", func(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pendingPolls := 0 + if r.URL.Path == "/api/accounts/deviceauth/token" { + pendingPolls = int(polls.Add(1)) + } + writeCodexDeviceServer(w, r, pendingPolls) + })) + defer server.Close() + + authn := NewCodexAuthenticator() + authn.httpClient = rewriteClient(server) + authn.pollIntervalOverride = time.Millisecond + record, err := authn.Login(context.Background(), &config.Config{}, &LoginOptions{ + NoBrowser: true, + Metadata: map[string]string{codexLoginModeMetadataKey: codexLoginModeDevice}, + }) + if err != nil { + t.Fatalf("Login: %v", err) + } + storage, ok := record.Storage.(*codexauth.CodexTokenStorage) + if !ok || storage.AccessToken != "access-1" || storage.Email != "user@openai.com" { + t.Fatalf("storage = %#v", record.Storage) + } + if polls.Load() < 2 { + t.Fatalf("token polls = %d, want pending then exchange", polls.Load()) + } + }) +} + +func TestPollDeviceFlowExchangesApprovedCode(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writeXAIDeviceServer(w, r, 2) + })) + defer server.Close() + authn := XAIAuthenticator{httpClient: rewriteClient(server)} + flow, err := authn.StartDeviceFlow(context.Background(), &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + polled, err := authn.PollDeviceFlow(context.Background(), &config.Config{}, flow) + if err != nil { + t.Fatalf("PollDeviceFlow: %v", err) + } + if polled == nil || polled.Pending || polled.Auth == nil || polled.Auth.Metadata["access_token"] != "access-1" { + t.Fatalf("poll = %+v", polled) + } + + denied := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/oauth2/token" { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{"error": "access_denied"}) + return + } + writeXAIDeviceServer(w, r, 0) + })) + defer denied.Close() + authn.httpClient = rewriteClient(denied) + flow, err = authn.StartDeviceFlow(context.Background(), &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + polled, err = authn.PollDeviceFlow(context.Background(), &config.Config{}, flow) + if err == nil || !strings.Contains(err.Error(), "authorization denied") { + t.Fatalf("PollDeviceFlow error = %v, want authorization denied", err) + } + if polled != nil && polled.Auth != nil { + t.Fatalf("denied poll returned auth %+v", polled.Auth) + } +} + +func TestCompleteDeviceFlowExchangesCodexCode(t *testing.T) { + var polls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + pendingPolls := 0 + if r.URL.Path == "/api/accounts/deviceauth/token" { + pendingPolls = int(polls.Add(1)) + } + writeCodexDeviceServer(w, r, pendingPolls) + })) + defer server.Close() + + authn := NewCodexAuthenticator() + authn.httpClient = rewriteClient(server) + authn.pollIntervalOverride = time.Millisecond + flow, err := authn.StartDeviceFlow(context.Background(), &config.Config{}) + if err != nil { + t.Fatalf("StartDeviceFlow: %v", err) + } + if polls.Load() != 0 { + t.Fatalf("start polled %d times", polls.Load()) + } + record, err := authn.CompleteDeviceFlow(context.Background(), &config.Config{}, flow) + if err != nil { + t.Fatalf("CompleteDeviceFlow: %v", err) + } + storage, ok := record.Storage.(*codexauth.CodexTokenStorage) + if !ok || storage.AccessToken != "access-1" { + t.Fatalf("storage = %#v", record.Storage) + } + if polls.Load() < 2 { + t.Fatalf("token polls = %d, want pending then exchange", polls.Load()) + } +} + +func writeXAIDeviceServer(w http.ResponseWriter, r *http.Request, pollNumber int) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/.well-known/openid-configuration": + _ = json.NewEncoder(w).Encode(map[string]string{ + "device_authorization_endpoint": "https://auth.x.ai/oauth2/device/code", + "token_endpoint": "https://auth.x.ai/oauth2/token", + }) + case "/oauth2/device/code": + _ = json.NewEncoder(w).Encode(map[string]any{ + "device_code": "device-abc", + "user_code": "ABCD-1234", + "verification_uri": "https://accounts.x.ai/oauth2/device", + "verification_uri_complete": "https://accounts.x.ai/oauth2/device?user_code=ABCD-1234", + "expires_in": 1800, + "interval": 5, + }) + case "/oauth2/token": + if pollNumber < 2 { + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{"error": "authorization_pending"}) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-1", + "refresh_token": "refresh-1", + "token_type": "Bearer", + "expires_in": 3600, + "id_token": testJWT(`{"email":"user@x.ai","sub":"sub-1"}`), + }) + default: + http.NotFound(w, r) + } +} + +func writeMetaDeviceServer(w http.ResponseWriter, r *http.Request, pollNumber int) { + w.Header().Set("Content-Type", "application/json") + switch { + case strings.Contains(r.URL.Path, "/oidc/device/authorization"): + _ = json.NewEncoder(w).Encode(map[string]any{ + "device_code": "meta-device", + "user_code": "TEST-1234", + "verification_uri": "https://auth.meta.com/device", + "verification_uri_complete": "https://auth.meta.com/device?user_code=TEST-1234", + "expires_in": 900, + "interval": 5, + }) + case strings.Contains(r.URL.Path, "/oidc/device/token"): + if pollNumber < 2 { + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{"error": "authorization_pending"}) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "dca-access", + "token_type": "Bearer", + "expires_in": 3600, + }) + case strings.Contains(r.URL.Path, "/muse-code/key"): + _ = json.NewEncoder(w).Encode(map[string]any{ + "api_key": "LLM|minted", + "user_email": "user@meta.com", + "user_full_name": "Meta User", + }) + default: + http.NotFound(w, r) + } +} + +func writeCodexDeviceServer(w http.ResponseWriter, r *http.Request, pollNumber int) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/api/accounts/deviceauth/usercode": + _ = json.NewEncoder(w).Encode(map[string]any{ + "device_auth_id": "auth-id-1", + "user_code": "CODE-1234", + "interval": 5, + }) + case "/api/accounts/deviceauth/token": + if pollNumber < 2 { + w.WriteHeader(http.StatusForbidden) + return + } + _ = json.NewEncoder(w).Encode(map[string]string{ + "authorization_code": "auth-code", + "code_verifier": "verifier", + "code_challenge": "challenge", + }) + case "/oauth/token": + _ = json.NewEncoder(w).Encode(map[string]any{ + "access_token": "access-1", + "refresh_token": "refresh-1", + "id_token": testJWT(`{"email":"user@openai.com","https://api.openai.com/auth":{"chatgpt_account_id":"acct-1","chatgpt_plan_type":"plus"}}`), + "expires_in": 3600, + }) + default: + http.NotFound(w, r) + } +} + +func testJWT(payload string) string { + header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"none","typ":"JWT"}`)) + body := base64.RawURLEncoding.EncodeToString([]byte(payload)) + return header + "." + body + ".sig" +} + +func rewriteClient(server *httptest.Server) *http.Client { + return &http.Client{ + Timeout: 5 * time.Second, + Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + clone := req.Clone(req.Context()) + clone.URL.Scheme = "http" + clone.URL.Host = server.Listener.Addr().String() + clone.Host = clone.URL.Host + return server.Client().Transport.RoundTrip(clone) + }), + } +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/sdk/auth/meta.go b/sdk/auth/meta.go index 7650cd15cb..cdfaa41c05 100644 --- a/sdk/auth/meta.go +++ b/sdk/auth/meta.go @@ -3,10 +3,10 @@ package auth import ( "context" "fmt" + "net/http" "strings" "time" - metaauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/meta" "github.com/router-for-me/CLIProxyAPI/v7/internal/browser" "github.com/router-for-me/CLIProxyAPI/v7/internal/config" coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" @@ -14,7 +14,12 @@ import ( ) // MetaAuthenticator implements the Meta Muse OAuth device-code flow. -type MetaAuthenticator struct{} +type MetaAuthenticator struct { + httpClient *http.Client + // pollIntervalOverride, when set, replaces the wait between device-flow polls. + // Production leaves it zero so the provider interval is used. + pollIntervalOverride time.Duration +} // NewMetaAuthenticator constructs a new Meta authenticator. func NewMetaAuthenticator() Authenticator { @@ -45,22 +50,20 @@ func (a MetaAuthenticator) Login(ctx context.Context, cfg *config.Config, opts * opts = &LoginOptions{} } - authSvc := metaauth.NewMetaAuth(cfg) - fmt.Println("Starting Meta (Muse) authentication...") - deviceCode, err := authSvc.StartDeviceFlow(ctx) + flow, err := a.StartDeviceFlow(ctx, cfg) if err != nil { - return nil, fmt.Errorf("meta: failed to start device flow: %w", err) + return nil, err } - verificationURL := strings.TrimSpace(deviceCode.VerificationURIComplete) + verificationURL := strings.TrimSpace(flow.VerificationURIComplete) if verificationURL == "" { - verificationURL = strings.TrimSpace(deviceCode.VerificationURI) + verificationURL = strings.TrimSpace(flow.VerificationURI) } fmt.Printf("\nTo authenticate, please visit:\n%s\n\n", verificationURL) - if deviceCode.UserCode != "" { - fmt.Printf("Then enter this code: %s\n\n", deviceCode.UserCode) + if flow.UserCode != "" { + fmt.Printf("Then enter this code: %s\n\n", flow.UserCode) } if !opts.NoBrowser { @@ -76,78 +79,13 @@ func (a MetaAuthenticator) Login(ctx context.Context, cfg *config.Config, opts * } fmt.Println("Waiting for authorization...") - if deviceCode.ExpiresIn > 0 { - fmt.Printf("(This will timeout in %d seconds if not authorized)\n", deviceCode.ExpiresIn) + if flow.ExpiresIn > 0 { + fmt.Printf("(This will timeout in %d seconds if not authorized)\n", int(flow.ExpiresIn/time.Second)) } - bundle, errWait := authSvc.WaitForAuthorization(ctx, deviceCode) + record, errWait := a.CompleteDeviceFlow(ctx, cfg, flow) if errWait != nil { return nil, fmt.Errorf("meta: %w", errWait) } - - tokenStorage := authSvc.CreateTokenStorage(bundle) - if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" { - return nil, fmt.Errorf("meta token storage missing access token") - } - - fileName := metaauth.CredentialFileName(tokenStorage.Email, tokenStorage.DCAToken) - label := strings.TrimSpace(tokenStorage.Email) - if label == "" { - label = "Meta" - } - - metadata := map[string]any{ - "type": "meta", - "access_token": tokenStorage.AccessToken, - "token_type": tokenStorage.TokenType, - "expires_in": tokenStorage.ExpiresIn, - "expired": tokenStorage.Expired, - "last_refresh": tokenStorage.LastRefresh, - "base_url": tokenStorage.BaseURL, - "auth_kind": "oauth", - } - if tokenStorage.DCAExpired != "" { - metadata["dca_expired"] = tokenStorage.DCAExpired - } - if tokenStorage.DCAExpiresAt > 0 { - metadata["dca_expires_at"] = tokenStorage.DCAExpiresAt - } - if tokenStorage.APIKey != "" { - metadata["api_key"] = tokenStorage.APIKey - } - if tokenStorage.DCAToken != "" { - metadata["dca_token"] = tokenStorage.DCAToken - } - if tokenStorage.Email != "" { - metadata["email"] = tokenStorage.Email - } - if tokenStorage.Name != "" { - metadata["name"] = tokenStorage.Name - } - - attrs := map[string]string{ - "auth_kind": "oauth", - "base_url": tokenStorage.BaseURL, - } - if tokenStorage.APIKey != "" { - attrs["api_key"] = tokenStorage.APIKey - } - if tokenStorage.DCAToken != "" { - attrs["dca_token"] = tokenStorage.DCAToken - } - if tokenStorage.Email != "" { - attrs["email"] = tokenStorage.Email - } - - fmt.Println("Meta authentication successful") - - return &coreauth.Auth{ - ID: fileName, - Provider: a.Provider(), - FileName: fileName, - Label: label, - Storage: tokenStorage, - Metadata: metadata, - Attributes: attrs, - }, nil + return record, nil } diff --git a/sdk/auth/meta_device.go b/sdk/auth/meta_device.go new file mode 100644 index 0000000000..4eaeda11d6 --- /dev/null +++ b/sdk/auth/meta_device.go @@ -0,0 +1,201 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + metaauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/meta" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func (a MetaAuthenticator) metaAuth(cfg *config.Config) *metaauth.MetaAuth { + svc := metaauth.NewMetaAuth(cfg) + if a.httpClient != nil { + svc.SetHTTPClient(a.httpClient) + } + return svc +} + +func (a MetaAuthenticator) effectiveInterval(d time.Duration) time.Duration { + if a.pollIntervalOverride > 0 { + return a.pollIntervalOverride + } + if d <= 0 { + return metaauth.DefaultPollInterval + } + return d +} + +// StartDeviceFlow requests a Meta device code and returns it without waiting +// for the user to authorize. +func (a MetaAuthenticator) StartDeviceFlow(ctx context.Context, cfg *config.Config) (*DeviceFlow, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if ctx == nil { + ctx = context.Background() + } + deviceCode, err := a.metaAuth(cfg).StartDeviceFlow(ctx) + if err != nil { + return nil, fmt.Errorf("meta: failed to start device flow: %w", err) + } + if deviceCode == nil || strings.TrimSpace(deviceCode.DeviceCode) == "" { + return nil, fmt.Errorf("meta: failed to start device flow: empty device code") + } + flow := &DeviceFlow{ + Provider: a.Provider(), + UserCode: strings.TrimSpace(deviceCode.UserCode), + DeviceCode: strings.TrimSpace(deviceCode.DeviceCode), + VerificationURI: strings.TrimSpace(deviceCode.VerificationURI), + VerificationURIComplete: strings.TrimSpace(deviceCode.VerificationURIComplete), + TokenEndpoint: strings.TrimSpace(deviceCode.TokenEndpoint), + } + if deviceCode.ExpiresIn > 0 { + flow.ExpiresIn = time.Duration(deviceCode.ExpiresIn) * time.Second + } + if deviceCode.Interval > 0 { + flow.Interval = time.Duration(deviceCode.Interval) * time.Second + } + return flow, nil +} + +// PollDeviceFlow exchanges a Meta device code once. +// A pending result means the user has not approved yet. +func (a MetaAuthenticator) PollDeviceFlow(ctx context.Context, cfg *config.Config, flow *DeviceFlow) (*DeviceFlowPoll, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if flow == nil || strings.TrimSpace(flow.DeviceCode) == "" { + return nil, fmt.Errorf("meta device flow is required") + } + if ctx == nil { + ctx = context.Background() + } + interval := a.effectiveInterval(flow.Interval) + svc := a.metaAuth(cfg) + bundle, next, err := svc.PollDeviceCodeOnce(ctx, &metaauth.DeviceCodeResponse{ + DeviceCode: flow.DeviceCode, + UserCode: flow.UserCode, + VerificationURI: flow.VerificationURI, + VerificationURIComplete: flow.VerificationURIComplete, + ExpiresIn: int(flow.ExpiresIn / time.Second), + Interval: int(interval / time.Second), + TokenEndpoint: flow.TokenEndpoint, + }, interval) + if err != nil { + return nil, err + } + if bundle == nil { + return &DeviceFlowPoll{Pending: true, Interval: a.effectiveInterval(next)}, nil + } + record, errRecord := a.authFromMetaBundle(svc, bundle) + if errRecord != nil { + return nil, errRecord + } + return &DeviceFlowPoll{Interval: a.effectiveInterval(next), Auth: record}, nil +} + +// CompleteDeviceFlow polls until the Meta device code is authorized or the flow fails. +func (a MetaAuthenticator) CompleteDeviceFlow(ctx context.Context, cfg *config.Config, flow *DeviceFlow) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if flow == nil || strings.TrimSpace(flow.DeviceCode) == "" { + return nil, fmt.Errorf("meta device flow is required") + } + if ctx == nil { + ctx = context.Background() + } + maxDuration := metaauth.MaxPollDuration + if flow.ExpiresIn > 0 && flow.ExpiresIn < maxDuration { + maxDuration = flow.ExpiresIn + } + ctx, cancel := context.WithTimeout(ctx, maxDuration) + defer cancel() + + record, err := pollUntilDeviceAuthorized(ctx, a.effectiveInterval(flow.Interval), time.Time{}, nil, func() (*DeviceFlowPoll, error) { + polled, errPoll := a.PollDeviceFlow(ctx, cfg, flow) + if errPoll != nil || polled == nil || polled.Auth != nil { + return polled, errPoll + } + polled.Interval = a.effectiveInterval(polled.Interval) + flow.Interval = polled.Interval + return polled, nil + }) + if err != nil && ctx.Err() != nil && (errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled)) { + return nil, fmt.Errorf("meta auth: authorization timed out or canceled: %w", ctx.Err()) + } + return record, err +} + +func (a MetaAuthenticator) authFromMetaBundle(authSvc *metaauth.MetaAuth, bundle *metaauth.MetaAuthBundle) (*coreauth.Auth, error) { + tokenStorage := authSvc.CreateTokenStorage(bundle) + if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" { + return nil, fmt.Errorf("meta token storage missing access token") + } + + fileName := metaauth.CredentialFileName(tokenStorage.Email, tokenStorage.DCAToken) + label := strings.TrimSpace(tokenStorage.Email) + if label == "" { + label = "Meta" + } + + metadata := map[string]any{ + "type": "meta", + "access_token": tokenStorage.AccessToken, + "token_type": tokenStorage.TokenType, + "expires_in": tokenStorage.ExpiresIn, + "expired": tokenStorage.Expired, + "last_refresh": tokenStorage.LastRefresh, + "base_url": tokenStorage.BaseURL, + "auth_kind": "oauth", + } + if tokenStorage.DCAExpired != "" { + metadata["dca_expired"] = tokenStorage.DCAExpired + } + if tokenStorage.DCAExpiresAt > 0 { + metadata["dca_expires_at"] = tokenStorage.DCAExpiresAt + } + if tokenStorage.APIKey != "" { + metadata["api_key"] = tokenStorage.APIKey + } + if tokenStorage.DCAToken != "" { + metadata["dca_token"] = tokenStorage.DCAToken + } + if tokenStorage.Email != "" { + metadata["email"] = tokenStorage.Email + } + if tokenStorage.Name != "" { + metadata["name"] = tokenStorage.Name + } + + attrs := map[string]string{ + "auth_kind": "oauth", + "base_url": tokenStorage.BaseURL, + } + if tokenStorage.APIKey != "" { + attrs["api_key"] = tokenStorage.APIKey + } + if tokenStorage.DCAToken != "" { + attrs["dca_token"] = tokenStorage.DCAToken + } + if tokenStorage.Email != "" { + attrs["email"] = tokenStorage.Email + } + + fmt.Println("Meta authentication successful") + + return &coreauth.Auth{ + ID: fileName, + Provider: a.Provider(), + FileName: fileName, + Label: label, + Storage: tokenStorage, + Metadata: metadata, + Attributes: attrs, + }, nil +} diff --git a/sdk/auth/xai.go b/sdk/auth/xai.go index 039878b243..bd2622c4ca 100644 --- a/sdk/auth/xai.go +++ b/sdk/auth/xai.go @@ -3,6 +3,7 @@ package auth import ( "context" "fmt" + "net/http" "strings" "time" @@ -14,7 +15,12 @@ import ( ) // XAIAuthenticator implements the xAI Grok OAuth device-code flow. -type XAIAuthenticator struct{} +type XAIAuthenticator struct { + httpClient *http.Client + // pollIntervalOverride, when set, replaces the wait between device-flow polls. + // Production leaves it zero so the provider interval is used. + pollIntervalOverride time.Duration +} // NewXAIAuthenticator constructs a new xAI authenticator. func NewXAIAuthenticator() Authenticator { @@ -44,22 +50,20 @@ func (a XAIAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *L opts = &LoginOptions{} } - authSvc := xaiauth.NewXAIAuth(cfg) - fmt.Println("Starting xAI authentication...") - deviceCode, err := authSvc.StartDeviceFlow(ctx) + flow, err := a.StartDeviceFlow(ctx, cfg) if err != nil { - return nil, fmt.Errorf("xai: failed to start device flow: %w", err) + return nil, err } - verificationURL := strings.TrimSpace(deviceCode.VerificationURIComplete) + verificationURL := strings.TrimSpace(flow.VerificationURIComplete) if verificationURL == "" { - verificationURL = strings.TrimSpace(deviceCode.VerificationURI) + verificationURL = strings.TrimSpace(flow.VerificationURI) } fmt.Printf("\nTo authenticate, please visit:\n%s\n\n", verificationURL) - if deviceCode.UserCode != "" { - fmt.Printf("Then enter this code: %s\n\n", deviceCode.UserCode) + if flow.UserCode != "" { + fmt.Printf("Then enter this code: %s\n\n", flow.UserCode) } if !opts.NoBrowser { @@ -75,58 +79,13 @@ func (a XAIAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *L } fmt.Println("Waiting for authorization...") - if deviceCode.ExpiresIn > 0 { - fmt.Printf("(This will timeout in %d seconds if not authorized)\n", deviceCode.ExpiresIn) + if flow.ExpiresIn > 0 { + fmt.Printf("(This will timeout in %d seconds if not authorized)\n", int(flow.ExpiresIn/time.Second)) } - bundle, errWait := authSvc.WaitForAuthorization(ctx, deviceCode) + record, errWait := a.CompleteDeviceFlow(ctx, cfg, flow) if errWait != nil { return nil, fmt.Errorf("xai: %w", errWait) } - - tokenStorage := authSvc.CreateTokenStorage(bundle) - if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" { - return nil, fmt.Errorf("xai token storage missing access token") - } - - fileName := xaiauth.CredentialFileName(tokenStorage.Email, tokenStorage.Subject) - label := strings.TrimSpace(tokenStorage.Email) - if label == "" { - label = "xAI" - } - - metadata := map[string]any{ - "type": "xai", - "access_token": tokenStorage.AccessToken, - "refresh_token": tokenStorage.RefreshToken, - "id_token": tokenStorage.IDToken, - "token_type": tokenStorage.TokenType, - "expires_in": tokenStorage.ExpiresIn, - "expired": tokenStorage.Expire, - "last_refresh": tokenStorage.LastRefresh, - "base_url": tokenStorage.BaseURL, - "token_endpoint": tokenStorage.TokenEndpoint, - "auth_kind": "oauth", - } - if tokenStorage.Email != "" { - metadata["email"] = tokenStorage.Email - } - if tokenStorage.Subject != "" { - metadata["sub"] = tokenStorage.Subject - } - - fmt.Println("xAI authentication successful") - - return &coreauth.Auth{ - ID: fileName, - Provider: a.Provider(), - FileName: fileName, - Label: label, - Storage: tokenStorage, - Metadata: metadata, - Attributes: map[string]string{ - "auth_kind": "oauth", - "base_url": tokenStorage.BaseURL, - }, - }, nil + return record, nil } diff --git a/sdk/auth/xai_device.go b/sdk/auth/xai_device.go new file mode 100644 index 0000000000..6df601a449 --- /dev/null +++ b/sdk/auth/xai_device.go @@ -0,0 +1,179 @@ +package auth + +import ( + "context" + "fmt" + "strings" + "time" + + xaiauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/xai" + "github.com/router-for-me/CLIProxyAPI/v7/internal/config" + coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" +) + +func (a XAIAuthenticator) xaiAuth(cfg *config.Config) *xaiauth.XAIAuth { + svc := xaiauth.NewXAIAuth(cfg) + if a.httpClient != nil { + svc.SetHTTPClient(a.httpClient) + } + return svc +} + +func (a XAIAuthenticator) effectiveInterval(d time.Duration) time.Duration { + if a.pollIntervalOverride > 0 { + return a.pollIntervalOverride + } + if d < 5*time.Second { + return 5 * time.Second + } + return d +} + +// StartDeviceFlow requests an xAI device code and returns it without waiting +// for the user to authorize. +func (a XAIAuthenticator) StartDeviceFlow(ctx context.Context, cfg *config.Config) (*DeviceFlow, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if ctx == nil { + ctx = context.Background() + } + deviceCode, err := a.xaiAuth(cfg).StartDeviceFlow(ctx) + if err != nil { + return nil, fmt.Errorf("xai: failed to start device flow: %w", err) + } + if deviceCode == nil || strings.TrimSpace(deviceCode.DeviceCode) == "" { + return nil, fmt.Errorf("xai: failed to start device flow: empty device code") + } + flow := &DeviceFlow{ + Provider: a.Provider(), + UserCode: strings.TrimSpace(deviceCode.UserCode), + DeviceCode: strings.TrimSpace(deviceCode.DeviceCode), + VerificationURI: strings.TrimSpace(deviceCode.VerificationURI), + VerificationURIComplete: strings.TrimSpace(deviceCode.VerificationURIComplete), + TokenEndpoint: strings.TrimSpace(deviceCode.TokenEndpoint), + } + if deviceCode.ExpiresIn > 0 { + flow.ExpiresIn = time.Duration(deviceCode.ExpiresIn) * time.Second + } + if deviceCode.Interval > 0 { + flow.Interval = time.Duration(deviceCode.Interval) * time.Second + } + return flow, nil +} + +// PollDeviceFlow exchanges an xAI device code once. +// A pending result means the user has not approved yet. +func (a XAIAuthenticator) PollDeviceFlow(ctx context.Context, cfg *config.Config, flow *DeviceFlow) (*DeviceFlowPoll, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if flow == nil || strings.TrimSpace(flow.DeviceCode) == "" { + return nil, fmt.Errorf("xai device flow is required") + } + if ctx == nil { + ctx = context.Background() + } + interval := a.effectiveInterval(flow.Interval) + svc := a.xaiAuth(cfg) + token, next, err := svc.PollDeviceCodeOnce(ctx, &xaiauth.DeviceCodeResponse{ + DeviceCode: flow.DeviceCode, + UserCode: flow.UserCode, + VerificationURI: flow.VerificationURI, + VerificationURIComplete: flow.VerificationURIComplete, + ExpiresIn: int(flow.ExpiresIn / time.Second), + Interval: int(interval / time.Second), + TokenEndpoint: flow.TokenEndpoint, + }, interval) + if err != nil { + return nil, err + } + if token == nil { + return &DeviceFlowPoll{Pending: true, Interval: a.effectiveInterval(next)}, nil + } + record, errRecord := a.authFromXAIBundle(svc, &xaiauth.AuthBundle{ + TokenData: *token, + LastRefresh: time.Now().UTC().Format(time.RFC3339), + BaseURL: xaiauth.DefaultAPIBaseURL, + TokenEndpoint: strings.TrimSpace(flow.TokenEndpoint), + }) + if errRecord != nil { + return nil, errRecord + } + return &DeviceFlowPoll{Interval: a.effectiveInterval(next), Auth: record}, nil +} + +// CompleteDeviceFlow polls until the xAI device code is authorized or the flow fails. +func (a XAIAuthenticator) CompleteDeviceFlow(ctx context.Context, cfg *config.Config, flow *DeviceFlow) (*coreauth.Auth, error) { + if cfg == nil { + return nil, fmt.Errorf("cliproxy auth: configuration is required") + } + if flow == nil || strings.TrimSpace(flow.DeviceCode) == "" { + return nil, fmt.Errorf("xai device flow is required") + } + interval := a.effectiveInterval(flow.Interval) + deadline := time.Now().Add(xaiauth.MaxPollDuration) + if flow.ExpiresIn > 0 { + expiresAt := time.Now().Add(flow.ExpiresIn) + if expiresAt.Before(deadline) { + deadline = expiresAt + } + } + return pollUntilDeviceAuthorized(ctx, interval, deadline, fmt.Errorf("xai device code expired"), func() (*DeviceFlowPoll, error) { + polled, err := a.PollDeviceFlow(ctx, cfg, flow) + if err != nil || polled == nil || polled.Auth != nil { + return polled, err + } + polled.Interval = a.effectiveInterval(polled.Interval) + flow.Interval = polled.Interval + return polled, nil + }) +} + +func (a XAIAuthenticator) authFromXAIBundle(authSvc *xaiauth.XAIAuth, bundle *xaiauth.AuthBundle) (*coreauth.Auth, error) { + tokenStorage := authSvc.CreateTokenStorage(bundle) + if tokenStorage == nil || strings.TrimSpace(tokenStorage.AccessToken) == "" { + return nil, fmt.Errorf("xai token storage missing access token") + } + + fileName := xaiauth.CredentialFileName(tokenStorage.Email, tokenStorage.Subject) + label := strings.TrimSpace(tokenStorage.Email) + if label == "" { + label = "xAI" + } + + metadata := map[string]any{ + "type": "xai", + "access_token": tokenStorage.AccessToken, + "refresh_token": tokenStorage.RefreshToken, + "id_token": tokenStorage.IDToken, + "token_type": tokenStorage.TokenType, + "expires_in": tokenStorage.ExpiresIn, + "expired": tokenStorage.Expire, + "last_refresh": tokenStorage.LastRefresh, + "base_url": tokenStorage.BaseURL, + "token_endpoint": tokenStorage.TokenEndpoint, + "auth_kind": "oauth", + } + if tokenStorage.Email != "" { + metadata["email"] = tokenStorage.Email + } + if tokenStorage.Subject != "" { + metadata["sub"] = tokenStorage.Subject + } + + fmt.Println("xAI authentication successful") + + return &coreauth.Auth{ + ID: fileName, + Provider: a.Provider(), + FileName: fileName, + Label: label, + Storage: tokenStorage, + Metadata: metadata, + Attributes: map[string]string{ + "auth_kind": "oauth", + "base_url": tokenStorage.BaseURL, + }, + }, nil +}