git.giacomolli.uk - onechat/blob - chatgpt_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 // chatgptTestAccount is the fake ChatGPT account ID in the fake tokens.
  22 const chatgptTestAccount = "account-test"
  23
  24 func TestLoginWritesTheTokenFile(t *testing.T) {
  25         server := newChatGPTFake(t)
  26         server.pending = 2
  27         auth := newTestChatGPTAuth(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, chatgptClientID, form.Get("client_id"))
  42         require.Equal(t, server.URL+"/deviceauth/callback", form.Get("redirect_uri"))
  43
  44         tokens := chatgptReadTokenFile(t, auth.path)
  45         require.Equal(t, "refresh-1", tokens.RefreshToken)
  46         require.Equal(t, chatgptTestAccount, 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 := newChatGPTFake(t)
  58         server.deny = true
  59         auth := newTestChatGPTAuth(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 := newChatGPTFake(t)
  68         auth := newTestChatGPTAuth(t, server)
  69         chatgptWriteTestTokens(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, chatgptTestAccount, 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 := newChatGPTFake(t)
 100                         server.newRefresh = tc.newRefresh
 101                         auth := newTestChatGPTAuth(t, server)
 102                         chatgptWriteTestTokens(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 := chatgptReadTokenFile(t, auth.path)
 111                         require.Equal(t, server.access[1], tokens.AccessToken)
 112                         require.Equal(t, tc.wantRefresh, tokens.RefreshToken)
 113                         require.Equal(t, chatgptTestAccount, 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 := newChatGPTFake(t)
 125         server.newRefresh = "refresh-new"
 126         auth := newTestChatGPTAuth(t, server)
 127         chatgptWriteTestTokens(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 := newChatGPTFake(t)
 148         auth := newTestChatGPTAuth(t, server)
 149
 150         err := auth.Authorize(t.Context(), http.Header{})
 151         require.ErrorIs(t, err, errChatGPTLoginNeeded)
 152
 153         server.failRefresh = true
 154         chatgptWriteTestTokens(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, errChatGPTLoginNeeded)
 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 := newChatGPTFake(t)
 169         auth := newTestChatGPTAuth(t, server)
 170         chatgptWriteTestTokens(t, auth, server.token(-time.Hour), "refresh-0")
 171
 172         responses := newModelFake(t)
 173         client := newModelClient(responses.URL, auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
 174         _, err := client.Create(t.Context(), modelTestRequest())
 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, chatgptTestAccount, call.header.Get("ChatGPT-Account-Id"))
 181         require.Equal(t, server.access[1], chatgptReadTokenFile(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", modelTestKey)
 189         address, auth, err := modelEndpoint(config, dir)
 190         require.NoError(t, err)
 191         require.Equal(t, modelChatGPTURL, address)
 192         require.IsType(t, &ChatGPTAuth{}, auth)
 193         require.Equal(t, filepath.Join(dir, chatgptTokenFileName), auth.(*ChatGPTAuth).path)
 194
 195         config.Endpoint = endpointAPI
 196         address, auth, err = modelEndpoint(config, dir)
 197         require.NoError(t, err)
 198         require.Equal(t, modelAPIURL, address)
 199         require.Equal(t, ModelAPIKey(modelTestKey), auth)
 200
 201         t.Setenv("OPENAI_API_KEY", "")
 202         _, _, err = modelEndpoint(config, dir)
 203         require.ErrorContains(t, err, "OPENAI_API_KEY")
 204 }
 205
 206 // ChatGPTFake 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 ChatGPTFake 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 // token returns a new access token that expires after ttl, for a token file
 224 // that a test writes.
 225 func (f *ChatGPTFake) token(ttl time.Duration) string {
 226         f.mutex.Lock()
 227         defer f.mutex.Unlock()
 228
 229         return f.newToken(ttl)
 230 }
 231
 232 // newToken returns a new fake JWT with the expiry and the account claim, and
 233 // keeps it in access. The caller holds the mutex.
 234 func (f *ChatGPTFake) newToken(ttl time.Duration) string {
 235         claims, err := json.Marshal(map[string]any{
 236                 "exp": time.Now().Add(ttl).Unix(),
 237                 chatgptAuthClaim: map[string]any{
 238                         "chatgpt_account_id": chatgptTestAccount,
 239                 },
 240                 "serial": len(f.access),
 241         })
 242         if err != nil {
 243                 panic(err)
 244         }
 245
 246         encode := base64.RawURLEncoding.EncodeToString
 247         token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
 248         f.access = append(f.access, token)
 249
 250         return token
 251 }
 252
 253 func (f *ChatGPTFake) userCode(w http.ResponseWriter, r *http.Request) {
 254         var body map[string]string
 255         err := json.NewDecoder(r.Body).Decode(&body)
 256         if err != nil || body["client_id"] != chatgptClientID {
 257                 http.Error(w, "bad client", http.StatusBadRequest)
 258                 return
 259         }
 260
 261         chatgptWriteJSON(w, map[string]any{
 262                 "device_auth_id": "device-1",
 263                 "user_code":      "CODE-1234",
 264                 "interval":       "0",
 265         })
 266 }
 267
 268 func (f *ChatGPTFake) deviceToken(w http.ResponseWriter, r *http.Request) {
 269         f.mutex.Lock()
 270         defer f.mutex.Unlock()
 271
 272         var body map[string]string
 273         err := json.NewDecoder(r.Body).Decode(&body)
 274         if err != nil || body["device_auth_id"] != "device-1" || body["user_code"] != "CODE-1234" {
 275                 http.Error(w, "bad device", http.StatusBadRequest)
 276                 return
 277         }
 278
 279         f.polls++
 280         if f.deny {
 281                 http.Error(w, "denied", http.StatusBadRequest)
 282                 return
 283         }
 284
 285         if f.polls <= f.pending {
 286                 http.Error(w, "pending", http.StatusForbidden)
 287                 return
 288         }
 289
 290         chatgptWriteJSON(w, map[string]any{
 291                 "authorization_code": "auth-code",
 292                 "code_challenge":     "challenge",
 293                 "code_verifier":      "verifier",
 294         })
 295 }
 296
 297 func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 298         f.mutex.Lock()
 299         defer f.mutex.Unlock()
 300
 301         if r.Header.Get("Content-Type") == "application/x-www-form-urlencoded" {
 302                 err := r.ParseForm()
 303                 if err != nil {
 304                         http.Error(w, err.Error(), http.StatusBadRequest)
 305                         return
 306                 }
 307
 308                 f.exchange = r.PostForm
 309                 chatgptWriteJSON(w, map[string]any{
 310                         "access_token":  f.newToken(time.Hour),
 311                         "refresh_token": "refresh-1",
 312                         "id_token":      f.newToken(time.Hour),
 313                 })
 314
 315                 return
 316         }
 317
 318         var body map[string]string
 319         err := json.NewDecoder(r.Body).Decode(&body)
 320         if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != chatgptClientID {
 321                 http.Error(w, "bad refresh", http.StatusBadRequest)
 322                 return
 323         }
 324
 325         f.refreshes++
 326         f.refreshed = body["refresh_token"]
 327         if f.failRefresh {
 328                 http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized)
 329                 return
 330         }
 331
 332         answer := map[string]any{
 333                 "access_token": f.newToken(time.Hour),
 334                 "id_token":     f.newToken(time.Hour),
 335         }
 336
 337         if f.newRefresh != "" {
 338                 answer["refresh_token"] = f.newRefresh
 339         }
 340
 341         chatgptWriteJSON(w, answer)
 342 }
 343
 344 func newChatGPTFake(t *testing.T) *ChatGPTFake {
 345         t.Helper()
 346
 347         f := &ChatGPTFake{}
 348         mux := http.NewServeMux()
 349         mux.HandleFunc("POST /api/accounts/deviceauth/usercode", f.userCode)
 350         mux.HandleFunc("POST /api/accounts/deviceauth/token", f.deviceToken)
 351         mux.HandleFunc("POST /oauth/token", f.oauthToken)
 352         f.Server = httptest.NewServer(mux)
 353         t.Cleanup(f.Close)
 354
 355         return f
 356 }
 357
 358 func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
 359         t.Helper()
 360
 361         return newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL)
 362 }
 363
 364 func chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile {
 365         t.Helper()
 366
 367         data, err := os.ReadFile(path)
 368         require.NoError(t, err)
 369
 370         var tokens ChatGPTTokenFile
 371         require.NoError(t, json.Unmarshal(data, &tokens))
 372
 373         return tokens
 374 }
 375
 376 func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) {
 377         t.Helper()
 378
 379         err := auth.write(ChatGPTTokenFile{
 380                 AccessToken:  access,
 381                 RefreshToken: refresh,
 382                 IDToken:      access,
 383                 AccountID:    chatgptTestAccount,
 384         })
 385         require.NoError(t, err)
 386 }
 387
 388 func chatgptWriteJSON(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 }