package codebuddy import ( "bytes" "context" "encoding/base64" "encoding/json" "fmt" "io" "net/http" "net/url" "strings" "time" "github.com/google/uuid" log "github.com/sirupsen/logrus" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" "github.com/router-for-me/CLIProxyAPI/v6/internal/util" ) const ( BaseURL = "https://copilot.tencent.com" DefaultDomain = "www.codebuddy.cn" UserAgent = "CLI/2.63.2 CodeBuddy/2.63.2" codeBuddyStatePath = "/v2/plugin/auth/state" codeBuddyTokenPath = "/v2/plugin/auth/token" codeBuddyRefreshPath = "/v2/plugin/auth/token/refresh" pollInterval = 5 * time.Second maxPollDuration = 5 * time.Minute codeLoginPending = 11217 codeSuccess = 0 ) type CodeBuddyAuth struct { httpClient *http.Client cfg *config.Config baseURL string } func NewCodeBuddyAuth(cfg *config.Config) *CodeBuddyAuth { httpClient := &http.Client{Timeout: 30 * time.Second} if cfg != nil { httpClient = util.SetProxy(&cfg.SDKConfig, httpClient) } return &CodeBuddyAuth{httpClient: httpClient, cfg: cfg, baseURL: BaseURL} } // AuthState holds the state and auth URL returned by the auth state API. type AuthState struct { State string AuthURL string } // FetchAuthState calls POST /v2/plugin/auth/state?platform=CLI to get the state and login URL. func (a *CodeBuddyAuth) FetchAuthState(ctx context.Context) (*AuthState, error) { stateURL := fmt.Sprintf("%s%s?platform=CLI", a.baseURL, codeBuddyStatePath) body := []byte("{}") req, err := http.NewRequestWithContext(ctx, http.MethodPost, stateURL, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("codebuddy: failed to create auth state request: %w", err) } requestID := uuid.NewString() req.Header.Set("Accept", "application/json, text/plain, */*") req.Header.Set("Content-Type", "application/json") req.Header.Set("X-Requested-With", "XMLHttpRequest") req.Header.Set("X-Domain", "copilot.tencent.com") req.Header.Set("X-No-Authorization", "true") req.Header.Set("X-No-User-Id", "true") req.Header.Set("X-No-Enterprise-Id", "true") req.Header.Set("X-No-Department-Info", "true") req.Header.Set("X-Product", "SaaS") req.Header.Set("User-Agent", UserAgent) req.Header.Set("X-Request-ID", requestID) resp, err := a.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("codebuddy: auth state request failed: %w", err) } defer func() { if errClose := resp.Body.Close(); errClose != nil { log.Errorf("codebuddy auth state: close body error: %v", errClose) } }() bodyBytes, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("codebuddy: failed to read auth state response: %w", err) } if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("codebuddy: auth state request returned status %d: %s", resp.StatusCode, string(bodyBytes)) } var result struct { Code int `json:"code"` Msg string `json:"msg"` Data *struct { State string `json:"state"` AuthURL string `json:"authUrl"` } `json:"data"` } if err = json.Unmarshal(bodyBytes, &result); err != nil { return nil, fmt.Errorf("codebuddy: failed to parse auth state response: %w", err) } if result.Code != codeSuccess { return nil, fmt.Errorf("codebuddy: auth state request failed with code %d: %s", result.Code, result.Msg) } if result.Data == nil || result.Data.State == "" || result.Data.AuthURL == "" { return nil, fmt.Errorf("codebuddy: auth state response missing state or authUrl") } return &AuthState{ State: result.Data.State, AuthURL: result.Data.AuthURL, }, nil } type pollResponse struct { Code int `json:"code"` Msg string `json:"msg"` RequestID string `json:"requestId"` Data *struct { AccessToken string `json:"accessToken"` RefreshToken string `json:"refreshToken"` ExpiresIn int64 `json:"expiresIn"` TokenType string `json:"tokenType"` Domain string `json:"domain"` } `json:"data"` } // doPollRequest performs a single polling request, safely reading and closing the response body func (a *CodeBuddyAuth) doPollRequest(ctx context.Context, pollURL string) ([]byte, int, error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, pollURL, nil) if err != nil { return nil, 0, fmt.Errorf("%w: %v", ErrTokenFetchFailed, err) } a.applyPollHeaders(req) resp, err := a.httpClient.Do(req) if err != nil { return nil, 0, err } defer func() { if errClose := resp.Body.Close(); errClose != nil { log.Errorf("codebuddy poll: close body error: %v", errClose) } }() body, err := io.ReadAll(resp.Body) if err != nil { return nil, resp.StatusCode, fmt.Errorf("codebuddy poll: failed to read response body: %w", err) } return body, resp.StatusCode, nil } // PollForToken polls until the user completes browser authorization and returns auth data. func (a *CodeBuddyAuth) PollForToken(ctx context.Context, state string) (*CodeBuddyTokenStorage, error) { deadline := time.Now().Add(maxPollDuration) pollURL := fmt.Sprintf("%s%s?state=%s", a.baseURL, codeBuddyTokenPath, url.QueryEscape(state)) for time.Now().Before(deadline) { select { case <-ctx.Done(): return nil, ctx.Err() case <-time.After(pollInterval): } body, statusCode, err := a.doPollRequest(ctx, pollURL) if err != nil { log.Debugf("codebuddy poll: request error: %v", err) continue } if statusCode != http.StatusOK { log.Debugf("codebuddy poll: unexpected status %d", statusCode) continue } var result pollResponse if err := json.Unmarshal(body, &result); err != nil { continue } switch result.Code { case codeSuccess: if result.Data == nil { return nil, fmt.Errorf("%w: empty data in response", ErrTokenFetchFailed) } userID, _ := a.DecodeUserID(result.Data.AccessToken) return &CodeBuddyTokenStorage{ AccessToken: result.Data.AccessToken, RefreshToken: result.Data.RefreshToken, ExpiresIn: result.Data.ExpiresIn, TokenType: result.Data.TokenType, Domain: result.Data.Domain, UserID: userID, Type: "codebuddy", }, nil case codeLoginPending: // continue polling default: // TODO: when the CodeBuddy API error code for user denial is known, // return ErrAccessDenied here instead of ErrTokenFetchFailed. return nil, fmt.Errorf("%w: server returned code %d: %s", ErrTokenFetchFailed, result.Code, result.Msg) } } return nil, ErrPollingTimeout } // DecodeUserID decodes the sub field from a JWT access token as the user ID. func (a *CodeBuddyAuth) DecodeUserID(accessToken string) (string, error) { parts := strings.Split(accessToken, ".") if len(parts) < 2 { return "", ErrJWTDecodeFailed } payload, err := base64.RawURLEncoding.DecodeString(parts[1]) if err != nil { return "", fmt.Errorf("%w: %v", ErrJWTDecodeFailed, err) } var claims struct { Sub string `json:"sub"` } if err := json.Unmarshal(payload, &claims); err != nil { return "", fmt.Errorf("%w: %v", ErrJWTDecodeFailed, err) } if claims.Sub == "" { return "", fmt.Errorf("%w: sub claim is empty", ErrJWTDecodeFailed) } return claims.Sub, nil } // RefreshToken exchanges a refresh token for a new access token. // It calls POST /v2/plugin/auth/token/refresh with the required headers. func (a *CodeBuddyAuth) RefreshToken(ctx context.Context, accessToken, refreshToken, userID, domain string) (*CodeBuddyTokenStorage, error) { if domain == "" { domain = DefaultDomain } refreshURL := fmt.Sprintf("%s%s", a.baseURL, codeBuddyRefreshPath) body := []byte("{}") req, err := http.NewRequestWithContext(ctx, http.MethodPost, refreshURL, bytes.NewReader(body)) if err != nil { return nil, fmt.Errorf("codebuddy: failed to create refresh request: %w", err) } requestID := strings.ReplaceAll(uuid.New().String(), "-", "") req.Header.Set("Accept", "application/json, text/plain, */*") req.Header.Set("Content-Type", "application/json") req.Header.Set("X-Requested-With", "XMLHttpRequest") req.Header.Set("X-Domain", domain) req.Header.Set("X-Refresh-Token", refreshToken) req.Header.Set("X-Auth-Refresh-Source", "plugin") req.Header.Set("X-Request-ID", requestID) req.Header.Set("Authorization", "Bearer "+accessToken) req.Header.Set("X-User-Id", userID) req.Header.Set("X-Product", "SaaS") req.Header.Set("User-Agent", UserAgent) resp, err := a.httpClient.Do(req) if err != nil { return nil, fmt.Errorf("codebuddy: refresh request failed: %w", err) } defer func() { if errClose := resp.Body.Close(); errClose != nil { log.Errorf("codebuddy refresh: close body error: %v", errClose) } }() bodyBytes, err := io.ReadAll(resp.Body) if err != nil { return nil, fmt.Errorf("codebuddy: failed to read refresh response: %w", err) } if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { return nil, fmt.Errorf("codebuddy: refresh token rejected (status %d)", resp.StatusCode) } if resp.StatusCode != http.StatusOK { return nil, fmt.Errorf("codebuddy: refresh failed with status %d: %s", resp.StatusCode, string(bodyBytes)) } var result struct { Code int `json:"code"` Msg string `json:"msg"` Data *struct { AccessToken string `json:"accessToken"` RefreshToken string `json:"refreshToken"` ExpiresIn int64 `json:"expiresIn"` RefreshExpiresIn int64 `json:"refreshExpiresIn"` TokenType string `json:"tokenType"` Domain string `json:"domain"` } `json:"data"` } if err = json.Unmarshal(bodyBytes, &result); err != nil { return nil, fmt.Errorf("codebuddy: failed to parse refresh response: %w", err) } if result.Code != codeSuccess { return nil, fmt.Errorf("codebuddy: refresh failed with code %d: %s", result.Code, result.Msg) } if result.Data == nil { return nil, fmt.Errorf("codebuddy: empty data in refresh response") } newUserID, _ := a.DecodeUserID(result.Data.AccessToken) if newUserID == "" { newUserID = userID } tokenDomain := result.Data.Domain if tokenDomain == "" { tokenDomain = domain } return &CodeBuddyTokenStorage{ AccessToken: result.Data.AccessToken, RefreshToken: result.Data.RefreshToken, ExpiresIn: result.Data.ExpiresIn, RefreshExpiresIn: result.Data.RefreshExpiresIn, TokenType: result.Data.TokenType, Domain: tokenDomain, UserID: newUserID, Type: "codebuddy", }, nil } func (a *CodeBuddyAuth) applyPollHeaders(req *http.Request) { req.Header.Set("Accept", "application/json, text/plain, */*") req.Header.Set("User-Agent", UserAgent) req.Header.Set("X-Requested-With", "XMLHttpRequest") req.Header.Set("X-No-Authorization", "true") req.Header.Set("X-No-User-Id", "true") req.Header.Set("X-No-Enterprise-Id", "true") req.Header.Set("X-No-Department-Info", "true") req.Header.Set("X-Product", "SaaS") }