diff --git a/README.md b/README.md index e3e680f9..09d69b39 100644 --- a/README.md +++ b/README.md @@ -133,7 +133,7 @@ Instead of writing `.proto` files by hand, each SDK ships a tool to generate pro ### Arrow Flight ingestion (Beta) -Available in the Rust, Python, Go, TypeScript, and Java SDKs starting from version 2.0.0, and in the C++ SDK from its initial `0.1.0` release. Currently in Beta — the API is stabilising but may still change before reaching GA. A third record format option alongside JSON and Protocol Buffers: send Apache Arrow `RecordBatch` data directly to Zerobus over the Arrow Flight protocol, on the same gRPC connection. Best fit when: +Available in the Rust, Python, Go, TypeScript, and Java SDKs starting from their 2.0.0 releases, and in the C++ SDK from its initial `0.1.0` release. Currently in Beta — the API is stabilising but may still change before reaching GA. A third record format option alongside JSON and Protocol Buffers: send Apache Arrow `RecordBatch` data directly to Zerobus over the Arrow Flight protocol, on the same gRPC connection. Best fit when: - Your workload is naturally columnar or batched — analytics pipelines, gateways aggregating short windows of rows, wide/numeric schemas where row-by-row serialization adds noticeable CPU overhead. - Your application already produces Arrow data — pyarrow, the [arrow-rs](https://github.com/apache/arrow-rs) crates, DataFusion, Polars, or other libraries built on Arrow. diff --git a/cpp/include/zerobus/arrow_stream.hpp b/cpp/include/zerobus/arrow_stream.hpp index 36de4e27..ac820101 100644 --- a/cpp/include/zerobus/arrow_stream.hpp +++ b/cpp/include/zerobus/arrow_stream.hpp @@ -64,10 +64,10 @@ class ArrowStream { /// Flush all pending batches and wait for their acknowledgment. void flush(); - /// Return all unacknowledged batches from a closed or failed stream, each as - /// a self-contained Arrow IPC stream (schema + one batch). Remains callable - /// after a failed `close()`, which keeps the handle alive so recovery is - /// possible. + /// Return all unacknowledged batches from a failed stream, each as a + /// self-contained Arrow IPC stream (schema + one batch). Callable after a + /// failed `close()` (the handle stays alive for recovery); after a successful + /// `close()` the handle is freed, so this throws instead. std::vector> get_unacked_batches(); /// Gracefully close the stream, flushing pending batches first. Idempotent. diff --git a/purego/README.md b/purego/README.md index e54a2d43..d3e2b219 100644 --- a/purego/README.md +++ b/purego/README.md @@ -20,9 +20,14 @@ Implemented packages: `RecordType`, descriptor requirement for `PROTO`) and sets auth metadata from `StreamParams.Token` (`"Bearer "` for bare tokens, verbatim when a known scheme — `Bearer`, `Basic`, or `DPoP` — is already present). - -Planned layers: OAuth/auth, ingest/ack state management, recovery, and the -public API. +- `internal/auth` — token providers for stream authentication. `OAuthTokenProvider` + implements the Unity Catalog OAuth 2.0 client-credentials flow with per-table + token caching and proactive refresh; `StaticTokenProvider` wraps a fixed token + for tests or externally managed lifecycles. The package also exposes + `HeadersProvider` variants (`OAuthHeadersProvider`, `StaticHeadersProvider`) + so transport open can authenticate via provider callbacks like other SDKs. + +Planned layers: ingest/ack state management, recovery, and the public API. ## Regenerating the protobuf bindings diff --git a/purego/internal/auth/oauth.go b/purego/internal/auth/oauth.go new file mode 100644 index 00000000..e9905216 --- /dev/null +++ b/purego/internal/auth/oauth.go @@ -0,0 +1,550 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "math" + "net" + "net/http" + "net/url" + "strings" + "time" + "unicode" +) + +// OAuthTokenProvider obtains Unity Catalog OAuth 2.0 tokens using the client +// credentials flow and caches them until they approach expiry. +// +// UC issues downscoped tokens — each token is minted with authorization_details +// that limit it to a specific table — so the table name is part of the cache +// key. A single OAuthTokenProvider can serve many tables; tokens are cached +// independently per table. +// +// Refresh is on-access, not background: a token is re-minted on the next Token +// call once it is within the 5-minute window before expiry. If that refresh +// fails transiently the still-valid cached token is served; a non-retryable +// failure (revoked credentials, 4xx) propagates. A token left untouched past +// expiry is not refreshed until the next call, which pays a full cold mint. +// +// OAuthTokenProvider is safe for concurrent use. +type OAuthTokenProvider struct { + clientID string + clientSecret string + workspaceID string + zerobusHost string + ucEndpoint string // e.g. "https://workspace.databricks.com" + + cache *tokenCache + client *http.Client + logger *slog.Logger + + // cacheOpts are applied to the provider's own default cache. They are + // discarded when a shared cache is installed via WithSharedTokenCache, since + // that cache owns its own configuration. + cacheOpts []CacheOption +} + +// OAuthOption configures an [OAuthTokenProvider]. +type OAuthOption func(*OAuthTokenProvider) + +// WithHTTPClient overrides the HTTP client used for token requests. The default +// is a client with a 30-second timeout. A nil client is ignored so the default +// is preserved. +func WithHTTPClient(c *http.Client) OAuthOption { + return func(p *OAuthTokenProvider) { + if c != nil { + p.client = c + } + } +} + +// SharedTokenCache is an OAuth token cache shared across multiple +// [OAuthTokenProvider] instances via [WithSharedTokenCache]. Obtain one with +// [NewSharedTokenCache]. It exposes no methods and is safe for concurrent use. +type SharedTokenCache = tokenCache + +// WithSharedTokenCache installs a shared token cache. Use this when multiple +// OAuthTokenProviders are created for the same credentials but different tables +// so they share one cache map rather than each allocating their own. +// +// Obtain a cache with [NewSharedTokenCache]. A nil cache is ignored so the +// provider's own default cache is preserved. +func WithSharedTokenCache(cache *SharedTokenCache) OAuthOption { + return func(p *OAuthTokenProvider) { + if cache != nil { + p.cache = cache + } + } +} + +// WithTokenCacheEnabled toggles token caching for the provider's own cache. When +// disabled, every [OAuthTokenProvider.Token] call mints a fresh token. Caching +// is enabled by default. Ignored when a shared cache is installed via +// [WithSharedTokenCache] — that cache carries its own configuration. +func WithTokenCacheEnabled(enabled bool) OAuthOption { + return func(p *OAuthTokenProvider) { + p.cacheOpts = append(p.cacheOpts, CacheEnabled(enabled)) + } +} + +// WithRefreshBuffer sets the lead time before expiry at which the provider's own +// cache proactively re-mints a token. Defaults to 5 minutes; a non-positive +// value is ignored. Ignored when a shared cache is installed via +// [WithSharedTokenCache]. +func WithRefreshBuffer(d time.Duration) OAuthOption { + return func(p *OAuthTokenProvider) { + p.cacheOpts = append(p.cacheOpts, CacheRefreshBuffer(d)) + } +} + +// WithLogger sets the [slog.Logger] used for token-mint observability (mint +// reason, latency, retryability). A nil logger is ignored so the default +// (a no-op logger) is preserved. +func WithLogger(l *slog.Logger) OAuthOption { + return func(p *OAuthTokenProvider) { + if l != nil { + p.logger = l + } + } +} + +// NewOAuthTokenProvider creates a provider that mints UC OAuth tokens using the +// client credentials flow. +// +// zerobusEndpoint is the Zerobus service URL. The workspace ID used in OAuth +// resource audience is derived from its host prefix (matching the Rust SDK). +// ucEndpoint is the workspace URL (e.g. "https://my-workspace.databricks.com"). +func NewOAuthTokenProvider( + clientID, clientSecret, zerobusEndpoint, ucEndpoint string, + opts ...OAuthOption, +) (*OAuthTokenProvider, error) { + ucEndpoint = strings.TrimRight(strings.TrimSpace(ucEndpoint), "/") + if ucEndpoint == "" { + return nil, fmt.Errorf("auth: oauth: ucEndpoint is required") + } + if err := validateEndpoint(ucEndpoint); err != nil { + return nil, fmt.Errorf("auth: oauth: %w", err) + } + if clientID == "" { + return nil, fmt.Errorf("auth: oauth: clientID is required") + } + if clientSecret == "" { + return nil, fmt.Errorf("auth: oauth: clientSecret is required") + } + workspaceID, zerobusHost, err := deriveWorkspaceIDFromEndpoint(zerobusEndpoint) + if err != nil { + return nil, fmt.Errorf("auth: oauth: %w", err) + } + + p := &OAuthTokenProvider{ + clientID: clientID, + clientSecret: clientSecret, + workspaceID: workspaceID, + zerobusHost: zerobusHost, + ucEndpoint: ucEndpoint, + client: &http.Client{Timeout: 30 * time.Second}, + logger: slog.New(slog.DiscardHandler), + } + for _, opt := range opts { + opt(p) + } + // A shared cache (set via WithSharedTokenCache) owns its own configuration; + // otherwise build the provider's own cache from any cache options collected. + if p.cache == nil { + p.cache = newTokenCache(p.cacheOpts...) + } + return p, nil +} + +// NewSharedTokenCache allocates a token cache that can be passed to multiple +// [OAuthTokenProvider] instances via [WithSharedTokenCache]. Sharing a cache +// ensures tokens for the same (clientID, secret, table) are reused across +// providers rather than each minting independently. +// +// Configure it with [CacheEnabled] and [CacheRefreshBuffer]. +func NewSharedTokenCache(opts ...CacheOption) *SharedTokenCache { + return newTokenCache(opts...) +} + +// Token returns a valid bearer token for tableName, minting a new one via +// Unity Catalog's OIDC token endpoint when the cache is empty or nearing +// expiry. +// +// ctx bounds the token request when a mint is required. Cancelling ctx during a +// cached hit is a no-op. +// +// A successful mint is cached only when UC reports a usable expires_in. If UC +// omits expires_in (or reports a non-positive value), the token is returned but +// not cached to match Rust SDK behavior. +func (p *OAuthTokenProvider) Token(ctx context.Context, tableName string) (string, error) { + tableName = strings.TrimSpace(tableName) + if err := validateTableName(tableName); err != nil { + return "", fmt.Errorf("auth: oauth: %w", err) + } + return p.cache.getOrFetch(ctx, p.clientID, p.clientSecret, tableName, + func(ctx context.Context, reason mintReason) (fetchedToken, error) { + return p.mint(ctx, tableName, reason) + }, + ) +} + +// FetchToken mints a token directly from Unity Catalog, bypassing the cache. It +// neither reads nor writes the cache, so callers get a guaranteed-fresh token; +// most callers should use [OAuthTokenProvider.Token] instead so tokens are +// reused across streams. +func (p *OAuthTokenProvider) FetchToken(ctx context.Context, tableName string) (string, error) { + tableName = strings.TrimSpace(tableName) + if err := validateTableName(tableName); err != nil { + return "", fmt.Errorf("auth: oauth: %w", err) + } + fetched, err := p.mint(ctx, tableName, mintReasonDirect) + if err != nil { + return "", err + } + return fetched.token, nil +} + +// mint fetches a token and emits structured observability for the outcome. It +// is the single mint entry point shared by the cached and direct paths. +func (p *OAuthTokenProvider) mint(ctx context.Context, tableName string, reason mintReason) (fetchedToken, error) { + started := time.Now() + fetched, err := p.fetchToken(ctx, tableName) + elapsed := time.Since(started) + + switch { + case err != nil: + p.logger.LogAttrs(ctx, slog.LevelWarn, "failed to mint UC OAuth token", + slog.String("table", tableName), + slog.String("reason", reason.String()), + slog.Bool("retryable", isRetryable(err)), + slog.Duration("elapsed", elapsed), + slog.String("error", err.Error()), + ) + case fetched.expiresIn == nil: + p.logger.LogAttrs(ctx, slog.LevelWarn, "minted UC OAuth token but UC returned no expires_in; token will not be cached", + slog.String("table", tableName), + slog.String("reason", reason.String()), + slog.Duration("elapsed", elapsed), + ) + default: + p.logger.LogAttrs(ctx, slog.LevelInfo, "minted UC OAuth token", + slog.String("table", tableName), + slog.String("reason", reason.String()), + slog.Duration("expires_in", *fetched.expiresIn), + slog.Duration("elapsed", elapsed), + ) + } + return fetched, err +} + +// Invalidate drops the cached token for this provider's table so the next +// Token call re-mints from Unity Catalog. Call this when the server rejects +// the token with an authentication error. +func (p *OAuthTokenProvider) Invalidate(_ context.Context, tableName string) { + tableName = strings.TrimSpace(tableName) + if err := validateTableName(tableName); err != nil { + return + } + p.cache.invalidate(p.clientID, p.clientSecret, tableName) +} + +// fetchToken performs the OAuth 2.0 client credentials request against Unity +// Catalog's OIDC token endpoint. +func (p *OAuthTokenProvider) fetchToken(ctx context.Context, tableName string) (fetchedToken, error) { + catalog, schema, _, err := parseTableName(tableName) + if err != nil { + return fetchedToken{}, err + } + + authDetails, err := buildAuthorizationDetails(catalog, schema, tableName) + if err != nil { + return fetchedToken{}, fmt.Errorf("auth: oauth: marshal authorization_details: %w", err) + } + + form := url.Values{ + "grant_type": {"client_credentials"}, + "scope": {"all-apis"}, + "resource": {fmt.Sprintf("api://databricks/workspaces/%s/zerobusDirectWriteApi", p.workspaceID)}, + "authorization_details": {authDetails}, + } + + tokenURL := p.ucEndpoint + "/oidc/v1/token" + req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(form.Encode())) + if err != nil { + return fetchedToken{}, &TokenError{msg: fmt.Sprintf("build token request: %v", err), retryable: false, cause: err} + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.SetBasicAuth(p.clientID, p.clientSecret) + + resp, err := p.client.Do(req) + if err != nil { + return fetchedToken{}, &TokenError{msg: fmt.Sprintf("token request: %v", err), retryable: isRetryableTransportError(err), cause: err} + } + // Drain any unread bytes before closing so the keep-alive connection can be + // reused: json.Decode stops at the end of the JSON value and classifyHTTPError + // caps its read, either of which can leave a trailing tail on the body. + defer func() { + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + }() + + if !isHTTPSuccess(resp.StatusCode) { + return fetchedToken{}, classifyHTTPError(resp) + } + + var body struct { + AccessToken string `json:"access_token"` + // RFC 6749 §4.2.2 defines expires_in as an integer number of seconds. + ExpiresIn int64 `json:"expires_in"` + } + if err := json.NewDecoder(resp.Body).Decode(&body); err != nil { + return fetchedToken{}, &TokenError{ + msg: fmt.Sprintf("parse token response: %v", err), + retryable: false, + cause: err, + } + } + if body.AccessToken == "" { + return fetchedToken{}, &TokenError{msg: "token response missing access_token", retryable: false} + } + if !isUsableAsHeader(body.AccessToken) { + return fetchedToken{}, &TokenError{msg: "access token contains invalid header characters", retryable: false} + } + + ft := fetchedToken{token: body.AccessToken} + if d, ok := secondsToDuration(body.ExpiresIn); ok { + ft.expiresIn = &d + } + return ft, nil +} + +// maxExpiresInSeconds bounds a server-reported expires_in so the seconds→ +// nanoseconds conversion can't overflow time.Duration (an int64 nanosecond +// count, ~292 years). A larger value is clamped to this ceiling rather than +// wrapping to a negative (past) TTL. +const maxExpiresInSeconds = int64(math.MaxInt64 / int64(time.Second)) + +// secondsToDuration converts a positive expires_in (seconds) to a Duration, +// clamped to maxExpiresInSeconds. It returns ok=false for a non-positive value, +// which signals "no usable TTL, don't cache". +func secondsToDuration(secs int64) (time.Duration, bool) { + if secs <= 0 { + return 0, false + } + if secs > maxExpiresInSeconds { + secs = maxExpiresInSeconds + } + return time.Duration(secs) * time.Second, true +} + +// authorizationDetailsEntry mirrors the JSON structure sent in the OAuth request. +type authorizationDetailsEntry struct { + Type string `json:"type"` + Privileges []string `json:"privileges"` + ObjectType string `json:"object_type"` + ObjectFullPath string `json:"object_full_path"` + Operations []string `json:"operations,omitempty"` +} + +func buildAuthorizationDetails(catalog, schema, fullTable string) (string, error) { + details := []authorizationDetailsEntry{ + { + Type: "unity_catalog_privileges", + Privileges: []string{"USE CATALOG"}, + ObjectType: "CATALOG", + ObjectFullPath: catalog, + }, + { + Type: "unity_catalog_privileges", + Privileges: []string{"USE SCHEMA"}, + ObjectType: "SCHEMA", + ObjectFullPath: catalog + "." + schema, + }, + { + Type: "unity_catalog_privileges", + Privileges: []string{"SELECT", "MODIFY"}, + ObjectType: "TABLE", + ObjectFullPath: fullTable, + Operations: []string{"zerobuswrite"}, + }, + } + b, err := json.Marshal(details) + if err != nil { + return "", err + } + return string(b), nil +} + +// TokenError is returned by [OAuthTokenProvider.Token] when a mint fails. +// It carries a retryability flag so the cache can decide whether to suppress +// the error and serve a still-valid cached token, and wraps the underlying +// cause (if any) so callers can inspect it with [errors.Is]/[errors.As]. +type TokenError struct { + msg string + retryable bool + cause error +} + +func (e *TokenError) Error() string { return "auth: oauth: " + e.msg } +func (e *TokenError) IsRetryable() bool { return e.retryable } + +// Unwrap exposes the cause so [errors.Is]/[errors.As] can inspect it. +func (e *TokenError) Unwrap() error { return e.cause } + +// isRetryableStatus reports whether an HTTP status is a transient failure worth +// suppressing when a cached token can be served. Only 5xx responses qualify; +// all 4xx (including 429 and 408) are non-retryable, matching the Rust SDK. +func isRetryableStatus(code int) bool { + return code >= 500 +} + +func classifyHTTPError(resp *http.Response) error { + // Best-effort read of the error payload, bounded so a misbehaving server + // can't stream an unbounded body into the error message. + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + detail := strings.TrimSpace(string(body)) + msg := fmt.Sprintf("HTTP %d", resp.StatusCode) + if detail != "" { + msg += ": " + detail + } + return &TokenError{msg: msg, retryable: isRetryableStatus(resp.StatusCode)} +} + +func isHTTPSuccess(code int) bool { return code >= 200 && code < 300 } + +// isRetryableTransportError reports whether a transport-level error from the +// token request is transient and safe to retry. Only timeouts and connection +// failures qualify; a cancelled context is the caller's own signal, and any +// other transport error is treated as non-retryable so a genuine failure isn't +// masked by a stale cached token. +func isRetryableTransportError(err error) bool { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + // Timeouts (client Timeout, dial/read deadlines) are transient. + var ne net.Error + if errors.As(err, &ne) && ne.Timeout() { + return true + } + // A failed dial/connection (refused, reset, unreachable) is transient too. + var oe *net.OpError + return errors.As(err, &oe) +} + +// isUsableAsHeader reports whether token can be sent as a gRPC/HTTP +// authorization header value. gRPC metadata rejects values with ASCII control +// characters (bytes 0–31 and DEL). +// +// Note: transport.isUsableAsHeader applies the same character check but +// intentionally treats the empty string as usable ("no header"); this copy +// returns false for empty because a minted token must be non-empty. Keep the +// character logic in sync if either is changed. +func isUsableAsHeader(token string) bool { + for _, r := range token { + if r > unicode.MaxASCII || unicode.IsControl(r) { + return false + } + } + return token != "" +} + +// validateEndpoint requires an https UC endpoint so client credentials never +// go over plaintext; plain http is allowed only for loopback hosts (local dev). +func validateEndpoint(endpoint string) error { + u, err := url.Parse(endpoint) + if err != nil { + return fmt.Errorf("ucEndpoint is not a valid URL: %w", err) + } + // A missing host (e.g. "https://") would otherwise pass scheme validation and + // only fail later at request time with an opaque transport error; reject it + // here so misconfiguration fails fast in the constructor. + if u.Hostname() == "" { + return fmt.Errorf("ucEndpoint has no host: %q", endpoint) + } + switch u.Scheme { + case "https": + return nil + case "http": + if isLoopbackHost(u.Hostname()) { + return nil + } + return fmt.Errorf("ucEndpoint must use https (got plaintext http for non-loopback host %q); "+ + "client credentials would be sent over an unencrypted connection", u.Host) + default: + return fmt.Errorf("ucEndpoint must be an https URL, got scheme %q", u.Scheme) + } +} + +// isLoopbackHost reports whether host is "localhost" or a loopback IP. +// The "localhost" match is case-insensitive since DNS hostnames are. +func isLoopbackHost(host string) bool { + if strings.EqualFold(host, "localhost") { + return true + } + if ip := net.ParseIP(host); ip != nil { + return ip.IsLoopback() + } + return false +} + +// deriveWorkspaceIDFromEndpoint extracts workspace ID from Zerobus endpoint +// host by taking the first DNS label, matching Rust SDK behavior. +func deriveWorkspaceIDFromEndpoint(endpoint string) (workspaceID, host string, err error) { + endpoint = strings.TrimSpace(endpoint) + if endpoint == "" { + return "", "", fmt.Errorf("zerobusEndpoint is required") + } + if !strings.HasPrefix(endpoint, "https://") && !strings.HasPrefix(endpoint, "http://") { + endpoint = "https://" + endpoint + } + u, err := url.Parse(endpoint) + if err != nil { + return "", "", fmt.Errorf("zerobusEndpoint is not a valid URL: %w", err) + } + host = u.Hostname() + if host == "" { + return "", "", fmt.Errorf("zerobusEndpoint has no host: %q", endpoint) + } + workspaceID, _, _ = strings.Cut(host, ".") + if workspaceID == "" { + return "", "", fmt.Errorf("failed to extract workspaceID from zerobusEndpoint host %q", host) + } + return workspaceID, host, nil +} + +// validateTableName checks that s is a non-empty three-part dotted name. +func validateTableName(s string) error { + _, _, _, err := parseTableName(s) + return err +} + +// parseTableName splits "catalog.schema.table" into its three components, +// returning an error if any part is empty or the format is wrong. +// +// Only unquoted three-part names are supported: the split is purely on ".", so +// a backtick-quoted identifier containing a dot (e.g. cat.sch.`weird.name`) is +// rejected. This matches the naming used for downscoped-token tables. +func parseTableName(s string) (catalog, schema, table string, err error) { + // SplitN with n=4 caps the slice at 4 parts: a valid name yields exactly 3, + // and anything with an extra dot yields 4 (rejected below) instead of + // splitting the whole string. The table part itself must not contain a dot. + parts := strings.SplitN(s, ".", 4) + if len(parts) != 3 { + return "", "", "", fmt.Errorf("table name must be catalog.schema.table, got %q", s) + } + catalog, schema, table = parts[0], parts[1], parts[2] + if catalog == "" { + return "", "", "", fmt.Errorf("catalog part of table name is empty in %q", s) + } + if schema == "" { + return "", "", "", fmt.Errorf("schema part of table name is empty in %q", s) + } + if table == "" { + return "", "", "", fmt.Errorf("table part of table name is empty in %q", s) + } + return catalog, schema, table, nil +} diff --git a/purego/internal/auth/oauth_test.go b/purego/internal/auth/oauth_test.go new file mode 100644 index 00000000..0de1825c --- /dev/null +++ b/purego/internal/auth/oauth_test.go @@ -0,0 +1,609 @@ +package auth + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +// tokenServer is a minimal OAuth 2.0 token endpoint stub. +type tokenServer struct { + accessToken string + expiresIn int // 0 means omit expires_in + statusCode int // 0 defaults to 200 + calls atomic.Int32 +} + +func (s *tokenServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + s.calls.Add(1) + code := s.statusCode + if code == 0 { + code = http.StatusOK + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + if code != http.StatusOK { + _ = json.NewEncoder(w).Encode(map[string]string{"error": "test_error"}) + return + } + resp := map[string]any{"access_token": s.accessToken} + if s.expiresIn > 0 { + resp["expires_in"] = s.expiresIn + } + _ = json.NewEncoder(w).Encode(resp) +} + +func newTestProvider(t *testing.T, srv *tokenServer, table string) (*OAuthTokenProvider, *httptest.Server) { + t.Helper() + ts := httptest.NewServer(srv) + p, err := NewOAuthTokenProvider("clientID", "clientSecret", "https://ws123.zerobus.databricks.com", ts.URL, + WithHTTPClient(ts.Client()), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + return p, ts +} + +func TestOAuthTokenProviderHappyPath(t *testing.T) { + srv := &tokenServer{accessToken: "eyJtb2NrfQ", expiresIn: 3600} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + tok, err := p.Token(context.Background(), "cat.sch.tbl") + if err != nil { + t.Fatalf("Token: %v", err) + } + if tok != "eyJtb2NrfQ" { + t.Fatalf("want %q, got %q", "eyJtb2NrfQ", tok) + } +} + +func TestOAuthTokenProviderCachesToken(t *testing.T) { + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("second Token: %v", err) + } + if got := srv.calls.Load(); got != 1 { + t.Fatalf("want 1 server call (cached), got %d", got) + } +} + +func TestOAuthTokenProviderInvalidateForcesRemint(t *testing.T) { + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("first Token: %v", err) + } + p.Invalidate(context.Background(), "cat.sch.tbl") + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("Token after Invalidate: %v", err) + } + if got := srv.calls.Load(); got != 2 { + t.Fatalf("want 2 server calls after invalidate, got %d", got) + } +} + +func TestOAuthTokenProvider5xxIsRetryable(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusInternalServerError} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background(), "cat.sch.tbl") + if err == nil { + t.Fatal("want error for 500, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || !te.IsRetryable() { + t.Fatalf("want retryable TokenError, got %T: %v", err, err) + } +} + +func TestOAuthTokenProvider4xxIsNonRetryable(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusUnauthorized} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background(), "cat.sch.tbl") + if err == nil { + t.Fatal("want error for 401, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || te.IsRetryable() { + t.Fatalf("want non-retryable TokenError, got %T (retryable=%v): %v", err, te != nil && te.IsRetryable(), err) + } +} + +func TestOAuthTokenProviderContextCancellation(t *testing.T) { + // unblock is closed when the test ends so the hanging handler exits before + // httptest.Server.Close() tries to drain in-flight requests. + unblock := make(chan struct{}) + hang := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-unblock: + case <-r.Context().Done(): + } + http.Error(w, "cancelled", http.StatusServiceUnavailable) + })) + defer hang.Close() // executed second: server is now idle + defer close(unblock) // executed first: unblocks the handler + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", hang.URL, + WithHTTPClient(hang.Client()), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err = p.Token(ctx, "c.s.t") + if err == nil { + t.Fatal("want error on ctx cancellation, got nil") + } + if elapsed := time.Since(start); elapsed > 1*time.Second { + t.Fatalf("Token took %v, expected it to return near 100ms deadline", elapsed) + } +} + +func TestOAuthTokenProviderConnectionRefusedIsRetryable(t *testing.T) { + // Point the provider at a closed port so Do() fails with a dial error, which + // must be classified as a retryable transport failure. + ts := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + client := ts.Client() + url := ts.URL + ts.Close() // close immediately so connections are refused + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", url, + WithHTTPClient(client), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + _, err = p.Token(context.Background(), "c.s.t") + if err == nil { + t.Fatal("want error for refused connection, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || !te.IsRetryable() { + t.Fatalf("want retryable TokenError for dial failure, got %T (retryable=%v): %v", + err, te != nil && te.IsRetryable(), err) + } +} + +func TestOAuthTokenProvider429IsNonRetryable(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusTooManyRequests} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background(), "cat.sch.tbl") + if err == nil { + t.Fatal("want error for 429, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || te.IsRetryable() { + t.Fatalf("want non-retryable TokenError for 429, got %T (retryable=%v): %v", + err, te != nil && te.IsRetryable(), err) + } +} + +func TestOAuthTokenProvider408IsNonRetryable(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusRequestTimeout} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background(), "cat.sch.tbl") + if err == nil { + t.Fatal("want error for 408, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || te.IsRetryable() { + t.Fatalf("want non-retryable TokenError for 408, got %T (retryable=%v): %v", + err, te != nil && te.IsRetryable(), err) + } +} + +func TestTokenErrorUnwrapsContextCancellation(t *testing.T) { + // A cancelled context must be visible via errors.Is on the returned error, + // and must NOT be classified as retryable (else a cancelled proactive + // refresh would silently serve a stale token). + unblock := make(chan struct{}) + hang := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-unblock: + case <-r.Context().Done(): + } + })) + defer hang.Close() + defer close(unblock) + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", hang.URL, + WithHTTPClient(hang.Client()), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + _, err = p.Token(ctx, "c.s.t") + if err == nil { + t.Fatal("want error on ctx cancellation, got nil") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("want errors.Is(err, DeadlineExceeded), got %v", err) + } + var te *TokenError + if asTokenError(err, &te) && te.IsRetryable() { + t.Fatal("context deadline must not be classified as retryable") + } +} + +func TestOAuthTokenProviderNilOptionsKeepDefaults(t *testing.T) { + // A nil client / cache passed via options must be ignored so the provider's + // defaults survive rather than being overwritten with nil (which would panic + // on the first Token call). + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + ts := httptest.NewServer(srv) + defer ts.Close() + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, + WithHTTPClient(nil), + WithSharedTokenCache(nil), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if p.client == nil { + t.Fatal("WithHTTPClient(nil) overwrote the default client") + } + if p.cache == nil { + t.Fatal("WithSharedTokenCache(nil) overwrote the default cache") + } + // It must not panic and must actually work end to end. + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("Token with nil options: %v", err) + } +} + +func TestClassifyHTTPErrorPreservesBody(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusUnauthorized} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background(), "cat.sch.tbl") + if err == nil { + t.Fatal("want error for 401, got nil") + } + // The server error payload (`{"error":"test_error"}`) must survive into the + // error message so on-call has something to debug with. + if !strings.Contains(err.Error(), "test_error") { + t.Fatalf("error message dropped server body: %q", err.Error()) + } + if !strings.Contains(err.Error(), "401") { + t.Fatalf("error message dropped status code: %q", err.Error()) + } +} + +func TestNewOAuthTokenProviderValidation(t *testing.T) { + cases := []struct { + name string + clientID, secret, zerobusEndpoint, ucEndpoint, table string + }{ + {"empty clientID", "", "s", "https://ws.zerobus.databricks.com", "https://host", "c.s.t"}, + {"empty secret", "id", "", "https://ws.zerobus.databricks.com", "https://host", "c.s.t"}, + {"empty zerobus endpoint", "id", "s", "", "https://host", "c.s.t"}, + {"empty ucEndpoint", "id", "s", "https://ws.zerobus.databricks.com", "", "c.s.t"}, + {"bad table (2 parts)", "id", "s", "https://ws.zerobus.databricks.com", "https://host", "c.s"}, + {"bad table (empty schema)", "id", "s", "https://ws.zerobus.databricks.com", "https://host", "c..t"}, + {"bad zerobus endpoint host", "id", "s", "https://", "https://host", "c.s.t"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + p, err := NewOAuthTokenProvider(tc.clientID, tc.secret, tc.zerobusEndpoint, tc.ucEndpoint) + if err == nil { + if _, tokErr := p.Token(context.Background(), tc.table); tokErr == nil { + t.Fatal("want constructor or token validation error, got nil") + } + } + }) + } +} + +func TestSecondsToDuration(t *testing.T) { + if d, ok := secondsToDuration(3600); !ok || d != time.Hour { + t.Fatalf("secondsToDuration(3600) = (%v, %v), want (1h, true)", d, ok) + } + if _, ok := secondsToDuration(0); ok { + t.Fatal("secondsToDuration(0) should report no usable TTL") + } + if _, ok := secondsToDuration(-5); ok { + t.Fatal("secondsToDuration(negative) should report no usable TTL") + } + // An absurd value must clamp to a positive duration, never wrap negative. + d, ok := secondsToDuration(1 << 60) + if !ok || d <= 0 { + t.Fatalf("secondsToDuration(1<<60) = (%v, %v), want a positive clamped duration", d, ok) + } +} + +func TestValidateEndpoint(t *testing.T) { + ok := []string{ + "https://workspace.databricks.com", + "http://localhost:8080", + "http://LOCALHOST:8080", // hostnames are case-insensitive + "http://127.0.0.1:1234", + "http://[::1]:9000", + } + for _, e := range ok { + if err := validateEndpoint(e); err != nil { + t.Errorf("validateEndpoint(%q): unexpected error: %v", e, err) + } + } + + bad := []string{ + "http://workspace.databricks.com", // plaintext to a remote host + "ftp://host", + "://nonsense", + "https://", // scheme but no host + "https:///", // no host, only a path + } + for _, e := range bad { + if err := validateEndpoint(e); err == nil { + t.Errorf("validateEndpoint(%q): want error, got nil", e) + } + } +} + +func TestOAuthTokenProviderCacheDisabledAlwaysMints(t *testing.T) { + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + ts := httptest.NewServer(srv) + defer ts.Close() + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, + WithHTTPClient(ts.Client()), + WithTokenCacheEnabled(false), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("second Token: %v", err) + } + if got := srv.calls.Load(); got != 2 { + t.Fatalf("want 2 server calls with caching disabled, got %d", got) + } +} + +func TestOAuthTokenProviderCustomRefreshBuffer(t *testing.T) { + // WithRefreshBuffer must reach the provider's own cache. (The refresh-timing + // behavior itself is covered by the tokenCache unit tests; a large buffer no + // longer forces a re-mint on every call because it is clamped to ttl/2.) + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", "https://host", + WithRefreshBuffer(90*time.Second), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if p.cache.refreshBuffer != 90*time.Second { + t.Fatalf("WithRefreshBuffer not applied: got %v", p.cache.refreshBuffer) + } +} + +func TestOAuthTokenProviderSharedCacheIgnoresProviderCacheOpts(t *testing.T) { + // A shared cache owns its own config; per-provider cache options are ignored + // so the shared cache still caches even though the provider asked to disable. + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + ts := httptest.NewServer(srv) + defer ts.Close() + + shared := NewSharedTokenCache() + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, + WithHTTPClient(ts.Client()), + WithSharedTokenCache(shared), + WithTokenCacheEnabled(false), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("second Token: %v", err) + } + if got := srv.calls.Load(); got != 1 { + t.Fatalf("shared cache should cache (1 call), got %d", got) + } +} + +func TestSharedTokenCacheDisabled(t *testing.T) { + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + ts := httptest.NewServer(srv) + defer ts.Close() + + shared := NewSharedTokenCache(CacheEnabled(false)) + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, + WithHTTPClient(ts.Client()), + WithSharedTokenCache(shared), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("second Token: %v", err) + } + if got := srv.calls.Load(); got != 2 { + t.Fatalf("want 2 calls with disabled shared cache, got %d", got) + } +} + +func TestOAuthTokenProviderFetchTokenBypassesCache(t *testing.T) { + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + // Warm the cache. + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("Token: %v", err) + } + // FetchToken must mint fresh regardless of the warm cache… + if _, err := p.FetchToken(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("FetchToken: %v", err) + } + if got := srv.calls.Load(); got != 2 { + t.Fatalf("FetchToken should bypass cache (2 calls), got %d", got) + } + // …and must not populate/replace the cached entry the next Token call uses. + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { + t.Fatalf("Token after FetchToken: %v", err) + } + if got := srv.calls.Load(); got != 2 { + t.Fatalf("Token after FetchToken should hit cache (still 2 calls), got %d", got) + } +} + +func TestOAuthTokenProviderLoggerReceivesMint(t *testing.T) { + var buf strings.Builder + logger := slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelInfo})) + + srv := &tokenServer{accessToken: "tok", expiresIn: 3600} + ts := httptest.NewServer(srv) + defer ts.Close() + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, + WithHTTPClient(ts.Client()), + WithLogger(logger), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { + t.Fatalf("Token: %v", err) + } + out := buf.String() + if !strings.Contains(out, "minted UC OAuth token") || !strings.Contains(out, "reason=cold_miss") { + t.Fatalf("logger did not receive expected mint line: %q", out) + } +} + +func TestAuthorizationDetailsContent(t *testing.T) { + var captured []byte + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := r.ParseForm(); err == nil { + captured = []byte(r.FormValue("authorization_details")) + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"access_token": "tok", "expires_in": 3600}) + })) + defer ts.Close() + + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, + WithHTTPClient(ts.Client()), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background(), "mycat.mysch.mytbl"); err != nil { + t.Fatalf("Token: %v", err) + } + + var details []authorizationDetailsEntry + if err := json.Unmarshal(captured, &details); err != nil { + t.Fatalf("parse authorization_details: %v", err) + } + if len(details) != 3 { + t.Fatalf("want 3 authorization_details entries, got %d", len(details)) + } + + catalog := details[0] + if catalog.ObjectFullPath != "mycat" { + t.Errorf("catalog entry: want ObjectFullPath %q, got %q", "mycat", catalog.ObjectFullPath) + } + + schema := details[1] + if schema.ObjectFullPath != "mycat.mysch" { + t.Errorf("schema entry: want ObjectFullPath %q, got %q", "mycat.mysch", schema.ObjectFullPath) + } + + table := details[2] + if table.ObjectFullPath != "mycat.mysch.mytbl" { + t.Errorf("table entry: want ObjectFullPath %q, got %q", "mycat.mysch.mytbl", table.ObjectFullPath) + } + if len(table.Operations) != 1 || table.Operations[0] != "zerobuswrite" { + t.Errorf("table entry: want Operations [zerobuswrite], got %v", table.Operations) + } +} + +func TestIsUsableAsHeader(t *testing.T) { + cases := []struct { + token string + want bool + }{ + {"eyJhbGciOiJSUzI1NiJ9.payload.sig", true}, + {"", false}, + {"bad\ntoken", false}, + {"bad\x00token", false}, + {strings.Repeat("a", 1000), true}, + } + for _, tc := range cases { + if got := isUsableAsHeader(tc.token); got != tc.want { + t.Errorf("isUsableAsHeader(%q) = %v, want %v", tc.token, got, tc.want) + } + } +} + +func TestParseTableName(t *testing.T) { + good := []struct{ in, c, s, tbl string }{ + {"a.b.c", "a", "b", "c"}, + {"cat.schema.table", "cat", "schema", "table"}, + } + for _, tc := range good { + c, s, tbl, err := parseTableName(tc.in) + if err != nil { + t.Errorf("parseTableName(%q): unexpected error: %v", tc.in, err) + continue + } + if c != tc.c || s != tc.s || tbl != tc.tbl { + t.Errorf("parseTableName(%q) = (%q,%q,%q), want (%q,%q,%q)", tc.in, c, s, tbl, tc.c, tc.s, tc.tbl) + } + } + + bad := []string{"", "a", "a.b", "a.b.c.d", ".b.c", "a..c", "a.b."} + for _, tc := range bad { + if _, _, _, err := parseTableName(tc); err == nil { + t.Errorf("parseTableName(%q): want error, got nil", tc) + } + } +} + +// asTokenError is a test helper to check if err is or wraps a *TokenError. +// TokenError implements Unwrap, so errors.As handles both single- and +// multi-error chains. +func asTokenError(err error, target **TokenError) bool { + return errors.As(err, target) +} diff --git a/purego/internal/auth/token_cache.go b/purego/internal/auth/token_cache.go new file mode 100644 index 00000000..ea4d5938 --- /dev/null +++ b/purego/internal/auth/token_cache.go @@ -0,0 +1,344 @@ +package auth + +import ( + "context" + "crypto/sha256" + "sync" + "time" +) + +// defaultRefreshBuffer is the lead time before a token's expiry at which the +// cache proactively re-mints; 5 minutes is the SDK-wide standard. +const defaultRefreshBuffer = 5 * time.Minute + +// fetchedToken is the raw result of a mint: the token string plus its +// server-reported lifetime (nil when the OAuth server omits expires_in). +type fetchedToken struct { + token string + expiresIn *time.Duration +} + +// mintReason is surfaced in log messages so operators can distinguish a cold +// start from a proactive refresh from caching being off. +type mintReason int + +const ( + mintReasonColdMiss mintReason = iota // no usable entry in cache + mintReasonRefresh // token is within the refresh buffer + mintReasonCacheDisabled // caching is off, so every call mints + mintReasonDirect // minted outside the cache via FetchToken +) + +func (r mintReason) String() string { + switch r { + case mintReasonColdMiss: + return "cold_miss" + case mintReasonRefresh: + return "refresh" + case mintReasonCacheDisabled: + return "cache_disabled" + case mintReasonDirect: + return "direct" + default: + return "unknown" + } +} + +// tokenKey identifies a cache entry. The client secret is stored as its +// SHA-256 digest so the raw credential is never kept in the map: distinct +// secrets still yield distinct keys (collision resistance), and a rotated +// secret gets a fresh entry. +type tokenKey struct { + clientID string + secretDigest [sha256.Size]byte + tableName string +} + +func newTokenKey(clientID, clientSecret, tableName string) tokenKey { + return tokenKey{ + clientID: clientID, + secretDigest: sha256.Sum256([]byte(clientSecret)), + tableName: tableName, + } +} + +// cachedToken is one live entry in the cache. refreshAt is precomputed from the +// TTL and the cache's refresh buffer so needsRefresh is a plain timestamp +// comparison; it never sits after expiresAt. +type cachedToken struct { + value string + expiresAt time.Time + refreshAt time.Time +} + +func (c *cachedToken) isExpired() bool { + return time.Now().After(c.expiresAt) +} + +// newCachedToken builds an entry for a token whose TTL started at mintedAt +// (response-receipt time, matching the Rust SDK). ttl must be positive. +// +// The effective refresh lead time is clamped to at most half the TTL: with the +// default 5-minute buffer a 10-minute token would otherwise be due for refresh +// the instant it is cached, collapsing to a re-mint on every call. Clamping +// guarantees at least ttl/2 of reuse before proactive refresh kicks in. +func newCachedToken(value string, ttl, refreshBuffer time.Duration, mintedAt time.Time) *cachedToken { + if maxBuffer := ttl / 2; refreshBuffer > maxBuffer { + refreshBuffer = maxBuffer + } + expiresAt := mintedAt.Add(ttl) + return &cachedToken{ + value: value, + expiresAt: expiresAt, + refreshAt: expiresAt.Add(-refreshBuffer), + } +} + +// tokenCacheEntry is the per-key slot. Its mutex is held only for short state +// transitions, never across the mint call, so callers single-flight the mint +// yet each waiter can still abandon on its own context. Different keys never +// block each other. +type tokenCacheEntry struct { + mu sync.Mutex + cached *cachedToken // nil when no valid entry has been stored yet + inflight *tokenFlight // non-nil while a mint is in progress +} + +// tokenFlight is a single in-progress mint. The leader records its resolved +// outcome (the same token or error it returns to its own caller) before closing +// done, so every waiter parked on the flight shares that outcome instead of +// re-minting — matching golang.org/x/sync/singleflight and bounding a failing +// token endpoint to one mint per burst rather than N serial attempts. +type tokenFlight struct { + done chan struct{} // closed when the mint completes + token string // resolved token (empty on error) + err error // resolved error (nil on success) +} + +// tokenCache caches OAuth tokens per (clientID, secret, tableName). +// +// It is safe for concurrent use. Construct one with [newTokenCache]; the +// methods do not guard against a nil receiver. +type tokenCache struct { + mu sync.Mutex + entries map[tokenKey]*tokenCacheEntry + refreshBuffer time.Duration + disabled bool // when true, every getOrFetch mints without caching +} + +// CacheOption configures a [SharedTokenCache] passed to [NewSharedTokenCache]. +// The same settings are available per-provider via [WithTokenCacheEnabled] and +// [WithRefreshBuffer]. +type CacheOption func(*tokenCache) + +// CacheEnabled toggles token caching. When disabled, every token request mints +// a fresh token instead of consulting the cache. Caching is enabled by default. +func CacheEnabled(enabled bool) CacheOption { + return func(c *tokenCache) { c.disabled = !enabled } +} + +// CacheRefreshBuffer sets the lead time before a token's expiry at which it is +// proactively re-minted; it defaults to 5 minutes. A non-positive value is +// ignored so the default holds. +func CacheRefreshBuffer(d time.Duration) CacheOption { + return func(c *tokenCache) { + if d > 0 { + c.refreshBuffer = d + } + } +} + +func newTokenCache(opts ...CacheOption) *tokenCache { + c := &tokenCache{ + entries: make(map[tokenKey]*tokenCacheEntry), + refreshBuffer: defaultRefreshBuffer, + } + for _, opt := range opts { + opt(c) + } + return c +} + +// getOrFetch returns a valid token for the given credentials and table, minting +// (via mint, bounded by ctx) only when the cache is empty or within the refresh +// window. A retryable refresh error falls back to the still-valid cached token; +// a non-retryable error propagates. +// +// mint runs at most once per key at a time; other callers wait for it and share +// the result, or return ctx.Err() if their own ctx expires while waiting. +func (c *tokenCache) getOrFetch( + ctx context.Context, + clientID, clientSecret, tableName string, + mint func(ctx context.Context, reason mintReason) (fetchedToken, error), +) (string, error) { + if c.disabled { + fetched, err := mint(ctx, mintReasonCacheDisabled) + if err != nil { + return "", err + } + return fetched.token, nil + } + + key := newTokenKey(clientID, clientSecret, tableName) + entry := c.slot(key) + + entry.mu.Lock() + + if entry.cached != nil && !c.needsRefresh(entry.cached) { + token := entry.cached.value + entry.mu.Unlock() + return token, nil + } + + // A mint is already in progress: wait for it (or our own ctx) and share the + // leader's resolved outcome rather than launching a second mint. This bounds + // a failing token endpoint to one mint per burst instead of N serial retries. + if entry.inflight != nil { + flight := entry.inflight + entry.mu.Unlock() + select { + case <-flight.done: + return flight.token, flight.err + case <-ctx.Done(): + return "", ctx.Err() + } + } + + // We are the leader: claim the mint. The reason is a refresh only when a + // still-valid token is being re-minted early; an expired entry is a cold + // miss operationally, not a refresh. + reason := mintReasonColdMiss + if entry.cached != nil && !entry.cached.isExpired() { + reason = mintReasonRefresh + } + flight := &tokenFlight{done: make(chan struct{})} + entry.inflight = flight + entry.mu.Unlock() + + fetched, err := mint(ctx, reason) + + entry.mu.Lock() + entry.inflight = nil + + if err != nil { + token, resErr := "", err + if entry.cached != nil && !entry.cached.isExpired() && isRetryable(err) { + // Proactive refresh failed transiently; serve the still-valid token. + token, resErr = entry.cached.value, nil + } + flight.token, flight.err = token, resErr + // Publish the outcome to waiters under entry.mu, then release them; the + // mutex (not the close/write ordering) is what prevents a stale read. + close(flight.done) + entry.mu.Unlock() + return token, resErr + } + + token := fetched.token + + // Cache only tokens with a usable TTL. If refresh returned no expires_in, + // keep any existing still-valid token rather than discarding it. + if fetched.expiresIn != nil && *fetched.expiresIn > 0 { + mintedAt := time.Now() + entry.cached = newCachedToken(token, *fetched.expiresIn, c.refreshBuffer, mintedAt) + } else { + keepExisting := entry.cached != nil && !entry.cached.isExpired() + if !keepExisting { + entry.cached = nil + } + } + flight.token, flight.err = token, nil + close(flight.done) + entry.mu.Unlock() + + return token, nil +} + +// invalidate drops the cached token for the given credentials and table. The +// next getOrFetch call will re-mint from scratch. +func (c *tokenCache) invalidate(clientID, clientSecret, tableName string) { + key := newTokenKey(clientID, clientSecret, tableName) + + c.mu.Lock() + entry, ok := c.entries[key] + c.mu.Unlock() + + if !ok { + return + } + entry.mu.Lock() + entry.cached = nil + entry.mu.Unlock() +} + +// slot returns the per-key entry, creating it on first access. The outer lock +// is held only for the map lookup/insert, keeping contention off the hot path. +func (c *tokenCache) slot(key tokenKey) *tokenCacheEntry { + c.mu.Lock() + defer c.mu.Unlock() + + if e, ok := c.entries[key]; ok { + return e + } + // Sweep expired entries on a miss to prevent unbounded map growth in + // clients that ingest into many tables over the lifetime of one process. + c.pruneExpiredLocked() + e := &tokenCacheEntry{} + c.entries[key] = e + return e +} + +// pruneExpiredLocked drops cache entries whose token is fully expired and that +// have no mint in progress. Must be called with c.mu held. +// +// A goroutine can still hold a *tokenCacheEntry pointer that this prunes if the +// entry's cached token is expired and no mint is in flight — a fresh slot() for +// the same key would then create a second entry and briefly duplicate a mint. +// The inflight guard keeps this from touching an in-progress mint, and the +// window only opens for entries with no usable cached token (i.e. no-TTL / fully +// expired), so in normal TTL'd operation a valid entry is never pruned. +func (c *tokenCache) pruneExpiredLocked() { + for k, e := range c.entries { + if e.mu.TryLock() { + // Never evict a mint in flight: its leader writes back to this entry. + expired := e.inflight == nil && (e.cached == nil || e.cached.isExpired()) + e.mu.Unlock() + if expired { + delete(c.entries, k) + } + } + // A locked entry is mid-transition; leave it alone. + } +} + +func (c *tokenCache) needsRefresh(cached *cachedToken) bool { + return time.Now().After(cached.refreshAt) +} + +// isRetryable reports whether the error tree holds a retryable [retryableError], +// walking both single- and multi-error (errors.Join) unwrap chains. +func isRetryable(err error) bool { + switch x := err.(type) { + case nil: + return false + case retryableError: + return x.IsRetryable() + case interface{ Unwrap() error }: + return isRetryable(x.Unwrap()) + case interface{ Unwrap() []error }: + for _, e := range x.Unwrap() { + if isRetryable(e) { + return true + } + } + return false + default: + return false + } +} + +// retryableError is implemented by errors from the mint path that are safe to +// suppress when a cached token is still valid. +type retryableError interface { + IsRetryable() bool +} diff --git a/purego/internal/auth/token_cache_test.go b/purego/internal/auth/token_cache_test.go new file mode 100644 index 00000000..f8d53522 --- /dev/null +++ b/purego/internal/auth/token_cache_test.go @@ -0,0 +1,455 @@ +package auth + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" +) + +// retryErr is a test error that reports itself as retryable. +type retryErr struct{ msg string } + +func (e *retryErr) Error() string { return e.msg } +func (e *retryErr) IsRetryable() bool { return true } + +// fatalErr is a test error that is non-retryable. +type fatalErr struct{ msg string } + +func (e *fatalErr) Error() string { return e.msg } +func (e *fatalErr) IsRetryable() bool { return false } + +func makeMint(token string, ttlSecs int) func(context.Context, mintReason) (fetchedToken, error) { + return func(_ context.Context, _ mintReason) (fetchedToken, error) { + ft := fetchedToken{token: token} + if ttlSecs > 0 { + d := time.Duration(ttlSecs) * time.Second + ft.expiresIn = &d + } + return ft, nil + } +} + +func getOrFetch(t *testing.T, c *tokenCache, clientID, secret, table string, + mint func(context.Context, mintReason) (fetchedToken, error), +) string { + t.Helper() + tok, err := c.getOrFetch(context.Background(), clientID, secret, table, mint) + if err != nil { + t.Fatalf("getOrFetch: %v", err) + } + return tok +} + +func TestTokenCacheCachesAcrossCalls(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + return fetchedToken{token: "tok", expiresIn: dur(3600)}, nil + } + + a := getOrFetch(t, c, "id", "secret", "c.s.t", mint) + b := getOrFetch(t, c, "id", "secret", "c.s.t", mint) + + if a != "tok" || b != "tok" { + t.Fatalf("want tok/tok, got %q/%q", a, b) + } + if n := calls.Load(); n != 1 { + t.Fatalf("want 1 mint, got %d", n) + } +} + +func TestTokenCacheSeparateTablesGetSeparateEntries(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { + n := calls.Add(1) + return fetchedToken{token: "tok" + string(rune('0'+n)), expiresIn: dur(3600)}, nil + } + + a := getOrFetch(t, c, "id", "secret", "c.s.t1", mint) + b := getOrFetch(t, c, "id", "secret", "c.s.t2", mint) + + if a == b { + t.Fatal("different tables must get different tokens") + } + if n := calls.Load(); n != 2 { + t.Fatalf("want 2 mints, got %d", n) + } +} + +func TestTokenCacheRotatedSecretGetsFreshEntry(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + return fetchedToken{token: "tok", expiresIn: dur(3600)}, nil + } + + getOrFetch(t, c, "id", "secret-v1", "c.s.t", mint) + getOrFetch(t, c, "id", "secret-v2", "c.s.t", mint) + + if n := calls.Load(); n != 2 { + t.Fatalf("want 2 mints for different secrets, got %d", n) + } +} + +func TestTokenCacheNoTTLIsNotCached(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + return fetchedToken{token: "tok"}, nil // no expiresIn + } + + a := getOrFetch(t, c, "id", "secret", "c.s.t", mint) + b := getOrFetch(t, c, "id", "secret", "c.s.t", mint) + + if a != "tok" || b != "tok" { + t.Fatalf("want tok/tok, got %q/%q", a, b) + } + if n := calls.Load(); n != 2 { + t.Fatalf("want 2 mints for no-TTL token (not cached), got %d", n) + } +} + +func TestTokenCacheInvalidateForcesMintOnNextCall(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + return fetchedToken{token: "tok", expiresIn: dur(3600)}, nil + } + + getOrFetch(t, c, "id", "secret", "c.s.t", mint) + c.invalidate("id", "secret", "c.s.t") + getOrFetch(t, c, "id", "secret", "c.s.t", mint) + + if n := calls.Load(); n != 2 { + t.Fatalf("want 2 mints after invalidate, got %d", n) + } +} + +func TestTokenCacheWithinRefreshBufferRemints(t *testing.T) { + c := newTokenCache() + // A cached token past its refresh point is re-minted on the next call, and + // the fresh long-TTL result is then cached and reused. + seedRefreshable(c, "id", "secret", "c.s.t", "stale") + + var calls atomic.Int64 + mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + return fetchedToken{token: "fresh", expiresIn: dur(3600)}, nil + } + + a := getOrFetch(t, c, "id", "secret", "c.s.t", mint) // refresh due -> re-mints + b := getOrFetch(t, c, "id", "secret", "c.s.t", mint) // fresh token cached -> hit + + if a != "fresh" || b != "fresh" { + t.Fatalf("want fresh/fresh, got %q/%q", a, b) + } + if n := calls.Load(); n != 1 { + t.Fatalf("want 1 mint (refresh, then cache hit), got %d", n) + } +} + +func TestTokenCacheRetryableRefreshFailureFallsBack(t *testing.T) { + c := newTokenCache() + // Seed a token that is valid but due for proactive refresh. + seedRefreshable(c, "id", "secret", "c.s.t", "valid") + + // Retryable refresh error: should fall back to the still-valid cached token. + tok, err := c.getOrFetch(context.Background(), "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + return fetchedToken{}, &retryErr{"transient"} + }, + ) + if err != nil { + t.Fatalf("want fallback to cached token, got error: %v", err) + } + if tok != "valid" { + t.Fatalf("want %q, got %q", "valid", tok) + } +} + +func TestTokenCacheNonRetryableRefreshErrorPropagates(t *testing.T) { + c := newTokenCache() + seedRefreshable(c, "id", "secret", "c.s.t", "valid") + + _, err := c.getOrFetch(context.Background(), "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + return fetchedToken{}, &fatalErr{"revoked"} + }, + ) + if err == nil { + t.Fatal("want error for non-retryable refresh failure, got nil") + } + var fe *fatalErr + if !errors.As(err, &fe) { + t.Fatalf("want fatalErr, got %T: %v", err, err) + } +} + +func TestTokenCacheNoTTLResponseKeepsExistingCachedToken(t *testing.T) { + c := newTokenCache() + // Seed a token that is due for refresh (refreshAt already in the past). + seedRefreshable(c, "id", "secret", "c.s.t", "valid") + + // A refresh returns a token with no TTL; caller gets the fresh token but the + // existing valid cached token is retained (Rust parity). + fresh := getOrFetch(t, c, "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + return fetchedToken{token: "nottl"}, nil + }, + ) + if fresh != "nottl" { + t.Fatalf("want %q, got %q", "nottl", fresh) + } + + // A subsequent retryable refresh failure should still fall back to the + // original valid cached token, proving it was retained. + tok, err := c.getOrFetch(context.Background(), "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + return fetchedToken{}, &retryErr{"blip"} + }, + ) + if err != nil { + t.Fatalf("want fallback to retained cached token, got error: %v", err) + } + if tok != "valid" { + t.Fatalf("want retained cached token %q, got %q", "valid", tok) + } +} + +func TestTokenCacheSingleFlightMintOnce(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + + const goroutines = 16 + var wg sync.WaitGroup + wg.Add(goroutines) + results := make([]string, goroutines) + + for i := range goroutines { + go func(i int) { + defer wg.Done() + tok, err := c.getOrFetch(context.Background(), "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + time.Sleep(20 * time.Millisecond) // hold lock so others queue up + return fetchedToken{token: "tok", expiresIn: dur(3600)}, nil + }, + ) + if err != nil { + t.Errorf("goroutine %d: %v", i, err) + return + } + results[i] = tok + }(i) + } + wg.Wait() + + for i, tok := range results { + if tok != "tok" { + t.Errorf("goroutine %d: want %q, got %q", i, "tok", tok) + } + } + if n := calls.Load(); n != 1 { + t.Fatalf("single-flight: want 1 mint, got %d", n) + } +} + +func TestTokenCacheConcurrentMintFailureSharesError(t *testing.T) { + c := newTokenCache() + var calls atomic.Int64 + + const goroutines = 16 + var wg sync.WaitGroup + wg.Add(goroutines) + errs := make([]error, goroutines) + + for i := range goroutines { + go func(i int) { + defer wg.Done() + _, errs[i] = c.getOrFetch(context.Background(), "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + calls.Add(1) + time.Sleep(20 * time.Millisecond) // hold the flight so others park on it + return fetchedToken{}, &fatalErr{"boom"} + }, + ) + }(i) + } + wg.Wait() + + // The leader's error is shared with every parked waiter: one mint, not N. + if n := calls.Load(); n != 1 { + t.Fatalf("concurrent failure should mint once, got %d mints", n) + } + for i, err := range errs { + var fe *fatalErr + if !errors.As(err, &fe) { + t.Errorf("goroutine %d: want shared fatalErr, got %T: %v", i, err, err) + } + } +} + +func TestTokenCacheContextDeadlineRespectedDuringMint(t *testing.T) { + c := newTokenCache() + started := make(chan struct{}) + release := make(chan struct{}) + defer close(release) + + // Leader holds the mint open until release is closed. + go func() { + _, _ = c.getOrFetch(context.Background(), "id", "s", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + close(started) + <-release + return fetchedToken{token: "tok", expiresIn: dur(3600)}, nil + }, + ) + }() + <-started + + // A second caller with a short deadline must not block on the leader's mint. + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + start := time.Now() + _, err := c.getOrFetch(ctx, "id", "s", "c.s.t", makeMint("tok2", 3600)) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("want DeadlineExceeded, got %v", err) + } + if elapsed := time.Since(start); elapsed > 1*time.Second { + t.Fatalf("waiter blocked %v on in-progress mint, want ~100ms", elapsed) + } +} + +func TestTokenCacheJoinedRetryableErrorFallsBack(t *testing.T) { + c := newTokenCache() + seedRefreshable(c, "id", "secret", "c.s.t", "valid") + + // A retryable error buried inside errors.Join must still be detected so the + // cache falls back to the still-valid token. + tok, err := c.getOrFetch(context.Background(), "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + joined := errors.Join(errors.New("context note"), &retryErr{"transient"}) + return fetchedToken{}, joined + }, + ) + if err != nil { + t.Fatalf("want fallback to cached token, got error: %v", err) + } + if tok != "valid" { + t.Fatalf("want %q, got %q", "valid", tok) + } +} + +func TestTokenCacheDisabledAlwaysMints(t *testing.T) { + c := newTokenCache(CacheEnabled(false)) + var calls atomic.Int64 + reasons := make(chan mintReason, 2) + mint := func(_ context.Context, r mintReason) (fetchedToken, error) { + calls.Add(1) + reasons <- r + return fetchedToken{token: "tok", expiresIn: dur(3600)}, nil + } + + getOrFetch(t, c, "id", "secret", "c.s.t", mint) + getOrFetch(t, c, "id", "secret", "c.s.t", mint) + + if n := calls.Load(); n != 2 { + t.Fatalf("want 2 mints with caching disabled, got %d", n) + } + close(reasons) + for r := range reasons { + if r != mintReasonCacheDisabled { + t.Fatalf("want mintReasonCacheDisabled, got %v", r) + } + } +} + +func TestTokenCacheCustomRefreshBuffer(t *testing.T) { + // A custom 20-minute buffer against a 1-hour token places the refresh point + // ~40 minutes out (20m < ttl/2, so no clamping) — proving the option took + // effect. Compare against the default 5-minute buffer, which would place it + // ~55 minutes out. + const buffer = 20 * time.Minute + c := newTokenCache(CacheRefreshBuffer(buffer)) + + before := time.Now() + getOrFetch(t, c, "id", "secret", "c.s.t", makeMint("tok", 3600)) + after := time.Now() + + entry := c.slot(newTokenKey("id", "secret", "c.s.t")) + entry.mu.Lock() + cached := entry.cached + entry.mu.Unlock() + if cached == nil { + t.Fatal("token should be cached") + } + + // refreshAt should sit buffer before expiresAt, i.e. ttl-buffer after mint. + wantMin := before.Add(time.Hour - buffer) + wantMax := after.Add(time.Hour - buffer) + if cached.refreshAt.Before(wantMin) || cached.refreshAt.After(wantMax) { + t.Fatalf("refreshAt %v outside expected window [%v, %v] for 20m buffer", + cached.refreshAt, wantMin, wantMax) + } +} + +func TestNewCachedToken(t *testing.T) { + now := time.Now() + + // Normal case: buffer smaller than ttl/2 is applied verbatim. + ct := newCachedToken("v", time.Hour, 5*time.Minute, now) + if !ct.expiresAt.Equal(now.Add(time.Hour)) { + t.Fatalf("expiresAt = %v, want +1h", ct.expiresAt) + } + if !ct.refreshAt.Equal(now.Add(time.Hour - 5*time.Minute)) { + t.Fatalf("refreshAt = %v, want +55m", ct.refreshAt) + } + + // Buffer larger than ttl/2 is clamped so at least half the TTL is reusable. + ct = newCachedToken("v", 10*time.Minute, time.Hour, now) + if !ct.refreshAt.Equal(now.Add(5 * time.Minute)) { + t.Fatalf("clamped refreshAt = %v, want +5m (ttl/2)", ct.refreshAt) + } + if !ct.refreshAt.After(now) { + t.Fatal("clamped refreshAt must stay in the future so the token is reusable") + } +} + +func TestTokenCacheNonPositiveRefreshBufferKeepsDefault(t *testing.T) { + c := newTokenCache(CacheRefreshBuffer(-1)) + if c.refreshBuffer != defaultRefreshBuffer { + t.Fatalf("non-positive buffer should be ignored, got %v", c.refreshBuffer) + } +} + +// dur returns a pointer to a Duration of d seconds, for test readability. +func dur(secs int) *time.Duration { + d := time.Duration(secs) * time.Second + return &d +} + +// seedRefreshable directly installs a cached token that is still valid but past +// its refresh point (refreshAt in the past, expiresAt in the future). This is +// the "due for proactive refresh yet safe to fall back to" state; seeding it +// explicitly avoids depending on TTL/buffer arithmetic that clamping changes. +func seedRefreshable(c *tokenCache, clientID, secret, table, value string) { + key := newTokenKey(clientID, secret, table) + entry := c.slot(key) + now := time.Now() + entry.mu.Lock() + entry.cached = &cachedToken{ + value: value, + expiresAt: now.Add(time.Hour), // still valid + refreshAt: now.Add(-time.Second), + } + entry.mu.Unlock() +} diff --git a/purego/internal/auth/token_provider.go b/purego/internal/auth/token_provider.go new file mode 100644 index 00000000..0f9fd11a --- /dev/null +++ b/purego/internal/auth/token_provider.go @@ -0,0 +1,149 @@ +// Package auth provides OAuth 2.0 and custom-header authentication for the +// Zerobus pure-Go SDK. +// +// It exposes two layers: +// - [TokenProvider], for obtaining bearer tokens. +// - [HeadersProvider], matching other SDKs' custom-header auth model. +// +// The package ships two token providers: +// - [OAuthTokenProvider] — Unity Catalog OAuth 2.0 client credentials flow +// with per-table token caching and proactive refresh. +// - [StaticTokenProvider] — wraps a fixed string for tests or externally +// managed token lifecycles. +// +// And two headers providers: +// - [OAuthHeadersProvider] — wraps an [OAuthTokenProvider] and emits the +// required Zerobus metadata headers. +// - [StaticHeadersProvider] — returns a fixed headers map. +package auth + +import ( + "context" + "fmt" + "strings" +) + +// TokenProvider returns a bearer token for authenticating a stream. +// +// Implementations must be safe for concurrent use. Token may be called on +// every stream open, so implementations should cache aggressively. +// +// The returned token must be a value acceptable as an HTTP/gRPC header value +// (ASCII, no control characters). It may carry a scheme prefix ("Bearer …", +// "Basic …") or be a bare token — [transport.StreamParams].Token normalises +// either form. +// +// tableName is the fully qualified target table (catalog.schema.table) used +// for downscoped-token minting. +// +// Invalidate is called when the server rejects the last token returned by Token +// with an authentication error, signalling that any cached credential for that +// table is stale and must be discarded so the next Token call re-mints. +// Implementations that hold no cache may make Invalidate a no-op. +type TokenProvider interface { + Token(ctx context.Context, tableName string) (string, error) + Invalidate(ctx context.Context, tableName string) +} + +// HeadersProvider provides gRPC metadata headers for stream authentication. +// +// GetHeaders returns headers for the given tableName. The transport layer will +// always enforce the authoritative table-name header from stream-open params. +// Implementations may include it for parity with other SDKs. +// +// Invalidate is called on server auth rejection so any cached credentials can +// be dropped before the next open attempt. +type HeadersProvider interface { + GetHeaders(ctx context.Context, tableName string) (map[string]string, error) + Invalidate(ctx context.Context, tableName string) +} + +// StaticTokenProvider returns the same fixed token on every call. It is +// intended for tests and for callers that manage token lifecycle externally. +// +// Invalidate is a no-op: the token is static and cannot be refreshed. +type StaticTokenProvider struct { + token string +} + +// NewStaticTokenProvider returns a [StaticTokenProvider] that always returns +// the given token. A bare token is accepted; the transport layer will prepend +// "Bearer " if needed. +func NewStaticTokenProvider(token string) *StaticTokenProvider { + return &StaticTokenProvider{token: strings.TrimSpace(token)} +} + +// Token returns the static token. +func (p *StaticTokenProvider) Token(_ context.Context, _ string) (string, error) { + if p.token == "" { + return "", fmt.Errorf("auth: static token is empty") + } + return p.token, nil +} + +// Invalidate is a no-op for a static token. +func (p *StaticTokenProvider) Invalidate(_ context.Context, _ string) {} + +// OAuthHeadersProvider bridges OAuth token minting into headers-provider shape. +type OAuthHeadersProvider struct { + tokenProvider *OAuthTokenProvider +} + +// NewOAuthHeadersProvider creates an OAuth-backed headers provider. +func NewOAuthHeadersProvider( + clientID, clientSecret, zerobusEndpoint, ucEndpoint string, + opts ...OAuthOption, +) (*OAuthHeadersProvider, error) { + p, err := NewOAuthTokenProvider(clientID, clientSecret, zerobusEndpoint, ucEndpoint, opts...) + if err != nil { + return nil, err + } + return &OAuthHeadersProvider{tokenProvider: p}, nil +} + +// GetHeaders returns Zerobus auth headers for tableName. +func (p *OAuthHeadersProvider) GetHeaders(ctx context.Context, tableName string) (map[string]string, error) { + token, err := p.tokenProvider.Token(ctx, tableName) + if err != nil { + return nil, err + } + return map[string]string{ + "authorization": "Bearer " + token, + "x-databricks-zerobus-table-name": tableName, + }, nil +} + +// Invalidate drops the cached token for tableName. +func (p *OAuthHeadersProvider) Invalidate(ctx context.Context, tableName string) { + p.tokenProvider.Invalidate(ctx, tableName) +} + +// StaticHeadersProvider returns a fixed header set. +type StaticHeadersProvider struct { + headers map[string]string +} + +// NewStaticHeadersProvider returns a provider that returns the same headers on +// every call. +func NewStaticHeadersProvider(headers map[string]string) *StaticHeadersProvider { + cloned := make(map[string]string, len(headers)) + for k, v := range headers { + cloned[k] = strings.TrimSpace(v) + } + return &StaticHeadersProvider{headers: cloned} +} + +// GetHeaders returns a copy of the configured headers. +func (p *StaticHeadersProvider) GetHeaders(_ context.Context, _ string) (map[string]string, error) { + if len(p.headers) == 0 { + return nil, fmt.Errorf("auth: static headers are empty") + } + out := make(map[string]string, len(p.headers)) + for k, v := range p.headers { + out[k] = v + } + return out, nil +} + +// Invalidate is a no-op for static headers. +func (p *StaticHeadersProvider) Invalidate(_ context.Context, _ string) {} diff --git a/purego/internal/transport/ephemeral_stream.go b/purego/internal/transport/ephemeral_stream.go index d4ef31af..1093356c 100644 --- a/purego/internal/transport/ephemeral_stream.go +++ b/purego/internal/transport/ephemeral_stream.go @@ -5,6 +5,8 @@ import ( "fmt" "strings" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "github.com/databricks/zerobus-sdk/purego/internal/zerobuspb" @@ -23,7 +25,17 @@ type StreamParams struct { // Token is the credential sent in the authorization header. A bare token is // prefixed with "Bearer "; a value that already carries a known scheme (e.g. // "Bearer ..." or "Basic ...") is sent verbatim. Empty means no header. + // Ignored when HeadersProvider is set. Token string + // HeadersProvider supplies custom auth headers, similar to other SDKs. + // When set, it is used instead of Token. + HeadersProvider HeadersProvider +} + +// HeadersProvider provides gRPC metadata headers for stream authentication. +type HeadersProvider interface { + GetHeaders(ctx context.Context, tableName string) (map[string]string, error) + Invalidate(ctx context.Context, tableName string) } // Stream is an open ephemeral ingestion stream over the proto/JSON @@ -63,6 +75,23 @@ func (c *Conn) Open(ctx context.Context, p StreamParams) (*Stream, error) { if p.RecordType == zerobuspb.RecordType_PROTO && len(p.DescriptorProto) == 0 { return nil, fmt.Errorf("transport: open %q: descriptor proto required for PROTO records", p.TableName) } + var headers map[string]string + if p.HeadersProvider != nil { + var err error + headers, err = p.HeadersProvider.GetHeaders(ctx, p.TableName) + if err != nil { + return nil, fmt.Errorf("transport: open %q: headers provider: %w", p.TableName, err) + } + for k, v := range headers { + if !isUsableAsHeader(v) { + return nil, fmt.Errorf("transport: open %q: header %q contains invalid value characters", p.TableName, strings.TrimSpace(k)) + } + } + } else if !isUsableAsHeader(p.Token) { + // Reject control chars at the wire boundary so every token source is + // covered, not just the OAuth mint path; gRPC would otherwise fail opaquely. + return nil, fmt.Errorf("transport: open %q: token contains invalid header characters", p.TableName) + } // openCtx bounds the open attempt: the caller's ctx, defaulted to // defaultHandshakeTimeout when it has no deadline. @@ -75,7 +104,7 @@ func (c *Conn) Open(ctx context.Context, p StreamParams) (*Stream, error) { // Detach the live stream from ctx's cancel/deadline (WithoutCancel keeps its // values, so caller metadata survives); Close is its only canceller. - streamCtx := withStreamMetadata(context.WithoutCancel(ctx), p.TableName, p.Token) + streamCtx := withStreamMetadataHeaders(context.WithoutCancel(ctx), p.TableName, headers, p.Token) streamCtx, cancelStream := context.WithCancel(streamCtx) // Until the handshake succeeds, bridge openCtx to cancelStream so a caller @@ -85,6 +114,9 @@ func (c *Conn) Open(ctx context.Context, p StreamParams) (*Stream, error) { stream, err := c.open(openCtx, streamCtx, cancelStream, p) if err != nil { + if p.HeadersProvider != nil && isAuthRejection(err) { + p.HeadersProvider.Invalidate(ctx, p.TableName) + } cancelStream() return nil, err } @@ -98,6 +130,11 @@ func (c *Conn) Open(ctx context.Context, p StreamParams) (*Stream, error) { return stream, nil } +func isAuthRejection(err error) bool { + code := status.Code(err) + return code == codes.Unauthenticated || code == codes.PermissionDenied +} + // open starts the RPC on streamCtx and runs the handshake bounded by hctx. // teardown cancels streamCtx so the handshake can reap its recv goroutine on // cancel/timeout. It supplies the proto-specific hooks and annotates their errors. @@ -168,11 +205,22 @@ func (s *Stream) Recv() (*zerobuspb.EphemeralStreamResponse, error) { return s.r // CloseSend signals that no more requests will be sent, half-closing the stream // while leaving Recv open to drain remaining responses. For a graceful shutdown, -// call CloseSend, read until io.EOF, then Close. +// prefer GracefulClose, which performs the CloseSend/drain/Close sequence for you. func (s *Stream) CloseSend() error { return s.closeSend() } +// GracefulClose half-closes the send side, drains remaining responses until the +// server ends the stream, then releases resources. Prefer it over Close when done +// sending: draining to end-of-stream lets the server see an orderly close rather +// than the abrupt reset Close produces. +// +// ctx bounds the drain; on expiry or a stream error it hard-aborts like Close and +// returns the cause, else nil. Must not be called concurrently with Recv, and no +// Send may follow (the send side is half-closed). A later Close is a no-op, since +// GracefulClose always releases resources before it returns. +func (s *Stream) GracefulClose(ctx context.Context) error { return s.gracefulClose(ctx) } + // Close aborts the stream and releases its resources. It is idempotent and safe -// to call after a graceful CloseSend/drain. Unlike CloseSend it does not wait -// for the server: any in-flight Send or Recv is unblocked with a cancellation -// error. +// to call after a graceful CloseSend/drain (or GracefulClose). Unlike GracefulClose +// it does not wait for the server: any in-flight Send or Recv is unblocked with a +// cancellation error. func (s *Stream) Close() { s.close() } diff --git a/purego/internal/transport/export_test.go b/purego/internal/transport/export_test.go index c636720e..4c843cd6 100644 --- a/purego/internal/transport/export_test.go +++ b/purego/internal/transport/export_test.go @@ -1,6 +1,10 @@ package transport -import "google.golang.org/grpc/credentials/insecure" +import ( + "time" + + "google.golang.org/grpc/credentials/insecure" +) // WithInsecure disables transport security. Auth tokens ride gRPC metadata, so // an insecure connection would send bearer tokens in plaintext; it lives in a @@ -9,3 +13,11 @@ import "google.golang.org/grpc/credentials/insecure" func WithInsecure() DialOption { return func(cfg *dialConfig) { cfg.creds = insecure.NewCredentials() } } + +// SetDefaultDrainTimeout overrides the no-deadline gracefulClose timeout and +// returns a restore func, so tests exercise the default-applied path quickly. +func SetDefaultDrainTimeout(d time.Duration) (restore func()) { + prev := defaultDrainTimeout + defaultDrainTimeout = d + return func() { defaultDrainTimeout = prev } +} diff --git a/purego/internal/transport/metadata.go b/purego/internal/transport/metadata.go index 699587f4..051f22ee 100644 --- a/purego/internal/transport/metadata.go +++ b/purego/internal/transport/metadata.go @@ -3,6 +3,7 @@ package transport import ( "context" "strings" + "unicode" "google.golang.org/grpc/metadata" ) @@ -23,14 +24,38 @@ var authSchemes = []string{"bearer", "basic", "dpop"} // are replaced (gRPC is first-value-wins, so a duplicate could mis-route or send // a stale token); unrelated caller metadata is preserved. func withStreamMetadata(ctx context.Context, tableName, token string) context.Context { + return withStreamMetadataHeaders(ctx, tableName, nil, token) +} + +// withStreamMetadataHeaders applies stream metadata using either headers from a +// provider or a direct token string. +func withStreamMetadataHeaders(ctx context.Context, tableName string, headers map[string]string, token string) context.Context { md, ok := metadata.FromOutgoingContext(ctx) if ok { md = md.Copy() } else { md = metadata.MD{} } + var authValue string + for key, value := range headers { + key = strings.ToLower(strings.TrimSpace(key)) + switch key { + case "": + continue + case mdTableName: + // table header is authoritative from stream-open params. + continue + case mdAuthorization: + authValue = value + default: + md.Set(key, strings.TrimSpace(value)) + } + } md.Set(mdTableName, tableName) - if v := authHeaderValue(token); v != "" { + if authValue == "" { + authValue = token + } + if v := authHeaderValue(authValue); v != "" { md.Set(mdAuthorization, v) } else { md.Delete(mdAuthorization) @@ -38,6 +63,18 @@ func withStreamMetadata(ctx context.Context, tableName, token string) context.Co return metadata.NewOutgoingContext(ctx, md) } +// isUsableAsHeader reports whether token is safe in a gRPC authorization +// header: no control or non-ASCII chars, which gRPC metadata rejects. An empty +// token is usable (it yields no header). +func isUsableAsHeader(token string) bool { + for _, r := range token { + if r > unicode.MaxASCII || unicode.IsControl(r) { + return false + } + } + return true +} + // authHeaderValue normalizes a token into an authorization header value: a value // already carrying a known scheme (Bearer/Basic/DPoP) is returned verbatim, a // bare token is prefixed with "Bearer ", and an empty token yields "". diff --git a/purego/internal/transport/rawstream.go b/purego/internal/transport/raw_stream.go similarity index 75% rename from purego/internal/transport/rawstream.go rename to purego/internal/transport/raw_stream.go index bde96a57..86fe713b 100644 --- a/purego/internal/transport/rawstream.go +++ b/purego/internal/transport/raw_stream.go @@ -13,6 +13,11 @@ import ( // no deadline, so Open can't hang if the server half-opens the stream. const defaultHandshakeTimeout = 30 * time.Second +// defaultDrainTimeout bounds gracefulClose when the caller's context has no +// deadline, so it can't hang on an unresponsive server. A var so tests can +// shrink it. +var defaultDrainTimeout = 30 * time.Second + // bidiRPC is the subset of a generated gRPC bidirectional streaming client that // rawStream needs. EphemeralStream satisfies it, as will Arrow Flight's DoPut. type bidiRPC[Req, Resp any] interface { @@ -94,6 +99,48 @@ func (s *rawStream[Req, Resp]) close() { }) } +// gracefulClose half-closes the send side, drains remaining responses to io.EOF, +// then releases resources. Draining to EOF lets the server see an orderly close; +// a bare close cancels the context and the server sees an abrupt reset instead. +// +// ctx bounds the drain: on ctx expiry or a non-EOF error it hard-aborts and +// returns the cause (ctx error preferred); a clean drain returns nil. Not safe to +// call concurrently with recv, and no send may follow (the send side is +// half-closed). Every return path calls close first, so a later close is a no-op. +func (s *rawStream[Req, Resp]) gracefulClose(ctx context.Context) error { + if _, ok := ctx.Deadline(); !ok { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, defaultDrainTimeout) + defer cancel() + } + + if err := s.closeSend(); err != nil { + s.close() + return err + } + + // Bridge ctx to close so a caller deadline unblocks the recv below, which + // otherwise waits on the stream's Close-only context. + stop := context.AfterFunc(ctx, s.close) + defer stop() + + for { + _, err := s.recv() + switch { + case err == io.EOF: + s.close() + return nil + case err != nil: + s.close() + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + return err + } + // Response arrived before EOF; discard and keep draining. + } +} + // handshake runs the create-stream exchange shared by every ingestion protocol: // send a setup message, await the readiness response, and validate it. Blocking // here surfaces setup failures (auth, schema, table access) at open time. diff --git a/purego/internal/transport/rawstream_test.go b/purego/internal/transport/raw_stream_test.go similarity index 83% rename from purego/internal/transport/rawstream_test.go rename to purego/internal/transport/raw_stream_test.go index 1bed944d..0aff0ea9 100644 --- a/purego/internal/transport/rawstream_test.go +++ b/purego/internal/transport/raw_stream_test.go @@ -152,4 +152,28 @@ func TestRawStreamRecvReturnsEOFUnwrapped(t *testing.T) { } } +// TestGracefulCloseCloseSendFailure: when CloseSend itself fails (e.g. the +// stream is already broken), gracefulClose must hard-abort and return the error +// rather than proceeding to drain. +func TestGracefulCloseCloseSendFailure(t *testing.T) { + closeSendBoom := errors.New("close-send boom") + var cancelled bool + s := &rawStream[string, string]{ + rpc: &fakeBidiRPC{closeErr: closeSendBoom}, + cancel: func() { cancelled = true }, + } + s.setID("test-stream") + + err := s.gracefulClose(context.Background()) + if err == nil { + t.Fatal("gracefulClose with failing CloseSend: got nil error, want failure") + } + if !strings.Contains(err.Error(), "test-stream") { + t.Errorf("error %q should contain stream name", err.Error()) + } + if !cancelled { + t.Error("gracefulClose did not call close (cancel) after CloseSend failure") + } +} + func strPtr(s string) *string { return &s } diff --git a/purego/internal/transport/transport_test.go b/purego/internal/transport/transport_test.go index 8ffa5052..f691ec9c 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -2,13 +2,17 @@ package transport_test import ( "context" + "errors" "io" "net" + "sync/atomic" "testing" "time" + "google.golang.org/grpc/codes" "google.golang.org/grpc" "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" "google.golang.org/grpc/test/bufconn" "google.golang.org/protobuf/proto" @@ -48,6 +52,14 @@ type fakeServer struct { // blocking until the stream context is cancelled. Used to exercise the // handshake deadline. hangHandshake bool + // hangDrain makes the server ignore the client's half-close and never end the + // stream, so GracefulClose can't drain to io.EOF and must hit its ctx deadline. + hangDrain bool + // authReject makes open fail with UNAUTHENTICATED after receiving create. + authReject bool + // drainGate, when non-nil, holds io.EOF back until closed, so a test can + // assert GracefulClose keeps draining rather than returning at the first ack. + drainGate chan struct{} } func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamServer) error { @@ -75,6 +87,9 @@ func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamSer <-stream.Context().Done() return stream.Context().Err() } + if f.authReject { + return status.Error(codes.Unauthenticated, "bad credentials") + } var resp *zerobuspb.EphemeralStreamResponse if f.badHandshake { @@ -99,6 +114,14 @@ func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamSer for { req, err := stream.Recv() if err == io.EOF { + if f.hangDrain { + <-stream.Context().Done() + return stream.Context().Err() + } + // Hold EOF until the test releases the gate. + if f.drainGate != nil { + <-f.drainGate + } return nil } if err != nil { @@ -120,6 +143,28 @@ func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamSer } } +type stubHeadersProvider struct { + headers map[string]string + calls atomic.Int32 + invalidateCalls atomic.Int32 + lastTable atomic.Value // string +} + +func (p *stubHeadersProvider) GetHeaders(_ context.Context, tableName string) (map[string]string, error) { + p.calls.Add(1) + p.lastTable.Store(tableName) + out := make(map[string]string, len(p.headers)) + for k, v := range p.headers { + out[k] = v + } + return out, nil +} + +func (p *stubHeadersProvider) Invalidate(_ context.Context, tableName string) { + p.invalidateCalls.Add(1) + p.lastTable.Store(tableName) +} + // firstMD returns the first metadata value for key, or "". func firstMD(md metadata.MD, key string) string { if vs := md.Get(key); len(vs) > 0 { @@ -230,6 +275,22 @@ func TestOpenPrefixesUnknownScheme(t *testing.T) { } } +func TestOpenRejectsTokenWithControlChars(t *testing.T) { + for _, tok := range []string{"tok\nen", "tok\x00en", "Bearer tok\ren"} { + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1)} + conn := dialFake(t, srv) + + _, err := conn.Open(context.Background(), transport.StreamParams{ + TableName: "c.s.t", + RecordType: zerobuspb.RecordType_JSON, + Token: tok, + }) + if err == nil { + t.Errorf("Open with token %q: got nil error, want rejection", tok) + } + } +} + func TestStreamSendRecv(t *testing.T) { srv := &fakeServer{streamID: "s1", seen: make(chan observed, 1)} conn := dialFake(t, srv) @@ -356,6 +417,64 @@ func TestOpenOmitsEmptyToken(t *testing.T) { } } +func TestOpenUsesHeadersProviderAndTableIsAuthoritative(t *testing.T) { + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1)} + conn := dialFake(t, srv) + + p := &stubHeadersProvider{ + headers: map[string]string{ + "authorization": "provider-token", + "x-databricks-zerobus-table-name": "wrong.table.name", + mdUserKey: "provider-md", + }, + } + _, err := conn.Open(context.Background(), transport.StreamParams{ + TableName: "c.s.t", + RecordType: zerobuspb.RecordType_JSON, + HeadersProvider: p, + }) + if err != nil { + t.Fatalf("Open: %v", err) + } + got := <-srv.seen + if got.tableName != "c.s.t" { + t.Fatalf("server saw table %q, want %q", got.tableName, "c.s.t") + } + if got.auth != "Bearer provider-token" { + t.Fatalf("server saw auth %q, want %q", got.auth, "Bearer provider-token") + } + if got.userMD != "provider-md" { + t.Fatalf("server saw custom metadata %q, want %q", got.userMD, "provider-md") + } + if p.calls.Load() != 1 { + t.Fatalf("GetHeaders calls = %d, want 1", p.calls.Load()) + } + last, _ := p.lastTable.Load().(string) + if last != "c.s.t" { + t.Fatalf("provider saw table %q, want %q", last, "c.s.t") + } +} + +func TestOpenAuthRejectionInvalidatesHeadersProvider(t *testing.T) { + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1), authReject: true} + conn := dialFake(t, srv) + + p := &stubHeadersProvider{ + headers: map[string]string{"authorization": "tok"}, + } + _, err := conn.Open(context.Background(), transport.StreamParams{ + TableName: "c.s.t", + RecordType: zerobuspb.RecordType_JSON, + HeadersProvider: p, + }) + if err == nil { + t.Fatal("Open with server auth rejection: got nil error") + } + if p.invalidateCalls.Load() != 1 { + t.Fatalf("Invalidate calls = %d, want 1", p.invalidateCalls.Load()) + } +} + func TestStreamCloseAbortsRecv(t *testing.T) { srv := &fakeServer{streamID: "s", seen: make(chan observed, 1)} conn := dialFake(t, srv) @@ -525,9 +644,35 @@ func TestOpenDeadlineDoesNotTearDownStream(t *testing.T) { } } -// TestStreamGracefulClose exercises the documented graceful shutdown sequence: -// CloseSend, drain to io.EOF, then Close. The server returns EOF once it sees -// the half-close, so Recv observes a clean end rather than a cancellation. +// TestStreamGracefulCloseDefaultsDeadline: with no deadline on the context, +// GracefulClose must apply defaultDrainTimeout and not hang on a stalled server. +// The default is shrunk to a few ms so the path runs quickly. +func TestStreamGracefulCloseDefaultsDeadline(t *testing.T) { + defer transport.SetDefaultDrainTimeout(50 * time.Millisecond)() + + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1), hangDrain: true} + conn := dialFake(t, srv) + + stream, err := conn.Open(context.Background(), transport.StreamParams{ + TableName: "c.s.t", + RecordType: zerobuspb.RecordType_JSON, + Token: "tok", + }) + if err != nil { + t.Fatalf("Open: %v", err) + } + <-srv.seen + + // No deadline on the context — the function must impose one itself and + // return an error rather than hanging forever. + err = stream.GracefulClose(context.Background()) + if err == nil { + t.Fatal("GracefulClose against a stalled server with no deadline: got nil, want timeout error") + } +} + +// TestStreamGracefulClose: the server ends the stream on half-close, so +// GracefulClose drains to a clean io.EOF and returns nil. func TestStreamGracefulClose(t *testing.T) { srv := &fakeServer{streamID: "s", seen: make(chan observed, 1)} conn := dialFake(t, srv) @@ -542,18 +687,94 @@ func TestStreamGracefulClose(t *testing.T) { } <-srv.seen - if err := stream.CloseSend(); err != nil { - t.Fatalf("CloseSend: %v", err) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := stream.GracefulClose(ctx); err != nil { + t.Fatalf("GracefulClose: %v", err) } - // Drain to the clean end-of-stream the server sends after the half-close. - for { - _, err := stream.Recv() - if err == io.EOF { - break - } + stream.Close() +} + +// TestStreamGracefulCloseDrainsPending: an in-flight ack precedes io.EOF, so +// GracefulClose must discard it and keep draining. drainGate holds EOF back to +// prove it's still draining after the ack, then returns nil once EOF arrives. +func TestStreamGracefulCloseDrainsPending(t *testing.T) { + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1), drainGate: make(chan struct{})} + conn := dialFake(t, srv) + + stream, err := conn.Open(context.Background(), transport.StreamParams{ + TableName: "c.s.t", + RecordType: zerobuspb.RecordType_JSON, + Token: "tok", + }) + if err != nil { + t.Fatalf("Open: %v", err) + } + <-srv.seen + + if err := stream.Send(&zerobuspb.EphemeralStreamRequest{ + Payload: &zerobuspb.EphemeralStreamRequest_IngestRecord{ + IngestRecord: &zerobuspb.IngestRecordRequest{ + OffsetId: proto.Int64(1), + Record: &zerobuspb.IngestRecordRequest_JsonRecord{JsonRecord: `{"id":1}`}, + }, + }, + }); err != nil { + t.Fatalf("Send: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + done := make(chan error, 1) + go func() { done <- stream.GracefulClose(ctx) }() + + // The server sends the ack, then blocks before EOF. A correct drain discards + // the ack and waits for EOF, so GracefulClose must not have returned yet. + select { + case err := <-done: + t.Fatalf("GracefulClose returned at the pending ack (err=%v); it must keep draining to io.EOF", err) + case <-time.After(200 * time.Millisecond): + } + + // Release EOF; the drain completes cleanly. + close(srv.drainGate) + select { + case err := <-done: if err != nil { - t.Fatalf("Recv during drain: got %v, want io.EOF", err) + t.Fatalf("GracefulClose after draining the pending ack: %v", err) } + case <-time.After(5 * time.Second): + t.Fatal("GracefulClose did not return after io.EOF was released") + } +} + +// TestStreamGracefulCloseHonorsDeadline: the server never ends the stream, so +// GracefulClose returns its ctx error promptly instead of blocking, and leaves +// the stream torn down. +func TestStreamGracefulCloseHonorsDeadline(t *testing.T) { + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1), hangDrain: true} + conn := dialFake(t, srv) + + stream, err := conn.Open(context.Background(), transport.StreamParams{ + TableName: "c.s.t", + RecordType: zerobuspb.RecordType_JSON, + Token: "tok", + }) + if err != nil { + t.Fatalf("Open: %v", err) + } + <-srv.seen + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + err = stream.GracefulClose(ctx) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("GracefulClose against a server that never ends the stream: got %v, want DeadlineExceeded", err) + } + // Bounded-out drain tears the stream down: Recv is unblocked, not hanging. + if _, err := stream.Recv(); err == nil { + t.Fatal("Recv after a deadline-bounded GracefulClose: got nil error, want cancellation") } - stream.Close() // releases resources; safe after a graceful drain }