sing-box/service/ocm/credential_oauth.go

package ocm

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"os"
	"os/user"
	"path/filepath"
	"strconv"
	"time"

	E "github.com/sagernet/sing/common/exceptions"
)

const (
	oauth2ClientID           = "app_EMoamEEZ73f0CkXaXp7hrann"
	oauth2TokenURL           = "https://auth.openai.com/oauth/token"
	openaiAPIBaseURL         = "https://api.openai.com"
	chatGPTBackendURL        = "https://chatgpt.com/backend-api/codex"
	tokenRefreshIntervalDays = 8
)

func getRealUser() (*user.User, error) {
	if sudoUser := os.Getenv("SUDO_USER"); sudoUser != "" {
		sudoUserInfo, err := user.Lookup(sudoUser)
		if err == nil {
			return sudoUserInfo, nil
		}
	}
	return user.Current()
}

func getDefaultCredentialsPath() (string, error) {
	if codexHome := os.Getenv("CODEX_HOME"); codexHome != "" {
		return filepath.Join(codexHome, "auth.json"), nil
	}
	userInfo, err := getRealUser()
	if err != nil {
		return "", err
	}
	return filepath.Join(userInfo.HomeDir, ".codex", "auth.json"), nil
}

func readCredentialsFromFile(path string) (*oauthCredentials, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var credentials oauthCredentials
	err = json.Unmarshal(data, &credentials)
	if err != nil {
		return nil, err
	}
	return &credentials, nil
}

func checkCredentialFileWritable(path string) error {
	file, err := os.OpenFile(path, os.O_WRONLY, 0)
	if err != nil {
		return err
	}
	return file.Close()
}

func writeCredentialsToFile(credentials *oauthCredentials, path string) error {
	data, err := json.MarshalIndent(credentials, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0o600)
}

type oauthCredentials struct {
	APIKey      string     `json:"OPENAI_API_KEY,omitempty"`
	Tokens      *tokenData `json:"tokens,omitempty"`
	LastRefresh *time.Time `json:"last_refresh,omitempty"`
}

type tokenData struct {
	IDToken      string `json:"id_token,omitempty"`
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
	AccountID    string `json:"account_id,omitempty"`
}

func (c *oauthCredentials) isAPIKeyMode() bool {
	return c.APIKey != ""
}

func (c *oauthCredentials) getAccessToken() string {
	if c.APIKey != "" {
		return c.APIKey
	}
	if c.Tokens != nil {
		return c.Tokens.AccessToken
	}
	return ""
}

func (c *oauthCredentials) getAccountID() string {
	if c.Tokens != nil {
		return c.Tokens.AccountID
	}
	return ""
}

func (c *oauthCredentials) needsRefresh() bool {
	if c.APIKey != "" {
		return false
	}
	if c.Tokens == nil || c.Tokens.RefreshToken == "" {
		return false
	}
	if c.LastRefresh == nil {
		return true
	}
	return time.Since(*c.LastRefresh) >= time.Duration(tokenRefreshIntervalDays)*24*time.Hour
}

func refreshToken(ctx context.Context, httpClient *http.Client, credentials *oauthCredentials) (*oauthCredentials, time.Duration, error) {
	if credentials.Tokens == nil || credentials.Tokens.RefreshToken == "" {
		return nil, 0, E.New("refresh token is empty")
	}

	requestBody, err := json.Marshal(map[string]string{
		"grant_type":    "refresh_token",
		"refresh_token": credentials.Tokens.RefreshToken,
		"client_id":     oauth2ClientID,
		"scope":         "openid profile email",
	})
	if err != nil {
		return nil, 0, E.Cause(err, "marshal request")
	}

	response, err := doHTTPWithRetry(ctx, httpClient, func() (*http.Request, error) {
		request, err := http.NewRequest("POST", oauth2TokenURL, bytes.NewReader(requestBody))
		if err != nil {
			return nil, err
		}
		request.Header.Set("Content-Type", "application/json")
		request.Header.Set("Accept", "application/json")
		return request, nil
	})
	if err != nil {
		return nil, 0, err
	}
	defer response.Body.Close()

	if response.StatusCode == http.StatusTooManyRequests {
		body, _ := io.ReadAll(response.Body)
		retryDelay := time.Duration(-1)
		if retryAfter := response.Header.Get("Retry-After"); retryAfter != "" {
			seconds, parseErr := strconv.ParseInt(retryAfter, 10, 64)
			if parseErr == nil && seconds > 0 {
				retryDelay = time.Duration(seconds) * time.Second
			}
		}
		return nil, retryDelay, E.New("refresh rate limited: ", response.Status, " ", string(body))
	}
	if response.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(response.Body)
		return nil, 0, E.New("refresh failed: ", response.Status, " ", string(body))
	}

	var tokenResponse struct {
		IDToken      string `json:"id_token"`
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
	}
	err = json.NewDecoder(response.Body).Decode(&tokenResponse)
	if err != nil {
		return nil, 0, E.Cause(err, "decode response")
	}

	newCredentials := *credentials
	if newCredentials.Tokens == nil {
		newCredentials.Tokens = &tokenData{}
	}
	if tokenResponse.IDToken != "" {
		newCredentials.Tokens.IDToken = tokenResponse.IDToken
	}
	if tokenResponse.AccessToken != "" {
		newCredentials.Tokens.AccessToken = tokenResponse.AccessToken
	}
	if tokenResponse.RefreshToken != "" {
		newCredentials.Tokens.RefreshToken = tokenResponse.RefreshToken
	}
	now := time.Now()
	newCredentials.LastRefresh = &now

	return &newCredentials, 0, nil
}

func cloneCredentials(credentials *oauthCredentials) *oauthCredentials {
	if credentials == nil {
		return nil
	}
	cloned := *credentials
	if credentials.Tokens != nil {
		clonedTokens := *credentials.Tokens
		cloned.Tokens = &clonedTokens
	}
	if credentials.LastRefresh != nil {
		lastRefresh := *credentials.LastRefresh
		cloned.LastRefresh = &lastRefresh
	}
	return &cloned
}

func credentialsEqual(left *oauthCredentials, right *oauthCredentials) bool {
	if left == nil || right == nil {
		return left == right
	}
	if left.APIKey != right.APIKey {
		return false
	}
	if (left.Tokens == nil) != (right.Tokens == nil) {
		return false
	}
	if left.Tokens != nil && *left.Tokens != *right.Tokens {
		return false
	}
	if (left.LastRefresh == nil) != (right.LastRefresh == nil) {
		return false
	}
	if left.LastRefresh != nil && !left.LastRefresh.Equal(*right.LastRefresh) {
		return false
	}
	return true
}