package main
import (
- "encoding/base64"
"encoding/json"
- "fmt"
"io"
"net/http"
- "net/http/httptest"
- "net/url"
"os"
"path/filepath"
"strings"
"github.com/stretchr/testify/require"
)
-// chatgptTestAccount is the fake ChatGPT account ID in the fake tokens.
-const chatgptTestAccount = "account-test"
-
func TestLoginWritesTheTokenFile(t *testing.T) {
server := newChatGPTFake(t)
server.pending = 2
auth := newTestChatGPTAuth(t, server)
var out strings.Builder
+
err := auth.Login(t.Context(), &out)
require.NoError(t, err)
chatgptWriteTestTokens(t, auth, server.token(time.Second), "refresh-0")
var group sync.WaitGroup
+
errs := make(chan error, 8)
for range 8 {
group.Go(func() {
group.Wait()
close(errs)
+
for err := range errs {
require.NoError(t, err)
}
server.failRefresh = true
chatgptWriteTestTokens(t, auth, server.token(-time.Minute), "refresh-0")
+
before, err := os.ReadFile(auth.path)
require.NoError(t, err)
require.NoError(t, err)
require.Equal(t, 1, server.refreshes)
+
call := responses.Call(0)
require.Equal(t, "Bearer "+server.access[1], call.header.Get("Authorization"))
require.Equal(t, chatgptTestAccount, call.header.Get("ChatGPT-Account-Id"))
config := defaultConfig()
t.Setenv("OPENAI_API_KEY", modelTestKey)
+
address, auth, err := modelEndpoint(config, dir)
require.NoError(t, err)
require.Equal(t, modelChatGPTURL, address)
require.Equal(t, ModelAPIKey(modelTestKey), auth)
t.Setenv("OPENAI_API_KEY", "")
+
_, _, err = modelEndpoint(config, dir)
require.ErrorContains(t, err, "OPENAI_API_KEY")
}
-// 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 chatgptReadTokenFile(t *testing.T, path string) ChatGPTTokenFile {
t.Helper()
require.NoError(t, err)
var tokens ChatGPTTokenFile
+
require.NoError(t, json.Unmarshal(data, &tokens))
return tokens
}
-
-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 / blobdiff