+package main
+
+import (
+ "encoding/base64"
+ "encoding/json"
+ "fmt"
+ "net/http"
+ "net/http/httptest"
+ "net/url"
+ "path/filepath"
+ "sync"
+ "testing"
+ "time"
+
+ "github.com/stretchr/testify/require"
+)
+
+// chatgptTestAccount is the fake ChatGPT account ID in the fake tokens.
+const chatgptTestAccount = "account-test"
+
+// ChatGPTFake is an httptest login service. The device token endpoint answers
+// 403 pending times before it grants the code. Each token answer has a new
+// access token, kept in access.
+type ChatGPTFake struct {
+ *httptest.Server
+ mutex sync.Mutex
+ pending int
+ deny bool
+ polls int
+ exchange url.Values
+ refreshes int
+ refreshed string
+ newRefresh string
+ failRefresh bool
+ access []string
+}
+
+// token returns a new access token that expires after ttl, for a token file
+// that a test writes.
+func (f *ChatGPTFake) token(ttl time.Duration) string {
+ f.mutex.Lock()
+ defer f.mutex.Unlock()
+
+ return f.newToken(ttl)
+}
+
+// newToken returns a new fake JWT with the expiry and the account claim, and
+// keeps it in access. The caller holds the mutex.
+func (f *ChatGPTFake) newToken(ttl time.Duration) string {
+ claims, err := json.Marshal(map[string]any{
+ "exp": time.Now().Add(ttl).Unix(),
+ chatgptAuthClaim: map[string]any{
+ "chatgpt_account_id": chatgptTestAccount,
+ },
+ "serial": len(f.access),
+ })
+ if err != nil {
+ panic(err)
+ }
+
+ encode := base64.RawURLEncoding.EncodeToString
+ token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
+ f.access = append(f.access, token)
+
+ return token
+}
+
+func (f *ChatGPTFake) userCode(w http.ResponseWriter, r *http.Request) {
+ var body map[string]string
+ err := json.NewDecoder(r.Body).Decode(&body)
+ if err != nil || body["client_id"] != chatgptClientID {
+ http.Error(w, "bad client", http.StatusBadRequest)
+ return
+ }
+
+ chatgptWriteJSON(w, map[string]any{
+ "device_auth_id": "device-1",
+ "user_code": "CODE-1234",
+ "interval": "0",
+ })
+}
+
+func (f *ChatGPTFake) deviceToken(w http.ResponseWriter, r *http.Request) {
+ f.mutex.Lock()
+ defer f.mutex.Unlock()
+
+ var body map[string]string
+ err := json.NewDecoder(r.Body).Decode(&body)
+ if err != nil || body["device_auth_id"] != "device-1" || body["user_code"] != "CODE-1234" {
+ http.Error(w, "bad device", http.StatusBadRequest)
+ return
+ }
+
+ f.polls++
+ if f.deny {
+ http.Error(w, "denied", http.StatusBadRequest)
+ return
+ }
+
+ if f.polls <= f.pending {
+ http.Error(w, "pending", http.StatusForbidden)
+ return
+ }
+
+ chatgptWriteJSON(w, map[string]any{
+ "authorization_code": "auth-code",
+ "code_challenge": "challenge",
+ "code_verifier": "verifier",
+ })
+}
+
+func (f *ChatGPTFake) oauthToken(w http.ResponseWriter, r *http.Request) {
+ f.mutex.Lock()
+ defer f.mutex.Unlock()
+
+ if r.Header.Get("Content-Type") == "application/x-www-form-urlencoded" {
+ err := r.ParseForm()
+ if err != nil {
+ http.Error(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ f.exchange = r.PostForm
+ chatgptWriteJSON(w, map[string]any{
+ "access_token": f.newToken(time.Hour),
+ "refresh_token": "refresh-1",
+ "id_token": f.newToken(time.Hour),
+ })
+
+ return
+ }
+
+ var body map[string]string
+ err := json.NewDecoder(r.Body).Decode(&body)
+ if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != chatgptClientID {
+ http.Error(w, "bad refresh", http.StatusBadRequest)
+ return
+ }
+
+ f.refreshes++
+ f.refreshed = body["refresh_token"]
+ if f.failRefresh {
+ http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized)
+ return
+ }
+
+ answer := map[string]any{
+ "access_token": f.newToken(time.Hour),
+ "id_token": f.newToken(time.Hour),
+ }
+
+ if f.newRefresh != "" {
+ answer["refresh_token"] = f.newRefresh
+ }
+
+ chatgptWriteJSON(w, answer)
+}
+
+func newChatGPTFake(t *testing.T) *ChatGPTFake {
+ t.Helper()
+
+ f := &ChatGPTFake{}
+ mux := http.NewServeMux()
+ mux.HandleFunc("POST /api/accounts/deviceauth/usercode", f.userCode)
+ mux.HandleFunc("POST /api/accounts/deviceauth/token", f.deviceToken)
+ mux.HandleFunc("POST /oauth/token", f.oauthToken)
+ f.Server = httptest.NewServer(mux)
+ t.Cleanup(f.Close)
+
+ return f
+}
+
+func newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
+ t.Helper()
+
+ return newChatGPTAuth(filepath.Join(t.TempDir(), chatgptTokenFileName), server.URL)
+}
+
+func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) {
+ t.Helper()
+
+ err := auth.write(ChatGPTTokenFile{
+ AccessToken: access,
+ RefreshToken: refresh,
+ IDToken: access,
+ AccountID: chatgptTestAccount,
+ })
+ require.NoError(t, err)
+}
+
+func chatgptWriteJSON(w http.ResponseWriter, value any) {
+ w.Header().Set("Content-Type", "application/json")
+ err := json.NewEncoder(w).Encode(value)
+ if err != nil {
+ panic(fmt.Sprint("encode answer: ", err))
+ }
+}
projects / onechat / commitdiff