git.giacomolli.uk - onechat/blob - chatgpt.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "cmp"
   6         "context"
   7         "encoding/base64"
   8         "encoding/json"
   9         "errors"
  10         "fmt"
  11         "io"
  12         "net/http"
  13         "net/url"
  14         "os"
  15         "path/filepath"
  16         "strconv"
  17         "strings"
  18         "sync"
  19         "time"
  20 )
  21
  22 const (
  23         // chatgptIssuer is the OpenAI login service.
  24         chatgptIssuer = "https://auth.openai.com"
  25
  26         // chatgptClientID is the OAuth client ID of the Codex CLI. The ChatGPT
  27         // backend takes the tokens of this client only.
  28         chatgptClientID = "app_EMoamEEZ73f0CkXaXp7hrann"
  29
  30         // chatgptAuthClaim is the JWT claim that holds the ChatGPT account ID.
  31         chatgptAuthClaim = "https://api.openai.com/auth"
  32
  33         // chatgptRefreshMargin is how long before its expiry an access token is
  34         // refreshed.
  35         chatgptRefreshMargin = 60 * time.Second
  36
  37         // chatgptLoginTimeout is how long a device code login waits for the
  38         // user.
  39         chatgptLoginTimeout = 15 * time.Minute
  40
  41         // chatgptTokenFileName is the name of the token file in the data
  42         // directory.
  43         chatgptTokenFileName = "auth.json"
  44
  45         // chatgptDefaultPollInterval is the poll interval of a device code
  46         // login when the login service gives none.
  47         chatgptDefaultPollInterval = 5 * time.Second
  48 )
  49
  50 // errChatGPTLoginNeeded tells that no valid token is available, so the user
  51 // must run "onechat login".
  52 var errChatGPTLoginNeeded = errors.New(`login needed: run "onechat login"`)
  53
  54 // ChatGPTAuth authorizes model calls to the ChatGPT backend with the tokens in
  55 // the token file. The tokens belong to OneChat only. The Codex CLI rotates the
  56 // refresh token on each refresh, so two programs with one token file log each
  57 // other out.
  58 //
  59 // Before each call it reads the expiry from the access token. Within
  60 // chatgptRefreshMargin of the expiry it gets new tokens with the refresh token
  61 // and replaces the file atomically. A mutex lets one refresh run at a time, and
  62 // a call that waits for it uses the new tokens.
  63 type ChatGPTAuth struct {
  64         mutex  sync.Mutex
  65         path   string
  66         issuer string
  67         http   *http.Client
  68 }
  69
  70 // ChatGPTDeviceCode is the answer of the user code endpoint. The interval is a
  71 // number or a string of seconds.
  72 type ChatGPTDeviceCode struct {
  73         DeviceAuthID string          `json:"device_auth_id"`
  74         UserCode     string          `json:"user_code"`
  75         UserCodeAlt  string          `json:"usercode"`
  76         Interval     json.RawMessage `json:"interval"`
  77 }
  78
  79 // ChatGPTTokenFile is the content of the token file.
  80 type ChatGPTTokenFile struct {
  81         AccessToken  string `json:"access_token"`
  82         RefreshToken string `json:"refresh_token"`
  83         IDToken      string `json:"id_token"`
  84         AccountID    string `json:"account_id"`
  85 }
  86
  87 // ChatGPTTokenResponse is the answer of the token endpoint.
  88 type ChatGPTTokenResponse struct {
  89         AccessToken  string `json:"access_token"`
  90         RefreshToken string `json:"refresh_token"`
  91         IDToken      string `json:"id_token"`
  92 }
  93
  94 // ChatGPTDeviceGrant is the answer of the device token endpoint after the user
  95 // approves the code.
  96 type ChatGPTDeviceGrant struct {
  97         AuthorizationCode string `json:"authorization_code"`
  98         CodeVerifier      string `json:"code_verifier"`
  99 }
 100
 101 // newChatGPTAuth returns the authorizer for the token file at path, with the
 102 // login service at issuer.
 103 func newChatGPTAuth(path, issuer string) *ChatGPTAuth {
 104         return &ChatGPTAuth{
 105                 path:   path,
 106                 issuer: issuer,
 107                 http:   &http.Client{},
 108         }
 109 }
 110
 111 // Authorize sets the access token and the account ID, after a refresh when the
 112 // token expires within chatgptRefreshMargin.
 113 func (a *ChatGPTAuth) Authorize(ctx context.Context, header http.Header) error {
 114         a.mutex.Lock()
 115         defer a.mutex.Unlock()
 116
 117         tokens, err := a.read()
 118         if err != nil {
 119                 return fmt.Errorf("read login tokens: %w", err)
 120         }
 121
 122         expiry, found := chatgptTokenExpiry(tokens.AccessToken)
 123         if !found || !time.Now().Before(expiry.Add(-chatgptRefreshMargin)) {
 124                 tokens, err = a.refresh(ctx, tokens)
 125                 if err != nil {
 126                         return fmt.Errorf("refresh login tokens: %w", err)
 127                 }
 128         }
 129
 130         header.Set("Authorization", "Bearer "+tokens.AccessToken)
 131         header.Set("ChatGPT-Account-Id", tokens.AccountID)
 132
 133         return nil
 134 }
 135
 136 // read reads the token file.
 137 func (a *ChatGPTAuth) read() (ChatGPTTokenFile, error) {
 138         data, err := os.ReadFile(a.path)
 139         if errors.Is(err, os.ErrNotExist) {
 140                 return ChatGPTTokenFile{}, errChatGPTLoginNeeded
 141         }
 142
 143         if err != nil {
 144                 return ChatGPTTokenFile{}, fmt.Errorf("read token file: %w", err)
 145         }
 146
 147         var tokens ChatGPTTokenFile
 148
 149         err = json.Unmarshal(data, &tokens)
 150         if err == nil && (tokens.AccessToken == "" || tokens.RefreshToken == "" || tokens.AccountID == "") {
 151                 err = errors.New("a token or the account ID is missing")
 152         }
 153
 154         if err != nil {
 155                 return ChatGPTTokenFile{}, fmt.Errorf("token file: %w: %w", err, errChatGPTLoginNeeded)
 156         }
 157
 158         return tokens, nil
 159 }
 160
 161 // refresh gets new tokens with the refresh token and replaces the token file.
 162 // It keeps the old refresh token and ID token when the answer has none.
 163 func (a *ChatGPTAuth) refresh(ctx context.Context, tokens ChatGPTTokenFile) (ChatGPTTokenFile, error) {
 164         var response ChatGPTTokenResponse
 165
 166         err := a.postJSON(ctx, "/oauth/token", map[string]string{
 167                 "client_id":     chatgptClientID,
 168                 "grant_type":    "refresh_token",
 169                 "refresh_token": tokens.RefreshToken,
 170                 "scope":         "openid profile email",
 171         }, &response)
 172         if err == nil && response.AccessToken == "" {
 173                 err = errors.New("refresh response has no access token")
 174         }
 175
 176         if err != nil {
 177                 return ChatGPTTokenFile{}, fmt.Errorf("refresh tokens: %w: %w", err, errChatGPTLoginNeeded)
 178         }
 179
 180         tokens.AccessToken = response.AccessToken
 181         tokens.RefreshToken = cmp.Or(response.RefreshToken, tokens.RefreshToken)
 182         tokens.IDToken = cmp.Or(response.IDToken, tokens.IDToken)
 183         tokens.AccountID = cmp.Or(chatgptAccountID(tokens.AccessToken), chatgptAccountID(tokens.IDToken), tokens.AccountID)
 184
 185         err = a.write(tokens)
 186         if err != nil {
 187                 return ChatGPTTokenFile{}, fmt.Errorf("save refreshed tokens: %w", err)
 188         }
 189
 190         return tokens, nil
 191 }
 192
 193 // postJSON posts body as JSON to the path of the login service and decodes the
 194 // answer into result.
 195 func (a *ChatGPTAuth) postJSON(ctx context.Context, path string, body, result any) error {
 196         data, err := json.Marshal(body)
 197         if err != nil {
 198                 return fmt.Errorf("encode request: %w", err)
 199         }
 200
 201         _, err = a.post(ctx, path, "application/json", data, result)
 202         if err != nil {
 203                 return fmt.Errorf("post JSON request: %w", err)
 204         }
 205
 206         return nil
 207 }
 208
 209 // write replaces the token file atomically. The file has mode 0600.
 210 func (a *ChatGPTAuth) write(tokens ChatGPTTokenFile) error {
 211         data, err := json.Marshal(tokens)
 212         if err != nil {
 213                 return fmt.Errorf("encode tokens: %w", err)
 214         }
 215
 216         err = writeFileAtomic(a.path, append(data, '\n'))
 217         if err != nil {
 218                 return fmt.Errorf("write token file: %w", err)
 219         }
 220
 221         return nil
 222 }
 223
 224 // post posts body to the path of the login service and decodes a 200 answer
 225 // into result. It returns the HTTP status, or 0 when no answer came. The error
 226 // holds the status, but not the body, which can hold a token.
 227 func (a *ChatGPTAuth) post(ctx context.Context, path, contentType string, body []byte, result any) (status int, err error) {
 228         request, err := http.NewRequestWithContext(ctx, http.MethodPost, a.issuer+path, bytes.NewReader(body))
 229         if err != nil {
 230                 return 0, fmt.Errorf("create request: %w", err)
 231         }
 232
 233         request.Header.Set("Content-Type", contentType)
 234
 235         response, err := a.http.Do(request)
 236         if err != nil {
 237                 return 0, fmt.Errorf("send request: %w", err)
 238         }
 239
 240         defer func() {
 241                 err = errors.Join(err, response.Body.Close())
 242         }()
 243
 244         if response.StatusCode != http.StatusOK {
 245                 return response.StatusCode, fmt.Errorf("HTTP %d", response.StatusCode)
 246         }
 247
 248         err = json.NewDecoder(response.Body).Decode(result)
 249         if err != nil {
 250                 return response.StatusCode, fmt.Errorf("decode answer: %w", err)
 251         }
 252
 253         return response.StatusCode, nil
 254 }
 255
 256 // Login runs the device code login. It writes the link and the code for the
 257 // user to out, waits until the user approves the code, and writes the tokens to
 258 // the token file.
 259 func (a *ChatGPTAuth) Login(ctx context.Context, out io.Writer) error {
 260         a.mutex.Lock()
 261         defer a.mutex.Unlock()
 262
 263         var code ChatGPTDeviceCode
 264
 265         err := a.postJSON(ctx, "/api/accounts/deviceauth/usercode", map[string]string{
 266                 "client_id": chatgptClientID,
 267         }, &code)
 268         if err != nil {
 269                 return fmt.Errorf("request device code (device code login must be allowed in the ChatGPT security settings): %w", err)
 270         }
 271
 272         userCode := cmp.Or(code.UserCode, code.UserCodeAlt)
 273         if code.DeviceAuthID == "" || userCode == "" {
 274                 return errors.New("device code response has no device ID or user code")
 275         }
 276
 277         _, 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)
 278         if err != nil {
 279                 return fmt.Errorf("write device code: %w", err)
 280         }
 281
 282         grant, err := a.poll(ctx, code.DeviceAuthID, userCode, chatgptPollInterval(code.Interval))
 283         if err != nil {
 284                 return fmt.Errorf("wait for device grant: %w", err)
 285         }
 286
 287         tokens, err := a.exchange(ctx, grant)
 288         if err != nil {
 289                 return fmt.Errorf("exchange device grant: %w", err)
 290         }
 291
 292         err = a.write(tokens)
 293         if err != nil {
 294                 return fmt.Errorf("save login tokens: %w", err)
 295         }
 296
 297         _, err = fmt.Fprintln(out, "Logged in.")
 298         if err != nil {
 299                 return fmt.Errorf("write login result: %w", err)
 300         }
 301
 302         return nil
 303 }
 304
 305 // poll asks for the grant of the device code until the user approves it or the
 306 // login times out.
 307 func (a *ChatGPTAuth) poll(ctx context.Context, deviceID, userCode string, interval time.Duration) (ChatGPTDeviceGrant, error) {
 308         ctx, cancel := context.WithTimeout(ctx, chatgptLoginTimeout)
 309         defer cancel()
 310
 311         body, err := json.Marshal(map[string]string{
 312                 "device_auth_id": deviceID,
 313                 "user_code":      userCode,
 314         })
 315         if err != nil {
 316                 return ChatGPTDeviceGrant{}, fmt.Errorf("encode poll request: %w", err)
 317         }
 318
 319         for {
 320                 var grant ChatGPTDeviceGrant
 321
 322                 status, err := a.post(ctx, "/api/accounts/deviceauth/token", "application/json", body, &grant)
 323                 if err == nil && (grant.AuthorizationCode == "" || grant.CodeVerifier == "") {
 324                         err = errors.New("device grant has no authorization code or code verifier")
 325                 }
 326
 327                 if err == nil {
 328                         return grant, nil
 329                 }
 330
 331                 if status != http.StatusForbidden && status != http.StatusNotFound {
 332                         return ChatGPTDeviceGrant{}, fmt.Errorf("poll device code: %w", err)
 333                 }
 334
 335                 timer := time.NewTimer(interval)
 336                 select {
 337                 case <-ctx.Done():
 338                         timer.Stop()
 339
 340                         return ChatGPTDeviceGrant{}, fmt.Errorf("wait for approval: %w", ctx.Err())
 341                 case <-timer.C:
 342                 }
 343         }
 344 }
 345
 346 // exchange trades the grant of the device code for tokens.
 347 func (a *ChatGPTAuth) exchange(ctx context.Context, grant ChatGPTDeviceGrant) (ChatGPTTokenFile, error) {
 348         form := url.Values{
 349                 "grant_type":    {"authorization_code"},
 350                 "code":          {grant.AuthorizationCode},
 351                 "redirect_uri":  {a.issuer + "/deviceauth/callback"},
 352                 "client_id":     {chatgptClientID},
 353                 "code_verifier": {grant.CodeVerifier},
 354         }
 355
 356         var response ChatGPTTokenResponse
 357
 358         _, err := a.post(ctx, "/oauth/token", "application/x-www-form-urlencoded", []byte(form.Encode()), &response)
 359         if err != nil {
 360                 return ChatGPTTokenFile{}, fmt.Errorf("exchange code for tokens: %w", err)
 361         }
 362
 363         if response.AccessToken == "" || response.RefreshToken == "" || response.IDToken == "" {
 364                 return ChatGPTTokenFile{}, errors.New("token response has no access, refresh, or ID token")
 365         }
 366
 367         tokens := ChatGPTTokenFile{
 368                 AccessToken:  response.AccessToken,
 369                 RefreshToken: response.RefreshToken,
 370                 IDToken:      response.IDToken,
 371         }
 372
 373         tokens.AccountID = cmp.Or(chatgptAccountID(tokens.AccessToken), chatgptAccountID(tokens.IDToken))
 374         if tokens.AccountID == "" {
 375                 return ChatGPTTokenFile{}, errors.New("tokens have no ChatGPT account ID")
 376         }
 377
 378         return tokens, nil
 379 }
 380
 381 // chatgptTokenExpiry returns the expiry of a JWT from its exp claim.
 382 func chatgptTokenExpiry(token string) (time.Time, bool) {
 383         var claims struct {
 384                 Exp *float64 `json:"exp"`
 385         }
 386
 387         if !chatgptDecodeClaims(token, &claims) || claims.Exp == nil {
 388                 return time.Time{}, false
 389         }
 390
 391         return time.Unix(int64(*claims.Exp), 0), true
 392 }
 393
 394 // chatgptAccountID returns the ChatGPT account ID of a JWT, or "" when it has
 395 // none.
 396 func chatgptAccountID(token string) string {
 397         var claims map[string]json.RawMessage
 398
 399         if !chatgptDecodeClaims(token, &claims) {
 400                 return ""
 401         }
 402
 403         var auth struct {
 404                 AccountID string `json:"chatgpt_account_id"`
 405         }
 406
 407         err := json.Unmarshal(claims[chatgptAuthClaim], &auth)
 408         if err != nil {
 409                 return ""
 410         }
 411
 412         return auth.AccountID
 413 }
 414
 415 // chatgptPollInterval returns the poll interval of a device code, given in
 416 // seconds as a number or a string.
 417 func chatgptPollInterval(raw json.RawMessage) time.Duration {
 418         seconds, err := strconv.ParseFloat(strings.Trim(string(raw), `"`), 64)
 419         if err != nil || seconds < 0 {
 420                 return chatgptDefaultPollInterval
 421         }
 422
 423         return time.Duration(seconds * float64(time.Second))
 424 }
 425
 426 // chatgptDecodeClaims decodes the payload of a JWT into claims, without a check
 427 // of the signature. The token comes from the login service over HTTPS, and the
 428 // client reads only the expiry and the account ID from it.
 429 func chatgptDecodeClaims(token string, claims any) bool {
 430         parts := strings.Split(token, ".")
 431         if len(parts) != 3 {
 432                 return false
 433         }
 434
 435         payload, err := base64.RawURLEncoding.DecodeString(strings.TrimRight(parts[1], "="))
 436         if err != nil {
 437                 return false
 438         }
 439
 440         return json.Unmarshal(payload, claims) == nil
 441 }
 442
 443 // chatgptLogin runs "onechat login": the device code login into the token file
 444 // of the data directory. It holds the lock of the token file, so it does not
 445 // run while the daemon runs.
 446 func chatgptLogin(ctx context.Context, out io.Writer) (err error) {
 447         dir, err := dataDir()
 448         if err != nil {
 449                 return fmt.Errorf("find data directory: %w", err)
 450         }
 451
 452         tokens, err := lockTokens(dir)
 453         if err != nil {
 454                 return fmt.Errorf("lock login tokens: %w", err)
 455         }
 456
 457         defer func() {
 458                 err = errors.Join(err, tokens.Close())
 459         }()
 460
 461         err = newChatGPTAuth(filepath.Join(dir, chatgptTokenFileName), chatgptIssuer).Login(ctx, out)
 462         if err != nil {
 463                 return fmt.Errorf("complete device login: %w", err)
 464         }
 465
 466         return nil
 467 }