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