Vectrum
Sign In
363 lines · 14.0 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 response is quoted back in an error. The body
// of a token-endpoint error is a place a submitted credential could plausibly
// be echoed, and a wrong endpoint can return a whole HTML page, so error text
// is clipped. This bounds the *message*, never the read: clipping the body we
// parse would truncate a perfectly good token response.
const maxErrorBody = 512

// maxBody bounds the read itself, so a misrouted request to something that
// streams forever cannot exhaust memory. A token response is a few hundred
// bytes; this is four orders of magnitude of headroom.
const maxBody = 1 << 20

// 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, or the server refused it. 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 — the server saying this grant is no good — wraps ErrExpired, so
// callers can classify it with errors.Is rather than by reading status codes
// back out of a string. Only 400/401/403 mean that.
//
// Everything else, 429 and 5xx included, is reported as an ordinary failure.
// The distinction is load-bearing: the caller remembers a refusal and stops
// asking, so treating a transient failure as one would disable refresh for the
// rest of the session on a login that is perfectly fine. 429 is exactly where
// that would bite, because it is the only response this endpoint has ever
// given us — and we still do not know whether it means "slow down" or "no".
// Being wrong in the direction of asking again is cheap; being wrong the other
// way costs the user their session.
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, maxBody))
	if err != nil {
		return nil, fmt.Errorf("claudecode: reading refresh response: %w", err)
	}
	switch resp.StatusCode {
	case http.StatusOK:
	case http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden:
		return nil, fmt.Errorf("%w (the server refused the refresh token: %d %s)",
			ErrExpired, resp.StatusCode, clip(raw))
	default:
		return nil, fmt.Errorf("claudecode: refresh failed (%d %s); if it keeps happening, run `claude` to log in again",
			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 != "" {
		next.RefreshToken = out.RefreshToken
	}
	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))
	// refreshTokenExpiresAt is deliberately not touched. We do not model it and
	// the exchange never states a rotated token's lifetime, so it survives as
	// one more key we preserve rather than one we guess at.

	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
	// refused holds the refusal the server gave us, if it gave one. There is no
	// stored refresh-token expiry to consult — the file's is not updated when a
	// token rotates, so trusting it would refuse a request the server would
	// have honoured. Asking once and remembering the answer is both more
	// accurate and cheaper than guessing from a timestamp. The error itself is
	// kept rather than a bare flag, so the reason survives to whoever asks
	// next instead of collapsing into an unexplained "expired".
	refused error

	// 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()

	if !s.creds.Expired(time.Now()) {
		return s.creds.AccessToken, nil
	}
	// Once the server has refused this refresh token, it will refuse it again.
	// Remembering that is what keeps a dead login from spending a round trip
	// on every subsequent request.
	if s.refused != nil {
		return "", s.refused
	}
	if s.creds.RefreshToken == "" {
		return "", ErrExpired
	}

	from := s.creds.RefreshToken
	next, err := Refresh(ctx, s.http, s.creds)
	if err != nil {
		// Only a refusal is remembered. A timeout, a 5xx, a 429 — those say
		// nothing about whether the token is good, and latching them would
		// turn a blip into a dead session.
		if errors.Is(err, ErrExpired) {
			s.refused = err
		}
		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.
		//
		// ErrStoreChanged is the exception, and it is not a failure: it means
		// something else wrote a newer credential while we ran. Nothing is
		// broken, nothing is lost, and warning about it would be alarming
		// about someone else's perfectly good login.
		err := Save(s.home, s.creds, from)
		if err != nil && !errors.Is(err, ErrStoreChanged) && s.OnWriteBackFailure != nil {
			s.OnWriteBackFailure(err)
		}
	}
	return s.creds.AccessToken, nil
}