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

   1 package main
   2
   3 import (
   4         "encoding/base64"
   5         "encoding/json"
   6         "fmt"
   7         "net/http"
   8         "net/http/httptest"
   9         "net/url"
  10         "path/filepath"
  11         "sync"
  12         "testing"
  13         "time"
  14
  15         "github.com/stretchr/testify/require"
  16 )
  17
  18 // chatgptTestAccount is the fake ChatGPT account ID in the fake tokens.
  19 const chatgptTestAccount = "account-test"
  20
  21 // ChatGPTFake is an httptest login service. The device token endpoint answers
  22 // 403 pending times before it grants the code. Each token answer has a new
  23 // access token, kept in access.
  24 type ChatGPTFake struct {
  25         *httptest.Server
  26         mutex       sync.Mutex
  27         pending     int
  28         deny        bool
  29         polls       int
  30         exchange    url.Values
  31         refreshes   int
  32         refreshed   string
  33         newRefresh  string
  34         failRefresh bool
  35         access      []string
  36 }
  37
  38 func newChatGPTFake(t *testing.T) *ChatGPTFake {
  39         t.Helper()
  40
  41         f := &ChatGPTFake{}
  42         mux := http.NewServeMux()
  43         mux.HandleFunc("POST /api/accounts/deviceauth/usercode", f.userCode)
  44         mux.HandleFunc("POST /api/accounts/deviceauth/token", f.deviceToken)
  45         mux.HandleFunc("POST /oauth/token", f.oauthToken)
  46
  47         f.Server = httptest.NewServer(mux)
  48         t.Cleanup(f.Close)
  49
  50         return f
  51 }
  52
  53 func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
  54         t.Helper()
  55
  56         return newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL)
  57 }
  58
  59 // token returns a new access token that expires after ttl, for a token file
  60 // that a test writes.
  61 func (f *ChatGPTFake) token(ttl time.Duration) string {
  62         f.mutex.Lock()
  63         defer f.mutex.Unlock()
  64
  65         return f.newToken(ttl)
  66 }
  67
  68 // newToken returns a new fake JWT with the expiry and the account claim, and
  69 // keeps it in access. The caller holds the mutex.
  70 func (f *ChatGPTFake) newToken(ttl time.Duration) string {
  71         claims, err := json.Marshal(map[string]any{
  72                 "exp": time.Now().Add(ttl).Unix(),
  73                 chatgptAuthClaim: map[string]any{
  74                         "chatgpt_account_id": chatgptTestAccount,
  75                 },
  76                 "serial": len(f.access),
  77         })
  78         if err != nil {
  79                 panic(err)
  80         }
  81
  82         encode := base64.RawURLEncoding.EncodeToString
  83         token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
  84         f.access = append(f.access, token)
  85
  86         return token
  87 }
  88
  89 func (f *ChatGPTFake) userCode(w http.ResponseWriter, r *http.Request) {
  90         var body map[string]string
  91
  92         err := json.NewDecoder(r.Body).Decode(&body)
  93         if err != nil || body["client_id"] != chatgptClientID {
  94                 http.Error(w, "bad client", http.StatusBadRequest)
  95
  96                 return
  97         }
  98
  99         chatgptWriteJSON(w, map[string]any{
 100                 "device_auth_id": "device-1",
 101                 "user_code":      "CODE-1234",
 102                 "interval":       "0",
 103         })
 104 }
 105
 106 func (f *ChatGPTFake) deviceToken(w http.ResponseWriter, r *http.Request) {
 107         f.mutex.Lock()
 108         defer f.mutex.Unlock()
 109
 110         var body map[string]string
 111
 112         err := json.NewDecoder(r.Body).Decode(&body)
 113         if err != nil || body["device_auth_id"] != "device-1" || body["user_code"] != "CODE-1234" {
 114                 http.Error(w, "bad device", http.StatusBadRequest)
 115
 116                 return
 117         }
 118
 119         f.polls++
 120         if f.deny {
 121                 http.Error(w, "denied", http.StatusBadRequest)
 122
 123                 return
 124         }
 125
 126         if f.polls <= f.pending {
 127                 http.Error(w, "pending", http.StatusForbidden)
 128
 129                 return
 130         }
 131
 132         chatgptWriteJSON(w, map[string]any{
 133                 "authorization_code": "auth-code",
 134                 "code_challenge":     "challenge",
 135                 "code_verifier":      "verifier",
 136         })
 137 }
 138
 139 func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
 140         f.mutex.Lock()
 141         defer f.mutex.Unlock()
 142
 143         if r.Header.Get("Content-Type") == "application/x-www-form-urlencoded" {
 144                 err := r.ParseForm()
 145                 if err != nil {
 146                         http.Error(w, err.Error(), http.StatusBadRequest)
 147
 148                         return
 149                 }
 150
 151                 f.exchange = r.PostForm
 152                 chatgptWriteJSON(w, map[string]any{
 153                         "access_token":  f.newToken(time.Hour),
 154                         "refresh_token": "refresh-1",
 155                         "id_token":      f.newToken(time.Hour),
 156                 })
 157
 158                 return
 159         }
 160
 161         var body map[string]string
 162
 163         err := json.NewDecoder(r.Body).Decode(&body)
 164         if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != chatgptClientID {
 165                 http.Error(w, "bad refresh", http.StatusBadRequest)
 166
 167                 return
 168         }
 169
 170         f.refreshes++
 171         f.refreshed = body["refresh_token"]
 172         if f.failRefresh {
 173                 http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized)
 174
 175                 return
 176         }
 177
 178         answer := map[string]any{
 179                 "access_token": f.newToken(time.Hour),
 180                 "id_token":     f.newToken(time.Hour),
 181         }
 182
 183         if f.newRefresh != "" {
 184                 answer["refresh_token"] = f.newRefresh
 185         }
 186
 187         chatgptWriteJSON(w, answer)
 188 }
 189
 190 func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) {
 191         t.Helper()
 192
 193         err := auth.write(ChatGPTTokenFile{
 194                 AccessToken:  access,
 195                 RefreshToken: refresh,
 196                 IDToken:      access,
 197                 AccountID:    chatgptTestAccount,
 198         })
 199         require.NoError(t, err)
 200 }
 201
 202 func chatgptWriteJSON(w http.ResponseWriter, value any) {
 203         w.Header().Set("Content-Type", "application/json")
 204
 205         err := json.NewEncoder(w).Encode(value)
 206         if err != nil {
 207                 panic(fmt.Sprint("encode answer: ", err))
 208         }
 209 }