Vectrum
Sign In
392 lines · 12.9 KB · Go
package claudecode

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

// liveCreds is a credential that has lapsed but whose refresh token has not,
// which is the only state in which a refresh should be attempted.
func liveCreds() *Credentials {
	return &Credentials{
		AccessToken:      "stale",
		RefreshToken:     "r1",
		ExpiresAt:        time.Now().Add(-time.Hour),
		RefreshExpiresAt: time.Now().Add(24 * time.Hour),
	}
}

func TestRefreshExchangesTheToken(t *testing.T) {
	var body map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewDecoder(r.Body).Decode(&body)
		io.WriteString(w, `{"access_token":"new-access","refresh_token":"new-refresh","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	got, err := Refresh(context.Background(), srv.Client(), &Credentials{AccessToken: "old", RefreshToken: "old-refresh"})
	if err != nil {
		t.Fatalf("Refresh: %v", err)
	}
	if got.AccessToken != "new-access" || got.RefreshToken != "new-refresh" {
		t.Errorf("got %q/%q, want new-access/new-refresh", got.AccessToken, got.RefreshToken)
	}
	if body["grant_type"] != "refresh_token" || body["refresh_token"] != "old-refresh" {
		t.Errorf("request body = %v", body)
	}
	if body["client_id"] != ClientID {
		t.Errorf("client_id = %v, want %s", body["client_id"], ClientID)
	}
	// expires_in is a duration; the stored form is an absolute instant.
	if !got.ExpiresAt.After(time.Now().Add(7 * time.Hour)) {
		t.Errorf("ExpiresAt = %s, want ~8h out", got.ExpiresAt)
	}
}

// A response that omits refresh_token means the old one stays valid. Blanking
// it would throw away the only thing that can refresh again.
func TestRefreshKeepsTheOldRefreshTokenWhenNoneIsReturned(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"access_token":"new-access","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	got, err := Refresh(context.Background(), srv.Client(), &Credentials{RefreshToken: "keep-me"})
	if err != nil {
		t.Fatalf("Refresh: %v", err)
	}
	if got.RefreshToken != "keep-me" {
		t.Errorf("RefreshToken = %q, want the original kept", got.RefreshToken)
	}
}

func TestRefreshRejectedSaysToLogInAgain(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadRequest)
		io.WriteString(w, `{"error":"invalid_grant"}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	_, err := Refresh(context.Background(), srv.Client(), &Credentials{RefreshToken: "dead"})
	if err == nil || !strings.Contains(err.Error(), "claude") {
		t.Fatalf("err = %v, want an error telling the user to run claude", err)
	}
}

// A 429 from these OAuth surfaces is not reliable evidence of a rate limit:
// the Messages API answers a request shape it declines the same way. Telling
// the user to wait would be advice that never comes good.
func TestRefreshRateLimitedDoesNotAdviseWaiting(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTooManyRequests)
		io.WriteString(w, `{"error":{"type":"rate_limit_error","message":"Rate limited. Please try again later."}}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	_, err := Refresh(context.Background(), srv.Client(), &Credentials{RefreshToken: "r"})
	if err == nil {
		t.Fatal("a 429 must be an error")
	}
	if !strings.Contains(err.Error(), "claude") {
		t.Errorf("a 429 should still point at `claude`, got %q", err)
	}
}

func TestRefreshTransportErrorIsReported(t *testing.T) {
	t.Setenv(tokenURLEnv, "http://127.0.0.1:1")
	_, err := Refresh(context.Background(), &http.Client{Timeout: time.Second}, &Credentials{RefreshToken: "r"})
	if err == nil {
		t.Fatal("a transport error must be reported")
	}
}

// The point of the feature: an expired token becomes a working one without the
// caller knowing anything happened, and only one exchange is spent doing it.
func TestSourceRefreshesOnExpiryAndCaches(t *testing.T) {
	var hits int
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits++
		io.WriteString(w, `{"access_token":"fresh","refresh_token":"r1","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	s := NewSource(liveCreds(), srv.Client(), "")
	for i := 0; i < 3; i++ {
		tok, err := s.Token(context.Background())
		if err != nil {
			t.Fatalf("Token: %v", err)
		}
		if tok != "fresh" {
			t.Fatalf("Token = %q, want fresh", tok)
		}
	}
	if hits != 1 {
		t.Errorf("refreshed %d times, want 1 — the result should be cached", hits)
	}
}

func TestSourceConcurrentTokensRefreshOnce(t *testing.T) {
	var mu sync.Mutex
	var hits int
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		hits++
		mu.Unlock()
		io.WriteString(w, `{"access_token":"fresh","refresh_token":"r1","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	s := NewSource(liveCreds(), srv.Client(), "")
	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if _, err := s.Token(context.Background()); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()
	if hits != 1 {
		t.Errorf("refreshed %d times under concurrency, want 1", hits)
	}
}

// An expired refresh token cannot be refreshed. Sending the request anyway
// spends a round trip to be told what we already knew.
func TestSourceExpiredRefreshTokenDoesNotCallOut(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("no request should be sent for a dead refresh token")
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	c := liveCreds()
	c.RefreshExpiresAt = time.Now().Add(-time.Hour)
	s := NewSource(c, srv.Client(), "")
	if _, err := s.Token(context.Background()); !errors.Is(err, ErrExpired) {
		t.Fatalf("err = %v, want ErrExpired", err)
	}
}

func TestSourceMissingRefreshTokenDoesNotCallOut(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("no request should be sent without a refresh token")
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	c := liveCreds()
	c.RefreshToken = ""
	s := NewSource(c, srv.Client(), "")
	if _, err := s.Token(context.Background()); !errors.Is(err, ErrExpired) {
		t.Fatalf("err = %v, want ErrExpired", err)
	}
}

// storeWith writes a credentials file carrying a key this package does not
// model, so a round trip through Save can be checked for dropping it.
func storeWith(t *testing.T, refresh string) string {
	t.Helper()
	home := t.TempDir()
	if err := os.MkdirAll(filepath.Join(home, ".claude"), 0o700); err != nil {
		t.Fatal(err)
	}
	body := `{"claudeAiOauth":{"accessToken":"stale","refreshToken":"` + refresh + `",
	  "expiresAt":1789668269460,"refreshTokenExpiresAt":1791404074460,
	  "scopes":["user:inference"],"subscriptionType":"team",
	  "rateLimitTier":"default_claude_ai"},"someOtherTool":{"keep":"me"}}`
	if err := os.WriteFile(filepath.Join(home, Filename), []byte(body), 0o600); err != nil {
		t.Fatal(err)
	}
	return home
}

func readStore(t *testing.T, home string) map[string]any {
	t.Helper()
	raw, err := os.ReadFile(filepath.Join(home, Filename))
	if err != nil {
		t.Fatal(err)
	}
	var doc map[string]any
	if err := json.Unmarshal(raw, &doc); err != nil {
		t.Fatalf("file is no longer valid JSON after Save: %v", err)
	}
	return doc
}

// Save rewrites a file another program owns, so it must put back everything it
// did not come to change — including keys this package has never heard of.
func TestSavePreservesUnknownKeys(t *testing.T) {
	home := storeWith(t, "r1")
	c, err := Load(home)
	if err != nil {
		t.Fatal(err)
	}
	c.AccessToken, c.RefreshToken = "new-access", "new-refresh"
	c.ExpiresAt = time.UnixMilli(1999999999000)
	if err := Save(home, c); err != nil {
		t.Fatalf("Save: %v", err)
	}

	doc := readStore(t, home)
	oauth, ok := doc["claudeAiOauth"].(map[string]any)
	if !ok {
		t.Fatal("claudeAiOauth is missing after Save")
	}
	if oauth["rateLimitTier"] != "default_claude_ai" {
		t.Errorf("rateLimitTier = %v, want it preserved — Save dropped a key it does not model", oauth["rateLimitTier"])
	}
	if oauth["subscriptionType"] != "team" {
		t.Errorf("subscriptionType = %v, want it preserved", oauth["subscriptionType"])
	}
	if _, ok := doc["someOtherTool"]; !ok {
		t.Error("Save dropped a whole top-level key it does not model")
	}
	if oauth["accessToken"] != "new-access" || oauth["refreshToken"] != "new-refresh" {
		t.Errorf("Save did not write the new tokens: %v", oauth)
	}
	// The stored form is epoch milliseconds, and it must stay an integer
	// literal rather than turning into 1.999999999e+12.
	if got := string(mustRaw(t, home)); !strings.Contains(got, "1999999999000") {
		t.Errorf("expiresAt was not written as an integer millisecond literal:\n%s", got)
	}
}

func mustRaw(t *testing.T, home string) []byte {
	t.Helper()
	raw, err := os.ReadFile(filepath.Join(home, Filename))
	if err != nil {
		t.Fatal(err)
	}
	return raw
}

func TestSaveKeepsTheFileOwnerOnly(t *testing.T) {
	home := storeWith(t, "r1")
	c, err := Load(home)
	if err != nil {
		t.Fatal(err)
	}
	if err := Save(home, c); err != nil {
		t.Fatalf("Save: %v", err)
	}
	info, err := os.Stat(filepath.Join(home, Filename))
	if err != nil {
		t.Fatal(err)
	}
	if perm := info.Mode().Perm(); perm != 0o600 {
		t.Errorf("mode = %o, want 600 — a credential must not become readable", perm)
	}
}

// The whole rotation policy, in two tests. When the server hands back the same
// refresh token, the file is Claude Code's business and we do not touch it.
func TestSourceDoesNotWriteBackWhenNotRotated(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"access_token":"fresh","refresh_token":"r1","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	home := storeWith(t, "r1")
	before := mustRaw(t, home)

	c := liveCreds()
	c.RefreshToken = "r1"
	s := NewSource(c, srv.Client(), home)
	if _, err := s.Token(context.Background()); err != nil {
		t.Fatalf("Token: %v", err)
	}
	if got := mustRaw(t, home); string(got) != string(before) {
		t.Errorf("the file was rewritten despite no rotation:\nbefore %s\nafter  %s", before, got)
	}
}

// When the server rotates, the stored refresh token is already dead. Leaving
// it there is what breaks the user's Claude Code login, so this is the one
// case where writing is the safe act rather than the risky one.
func TestSourceWritesBackWhenRotated(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"access_token":"fresh","refresh_token":"r2-rotated","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	home := storeWith(t, "r1")
	c := liveCreds()
	c.RefreshToken = "r1"
	s := NewSource(c, srv.Client(), home)
	if _, err := s.Token(context.Background()); err != nil {
		t.Fatalf("Token: %v", err)
	}

	oauth := readStore(t, home)["claudeAiOauth"].(map[string]any)
	if oauth["refreshToken"] != "r2-rotated" {
		t.Errorf("refreshToken = %v, want the rotated one written back", oauth["refreshToken"])
	}
	if oauth["accessToken"] != "fresh" {
		t.Errorf("accessToken = %v, want it written back alongside", oauth["accessToken"])
	}
	if oauth["rateLimitTier"] != "default_claude_ai" {
		t.Error("write-back dropped an unmodelled key")
	}
}

// Write-back is best-effort: a read-only store is a worse session, not a
// failed one, and failing the request would be a strictly worse outcome than
// running on the token we just successfully obtained.
func TestSourceSurvivesAFailedWriteBack(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"access_token":"fresh","refresh_token":"r2-rotated","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	c := liveCreds()
	c.RefreshToken = "r1"
	// A home with no credentials file at all: Save cannot succeed.
	s := NewSource(c, srv.Client(), t.TempDir())
	tok, err := s.Token(context.Background())
	if err != nil {
		t.Fatalf("a failed write-back must not fail the request: %v", err)
	}
	if tok != "fresh" {
		t.Errorf("Token = %q, want the refreshed token", tok)
	}
}

// With no home, write-back is disabled outright and refresh still works.
func TestSourceWithoutAHomeStillRefreshes(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"access_token":"fresh","refresh_token":"r2","expires_in":28800}`)
	}))
	defer srv.Close()
	t.Setenv(tokenURLEnv, srv.URL)

	s := NewSource(liveCreds(), srv.Client(), "")
	tok, err := s.Token(context.Background())
	if err != nil {
		t.Fatalf("Token: %v", err)
	}
	if tok != "fresh" {
		t.Errorf("Token = %q, want fresh", tok)
	}
}