at-container-registry/pkg/auth/oauth/transport.go

package oauth

import (
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"sync"
	"time"
)

// refreshDetachTimeout bounds a detached token-refresh POST. Refreshes are
// normally sub-second; this only exists so a hung auth server can't pin the
// per-DID session lock forever.
//
// Note this is per-POST, not per-refresh: indigo's postToAuthServer retries once
// on a use_dpop_nonce challenge, so a single refresh can issue two detached
// POSTs and hold the per-DID lock for up to 2x this value. Keep that product
// comfortably under the callers' own deadlines (Docker gives up on /auth/token
// well before it).
const refreshDetachTimeout = 30 * time.Second

// badTokenEndpointOnce dedupes the "couldn't parse the token endpoint" warning
// per endpoint value rather than once per process. That path fails open — it
// returns the client unwrapped, silently reinstating the refresh-burn bug this
// file exists to fix — so it must stay visible for every distinct auth server
// it affects, not just the first one after boot.
var badTokenEndpointOnce sync.Map

// refreshDetachTransport detaches the request context for OAuth token-refresh
// POSTs. A refresh is non-idempotent: the auth server rotates the refresh
// token as soon as it processes the request, whether or not we stick around
// for the response. If the inbound request context is canceled mid-refresh
// (e.g. Docker gives up on a slow /auth/token), aborting the POST strands the
// rotated token server-side; our stored refresh token is then already
// consumed, the next refresh fails with invalid_grant "Refresh token
// replayed", and the whole session gets deleted. Once a refresh starts it
// must run to completion, capped by its own timeout (same rationale as the
// upload finalization in pkg/hold/oci/xrpc.go).
//
// Only POSTs to the session's auth-server token endpoint are detached; every
// other request keeps normal cancellation semantics.
type refreshDetachTransport struct {
	base          http.RoundTripper
	tokenEndpoint *url.URL
	timeout       time.Duration
}

func (t *refreshDetachTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if req.Method != http.MethodPost || !sameEndpoint(req.URL, t.tokenEndpoint) {
		return t.base.RoundTrip(req)
	}

	// context.WithoutCancel keeps values (trace/log metadata) but drops
	// cancellation and deadline. The cancel func must outlive RoundTrip: the
	// body is read by the caller, so it is released on Body.Close() instead
	// of a defer here.
	ctx, cancel := context.WithTimeout(context.WithoutCancel(req.Context()), t.timeout)
	resp, err := t.base.RoundTrip(req.Clone(ctx))
	if err != nil {
		cancel()
		return nil, err
	}
	resp.Body = &cancelOnClose{ReadCloser: resp.Body, cancel: cancel}
	return resp, nil
}

// cancelOnClose ties a context.CancelFunc to response-body Close so the
// detached timeout context stays alive for the full body read.
type cancelOnClose struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (c *cancelOnClose) Close() error {
	c.cancel()
	return c.ReadCloser.Close()
}

// newRefreshDetachClient wraps inner so that token-refresh POSTs to
// tokenEndpoint survive cancellation of the inbound request context. If the
// endpoint can't be parsed (corrupt session data), the client is returned
// unwrapped, preserving the old behavior.
func newRefreshDetachClient(inner *http.Client, tokenEndpoint string) *http.Client {
	endpoint, err := url.Parse(tokenEndpoint)
	if err != nil || endpoint.Host == "" {
		if _, seen := badTokenEndpointOnce.LoadOrStore(tokenEndpoint, struct{}{}); !seen {
			slog.Warn("Not detaching token-refresh context: unparseable auth server token endpoint",
				"component", "oauth/refresher",
				"tokenEndpoint", tokenEndpoint,
				"error", err)
		}
		return inner
	}

	if inner == nil {
		inner = http.DefaultClient
	}
	base := inner.Transport
	if base == nil {
		base = http.DefaultTransport
	}

	wrapped := *inner // shallow copy: keep Jar, Timeout, redirect policy
	wrapped.Transport = &refreshDetachTransport{base: base, tokenEndpoint: endpoint, timeout: refreshDetachTimeout}
	return &wrapped
}

// sameEndpoint reports whether two URLs address the same endpoint, comparing
// scheme, host, and path rather than raw strings so that canonicalization
// differences (default ports, escaping) don't cause a miss.
func sameEndpoint(a, b *url.URL) bool {
	if a == nil || b == nil {
		return false
	}
	return a.Scheme == b.Scheme && a.Host == b.Host && a.Path == b.Path
}