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

   1 package main
   2
   3 import (
   4         "encoding/base64"
   5         "encoding/json"
   6         "fmt"
   7         "io"
   8         "net/http"
   9         "net/http/httptest"
  10         "net/url"
  11         "os"
  12         "path/filepath"
  13         "strings"
  14         "sync"
  15         "testing"
  16         "time"
  17
  18         "github.com/stretchr/testify/require"
  19 )
  20
  21 // testAccount is the fake ChatGPT account ID in the fake tokens.
  22 const testAccount = "account-test"
  23
  24 func TestLoginWritesTheTokenFile(t *testing.T) {
  25         server := newFakeAuth(t)
  26         server.pending = 2
  27         auth := newTestAuth(t, server)
  28
  29         var out strings.Builder
  30         err := auth.Login(t.Context(), &out)
  31         require.NoError(t, err)
  32
  33         require.Contains(t, out.String(), server.URL+"/codex/device")
  34         require.Contains(t, out.String(), "CODE-1234")
  35         require.Equal(t, 3, server.polls)
  36
  37         form := server.exchange
  38         require.Equal(t, "authorization_code", form.Get("grant_type"))
  39         require.Equal(t, "auth-code", form.Get("code"))
  40         require.Equal(t, "verifier", form.Get("code_verifier"))
  41         require.Equal(t, codexClientID, form.Get("client_id"))
  42         require.Equal(t, server.URL+"/deviceauth/callback", form.Get("redirect_uri"))
  43
  44         tokens := readTokenFile(t, auth.path)
  45         require.Equal(t, "refresh-1", tokens.RefreshToken)
  46         require.Equal(t, testAccount, tokens.AccountID)
  47         require.Equal(t, server.access[0], tokens.AccessToken)
  48         require.NotContains(t, out.String(), tokens.AccessToken)
  49         require.NotContains(t, out.String(), tokens.RefreshToken)
  50
  51         info, err := os.Stat(auth.path)
  52         require.NoError(t, err)
  53         require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
  54 }
  55
  56 func TestLoginFailsWhenTheCodeIsDenied(t *testing.T) {
  57         server := newFakeAuth(t)
  58         server.deny = true
  59         auth := newTestAuth(t, server)
  60
  61         err := auth.Login(t.Context(), io.Discard)
  62         require.ErrorContains(t, err, "HTTP 400")
  63         require.NoFileExists(t, auth.path)
  64 }
  65
  66 func TestAuthorizeUsesAFreshTokenWithoutRefresh(t *testing.T) {
  67         server := newFakeAuth(t)
  68         auth := newTestAuth(t, server)
  69         writeTestTokens(t, auth, server.token(time.Hour), "refresh-0")
  70
  71         header := http.Header{}
  72         err := auth.Authorize(t.Context(), header)
  73         require.NoError(t, err)
  74         require.Equal(t, "Bearer "+server.access[0], header.Get("Authorization"))
  75         require.Equal(t, testAccount, header.Get("ChatGPT-Account-Id"))
  76         require.Equal(t, 0, server.refreshes)
  77 }
  78
  79 func TestAuthorizeRefreshesNearTheExpiry(t *testing.T) {
  80         testCases := []struct {
  81                 name        string
  82                 newRefresh  string
  83                 wantRefresh string
  84         }{
  85                 {
  86                         name:        "new refresh token kept",
  87                         newRefresh:  "refresh-new",
  88                         wantRefresh: "refresh-new",
  89                 },
  90                 {
  91                         name:        "old refresh token kept without a new one",
  92                         newRefresh:  "",
  93                         wantRefresh: "refresh-0",
  94                 },
  95         }
  96
  97         for _, tc := range testCases {
  98                 t.Run(tc.name, func(t *testing.T) {
  99                         server := newFakeAuth(t)
 100                         server.newRefresh = tc.newRefresh
 101                         auth := newTestAuth(t, server)
 102                         writeTestTokens(t, auth, server.token(59*time.Second), "refresh-0")
 103
 104                         header := http.Header{}
 105                         err := auth.Authorize(t.Context(), header)
 106                         require.NoError(t, err)
 107                         require.Equal(t, 1, server.refreshes)
 108                         require.Equal(t, "refresh-0", server.refreshed)
 109
 110                         tokens := readTokenFile(t, auth.path)
 111                         require.Equal(t, server.access[1], tokens.AccessToken)
 112                         require.Equal(t, tc.wantRefresh, tokens.RefreshToken)
 113                         require.Equal(t, testAccount, tokens.AccountID)
 114                         require.Equal(t, "Bearer "+server.access[1], header.Get("Authorization"))
 115
 116                         err = auth.Authorize(t.Context(), header)
 117                         require.NoError(t, err)
 118                         require.Equal(t, 1, server.refreshes)
 119                 })
 120         }
 121 }
 122
 123 func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {
 124         server := newFakeAuth(t)
 125         server.newRefresh = "refresh-new"
 126         auth := newTestAuth(t, server)
 127         writeTestTokens(t, auth, server.token(time.Second), "refresh-0")
 128
 129         var group sync.WaitGroup
 130         errs := make(chan error, 8)
 131         for range 8 {
 132                 group.Go(func() {
 133                         errs <- auth.Authorize(t.Context(), http.Header{})
 134                 })
 135         }
 136
 137         group.Wait()
 138         close(errs)
 139         for err := range errs {
 140                 require.NoError(t, err)
 141         }
 142
 143         require.Equal(t, 1, server.refreshes)
 144 }
 145
 146 func TestAuthorizeNeedsALogin(t *testing.T) {
 147         server := newFakeAuth(t)
 148         auth := newTestAuth(t, server)
 149
 150         err := auth.Authorize(t.Context(), http.Header{})
 151         require.ErrorIs(t, err, errLoginNeeded)
 152
 153         server.failRefresh = true
 154         writeTestTokens(t, auth, server.token(-time.Minute), "refresh-0")
 155         before, err := os.ReadFile(auth.path)
 156         require.NoError(t, err)
 157
 158         err = auth.Authorize(t.Context(), http.Header{})
 159         require.ErrorIs(t, err, errLoginNeeded)
 160         require.NotContains(t, err.Error(), "refresh-0")
 161
 162         after, err := os.ReadFile(auth.path)
 163         require.NoError(t, err)
 164         require.Equal(t, before, after)
 165 }
 166
 167 func TestForcedRefreshLetsTheNextCallWork(t *testing.T) {
 168         server := newFakeAuth(t)
 169         auth := newTestAuth(t, server)
 170         writeTestTokens(t, auth, server.token(-time.Hour), "refresh-0")
 171
 172         responses := newFakeResponses(t)
 173         client := newClient(responses.URL, auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 174         _, err := client.Create(t.Context(), testRequest())
 175         require.NoError(t, err)
 176
 177         require.Equal(t, 1, server.refreshes)
 178         call := responses.Call(0)
 179         require.Equal(t, "Bearer "+server.access[1], call.header.Get("Authorization"))
 180         require.Equal(t, testAccount, call.header.Get("ChatGPT-Account-Id"))
 181         require.Equal(t, server.access[1], readTokenFile(t, auth.path).AccessToken)
 182 }
 183
 184 func TestModelEndpointFollowsTheConfig(t *testing.T) {
 185         dir := t.TempDir()
 186         config := defaultConfig()
 187
 188         t.Setenv("OPENAI_API_KEY", testKey)
 189         address, auth, err := modelEndpoint(config, dir)
 190         require.NoError(t, err)
 191         require.Equal(t, chatgptURL, address)
 192         require.IsType(t, &chatgptAuth{}, auth)
 193         require.Equal(t, filepath.Join(dir, tokenFileName), auth.(*chatgptAuth).path)
 194
 195         config.Endpoint = endpointAPI
 196         address, auth, err = modelEndpoint(config, dir)
 197         require.NoError(t, err)
 198         require.Equal(t, apiURL, address)
 199         require.Equal(t, apiKey(testKey), auth)
 200
 201         t.Setenv("OPENAI_API_KEY", "")
 202         _, _, err = modelEndpoint(config, dir)
 203         require.ErrorContains(t, err, "OPENAI_API_KEY")
 204 }
 205
 206 // fakeAuth is an httptest login service. The device token endpoint answers
 207 // 403 pending times before it grants the code. Each token answer has a new
 208 // access token, kept in access.
 209 type fakeAuth struct {
 210         *httptest.Server
 211         mutex       sync.Mutex
 212         pending     int
 213         deny        bool
 214         polls       int
 215         exchange    url.Values
 216         refreshes   int
 217         refreshed   string
 218         newRefresh  string
 219         failRefresh bool
 220         access      []string
 221 }
 222
 223 func newFakeAuth(t *testing.T) *fakeAuth {
 224         t.Helper()
 225
 226         f := &fakeAuth{}
 227         mux := http.NewServeMux()
 228         mux.HandleFunc("POST /api/accounts/deviceauth/usercode", f.userCode)
 229         mux.HandleFunc("POST /api/accounts/deviceauth/token", f.deviceToken)
 230         mux.HandleFunc("POST /oauth/token", f.oauthToken)
 231         f.Server = httptest.NewServer(mux)
 232         t.Cleanup(f.Close)
 233
 234         return f
 235 }
 236
 237 func (f *fakeAuth) userCode(w http.ResponseWriter, r *http.Request) {
 238         var body map[string]string
 239         err := json.NewDecoder(r.Body).Decode(&body)
 240         if err != nil || body["client_id"] != codexClientID {
 241                 http.Error(w, "bad client", http.StatusBadRequest)
 242                 return
 243         }
 244
 245         writeJSON(w, map[string]any{
 246                 "device_auth_id": "device-1",
 247                 "user_code":      "CODE-1234",
 248                 "interval":       "0",
 249         })
 250 }
 251
 252 func (f *fakeAuth) deviceToken(w http.ResponseWriter, r *http.Request) {
 253         f.mutex.Lock()
 254         defer f.mutex.Unlock()
 255
 256         var body map[string]string
 257         err := json.NewDecoder(r.Body).Decode(&body)
 258         if err != nil || body["device_auth_id"] != "device-1" || body["user_code"] != "CODE-1234" {
 259                 http.Error(w, "bad device", http.StatusBadRequest)
 260                 return
 261         }
 262
 263         f.polls++
 264         if f.deny {
 265                 http.Error(w, "denied", http.StatusBadRequest)
 266                 return
 267         }
 268
 269         if f.polls <= f.pending {
 270                 http.Error(w, "pending", http.StatusForbidden)
 271                 return
 272         }
 273
 274         writeJSON(w, map[string]any{
 275                 "authorization_code": "auth-code",
 276                 "code_challenge":     "challenge",
 277                 "code_verifier":      "verifier",
 278         })
 279 }
 280
 281 func (f *fakeAuth) oauthToken(w http.ResponseWriter, r *http.Request) {
 282         f.mutex.Lock()
 283         defer f.mutex.Unlock()
 284
 285         if r.Header.Get("Content-Type") == "application/x-www-form-urlencoded" {
 286                 err := r.ParseForm()
 287                 if err != nil {
 288                         http.Error(w, err.Error(), http.StatusBadRequest)
 289                         return
 290                 }
 291
 292                 f.exchange = r.PostForm
 293                 writeJSON(w, map[string]any{
 294                         "access_token":  f.newToken(time.Hour),
 295                         "refresh_token": "refresh-1",
 296                         "id_token":      f.newToken(time.Hour),
 297                 })
 298
 299                 return
 300         }
 301
 302         var body map[string]string
 303         err := json.NewDecoder(r.Body).Decode(&body)
 304         if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != codexClientID {
 305                 http.Error(w, "bad refresh", http.StatusBadRequest)
 306                 return
 307         }
 308
 309         f.refreshes++
 310         f.refreshed = body["refresh_token"]
 311         if f.failRefresh {
 312                 http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized)
 313                 return
 314         }
 315
 316         answer := map[string]any{
 317                 "access_token": f.newToken(time.Hour),
 318                 "id_token":     f.newToken(time.Hour),
 319         }
 320
 321         if f.newRefresh != "" {
 322                 answer["refresh_token"] = f.newRefresh
 323         }
 324
 325         writeJSON(w, answer)
 326 }
 327
 328 // token returns a new access token that expires after ttl, for a token file
 329 // that a test writes.
 330 func (f *fakeAuth) token(ttl time.Duration) string {
 331         f.mutex.Lock()
 332         defer f.mutex.Unlock()
 333
 334         return f.newToken(ttl)
 335 }
 336
 337 // newToken returns a new fake JWT with the expiry and the account claim, and
 338 // keeps it in access. The caller holds the mutex.
 339 func (f *fakeAuth) newToken(ttl time.Duration) string {
 340         claims, err := json.Marshal(map[string]any{
 341                 "exp": time.Now().Add(ttl).Unix(),
 342                 authClaim: map[string]any{
 343                         "chatgpt_account_id": testAccount,
 344                 },
 345                 "serial": len(f.access),
 346         })
 347         if err != nil {
 348                 panic(err)
 349         }
 350
 351         encode := base64.RawURLEncoding.EncodeToString
 352         token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
 353         f.access = append(f.access, token)
 354
 355         return token
 356 }
 357
 358 func newTestAuth(t *testing.T, server *fakeAuth) *chatgptAuth {
 359         t.Helper()
 360
 361         return newChatGPTAuth(filepath.Join(t.TempDir(), tokenFileName), server.URL)
 362 }
 363
 364 func writeTestTokens(t *testing.T, auth *chatgptAuth, access, refresh string) {
 365         t.Helper()
 366
 367         err := auth.write(tokenFile{
 368                 AccessToken:  access,
 369                 RefreshToken: refresh,
 370                 IDToken:      access,
 371                 AccountID:    testAccount,
 372         })
 373         require.NoError(t, err)
 374 }
 375
 376 func readTokenFile(t *testing.T, path string) tokenFile {
 377         t.Helper()
 378
 379         data, err := os.ReadFile(path)
 380         require.NoError(t, err)
 381
 382         var tokens tokenFile
 383         require.NoError(t, json.Unmarshal(data, &tokens))
 384
 385         return tokens
 386 }
 387
 388 func writeJSON(w http.ResponseWriter, value any) {
 389         w.Header().Set("Content-Type", "application/json")
 390         err := json.NewEncoder(w).Encode(value)
 391         if err != nil {
 392                 panic(fmt.Sprint("encode answer: ", err))
 393         }
 394 }