[onechat] / chatgpt.go

package main

import (
        "bytes"
        "cmp"
        "context"
        "encoding/base64"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net/http"
        "net/url"
        "os"
        "path/filepath"
        "strconv"
        "strings"
        "sync"
        "time"
)

const (
        // chatgptIssuer is the OpenAI login service.
        chatgptIssuer = "https://auth.openai.com"

        // chatgptClientID is the OAuth client ID of the Codex CLI. The ChatGPT
        // backend takes the tokens of this client only.
        chatgptClientID = "app_EMoamEEZ73f0CkXaXp7hrann"

        // chatgptAuthClaim is the JWT claim that holds the ChatGPT account ID.
        chatgptAuthClaim = "https://api.openai.com/auth"

        // chatgptRefreshMargin is how long before its expiry an access token is
        // refreshed.
        chatgptRefreshMargin = 60 * time.Second

        // chatgptLoginTimeout is how long a device code login waits for the
        // user.
        chatgptLoginTimeout = 15 * time.Minute

        // chatgptTokenFileName is the name of the token file in the data
        // directory.
        chatgptTokenFileName = "auth.json"

        // chatgptDefaultPollInterval is the poll interval of a device code
        // login when the login service gives none.
        chatgptDefaultPollInterval = 5 * time.Second
)

// errChatGPTLoginNeeded tells that no valid token is available, so the user
// must run "onechat login".
var errChatGPTLoginNeeded = errors.New(`login needed: run "onechat login"`)

// chatgptLogin runs "onechat login": the device code login into the token file
// of the data directory. It holds the lock of the token file, so it does not
// run while the daemon runs.
func chatgptLogin(ctx context.Context, out io.Writer) (err error) {
        dir, err := dataDir()
        if err != nil {
                return err
        }

        tokens, err := lockTokens(dir)
        if err != nil {
                return err
        }

        defer func() {
                err = errors.Join(err, tokens.Close())
        }()

        return newChatGPTAuth(filepath.Join(dir, chatgptTokenFileName), chatgptIssuer).Login(ctx, out)
}

// newChatGPTAuth returns the authorizer for the token file at path, with the
// login service at issuer.
func newChatGPTAuth(path, issuer string) *ChatGPTAuth {
        return &ChatGPTAuth{
                path:   path,
                issuer: issuer,
                http:   &http.Client{},
        }
}

// ChatGPTAuth authorizes model calls to the ChatGPT backend with the tokens in
// the token file. The tokens belong to OneChat only. The Codex CLI rotates the
// refresh token on each refresh, so two programs with one token file log each
// other out.
//
// Before each call it reads the expiry from the access token. Within
// chatgptRefreshMargin of the expiry it gets new tokens with the refresh token
// and replaces the file atomically. A mutex lets one refresh run at a time, and
// a call that waits for it uses the new tokens.
type ChatGPTAuth struct {
        mutex  sync.Mutex
        path   string
        issuer string
        http   *http.Client
}

// Authorize sets the access token and the account ID, after a refresh when the
// token expires within chatgptRefreshMargin.
func (a *ChatGPTAuth) Authorize(ctx context.Context, header http.Header) error {
        a.mutex.Lock()
        defer a.mutex.Unlock()

        tokens, err := a.read()
        if err != nil {
                return err
        }

        expiry, found := chatgptTokenExpiry(tokens.AccessToken)
        if !found || !time.Now().Before(expiry.Add(-chatgptRefreshMargin)) {
                tokens, err = a.refresh(ctx, tokens)
                if err != nil {
                        return err
                }
        }

        header.Set("Authorization", "Bearer "+tokens.AccessToken)
        header.Set("ChatGPT-Account-Id", tokens.AccountID)

        return nil
}

// Login runs the device code login. It writes the link and the code for the
// user to out, waits until the user approves the code, and writes the tokens
// to the token file.
func (a *ChatGPTAuth) Login(ctx context.Context, out io.Writer) error {
        a.mutex.Lock()
        defer a.mutex.Unlock()

        var code ChatGPTDeviceCode
        err := a.postJSON(ctx, "/api/accounts/deviceauth/usercode", map[string]string{
                "client_id": chatgptClientID,
        }, &code)
        if err != nil {
                return fmt.Errorf("request device code (device code login must be allowed in the ChatGPT security settings): %w", err)
        }

        userCode := cmp.Or(code.UserCode, code.UserCodeAlt)
        if code.DeviceAuthID == "" || userCode == "" {
                return errors.New("device code response has no device ID or user code")
        }

        _, err = fmt.Fprintf(out, "Open %s/codex/device on any device and enter this code:\n\n    %s\n\nWaiting for the approval...\n", a.issuer, userCode)
        if err != nil {
                return fmt.Errorf("write device code: %w", err)
        }

        grant, err := a.poll(ctx, code.DeviceAuthID, userCode, chatgptPollInterval(code.Interval))
        if err != nil {
                return err
        }

        tokens, err := a.exchange(ctx, grant)
        if err != nil {
                return err
        }

        err = a.write(tokens)
        if err != nil {
                return err
        }

        _, err = fmt.Fprintln(out, "Logged in.")
        if err != nil {
                return fmt.Errorf("write login result: %w", err)
        }

        return nil
}

// read reads the token file.
func (a *ChatGPTAuth) read() (ChatGPTTokenFile, error) {
        data, err := os.ReadFile(a.path)
        if errors.Is(err, os.ErrNotExist) {
                return ChatGPTTokenFile{}, errChatGPTLoginNeeded
        }

        if err != nil {
                return ChatGPTTokenFile{}, fmt.Errorf("read token file: %w", err)
        }

        var tokens ChatGPTTokenFile
        err = json.Unmarshal(data, &tokens)
        if err == nil && (tokens.AccessToken == "" || tokens.RefreshToken == "" || tokens.AccountID == "") {
                err = errors.New("a token or the account ID is missing")
        }

        if err != nil {
                return ChatGPTTokenFile{}, fmt.Errorf("token file: %w: %w", err, errChatGPTLoginNeeded)
        }

        return tokens, nil
}

// refresh gets new tokens with the refresh token and replaces the token file.
// It keeps the old refresh token and ID token when the answer has none.
func (a *ChatGPTAuth) refresh(ctx context.Context, tokens ChatGPTTokenFile) (ChatGPTTokenFile, error) {
        var response ChatGPTTokenResponse
        err := a.postJSON(ctx, "/oauth/token", map[string]string{
                "client_id":     chatgptClientID,
                "grant_type":    "refresh_token",
                "refresh_token": tokens.RefreshToken,
                "scope":         "openid profile email",
        }, &response)
        if err == nil && response.AccessToken == "" {
                err = errors.New("refresh response has no access token")
        }

        if err != nil {
                return ChatGPTTokenFile{}, fmt.Errorf("refresh tokens: %w: %w", err, errChatGPTLoginNeeded)
        }

        tokens.AccessToken = response.AccessToken
        tokens.RefreshToken = cmp.Or(response.RefreshToken, tokens.RefreshToken)
        tokens.IDToken = cmp.Or(response.IDToken, tokens.IDToken)
        tokens.AccountID = cmp.Or(chatgptAccountID(tokens.AccessToken), chatgptAccountID(tokens.IDToken), tokens.AccountID)

        err = a.write(tokens)
        if err != nil {
                return ChatGPTTokenFile{}, err
        }

        return tokens, nil
}

// postJSON posts body as JSON to the path of the login service and decodes the
// answer into result.
func (a *ChatGPTAuth) postJSON(ctx context.Context, path string, body, result any) error {
        data, err := json.Marshal(body)
        if err != nil {
                return fmt.Errorf("encode request: %w", err)
        }

        _, err = a.post(ctx, path, "application/json", data, result)

        return err
}

// poll asks for the grant of the device code until the user approves it or the
// login times out.
func (a *ChatGPTAuth) poll(ctx context.Context, deviceID, userCode string, interval time.Duration) (ChatGPTDeviceGrant, error) {
        ctx, cancel := context.WithTimeout(ctx, chatgptLoginTimeout)
        defer cancel()

        body, err := json.Marshal(map[string]string{
                "device_auth_id": deviceID,
                "user_code":      userCode,
        })
        if err != nil {
                return ChatGPTDeviceGrant{}, fmt.Errorf("encode poll request: %w", err)
        }

        for {
                var grant ChatGPTDeviceGrant
                status, err := a.post(ctx, "/api/accounts/deviceauth/token", "application/json", body, &grant)
                if err == nil && (grant.AuthorizationCode == "" || grant.CodeVerifier == "") {
                        err = errors.New("device grant has no authorization code or code verifier")
                }

                if err == nil {
                        return grant, nil
                }

                if status != http.StatusForbidden && status != http.StatusNotFound {
                        return ChatGPTDeviceGrant{}, fmt.Errorf("poll device code: %w", err)
                }

                timer := time.NewTimer(interval)
                select {
                case <-ctx.Done():
                        timer.Stop()
                        return ChatGPTDeviceGrant{}, fmt.Errorf("wait for approval: %w", ctx.Err())
                case <-timer.C:
                }
        }
}

// exchange trades the grant of the device code for tokens.
func (a *ChatGPTAuth) exchange(ctx context.Context, grant ChatGPTDeviceGrant) (ChatGPTTokenFile, error) {
        form := url.Values{
                "grant_type":    {"authorization_code"},
                "code":          {grant.AuthorizationCode},
                "redirect_uri":  {a.issuer + "/deviceauth/callback"},
                "client_id":     {chatgptClientID},
                "code_verifier": {grant.CodeVerifier},
        }

        var response ChatGPTTokenResponse
        _, err := a.post(ctx, "/oauth/token", "application/x-www-form-urlencoded", []byte(form.Encode()), &response)
        if err != nil {
                return ChatGPTTokenFile{}, fmt.Errorf("exchange code for tokens: %w", err)
        }

        if response.AccessToken == "" || response.RefreshToken == "" || response.IDToken == "" {
                return ChatGPTTokenFile{}, errors.New("token response has no access, refresh, or ID token")
        }

        tokens := ChatGPTTokenFile{
                AccessToken:  response.AccessToken,
                RefreshToken: response.RefreshToken,
                IDToken:      response.IDToken,
        }

        tokens.AccountID = cmp.Or(chatgptAccountID(tokens.AccessToken), chatgptAccountID(tokens.IDToken))
        if tokens.AccountID == "" {
                return ChatGPTTokenFile{}, errors.New("tokens have no ChatGPT account ID")
        }

        return tokens, nil
}

// write replaces the token file atomically. The file has mode 0600.
func (a *ChatGPTAuth) write(tokens ChatGPTTokenFile) error {
        data, err := json.Marshal(tokens)
        if err != nil {
                return fmt.Errorf("encode tokens: %w", err)
        }

        err = writeFileAtomic(a.path, append(data, '\n'))
        if err != nil {
                return fmt.Errorf("write token file: %w", err)
        }

        return nil
}

// post posts body to the path of the login service and decodes a 200 answer
// into result. It returns the HTTP status, or 0 when no answer came. The error
// holds the status, but not the body, which can hold a token.
func (a *ChatGPTAuth) post(ctx context.Context, path, contentType string, body []byte, result any) (status int, err error) {
        request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.issuer+path, bytes.NewReader(body))
        if err != nil {
                return 0, fmt.Errorf("create request: %w", err)
        }

        request.Header.Set("Content-Type", contentType)
        response, err := a.http.Do(request)
        if err != nil {
                return 0, fmt.Errorf("send request: %w", err)
        }

        defer func() {
                err = errors.Join(err, response.Body.Close())
        }()

        if response.StatusCode != http.StatusOK {
                return response.StatusCode, fmt.Errorf("HTTP %d", response.StatusCode)
        }

        err = json.NewDecoder(response.Body).Decode(result)
        if err != nil {
                return response.StatusCode, fmt.Errorf("decode answer: %w", err)
        }

        return response.StatusCode, nil
}

// ChatGPTDeviceCode is the answer of the user code endpoint. The interval is a
// number or a string of seconds.
type ChatGPTDeviceCode struct {
        DeviceAuthID string          `json:"device_auth_id"`
        UserCode     string          `json:"user_code"`
        UserCodeAlt  string          `json:"usercode"`
        Interval     json.RawMessage `json:"interval"`
}

// ChatGPTTokenFile is the content of the token file.
type ChatGPTTokenFile struct {
        AccessToken  string `json:"access_token"`
        RefreshToken string `json:"refresh_token"`
        IDToken      string `json:"id_token"`
        AccountID    string `json:"account_id"`
}

// ChatGPTTokenResponse is the answer of the token endpoint.
type ChatGPTTokenResponse struct {
        AccessToken  string `json:"access_token"`
        RefreshToken string `json:"refresh_token"`
        IDToken      string `json:"id_token"`
}

// ChatGPTDeviceGrant is the answer of the device token endpoint after the user
// approves the code.
type ChatGPTDeviceGrant struct {
        AuthorizationCode string `json:"authorization_code"`
        CodeVerifier      string `json:"code_verifier"`
}

// chatgptTokenExpiry returns the expiry of a JWT from its exp claim.
func chatgptTokenExpiry(token string) (time.Time, bool) {
        var claims struct {
                Exp *float64 `json:"exp"`
        }

        if !chatgptDecodeClaims(token, &claims) || claims.Exp == nil {
                return time.Time{}, false
        }

        return time.Unix(int64(*claims.Exp), 0), true
}

// chatgptPollInterval returns the poll interval of a device code, given in
// seconds as a number or a string.
func chatgptPollInterval(raw json.RawMessage) time.Duration {
        seconds, err := strconv.ParseFloat(strings.Trim(string(raw), `"`), 64)
        if err != nil || seconds < 0 {
                return chatgptDefaultPollInterval
        }

        return time.Duration(seconds * float64(time.Second))
}

// chatgptAccountID returns the ChatGPT account ID of a JWT, or "" when it has
// none.
func chatgptAccountID(token string) string {
        var claims map[string]json.RawMessage
        if !chatgptDecodeClaims(token, &claims) {
                return ""
        }

        var auth struct {
                AccountID string `json:"chatgpt_account_id"`
        }

        err := json.Unmarshal(claims[chatgptAuthClaim], &auth)
        if err != nil {
                return ""
        }

        return auth.AccountID
}

// chatgptDecodeClaims decodes the payload of a JWT into claims, without a check
// of the signature. The token comes from the login service over HTTPS, and the
// client reads only the expiry and the account ID from it.
func chatgptDecodeClaims(token string, claims any) bool {
        parts := strings.Split(token, ".")
        if len(parts) != 3 {
                return false
        }

        payload, err := base64.RawURLEncoding.DecodeString(strings.TrimRight(parts[1], "="))
        if err != nil {
                return false
        }

        return json.Unmarshal(payload, claims) == nil
}