CLIProxyAPIPlus/internal/auth/codebuddy/codebuddy_auth.go

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")
}