Vectrum
Sign In
331 lines · 12.6 KB · Go
package claudecode

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"
)

// ClientID identifies Claude Code to the token endpoint.
const ClientID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"

// tokenURLEnv overrides the token endpoint. It exists so tests can point at an
// httptest server; nothing in normal operation sets it.
const tokenURLEnv = "CLAUDECODE_TOKEN_URL"

const defaultTokenURL = "https://console.anthropic.com/v1/oauth/token"

// refreshTimeout bounds a token exchange. The harness's own HTTP client allows
// ten minutes, which is right for a model response and badly wrong here: the
// refresh runs under a lock every other request is queued behind.
const refreshTimeout = 30 * time.Second

// maxErrorBody is how much of a failed response is quoted back. Unbounded, a
// wrong endpoint returning a large page becomes the error message; and the
// body of a token-endpoint error is a place a submitted credential could
// plausibly be echoed, so it is truncated as well as capped.
const maxErrorBody = 512

// tokenURL is the endpoint a refresh is sent to.
func tokenURL() string {
	if u := os.Getenv(tokenURLEnv); u != "" {
		return u
	}
	return defaultTokenURL
}

// ErrExpired means the credential has lapsed and cannot be renewed here —
// the refresh token is missing, itself expired, or was refused. Only logging
// in again fixes it, and the tool that owns the credential is the one to do it.
var ErrExpired = errors.New("claudecode: your Claude Code login has expired, run `claude` to refresh it")

// Refresh exchanges the refresh token for a new access token. It returns new
// credentials and writes nothing — persisting is a separate decision, made by
// Source, because whether it is safe depends on whether the server rotated.
//
// A refusal wraps ErrExpired, so callers can classify it with errors.Is rather
// than by reading status codes. That includes a 429: on these OAuth surfaces a
// 429 is not reliable evidence of a rate limit — the Messages API answers a
// request shape it declines the same way, and every probe of this endpoint has
// returned one without ever completing an exchange — so advising someone to
// wait would be advice that may never come good.
func Refresh(ctx context.Context, h *http.Client, c *Credentials) (*Credentials, error) {
	ctx, cancel := context.WithTimeout(ctx, refreshTimeout)
	defer cancel()

	body, err := json.Marshal(map[string]string{
		"grant_type":    "refresh_token",
		"refresh_token": c.RefreshToken,
		"client_id":     ClientID,
	})
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL(), bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("content-type", "application/json")

	resp, err := h.Do(req)
	if err != nil {
		return nil, fmt.Errorf("claudecode: refreshing: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxErrorBody+1))
	if err != nil {
		return nil, fmt.Errorf("claudecode: reading refresh response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%w (refresh refused: %d %s)", ErrExpired, resp.StatusCode, clip(raw))
	}

	var out struct {
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		ExpiresIn    int64  `json:"expires_in"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		// A 200 carrying HTML — a captive portal, a proxy interstitial — lands
		// here, so say what arrived rather than only that it would not parse.
		return nil, fmt.Errorf("claudecode: refresh returned something that is not a token response (%s): %w", clip(raw), err)
	}
	if out.AccessToken == "" {
		return nil, errors.New("claudecode: refresh returned no access token")
	}
	// Without a lifetime the new token would be born expired — Expired() would
	// report true immediately and every single request would refresh again,
	// hammering the endpoint on a loop of our own making.
	if out.ExpiresIn <= 0 {
		return nil, fmt.Errorf("claudecode: refresh returned no usable expires_in (%d)", out.ExpiresIn)
	}

	next := *c
	next.AccessToken = out.AccessToken
	next.ExpiresAt = time.Now().Add(time.Duration(out.ExpiresIn) * time.Second)
	// An omitted refresh_token means the old one is still good. Blanking it
	// would throw away the only thing that can refresh again.
	if out.RefreshToken != "" && out.RefreshToken != c.RefreshToken {
		next.RefreshToken = out.RefreshToken
		// The response says nothing about how long the NEW refresh token
		// lives, and carrying the old one's expiry forward would eventually
		// make a perfectly good token look dead. Unknown is the truth.
		next.RefreshExpiresAt = time.Time{}
	}
	return &next, nil
}

// clip bounds untrusted response text before it reaches an error message.
func clip(b []byte) string {
	b = bytes.TrimSpace(b)
	if len(b) > maxErrorBody {
		return string(b[:maxErrorBody]) + "…"
	}
	return string(b)
}

// ErrStoreChanged means the credential store no longer holds the refresh token
// we refreshed from: something else — Claude Code itself, most likely — logged
// in while we were running. Its credential is newer than ours, so we leave it
// alone rather than overwrite a fresh login with our older chain.
var ErrStoreChanged = errors.New("claudecode: credential store changed underneath us")

// Save writes credentials back to the store, provided the refresh token on
// disk is still expectRefresh. Passing "" skips that check.
//
// This rewrites a file Claude Code owns, so it is a read-modify-write: the
// document is parsed, only the fields we came to change are replaced, and
// everything else — keys this package does not model, keys it has never heard
// of — goes back exactly as it was found. Numbers are decoded as literals
// rather than float64s, or the millisecond timestamps would come back as
// 1.78966826946e+12 and Claude Code would be reading a file we corrupted.
//
// The write is a temp file in the same directory, fsynced, then renamed, and
// the directory is fsynced after. That survives both a process crash and a
// machine losing power: the rename is atomic, and the sync before it is what
// stops the rename from publishing a file whose contents have not landed yet.
func Save(home string, c *Credentials, expectRefresh string) error {
	path := filepath.Join(home, Filename)
	raw, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("claudecode: reading %s: %w", path, err)
	}

	var doc map[string]json.RawMessage
	dec := json.NewDecoder(bytes.NewReader(raw))
	dec.UseNumber()
	if err := dec.Decode(&doc); err != nil {
		return fmt.Errorf("claudecode: parsing %s: %w", path, err)
	}
	blob, ok := doc["claudeAiOauth"]
	if !ok {
		return fmt.Errorf("claudecode: %s has no claudeAiOauth object", path)
	}
	var oauth map[string]any
	dec = json.NewDecoder(bytes.NewReader(blob))
	dec.UseNumber()
	if err := dec.Decode(&oauth); err != nil {
		return fmt.Errorf("claudecode: parsing claudeAiOauth in %s: %w", path, err)
	}

	// Compare-and-swap against the owner of this file. Re-reading is not
	// enough on its own: between our Load at startup and this write, the user
	// may have run `claude` and logged in afresh, and that credential is
	// strictly better than ours.
	if expectRefresh != "" {
		if cur, _ := oauth["refreshToken"].(string); cur != expectRefresh {
			return ErrStoreChanged
		}
	}

	oauth["accessToken"] = c.AccessToken
	oauth["refreshToken"] = c.RefreshToken
	oauth["expiresAt"] = json.Number(strconv.FormatInt(c.ExpiresAt.UnixMilli(), 10))
	// A rotated refresh token has an unknown lifetime: the exchange does not
	// say. In memory that is the zero time, but the stale key is left in place
	// here rather than deleted. Deleting would be more honest and is the wrong
	// trade — Claude Code reads this file and we cannot see what it does with
	// a missing key, whereas a too-early expiry only makes it renew or ask for
	// a login sooner than it needed to. Between annoying a tool we do not own
	// and possibly breaking it, annoying wins.
	if !c.RefreshExpiresAt.IsZero() {
		oauth["refreshTokenExpiresAt"] = json.Number(strconv.FormatInt(c.RefreshExpiresAt.UnixMilli(), 10))
	}

	if blob, err = json.Marshal(oauth); err != nil {
		return fmt.Errorf("claudecode: encoding claudeAiOauth: %w", err)
	}
	doc["claudeAiOauth"] = blob
	out, err := json.Marshal(doc)
	if err != nil {
		return fmt.Errorf("claudecode: encoding %s: %w", path, err)
	}
	return writeAtomic(path, out)
}

// writeAtomic replaces path with data, or leaves it untouched.
func writeAtomic(path string, data []byte) error {
	dir := filepath.Dir(path)
	tmp, err := os.CreateTemp(dir, ".credentials-*.tmp")
	if err != nil {
		return fmt.Errorf("claudecode: writing %s: %w", path, err)
	}
	name := tmp.Name()
	// Removing by the captured name, and only while the temp file still
	// exists under it: after a successful rename this path is gone, and
	// os.Remove on a missing file is a no-op rather than a hazard.
	defer os.Remove(name)

	if err := tmp.Chmod(0o600); err != nil {
		tmp.Close()
		return fmt.Errorf("claudecode: securing %s: %w", name, err)
	}
	if _, err := tmp.Write(data); err != nil {
		tmp.Close()
		return fmt.Errorf("claudecode: writing %s: %w", name, err)
	}
	// Without this the rename can be durable while the contents are not, and
	// a machine that loses power here comes back to an empty credential file.
	if err := tmp.Sync(); err != nil {
		tmp.Close()
		return fmt.Errorf("claudecode: syncing %s: %w", name, err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("claudecode: closing %s: %w", name, err)
	}
	if err := os.Rename(name, path); err != nil {
		return fmt.Errorf("claudecode: replacing %s: %w", path, err)
	}
	// Fsync the directory so the rename itself survives a power loss. Failing
	// here means the new file is written and probably linked; reporting it as
	// a save failure would be more alarming than useful.
	if d, err := os.Open(dir); err == nil {
		d.Sync()
		d.Close()
	}
	return nil
}

// Source hands out an access token, refreshing it when it lapses. It satisfies
// the shape llm.WithTokenSource wants.
type Source struct {
	mu    sync.Mutex
	creds *Credentials
	http  *http.Client
	// home is where a rotated credential is written back. Empty disables
	// write-back entirely.
	home string

	// OnWriteBackFailure is called when the server rotated the refresh token
	// but the new one could not be stored. That combination is the one this
	// whole design exists to avoid — the token on disk is dead and its
	// replacement is about to be lost when this process exits — so it must not
	// pass silently. Leave it nil only if something else is watching.
	OnWriteBackFailure func(error)
}

// NewSource wraps credentials in a token source. h is the client refreshes are
// sent with; home is the directory holding the credential store, or "" to
// refresh without ever writing.
func NewSource(c *Credentials, h *http.Client, home string) *Source {
	if h == nil {
		h = &http.Client{Timeout: refreshTimeout}
	}
	return &Source{creds: c, http: h, home: home}
}

// Token returns a valid access token, refreshing first if the current one has
// lapsed.
//
// The refreshed credential is kept in memory. It is written back only when the
// server rotated the refresh token, and that is the point: if rotation
// happened, the token still on disk is already dead, so leaving it there is
// what breaks the user's Claude Code login. Writing is the safe act in that
// case and the risky one in every other, which is why it is conditional rather
// than routine.
//
// The lock means concurrent requests — streaming and buffered can be in flight
// together — produce one refresh rather than a thundering herd. Refresh bounds
// its own time, so the wait behind that lock is seconds, not the harness
// client's ten-minute model-response budget.
func (s *Source) Token(ctx context.Context) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	if !s.creds.Expired(now) {
		return s.creds.AccessToken, nil
	}
	if s.creds.RefreshToken == "" || s.creds.RefreshExpired(now) {
		return "", ErrExpired
	}

	from := s.creds.RefreshToken
	next, err := Refresh(ctx, s.http, s.creds)
	if err != nil {
		return "", err
	}
	rotated := next.RefreshToken != from
	s.creds = next

	if rotated && s.home != "" {
		// Best effort on the request, but never silent: we hold a token that
		// works, and failing the request would be strictly worse than a stale
		// file. What would be worse still is saying nothing, because the
		// user's Claude Code login is now broken in a way only this message
		// explains.
		if err := Save(s.home, s.creds, from); err != nil && s.OnWriteBackFailure != nil {
			s.OnWriteBackFailure(err)
		}
	}
	return s.creds.AccessToken, nil
}