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