From 8bc3676b961d620de0c09c9aaaeb29d4908a0901 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Thu, 9 Jul 2026 09:26:46 +0000 Subject: [PATCH 01/11] [PureGo] Add GracefulClose to transport Stream Encapsulate the clean-shutdown sequence (CloseSend, drain to io.EOF, Close) on rawStream so embedding streams don't re-implement it and get the order wrong. A bare Close cancels the context and the server sees an abrupt reset; draining to EOF lets it see an orderly close. ctx bounds the drain: on expiry or a stream error it hard-aborts and returns the cause, else nil. Close stays as the hard-abort escape hatch. Closes #487. Signed-off-by: Zlata Stefanovic --- purego/internal/transport/ephemeral_stream.go | 18 +++- purego/internal/transport/rawstream.go | 35 +++++++ purego/internal/transport/transport_test.go | 97 ++++++++++++++++--- 3 files changed, 131 insertions(+), 19 deletions(-) diff --git a/purego/internal/transport/ephemeral_stream.go b/purego/internal/transport/ephemeral_stream.go index d4ef31af..84d12ab5 100644 --- a/purego/internal/transport/ephemeral_stream.go +++ b/purego/internal/transport/ephemeral_stream.go @@ -168,11 +168,21 @@ 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; a later +// Close is a no-op. +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/rawstream.go b/purego/internal/transport/rawstream.go index bde96a57..bee9593f 100644 --- a/purego/internal/transport/rawstream.go +++ b/purego/internal/transport/rawstream.go @@ -94,6 +94,41 @@ 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. Safe to call once; a later close is a no-op. +func (s *rawStream[Req, Resp]) gracefulClose(ctx context.Context) error { + if err := s.closeSend(); err != nil { + s.close() // can't half-close cleanly; hard-abort + 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/transport_test.go b/purego/internal/transport/transport_test.go index 8ffa5052..fb2c2b1d 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -48,6 +48,9 @@ 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 } func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamServer) error { @@ -99,6 +102,10 @@ func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamSer for { req, err := stream.Recv() if err == io.EOF { + if f.hangDrain { + <-stream.Context().Done() // withhold the clean end + return stream.Context().Err() + } return nil } if err != nil { @@ -525,9 +532,8 @@ 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. +// 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 +548,79 @@ 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 - } - if err != nil { - t.Fatalf("Recv during drain: got %v, want io.EOF", err) - } + stream.Close() // idempotent no-op after a graceful shutdown +} + +// TestStreamGracefulCloseDrainsPending: an in-flight ack precedes io.EOF, so +// GracefulClose must discard it and keep draining rather than stop at the first +// response. +func TestStreamGracefulCloseDrainsPending(t *testing.T) { + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1)} + 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() + if err := stream.GracefulClose(ctx); err != nil { + t.Fatalf("GracefulClose with a pending response: %v", err) + } +} + +// 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() + + start := time.Now() + err = stream.GracefulClose(ctx) + if err == nil { + t.Fatal("GracefulClose against a server that never ends the stream: got nil, want deadline error") + } + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("GracefulClose took %v, expected it to return near the 200ms deadline", elapsed) + } + // 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 } From 139da7b665bf8fd44cb0211dde51d47f3780cc4e Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Thu, 9 Jul 2026 10:22:27 +0000 Subject: [PATCH 02/11] [PureGo] Clarify GracefulClose docs: no Send after, close-before-return Note the send side is half-closed (no Send may follow) and make the 'later Close is a no-op' guarantee explicit: every return path releases resources first. Signed-off-by: Zlata Stefanovic --- purego/internal/transport/ephemeral_stream.go | 5 +++-- purego/internal/transport/rawstream.go | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/purego/internal/transport/ephemeral_stream.go b/purego/internal/transport/ephemeral_stream.go index 84d12ab5..e02759ea 100644 --- a/purego/internal/transport/ephemeral_stream.go +++ b/purego/internal/transport/ephemeral_stream.go @@ -177,8 +177,9 @@ func (s *Stream) CloseSend() error { return s.closeSend() } // 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; a later -// Close is a no-op. +// 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 diff --git a/purego/internal/transport/rawstream.go b/purego/internal/transport/rawstream.go index bee9593f..7e239b47 100644 --- a/purego/internal/transport/rawstream.go +++ b/purego/internal/transport/rawstream.go @@ -100,7 +100,8 @@ func (s *rawStream[Req, Resp]) close() { // // 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. Safe to call once; a later close is a no-op. +// 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 err := s.closeSend(); err != nil { s.close() // can't half-close cleanly; hard-abort From 4ae878acd8979e2aa18fd3cc78c6bec0364300cb Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 10:41:57 +0000 Subject: [PATCH 03/11] [PureGo] Apply default deadline in GracefulClose when ctx has none context.Background() passed to GracefulClose would hang forever if the server stalled during the drain, since the only exit was ctx expiry. Apply defaultHandshakeTimeout internally when no deadline is present, mirroring the same guard in Open. Add TestStreamGracefulCloseDefaultsDeadline to cover the no-deadline path (marked -short-skippable as it takes ~30s to fire). Signed-off-by: Zlata Stefanovic --- purego/internal/transport/rawstream.go | 6 +++++ purego/internal/transport/transport_test.go | 29 +++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/purego/internal/transport/rawstream.go b/purego/internal/transport/rawstream.go index 7e239b47..957bcca9 100644 --- a/purego/internal/transport/rawstream.go +++ b/purego/internal/transport/rawstream.go @@ -103,6 +103,12 @@ func (s *rawStream[Req, Resp]) close() { // 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, defaultHandshakeTimeout) + defer cancel() + } + if err := s.closeSend(); err != nil { s.close() // can't half-close cleanly; hard-abort return err diff --git a/purego/internal/transport/transport_test.go b/purego/internal/transport/transport_test.go index fb2c2b1d..293840d2 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -532,6 +532,35 @@ func TestOpenDeadlineDoesNotTearDownStream(t *testing.T) { } } +// TestStreamGracefulCloseDefaultsDeadline: passing a context with no deadline +// (e.g. context.Background()) must not hang when the server stalls — the +// function applies defaultHandshakeTimeout internally, mirroring Open. +// Runs under -short but takes ~defaultHandshakeTimeout (30s) to complete. +func TestStreamGracefulCloseDefaultsDeadline(t *testing.T) { + if testing.Short() { + t.Skip("skipping: takes defaultHandshakeTimeout (30s) to exercise the no-deadline path") + } + 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) { From 4e99e19da4c7b833780fb43268501b5427390e09 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 10:45:13 +0000 Subject: [PATCH 04/11] [PureGo] Address review nits: CloseSend failure test, tighter timing bound MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add TestGracefulCloseCloseSendFailure to rawstream_test.go covering the path where CloseSend itself fails — verifies the error is returned and close() is called to hard-abort the stream. Tighten the elapsed-time assertion in TestStreamGracefulCloseHonorsDeadline from a fixed 5s bound to 5x the actual deadline (5*200ms = 1s), so the test catches genuine slowness rather than accepting a 4.9s return when the deadline was 200ms. Signed-off-by: Zlata Stefanovic --- purego/internal/transport/rawstream_test.go | 24 +++++++++++++++++++++ purego/internal/transport/transport_test.go | 5 +++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/purego/internal/transport/rawstream_test.go b/purego/internal/transport/rawstream_test.go index 1bed944d..0aff0ea9 100644 --- a/purego/internal/transport/rawstream_test.go +++ b/purego/internal/transport/rawstream_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 293840d2..bec7e42f 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -640,13 +640,14 @@ func TestStreamGracefulCloseHonorsDeadline(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() + const deadline = 200 * time.Millisecond start := time.Now() err = stream.GracefulClose(ctx) if err == nil { t.Fatal("GracefulClose against a server that never ends the stream: got nil, want deadline error") } - if elapsed := time.Since(start); elapsed > 5*time.Second { - t.Fatalf("GracefulClose took %v, expected it to return near the 200ms deadline", elapsed) + if elapsed := time.Since(start); elapsed > 5*deadline { + t.Fatalf("GracefulClose took %v, expected it to return near the %v deadline", elapsed, deadline) } // Bounded-out drain tears the stream down: Recv is unblocked, not hanging. if _, err := stream.Recv(); err == nil { From 734fea6fcc7a5274979771dd1dfcc1017966914b Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 10:50:57 +0000 Subject: [PATCH 05/11] [PureGo] Address review nits: CloseSend failure test, tighter timing bound MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add TestGracefulCloseCloseSendFailure to rawstream_test.go covering the path where CloseSend itself fails — verifies the error is returned and close() is called to hard-abort the stream. Replace the wall-clock elapsed check in TestStreamGracefulCloseHonorsDeadline with errors.Is(err, context.DeadlineExceeded): asserts what the deadline mechanism actually returned rather than inferring it from timing, removing the flakiness entirely. Signed-off-by: Zlata Stefanovic --- purego/internal/transport/transport_test.go | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/purego/internal/transport/transport_test.go b/purego/internal/transport/transport_test.go index bec7e42f..eee7fd7a 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -2,6 +2,7 @@ package transport_test import ( "context" + "errors" "io" "net" "testing" @@ -640,14 +641,9 @@ func TestStreamGracefulCloseHonorsDeadline(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() - const deadline = 200 * time.Millisecond - start := time.Now() err = stream.GracefulClose(ctx) - if err == nil { - t.Fatal("GracefulClose against a server that never ends the stream: got nil, want deadline error") - } - if elapsed := time.Since(start); elapsed > 5*deadline { - t.Fatalf("GracefulClose took %v, expected it to return near the %v deadline", elapsed, deadline) + 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 { From be8c9216795482f53393c55fbce31feb6b3aa134 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 11:14:03 +0000 Subject: [PATCH 06/11] [PureGo] Remove inline comments flagged in review Signed-off-by: Zlata Stefanovic --- purego/internal/transport/rawstream.go | 2 +- purego/internal/transport/transport_test.go | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/purego/internal/transport/rawstream.go b/purego/internal/transport/rawstream.go index 957bcca9..ce4f064c 100644 --- a/purego/internal/transport/rawstream.go +++ b/purego/internal/transport/rawstream.go @@ -110,7 +110,7 @@ func (s *rawStream[Req, Resp]) gracefulClose(ctx context.Context) error { } if err := s.closeSend(); err != nil { - s.close() // can't half-close cleanly; hard-abort + s.close() return err } diff --git a/purego/internal/transport/transport_test.go b/purego/internal/transport/transport_test.go index eee7fd7a..e1cf94c7 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -104,7 +104,7 @@ func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamSer req, err := stream.Recv() if err == io.EOF { if f.hangDrain { - <-stream.Context().Done() // withhold the clean end + <-stream.Context().Done() return stream.Context().Err() } return nil @@ -583,7 +583,7 @@ func TestStreamGracefulClose(t *testing.T) { if err := stream.GracefulClose(ctx); err != nil { t.Fatalf("GracefulClose: %v", err) } - stream.Close() // idempotent no-op after a graceful shutdown + stream.Close() } // TestStreamGracefulCloseDrainsPending: an in-flight ack precedes io.EOF, so From 7dd3703e6ee7d5f033263f53fdf6461efd8f812f Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 11:34:09 +0000 Subject: [PATCH 07/11] [PureGo] Rename rawstream.go to raw_stream.go Signed-off-by: Zlata Stefanovic --- purego/internal/transport/{rawstream.go => raw_stream.go} | 0 .../internal/transport/{rawstream_test.go => raw_stream_test.go} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename purego/internal/transport/{rawstream.go => raw_stream.go} (100%) rename purego/internal/transport/{rawstream_test.go => raw_stream_test.go} (100%) diff --git a/purego/internal/transport/rawstream.go b/purego/internal/transport/raw_stream.go similarity index 100% rename from purego/internal/transport/rawstream.go rename to purego/internal/transport/raw_stream.go diff --git a/purego/internal/transport/rawstream_test.go b/purego/internal/transport/raw_stream_test.go similarity index 100% rename from purego/internal/transport/rawstream_test.go rename to purego/internal/transport/raw_stream_test.go From 5ef98cfae39f9dc8f93ab6db7e6626ec08448be4 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 14:14:54 +0000 Subject: [PATCH 08/11] [PureGo] Add UC OAuth token provider with per-table caching Add internal/auth: OAuthTokenProvider implements the Unity Catalog OAuth 2.0 client-credentials flow, minting downscoped per-table tokens and caching them with single-flight and on-access proactive refresh. Also adds StaticTokenProvider and wires wire-boundary header validation into the transport Open path so every provider is covered. Co-authored-by: Isaac Signed-off-by: Zlata Stefanovic --- purego/README.md | 8 +- purego/internal/auth/oauth.go | 519 +++++++++++++++ purego/internal/auth/oauth_test.go | 608 ++++++++++++++++++ purego/internal/auth/token_cache.go | 359 +++++++++++ purego/internal/auth/token_cache_test.go | 464 +++++++++++++ purego/internal/auth/token_provider.go | 64 ++ purego/internal/transport/ephemeral_stream.go | 5 + purego/internal/transport/metadata.go | 13 + purego/internal/transport/transport_test.go | 16 + 9 files changed, 2054 insertions(+), 2 deletions(-) create mode 100644 purego/internal/auth/oauth.go create mode 100644 purego/internal/auth/oauth_test.go create mode 100644 purego/internal/auth/token_cache.go create mode 100644 purego/internal/auth/token_cache_test.go create mode 100644 purego/internal/auth/token_provider.go diff --git a/purego/README.md b/purego/README.md index e54a2d43..9ee5b323 100644 --- a/purego/README.md +++ b/purego/README.md @@ -20,9 +20,13 @@ 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). +- `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. Obtain a token with `Token(ctx)` + and pass it as `StreamParams.Token`. -Planned layers: OAuth/auth, ingest/ack state management, recovery, and the -public API. +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..e06b2e4e --- /dev/null +++ b/purego/internal/auth/oauth.go @@ -0,0 +1,519 @@ +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 + ucEndpoint string // e.g. "https://workspace.databricks.com" + tableName string + + 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 for +// tableName using the client credentials flow. +// +// ucEndpoint is the workspace URL (e.g. "https://my-workspace.databricks.com"). +// workspaceID is the numeric Databricks workspace ID. +func NewOAuthTokenProvider( + clientID, clientSecret, workspaceID, ucEndpoint, tableName string, + opts ...OAuthOption, +) (*OAuthTokenProvider, error) { + tableName = strings.TrimSpace(tableName) + if err := validateTableName(tableName); err != nil { + return nil, fmt.Errorf("auth: oauth: %w", err) + } + 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") + } + if workspaceID == "" { + return nil, fmt.Errorf("auth: oauth: workspaceID is required") + } + + p := &OAuthTokenProvider{ + clientID: clientID, + clientSecret: clientSecret, + workspaceID: workspaceID, + ucEndpoint: ucEndpoint, + tableName: tableName, + 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 the provider's table, 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. +// +// Every successful mint is cached. If UC reports an expires_in it is honored; if +// it omits one the token is cached under a bounded default lifetime instead of +// being re-minted on every call. The proactive-refresh lead time is clamped to +// at most half the token's TTL, so even a short-lived token gets some reuse +// before it is refreshed. +func (p *OAuthTokenProvider) Token(ctx context.Context) (string, error) { + return p.cache.getOrFetch(ctx, p.clientID, p.clientSecret, p.tableName, p.mint) +} + +// 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) (string, error) { + fetched, err := p.mint(ctx, 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, reason mintReason) (fetchedToken, error) { + started := time.Now() + fetched, err := p.fetchToken(ctx) + elapsed := time.Since(started) + + switch { + case err != nil: + p.logger.LogAttrs(ctx, slog.LevelWarn, "failed to mint UC OAuth token", + slog.String("table", p.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; caching under a conservative default lifetime", + slog.String("table", p.tableName), + slog.String("reason", reason.String()), + slog.Duration("elapsed", elapsed), + ) + default: + p.logger.LogAttrs(ctx, slog.LevelInfo, "minted UC OAuth token", + slog.String("table", p.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) { + p.cache.invalidate(p.clientID, p.clientSecret, p.tableName) +} + +// fetchToken performs the OAuth 2.0 client credentials request against Unity +// Catalog's OIDC token endpoint. +func (p *OAuthTokenProvider) fetchToken(ctx context.Context) (fetchedToken, error) { + catalog, schema, _, err := parseTableName(p.tableName) + if err != nil { + return fetchedToken{}, err + } + + authDetails, err := buildAuthorizationDetails(catalog, schema, p.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: any 5xx, plus 429 and 408. +func isRetryableStatus(code int) bool { + switch { + case code >= 500: + return true + case code == http.StatusTooManyRequests, code == http.StatusRequestTimeout: + return true + default: + return false + } +} + +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 +} + +// 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..c0be4052 --- /dev/null +++ b/purego/internal/auth/oauth_test.go @@ -0,0 +1,608 @@ +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", "ws123", ts.URL, table, + 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()) + 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()); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background()); 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()); err != nil { + t.Fatalf("first Token: %v", err) + } + p.Invalidate(context.Background()) + if _, err := p.Token(context.Background()); 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()) + 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()) + 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", "ws", hang.URL, "c.s.t", + 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) + 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", "ws", url, "c.s.t", + WithHTTPClient(client), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + _, err = p.Token(context.Background()) + 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 TestOAuthTokenProvider429IsRetryable(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusTooManyRequests} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background()) + if err == nil { + t.Fatal("want error for 429, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || !te.IsRetryable() { + t.Fatalf("want retryable TokenError for 429, got %T (retryable=%v): %v", + err, te != nil && te.IsRetryable(), err) + } +} + +func TestOAuthTokenProvider408IsRetryable(t *testing.T) { + srv := &tokenServer{statusCode: http.StatusRequestTimeout} + p, ts := newTestProvider(t, srv, "cat.sch.tbl") + defer ts.Close() + + _, err := p.Token(context.Background()) + if err == nil { + t.Fatal("want error for 408, got nil") + } + var te *TokenError + if !asTokenError(err, &te) || !te.IsRetryable() { + t.Fatalf("want 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", "ws", hang.URL, "c.s.t", + 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) + 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", "ws", ts.URL, "c.s.t", + 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()); 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()) + 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, wsID, ucEndpoint, table string + }{ + {"empty clientID", "", "s", "ws", "https://host", "c.s.t"}, + {"empty secret", "id", "", "ws", "https://host", "c.s.t"}, + {"empty wsID", "id", "s", "", "https://host", "c.s.t"}, + {"empty ucEndpoint", "id", "s", "ws", "", "c.s.t"}, + {"bad table (2 parts)", "id", "s", "ws", "https://host", "c.s"}, + {"bad table (empty schema)", "id", "s", "ws", "https://host", "c..t"}, + {"plaintext http endpoint", "id", "s", "ws", "http://host.databricks.com", "c.s.t"}, + {"non-http scheme", "id", "s", "ws", "ftp://host", "c.s.t"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + _, err := NewOAuthTokenProvider(tc.clientID, tc.secret, tc.wsID, tc.ucEndpoint, tc.table) + if err == nil { + t.Fatal("want 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", "ws", ts.URL, "c.s.t", + WithHTTPClient(ts.Client()), + WithTokenCacheEnabled(false), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background()); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background()); 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", "ws", "https://host", "c.s.t", + 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", "ws", ts.URL, "c.s.t", + WithHTTPClient(ts.Client()), + WithSharedTokenCache(shared), + WithTokenCacheEnabled(false), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background()); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background()); 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", "ws", ts.URL, "c.s.t", + WithHTTPClient(ts.Client()), + WithSharedTokenCache(shared), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background()); err != nil { + t.Fatalf("first Token: %v", err) + } + if _, err := p.Token(context.Background()); 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()); err != nil { + t.Fatalf("Token: %v", err) + } + // FetchToken must mint fresh regardless of the warm cache… + if _, err := p.FetchToken(context.Background()); 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()); 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", "ws", ts.URL, "c.s.t", + WithHTTPClient(ts.Client()), + WithLogger(logger), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background()); 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", "ws", ts.URL, "mycat.mysch.mytbl", + WithHTTPClient(ts.Client()), + ) + if err != nil { + t.Fatalf("NewOAuthTokenProvider: %v", err) + } + if _, err := p.Token(context.Background()); 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..cd0ee21a --- /dev/null +++ b/purego/internal/auth/token_cache.go @@ -0,0 +1,359 @@ +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 + +// defaultNoTTLLifetime is the lifetime assigned to a token whose response omits +// (or reports a non-positive) expires_in. Without it such a token would never +// be cached and every Token call would pay a full HTTP round-trip; a bounded, +// conservative lifetime instead lets it be reused for a short window. UC tokens +// carry an expires_in in practice, so this is a safety net, not the norm. +const defaultNoTTLLifetime = 10 * 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 +// (request-dispatch time, so latency is charged against the token rather than +// overestimating its validity). A non-positive ttl falls back to +// defaultNoTTLLifetime so a token that the server didn't tag with an expires_in +// still gets a bounded cache window. +// +// 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 ttl <= 0 { + ttl = defaultNoTTLLifetime + } + 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() + + // Anchor the TTL at request dispatch, not response receipt: the server starts + // the clock when it issues the token, so using receipt time would overestimate + // its remaining validity by the round-trip latency. + mintedAt := time.Now() + 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 + + // Always cache a successful mint. When the server reports a usable TTL we + // honor it; when it omits expires_in newCachedToken falls back to a bounded + // default lifetime so the token is still reused for a short window rather + // than re-minted on every call. + ttl := time.Duration(0) + if fetched.expiresIn != nil { + ttl = *fetched.expiresIn + } + entry.cached = newCachedToken(token, ttl, c.refreshBuffer, mintedAt) + 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..54bc2ec0 --- /dev/null +++ b/purego/internal/auth/token_cache_test.go @@ -0,0 +1,464 @@ +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 TestTokenCacheNoTTLIsCachedUnderDefaultLifetime(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) + + // A token with no expires_in is now cached under the default lifetime and + // reused, rather than re-minted on every call. + 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 for no-TTL token (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 TestTokenCacheNoTTLResponseReplacesCachedToken(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; the caller gets the new token and it + // is now cached under the default lifetime (no longer discarded). + 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) + } + + // The next call reuses the freshly cached no-TTL token rather than re-minting. + var minted atomic.Bool + tok := getOrFetch(t, c, "id", "secret", "c.s.t", + func(_ context.Context, _ mintReason) (fetchedToken, error) { + minted.Store(true) + return fetchedToken{token: "unexpected"}, nil + }, + ) + if tok != "nottl" { + t.Fatalf("want cached %q, got %q", "nottl", tok) + } + if minted.Load() { + t.Fatal("second call should hit cache, not re-mint") + } +} + +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) + } + + // No TTL: falls back to the default lifetime. + ct = newCachedToken("v", 0, 5*time.Minute, now) + if !ct.expiresAt.Equal(now.Add(defaultNoTTLLifetime)) { + t.Fatalf("no-TTL expiresAt = %v, want +%v", ct.expiresAt, defaultNoTTLLifetime) + } + + // 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..5a1a08de --- /dev/null +++ b/purego/internal/auth/token_provider.go @@ -0,0 +1,64 @@ +// Package auth provides OAuth 2.0 authentication for the Zerobus pure-Go SDK. +// +// The central abstraction is [TokenProvider], an interface for obtaining a +// bearer token on demand. Call its Token method before opening a stream and +// pass the result as [transport.StreamParams].Token. +// +// The package ships two implementations: +// - [OAuthTokenProvider] — Unity Catalog OAuth 2.0 client credentials flow, +// with per-table token caching and proactive refresh. This is the default +// for production use. +// - [StaticTokenProvider] — wraps a fixed string; useful for tests or when +// the caller manages token lifecycle outside the SDK. +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. +// +// Invalidate is called when the server rejects the last token returned by +// Token with an authentication error, signalling that any cached credential 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) (string, error) + Invalidate(ctx context.Context) +} + +// 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, 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) {} diff --git a/purego/internal/transport/ephemeral_stream.go b/purego/internal/transport/ephemeral_stream.go index e02759ea..db919be3 100644 --- a/purego/internal/transport/ephemeral_stream.go +++ b/purego/internal/transport/ephemeral_stream.go @@ -63,6 +63,11 @@ 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) } + // Reject control chars at the wire boundary so every TokenProvider is + // covered, not just the OAuth mint path; gRPC would otherwise fail opaquely. + if !isUsableAsHeader(p.Token) { + 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. diff --git a/purego/internal/transport/metadata.go b/purego/internal/transport/metadata.go index 699587f4..51f4de5b 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" ) @@ -38,6 +39,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/transport_test.go b/purego/internal/transport/transport_test.go index fb2c2b1d..b13584c5 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -237,6 +237,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) From 3050ebc4f24a2d561b381e7513adc35adbc5d044 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 14:31:23 +0000 Subject: [PATCH 09/11] [PureGo] Fix 30s CI hang in graceful-close test; doc/test nits Signed-off-by: Zlata Stefanovic --- README.md | 2 +- cpp/include/zerobus/arrow_stream.hpp | 8 ++-- purego/internal/transport/export_test.go | 14 +++++- purego/internal/transport/raw_stream.go | 7 ++- purego/internal/transport/transport_test.go | 48 +++++++++++++++------ 5 files changed, 60 insertions(+), 19 deletions(-) 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/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/raw_stream.go b/purego/internal/transport/raw_stream.go index ce4f064c..86fe713b 100644 --- a/purego/internal/transport/raw_stream.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 { @@ -105,7 +110,7 @@ func (s *rawStream[Req, Resp]) close() { func (s *rawStream[Req, Resp]) gracefulClose(ctx context.Context) error { if _, ok := ctx.Deadline(); !ok { var cancel context.CancelFunc - ctx, cancel = context.WithTimeout(ctx, defaultHandshakeTimeout) + ctx, cancel = context.WithTimeout(ctx, defaultDrainTimeout) defer cancel() } diff --git a/purego/internal/transport/transport_test.go b/purego/internal/transport/transport_test.go index e1cf94c7..d8de483c 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -52,6 +52,9 @@ type fakeServer struct { // 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 + // 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 { @@ -107,6 +110,10 @@ func (f *fakeServer) EphemeralStream(stream zerobuspb.Zerobus_EphemeralStreamSer <-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 { @@ -533,14 +540,12 @@ func TestOpenDeadlineDoesNotTearDownStream(t *testing.T) { } } -// TestStreamGracefulCloseDefaultsDeadline: passing a context with no deadline -// (e.g. context.Background()) must not hang when the server stalls — the -// function applies defaultHandshakeTimeout internally, mirroring Open. -// Runs under -short but takes ~defaultHandshakeTimeout (30s) to complete. +// 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) { - if testing.Short() { - t.Skip("skipping: takes defaultHandshakeTimeout (30s) to exercise the no-deadline path") - } + defer transport.SetDefaultDrainTimeout(50 * time.Millisecond)() + srv := &fakeServer{streamID: "s", seen: make(chan observed, 1), hangDrain: true} conn := dialFake(t, srv) @@ -587,10 +592,10 @@ func TestStreamGracefulClose(t *testing.T) { } // TestStreamGracefulCloseDrainsPending: an in-flight ack precedes io.EOF, so -// GracefulClose must discard it and keep draining rather than stop at the first -// response. +// 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)} + 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{ @@ -616,8 +621,27 @@ func TestStreamGracefulCloseDrainsPending(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - if err := stream.GracefulClose(ctx); err != nil { - t.Fatalf("GracefulClose with a pending response: %v", err) + + 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("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") } } From a5e15c38d4c428547939ac1d22a73cd3a640d585 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Fri, 10 Jul 2026 15:34:43 +0000 Subject: [PATCH 10/11] [PureGo] Align OAuth retry and TTL anchoring with Rust Treat all 4xx token responses (including 429 and 408) as non-retryable, and anchor cached token TTL at response receipt to match the Rust SDK. Signed-off-by: Zlata Stefanovic --- purego/internal/auth/oauth.go | 12 +++--------- purego/internal/auth/oauth_test.go | 12 ++++++------ purego/internal/auth/token_cache.go | 8 ++------ 3 files changed, 11 insertions(+), 21 deletions(-) diff --git a/purego/internal/auth/oauth.go b/purego/internal/auth/oauth.go index e06b2e4e..578aa3e7 100644 --- a/purego/internal/auth/oauth.go +++ b/purego/internal/auth/oauth.go @@ -384,16 +384,10 @@ func (e *TokenError) IsRetryable() bool { return e.retryable } 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: any 5xx, plus 429 and 408. +// 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 { - switch { - case code >= 500: - return true - case code == http.StatusTooManyRequests, code == http.StatusRequestTimeout: - return true - default: - return false - } + return code >= 500 } func classifyHTTPError(resp *http.Response) error { diff --git a/purego/internal/auth/oauth_test.go b/purego/internal/auth/oauth_test.go index c0be4052..c552f1cc 100644 --- a/purego/internal/auth/oauth_test.go +++ b/purego/internal/auth/oauth_test.go @@ -188,7 +188,7 @@ func TestOAuthTokenProviderConnectionRefusedIsRetryable(t *testing.T) { } } -func TestOAuthTokenProvider429IsRetryable(t *testing.T) { +func TestOAuthTokenProvider429IsNonRetryable(t *testing.T) { srv := &tokenServer{statusCode: http.StatusTooManyRequests} p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() @@ -198,13 +198,13 @@ func TestOAuthTokenProvider429IsRetryable(t *testing.T) { t.Fatal("want error for 429, got nil") } var te *TokenError - if !asTokenError(err, &te) || !te.IsRetryable() { - t.Fatalf("want retryable TokenError for 429, got %T (retryable=%v): %v", + 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 TestOAuthTokenProvider408IsRetryable(t *testing.T) { +func TestOAuthTokenProvider408IsNonRetryable(t *testing.T) { srv := &tokenServer{statusCode: http.StatusRequestTimeout} p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() @@ -214,8 +214,8 @@ func TestOAuthTokenProvider408IsRetryable(t *testing.T) { t.Fatal("want error for 408, got nil") } var te *TokenError - if !asTokenError(err, &te) || !te.IsRetryable() { - t.Fatalf("want retryable TokenError for 408, got %T (retryable=%v): %v", + 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) } } diff --git a/purego/internal/auth/token_cache.go b/purego/internal/auth/token_cache.go index cd0ee21a..f16a4aee 100644 --- a/purego/internal/auth/token_cache.go +++ b/purego/internal/auth/token_cache.go @@ -83,8 +83,7 @@ func (c *cachedToken) isExpired() bool { } // newCachedToken builds an entry for a token whose TTL started at mintedAt -// (request-dispatch time, so latency is charged against the token rather than -// overestimating its validity). A non-positive ttl falls back to +// (response-receipt time, matching the Rust SDK). A non-positive ttl falls back to // defaultNoTTLLifetime so a token that the server didn't tag with an expires_in // still gets a bounded cache window. // @@ -228,10 +227,6 @@ func (c *tokenCache) getOrFetch( entry.inflight = flight entry.mu.Unlock() - // Anchor the TTL at request dispatch, not response receipt: the server starts - // the clock when it issues the token, so using receipt time would overestimate - // its remaining validity by the round-trip latency. - mintedAt := time.Now() fetched, err := mint(ctx, reason) entry.mu.Lock() @@ -261,6 +256,7 @@ func (c *tokenCache) getOrFetch( if fetched.expiresIn != nil { ttl = *fetched.expiresIn } + mintedAt := time.Now() entry.cached = newCachedToken(token, ttl, c.refreshBuffer, mintedAt) flight.token, flight.err = token, nil close(flight.done) From 0b34456f8e26de9da53b9dfa1354f262136fbd65 Mon Sep 17 00:00:00 2001 From: Zlata Stefanovic Date: Mon, 13 Jul 2026 08:44:01 +0000 Subject: [PATCH 11/11] [Go] align purego auth with Rust behavior Signed-off-by: Zlata Stefanovic --- purego/README.md | 5 +- purego/internal/auth/oauth.go | 107 +++++++++++----- purego/internal/auth/oauth_test.go | 99 +++++++------- purego/internal/auth/token_cache.go | 33 ++--- purego/internal/auth/token_cache_test.go | 37 ++---- purego/internal/auth/token_provider.go | 121 +++++++++++++++--- purego/internal/transport/ephemeral_stream.go | 40 +++++- purego/internal/transport/metadata.go | 26 +++- purego/internal/transport/transport_test.go | 88 +++++++++++++ 9 files changed, 402 insertions(+), 154 deletions(-) diff --git a/purego/README.md b/purego/README.md index 9ee5b323..d3e2b219 100644 --- a/purego/README.md +++ b/purego/README.md @@ -23,8 +23,9 @@ Implemented packages: - `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. Obtain a token with `Token(ctx)` - and pass it as `StreamParams.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. diff --git a/purego/internal/auth/oauth.go b/purego/internal/auth/oauth.go index 578aa3e7..e9905216 100644 --- a/purego/internal/auth/oauth.go +++ b/purego/internal/auth/oauth.go @@ -35,8 +35,8 @@ type OAuthTokenProvider struct { clientID string clientSecret string workspaceID string + zerobusHost string ucEndpoint string // e.g. "https://workspace.databricks.com" - tableName string cache *tokenCache client *http.Client @@ -112,19 +112,16 @@ func WithLogger(l *slog.Logger) OAuthOption { } } -// NewOAuthTokenProvider creates a provider that mints UC OAuth tokens for -// tableName using the client credentials flow. +// 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"). -// workspaceID is the numeric Databricks workspace ID. func NewOAuthTokenProvider( - clientID, clientSecret, workspaceID, ucEndpoint, tableName string, + clientID, clientSecret, zerobusEndpoint, ucEndpoint string, opts ...OAuthOption, ) (*OAuthTokenProvider, error) { - tableName = strings.TrimSpace(tableName) - if err := validateTableName(tableName); err != nil { - return nil, fmt.Errorf("auth: oauth: %w", err) - } ucEndpoint = strings.TrimRight(strings.TrimSpace(ucEndpoint), "/") if ucEndpoint == "" { return nil, fmt.Errorf("auth: oauth: ucEndpoint is required") @@ -138,16 +135,17 @@ func NewOAuthTokenProvider( if clientSecret == "" { return nil, fmt.Errorf("auth: oauth: clientSecret is required") } - if workspaceID == "" { - return nil, fmt.Errorf("auth: oauth: workspaceID 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, - tableName: tableName, client: &http.Client{Timeout: 30 * time.Second}, logger: slog.New(slog.DiscardHandler), } @@ -172,28 +170,38 @@ func NewSharedTokenCache(opts ...CacheOption) *SharedTokenCache { return newTokenCache(opts...) } -// Token returns a valid bearer token for the provider's table, minting a new -// one via Unity Catalog's OIDC token endpoint when the cache is empty or -// nearing expiry. +// 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. // -// Every successful mint is cached. If UC reports an expires_in it is honored; if -// it omits one the token is cached under a bounded default lifetime instead of -// being re-minted on every call. The proactive-refresh lead time is clamped to -// at most half the token's TTL, so even a short-lived token gets some reuse -// before it is refreshed. -func (p *OAuthTokenProvider) Token(ctx context.Context) (string, error) { - return p.cache.getOrFetch(ctx, p.clientID, p.clientSecret, p.tableName, p.mint) +// 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) (string, error) { - fetched, err := p.mint(ctx, mintReasonDirect) +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 } @@ -202,29 +210,29 @@ func (p *OAuthTokenProvider) FetchToken(ctx context.Context) (string, error) { // 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, reason mintReason) (fetchedToken, error) { +func (p *OAuthTokenProvider) mint(ctx context.Context, tableName string, reason mintReason) (fetchedToken, error) { started := time.Now() - fetched, err := p.fetchToken(ctx) + 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", p.tableName), + 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; caching under a conservative default lifetime", - slog.String("table", p.tableName), + 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", p.tableName), + slog.String("table", tableName), slog.String("reason", reason.String()), slog.Duration("expires_in", *fetched.expiresIn), slog.Duration("elapsed", elapsed), @@ -236,19 +244,23 @@ func (p *OAuthTokenProvider) mint(ctx context.Context, reason mintReason) (fetch // 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) { - p.cache.invalidate(p.clientID, p.clientSecret, p.tableName) +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) (fetchedToken, error) { - catalog, schema, _, err := parseTableName(p.tableName) +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, p.tableName) + authDetails, err := buildAuthorizationDetails(catalog, schema, tableName) if err != nil { return fetchedToken{}, fmt.Errorf("auth: oauth: marshal authorization_details: %w", err) } @@ -479,6 +491,31 @@ func isLoopbackHost(host string) bool { 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) diff --git a/purego/internal/auth/oauth_test.go b/purego/internal/auth/oauth_test.go index c552f1cc..0de1825c 100644 --- a/purego/internal/auth/oauth_test.go +++ b/purego/internal/auth/oauth_test.go @@ -43,7 +43,7 @@ func (s *tokenServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { func newTestProvider(t *testing.T, srv *tokenServer, table string) (*OAuthTokenProvider, *httptest.Server) { t.Helper() ts := httptest.NewServer(srv) - p, err := NewOAuthTokenProvider("clientID", "clientSecret", "ws123", ts.URL, table, + p, err := NewOAuthTokenProvider("clientID", "clientSecret", "https://ws123.zerobus.databricks.com", ts.URL, WithHTTPClient(ts.Client()), ) if err != nil { @@ -57,7 +57,7 @@ func TestOAuthTokenProviderHappyPath(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - tok, err := p.Token(context.Background()) + tok, err := p.Token(context.Background(), "cat.sch.tbl") if err != nil { t.Fatalf("Token: %v", err) } @@ -71,10 +71,10 @@ func TestOAuthTokenProviderCachesToken(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { t.Fatalf("first Token: %v", err) } - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { t.Fatalf("second Token: %v", err) } if got := srv.calls.Load(); got != 1 { @@ -87,11 +87,11 @@ func TestOAuthTokenProviderInvalidateForcesRemint(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "cat.sch.tbl"); err != nil { t.Fatalf("first Token: %v", err) } - p.Invalidate(context.Background()) - if _, err := p.Token(context.Background()); err != nil { + 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 { @@ -104,7 +104,7 @@ func TestOAuthTokenProvider5xxIsRetryable(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - _, err := p.Token(context.Background()) + _, err := p.Token(context.Background(), "cat.sch.tbl") if err == nil { t.Fatal("want error for 500, got nil") } @@ -119,7 +119,7 @@ func TestOAuthTokenProvider4xxIsNonRetryable(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - _, err := p.Token(context.Background()) + _, err := p.Token(context.Background(), "cat.sch.tbl") if err == nil { t.Fatal("want error for 401, got nil") } @@ -143,7 +143,7 @@ func TestOAuthTokenProviderContextCancellation(t *testing.T) { defer hang.Close() // executed second: server is now idle defer close(unblock) // executed first: unblocks the handler - p, err := NewOAuthTokenProvider("id", "secret", "ws", hang.URL, "c.s.t", + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", hang.URL, WithHTTPClient(hang.Client()), ) if err != nil { @@ -154,7 +154,7 @@ func TestOAuthTokenProviderContextCancellation(t *testing.T) { defer cancel() start := time.Now() - _, err = p.Token(ctx) + _, err = p.Token(ctx, "c.s.t") if err == nil { t.Fatal("want error on ctx cancellation, got nil") } @@ -171,13 +171,13 @@ func TestOAuthTokenProviderConnectionRefusedIsRetryable(t *testing.T) { url := ts.URL ts.Close() // close immediately so connections are refused - p, err := NewOAuthTokenProvider("id", "secret", "ws", url, "c.s.t", + 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()) + _, err = p.Token(context.Background(), "c.s.t") if err == nil { t.Fatal("want error for refused connection, got nil") } @@ -193,7 +193,7 @@ func TestOAuthTokenProvider429IsNonRetryable(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - _, err := p.Token(context.Background()) + _, err := p.Token(context.Background(), "cat.sch.tbl") if err == nil { t.Fatal("want error for 429, got nil") } @@ -209,7 +209,7 @@ func TestOAuthTokenProvider408IsNonRetryable(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - _, err := p.Token(context.Background()) + _, err := p.Token(context.Background(), "cat.sch.tbl") if err == nil { t.Fatal("want error for 408, got nil") } @@ -234,7 +234,7 @@ func TestTokenErrorUnwrapsContextCancellation(t *testing.T) { defer hang.Close() defer close(unblock) - p, err := NewOAuthTokenProvider("id", "secret", "ws", hang.URL, "c.s.t", + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", hang.URL, WithHTTPClient(hang.Client()), ) if err != nil { @@ -244,7 +244,7 @@ func TestTokenErrorUnwrapsContextCancellation(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() - _, err = p.Token(ctx) + _, err = p.Token(ctx, "c.s.t") if err == nil { t.Fatal("want error on ctx cancellation, got nil") } @@ -265,7 +265,7 @@ func TestOAuthTokenProviderNilOptionsKeepDefaults(t *testing.T) { ts := httptest.NewServer(srv) defer ts.Close() - p, err := NewOAuthTokenProvider("id", "secret", "ws", ts.URL, "c.s.t", + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, WithHTTPClient(nil), WithSharedTokenCache(nil), ) @@ -279,7 +279,7 @@ func TestOAuthTokenProviderNilOptionsKeepDefaults(t *testing.T) { 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()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("Token with nil options: %v", err) } } @@ -289,7 +289,7 @@ func TestClassifyHTTPErrorPreservesBody(t *testing.T) { p, ts := newTestProvider(t, srv, "cat.sch.tbl") defer ts.Close() - _, err := p.Token(context.Background()) + _, err := p.Token(context.Background(), "cat.sch.tbl") if err == nil { t.Fatal("want error for 401, got nil") } @@ -305,23 +305,24 @@ func TestClassifyHTTPErrorPreservesBody(t *testing.T) { func TestNewOAuthTokenProviderValidation(t *testing.T) { cases := []struct { - name string - clientID, secret, wsID, ucEndpoint, table string + name string + clientID, secret, zerobusEndpoint, ucEndpoint, table string }{ - {"empty clientID", "", "s", "ws", "https://host", "c.s.t"}, - {"empty secret", "id", "", "ws", "https://host", "c.s.t"}, - {"empty wsID", "id", "s", "", "https://host", "c.s.t"}, - {"empty ucEndpoint", "id", "s", "ws", "", "c.s.t"}, - {"bad table (2 parts)", "id", "s", "ws", "https://host", "c.s"}, - {"bad table (empty schema)", "id", "s", "ws", "https://host", "c..t"}, - {"plaintext http endpoint", "id", "s", "ws", "http://host.databricks.com", "c.s.t"}, - {"non-http scheme", "id", "s", "ws", "ftp://host", "c.s.t"}, + {"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) { - _, err := NewOAuthTokenProvider(tc.clientID, tc.secret, tc.wsID, tc.ucEndpoint, tc.table) + p, err := NewOAuthTokenProvider(tc.clientID, tc.secret, tc.zerobusEndpoint, tc.ucEndpoint) if err == nil { - t.Fatal("want error, got nil") + if _, tokErr := p.Token(context.Background(), tc.table); tokErr == nil { + t.Fatal("want constructor or token validation error, got nil") + } } }) } @@ -377,17 +378,17 @@ func TestOAuthTokenProviderCacheDisabledAlwaysMints(t *testing.T) { ts := httptest.NewServer(srv) defer ts.Close() - p, err := NewOAuthTokenProvider("id", "secret", "ws", ts.URL, "c.s.t", + 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()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("first Token: %v", err) } - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("second Token: %v", err) } if got := srv.calls.Load(); got != 2 { @@ -399,7 +400,7 @@ 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", "ws", "https://host", "c.s.t", + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", "https://host", WithRefreshBuffer(90*time.Second), ) if err != nil { @@ -418,7 +419,7 @@ func TestOAuthTokenProviderSharedCacheIgnoresProviderCacheOpts(t *testing.T) { defer ts.Close() shared := NewSharedTokenCache() - p, err := NewOAuthTokenProvider("id", "secret", "ws", ts.URL, "c.s.t", + p, err := NewOAuthTokenProvider("id", "secret", "https://ws.zerobus.databricks.com", ts.URL, WithHTTPClient(ts.Client()), WithSharedTokenCache(shared), WithTokenCacheEnabled(false), @@ -426,10 +427,10 @@ func TestOAuthTokenProviderSharedCacheIgnoresProviderCacheOpts(t *testing.T) { if err != nil { t.Fatalf("NewOAuthTokenProvider: %v", err) } - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("first Token: %v", err) } - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("second Token: %v", err) } if got := srv.calls.Load(); got != 1 { @@ -443,17 +444,17 @@ func TestSharedTokenCacheDisabled(t *testing.T) { defer ts.Close() shared := NewSharedTokenCache(CacheEnabled(false)) - p, err := NewOAuthTokenProvider("id", "secret", "ws", ts.URL, "c.s.t", + 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()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("first Token: %v", err) } - if _, err := p.Token(context.Background()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("second Token: %v", err) } if got := srv.calls.Load(); got != 2 { @@ -467,18 +468,18 @@ func TestOAuthTokenProviderFetchTokenBypassesCache(t *testing.T) { defer ts.Close() // Warm the cache. - if _, err := p.Token(context.Background()); err != nil { + 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()); err != nil { + 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()); err != nil { + 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 { @@ -494,14 +495,14 @@ func TestOAuthTokenProviderLoggerReceivesMint(t *testing.T) { ts := httptest.NewServer(srv) defer ts.Close() - p, err := NewOAuthTokenProvider("id", "secret", "ws", ts.URL, "c.s.t", + 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()); err != nil { + if _, err := p.Token(context.Background(), "c.s.t"); err != nil { t.Fatalf("Token: %v", err) } out := buf.String() @@ -521,13 +522,13 @@ func TestAuthorizationDetailsContent(t *testing.T) { })) defer ts.Close() - p, err := NewOAuthTokenProvider("id", "secret", "ws", ts.URL, "mycat.mysch.mytbl", + 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()); err != nil { + if _, err := p.Token(context.Background(), "mycat.mysch.mytbl"); err != nil { t.Fatalf("Token: %v", err) } diff --git a/purego/internal/auth/token_cache.go b/purego/internal/auth/token_cache.go index f16a4aee..ea4d5938 100644 --- a/purego/internal/auth/token_cache.go +++ b/purego/internal/auth/token_cache.go @@ -11,13 +11,6 @@ import ( // cache proactively re-mints; 5 minutes is the SDK-wide standard. const defaultRefreshBuffer = 5 * time.Minute -// defaultNoTTLLifetime is the lifetime assigned to a token whose response omits -// (or reports a non-positive) expires_in. Without it such a token would never -// be cached and every Token call would pay a full HTTP round-trip; a bounded, -// conservative lifetime instead lets it be reused for a short window. UC tokens -// carry an expires_in in practice, so this is a safety net, not the norm. -const defaultNoTTLLifetime = 10 * 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 { @@ -83,18 +76,13 @@ func (c *cachedToken) isExpired() bool { } // newCachedToken builds an entry for a token whose TTL started at mintedAt -// (response-receipt time, matching the Rust SDK). A non-positive ttl falls back to -// defaultNoTTLLifetime so a token that the server didn't tag with an expires_in -// still gets a bounded cache window. +// (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 ttl <= 0 { - ttl = defaultNoTTLLifetime - } if maxBuffer := ttl / 2; refreshBuffer > maxBuffer { refreshBuffer = maxBuffer } @@ -248,16 +236,17 @@ func (c *tokenCache) getOrFetch( token := fetched.token - // Always cache a successful mint. When the server reports a usable TTL we - // honor it; when it omits expires_in newCachedToken falls back to a bounded - // default lifetime so the token is still reused for a short window rather - // than re-minted on every call. - ttl := time.Duration(0) - if fetched.expiresIn != nil { - ttl = *fetched.expiresIn + // 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 + } } - mintedAt := time.Now() - entry.cached = newCachedToken(token, ttl, c.refreshBuffer, mintedAt) flight.token, flight.err = token, nil close(flight.done) entry.mu.Unlock() diff --git a/purego/internal/auth/token_cache_test.go b/purego/internal/auth/token_cache_test.go index 54bc2ec0..f8d53522 100644 --- a/purego/internal/auth/token_cache_test.go +++ b/purego/internal/auth/token_cache_test.go @@ -97,7 +97,7 @@ func TestTokenCacheRotatedSecretGetsFreshEntry(t *testing.T) { } } -func TestTokenCacheNoTTLIsCachedUnderDefaultLifetime(t *testing.T) { +func TestTokenCacheNoTTLIsNotCached(t *testing.T) { c := newTokenCache() var calls atomic.Int64 mint := func(_ context.Context, _ mintReason) (fetchedToken, error) { @@ -108,13 +108,11 @@ func TestTokenCacheNoTTLIsCachedUnderDefaultLifetime(t *testing.T) { a := getOrFetch(t, c, "id", "secret", "c.s.t", mint) b := getOrFetch(t, c, "id", "secret", "c.s.t", mint) - // A token with no expires_in is now cached under the default lifetime and - // reused, rather than re-minted on every call. 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 for no-TTL token (cached), got %d", n) + if n := calls.Load(); n != 2 { + t.Fatalf("want 2 mints for no-TTL token (not cached), got %d", n) } } @@ -195,13 +193,13 @@ func TestTokenCacheNonRetryableRefreshErrorPropagates(t *testing.T) { } } -func TestTokenCacheNoTTLResponseReplacesCachedToken(t *testing.T) { +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; the caller gets the new token and it - // is now cached under the default lifetime (no longer discarded). + // 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 @@ -211,19 +209,18 @@ func TestTokenCacheNoTTLResponseReplacesCachedToken(t *testing.T) { t.Fatalf("want %q, got %q", "nottl", fresh) } - // The next call reuses the freshly cached no-TTL token rather than re-minting. - var minted atomic.Bool - tok := getOrFetch(t, c, "id", "secret", "c.s.t", + // 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) { - minted.Store(true) - return fetchedToken{token: "unexpected"}, nil + return fetchedToken{}, &retryErr{"blip"} }, ) - if tok != "nottl" { - t.Fatalf("want cached %q, got %q", "nottl", tok) + if err != nil { + t.Fatalf("want fallback to retained cached token, got error: %v", err) } - if minted.Load() { - t.Fatal("second call should hit cache, not re-mint") + if tok != "valid" { + t.Fatalf("want retained cached token %q, got %q", "valid", tok) } } @@ -417,12 +414,6 @@ func TestNewCachedToken(t *testing.T) { t.Fatalf("refreshAt = %v, want +55m", ct.refreshAt) } - // No TTL: falls back to the default lifetime. - ct = newCachedToken("v", 0, 5*time.Minute, now) - if !ct.expiresAt.Equal(now.Add(defaultNoTTLLifetime)) { - t.Fatalf("no-TTL expiresAt = %v, want +%v", ct.expiresAt, defaultNoTTLLifetime) - } - // 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)) { diff --git a/purego/internal/auth/token_provider.go b/purego/internal/auth/token_provider.go index 5a1a08de..0f9fd11a 100644 --- a/purego/internal/auth/token_provider.go +++ b/purego/internal/auth/token_provider.go @@ -1,15 +1,20 @@ -// Package auth provides OAuth 2.0 authentication for the Zerobus pure-Go SDK. +// Package auth provides OAuth 2.0 and custom-header authentication for the +// Zerobus pure-Go SDK. // -// The central abstraction is [TokenProvider], an interface for obtaining a -// bearer token on demand. Call its Token method before opening a stream and -// pass the result as [transport.StreamParams].Token. +// It exposes two layers: +// - [TokenProvider], for obtaining bearer tokens. +// - [HeadersProvider], matching other SDKs' custom-header auth model. // -// The package ships two implementations: -// - [OAuthTokenProvider] — Unity Catalog OAuth 2.0 client credentials flow, -// with per-table token caching and proactive refresh. This is the default -// for production use. -// - [StaticTokenProvider] — wraps a fixed string; useful for tests or when -// the caller manages token lifecycle outside the SDK. +// 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 ( @@ -28,13 +33,29 @@ import ( // "Basic …") or be a bare token — [transport.StreamParams].Token normalises // either form. // -// Invalidate is called when the server rejects the last token returned by -// Token with an authentication error, signalling that any cached credential is -// stale and must be discarded so the next Token call re-mints. Implementations -// that hold no cache may make Invalidate a no-op. +// 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) (string, error) - Invalidate(ctx context.Context) + 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 @@ -53,7 +74,7 @@ func NewStaticTokenProvider(token string) *StaticTokenProvider { } // Token returns the static token. -func (p *StaticTokenProvider) Token(_ context.Context) (string, error) { +func (p *StaticTokenProvider) Token(_ context.Context, _ string) (string, error) { if p.token == "" { return "", fmt.Errorf("auth: static token is empty") } @@ -61,4 +82,68 @@ func (p *StaticTokenProvider) Token(_ context.Context) (string, error) { } // Invalidate is a no-op for a static token. -func (p *StaticTokenProvider) Invalidate(_ context.Context) {} +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 db919be3..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,9 +75,21 @@ 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) } - // Reject control chars at the wire boundary so every TokenProvider is - // covered, not just the OAuth mint path; gRPC would otherwise fail opaquely. - if !isUsableAsHeader(p.Token) { + 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) } @@ -80,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 @@ -90,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 } @@ -103,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. diff --git a/purego/internal/transport/metadata.go b/purego/internal/transport/metadata.go index 51f4de5b..051f22ee 100644 --- a/purego/internal/transport/metadata.go +++ b/purego/internal/transport/metadata.go @@ -24,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) diff --git a/purego/internal/transport/transport_test.go b/purego/internal/transport/transport_test.go index b0041908..f691ec9c 100644 --- a/purego/internal/transport/transport_test.go +++ b/purego/internal/transport/transport_test.go @@ -5,11 +5,14 @@ import ( "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" @@ -52,6 +55,8 @@ type fakeServer struct { // 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{} @@ -82,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 { @@ -135,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 { @@ -387,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)