From acb86eabfa81eaf0c51d1d459ee47b0268999817 Mon Sep 17 00:00:00 2001 From: Jacki Ng Date: Wed, 23 Sep 2026 19:35:28 +0800 Subject: [PATCH] feat(auth): export split-phase device flow for Codex, xAI, and Meta StartDeviceFlow returns the device code without waiting for approval. PollDeviceFlow exchanges it once and CompleteDeviceFlow repeats until the user authorizes. Login calls those phases. --- internal/auth/codex/openai_auth.go | 9 + internal/auth/meta/meta.go | 179 +++++++----- internal/auth/xai/xai.go | 52 +++- sdk/auth/codex.go | 4 + sdk/auth/codex_device.go | 230 ++++++++++----- sdk/auth/device_flow.go | 77 +++++ sdk/auth/device_flow_test.go | 451 +++++++++++++++++++++++++++++ sdk/auth/meta.go | 96 ++---- sdk/auth/meta_device.go | 201 +++++++++++++ sdk/auth/xai.go | 75 ++--- sdk/auth/xai_device.go | 179 ++++++++++++ 11 files changed, 1270 insertions(+), 283 deletions(-) create mode 100644 sdk/auth/device_flow.go create mode 100644 sdk/auth/device_flow_test.go create mode 100644 sdk/auth/meta_device.go create mode 100644 sdk/auth/xai_device.go diff --git a/internal/auth/codex/openai_auth.go b/internal/auth/codex/openai_auth.go index 2c1eac086f4..fad2ce577e6 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 3eb35541922..495286dd06f 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 a0b76ea838f..86fbdaf72cc 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 79448991b8d..a436364f7d5 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 d7ea4e1fe93..b412c4169ba 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 00000000000..83140ab1aa6 --- /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 00000000000..489e5f10307 --- /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 7650cd15cbd..cdfaa41c05d 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 00000000000..4eaeda11d62 --- /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 039878b2432..bd2622c4ca2 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 00000000000..6df601a449d --- /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 +}