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