This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new 776854007 fix(catalog/rest)!: make OAuth token refresh context-aware
and singleflight (#2265)
776854007 is described below
commit 7768540070d5702bde7b80f2d7fd57de591f3a3e
Author: Ilia Demianenko <[email protected]>
AuthorDate: Thu Oct 22 14:25:47 2026 -0600
fix(catalog/rest)!: make OAuth token refresh context-aware and singleflight
(#2265)
AuthHeaderWithContext checked ctx once and then called Token()
synchronously, so a caller blocked on a stalled token endpoint until the
refresh client's Timeout regardless of its own deadline. Concurrent
callers each ran their own refresh in turn behind the token source's
mutex.
Share one in-flight Token call per manager with a singleflight.Group and
let the caller wait for it only until ctx is done. The refresh keeps
running, bounded by the refresh client's Timeout, and its result serves
later callers. AuthHeader goes through the same group.
BREAKING CHANGE: Oauth2AuthManager is no longer comparable. The type has
only unexported fields and pointer-receiver methods, so no caller can
hold two meaningful values to compare.
Signed-off-by: Ilia Demianenko <[email protected]>
---
catalog/rest/auth.go | 37 +++++++++++++++++-------
catalog/rest/auth_test.go | 73 +++++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 99 insertions(+), 11 deletions(-)
diff --git a/catalog/rest/auth.go b/catalog/rest/auth.go
index 79db6534e..4789eb78d 100644
--- a/catalog/rest/auth.go
+++ b/catalog/rest/auth.go
@@ -23,6 +23,7 @@ import (
"fmt"
"golang.org/x/oauth2"
+ "golang.org/x/sync/singleflight"
)
// AuthManager is an interface for providing custom authorization headers.
@@ -49,32 +50,44 @@ type ContextAuthManager interface {
// handles caching, thread-safe refresh, and expiry management.
type Oauth2AuthManager struct {
tokenSource oauth2.TokenSource
+ // refresh shares one in-flight Token call, so a caller that stops
waiting
+ // leaves nothing queued behind it.
+ refresh singleflight.Group
}
var _ ContextAuthManager = (*Oauth2AuthManager)(nil)
// AuthHeader returns the authorization header with the bearer token.
func (o *Oauth2AuthManager) AuthHeader() (string, string, error) {
- return o.authHeader()
+ tok, err, _ := o.refresh.Do("token", o.fetchToken)
+
+ return o.authHeader(tok, err)
}
// AuthHeaderWithContext returns the authorization header, honoring ctx. A
-// cached, unexpired token is returned without any network call, so a bounded
-// caller (such as the metrics dispatcher) pays no auth latency in the common
-// case. When a refresh is required it goes through the token-refresh HTTP
-// client, whose Timeout bounds how long a stalled token endpoint can block;
the
-// ctx check here additionally short-circuits work for a caller whose deadline
-// has already elapsed.
+// cached, unexpired token is returned without a network call. When a refresh
+// is required, the caller waits for it only until ctx is done; the refresh
+// itself keeps running, bounded by the token-refresh client's Timeout, and its
+// result serves later callers.
func (o *Oauth2AuthManager) AuthHeaderWithContext(ctx context.Context)
(string, string, error) {
if err := ctx.Err(); err != nil {
return "", "", err
}
- return o.authHeader()
+ select {
+ case <-ctx.Done():
+ return "", "", ctx.Err()
+ case result := <-o.refresh.DoChan("token", o.fetchToken):
+ return o.authHeader(result.Val, result.Err)
+ }
}
-func (o *Oauth2AuthManager) authHeader() (string, string, error) {
- tok, err := o.tokenSource.Token()
+// fetchToken adapts the token source to singleflight's untyped callback.
+func (o *Oauth2AuthManager) fetchToken() (any, error) {
+ return o.tokenSource.Token()
+}
+
+func (o *Oauth2AuthManager) authHeader(tok any, err error) (string, string,
error) {
if err != nil {
if re, ok := errors.AsType[*oauth2.RetrieveError](err); ok {
return "", "", oauthError{
@@ -87,7 +100,9 @@ func (o *Oauth2AuthManager) authHeader() (string, string,
error) {
return "", "", fmt.Errorf("%w: %s", ErrOAuthError, err)
}
- return "Authorization", tok.Type() + " " + tok.AccessToken, nil
+ token := tok.(*oauth2.Token)
+
+ return "Authorization", token.Type() + " " + token.AccessToken, nil
}
// oauthError wraps OAuth2 error details and implements the error chain
diff --git a/catalog/rest/auth_test.go b/catalog/rest/auth_test.go
index f7bb0534b..165be30bc 100644
--- a/catalog/rest/auth_test.go
+++ b/catalog/rest/auth_test.go
@@ -23,7 +23,10 @@ import (
"errors"
"net/http"
"net/http/httptest"
+ "sync"
+ "sync/atomic"
"testing"
+ "time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -117,3 +120,73 @@ func
TestOauth2AuthManager_AuthHeader_FetchToken_ErrorResponse(t *testing.T) {
assert.Contains(t, err.Error(), "invalid_client")
assert.Contains(t, err.Error(), "Invalid client credentials")
}
+
+func TestOauth2AuthManager_AuthHeaderWithContext_CallerDeadline(t *testing.T) {
+ mux := http.NewServeMux()
+ server := httptest.NewServer(mux)
+ defer server.Close()
+
+ // The token endpoint stalls until the test ends. Nothing cancels the
refresh
+ // request itself, so the release must come before server.Close waits
on it.
+ release := make(chan struct{})
+ defer close(release)
+ mux.HandleFunc("/oauth/token", func(w http.ResponseWriter, r
*http.Request) {
+ <-release
+ })
+
+ cfg := &clientcredentials.Config{
+ ClientID: "client",
+ ClientSecret: "secret",
+ TokenURL: server.URL + "/oauth/token",
+ AuthStyle: oauth2.AuthStyleInParams,
+ }
+
+ tokenCtx := context.WithValue(context.Background(), oauth2.HTTPClient,
server.Client())
+ manager := &Oauth2AuthManager{
+ tokenSource: cfg.TokenSource(tokenCtx),
+ }
+
+ ctx, cancel := context.WithTimeout(context.Background(),
50*time.Millisecond)
+ defer cancel()
+ _, _, err := manager.AuthHeaderWithContext(ctx)
+
+ require.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
+// blockingTokenSource counts Token calls and holds each until released.
+type blockingTokenSource struct {
+ calls atomic.Int32
+ release chan struct{}
+}
+
+func (s *blockingTokenSource) Token() (*oauth2.Token, error) {
+ s.calls.Add(1)
+ <-s.release
+
+ return &oauth2.Token{AccessToken: "shared_token", TokenType: "Bearer"},
nil
+}
+
+func TestOauth2AuthManager_AuthHeaderWithContext_SharesRefresh(t *testing.T) {
+ source := &blockingTokenSource{release: make(chan struct{})}
+ manager := &Oauth2AuthManager{tokenSource: source}
+
+ // Concurrent callers give up before the refresh completes.
+ var callers sync.WaitGroup
+ for range 5 {
+ callers.Go(func() {
+ ctx, cancel :=
context.WithTimeout(context.Background(), 50*time.Millisecond)
+ defer cancel()
+ _, _, err := manager.AuthHeaderWithContext(ctx)
+ assert.ErrorIs(t, err, context.DeadlineExceeded)
+ })
+ }
+ callers.Wait()
+ assert.EqualValues(t, 1, source.calls.Load(), "waiting callers must
share one refresh")
+
+ // The shared refresh completes in the background and serves later
callers.
+ close(source.release)
+ key, value, err := manager.AuthHeader()
+ require.NoError(t, err)
+ assert.Equal(t, "Authorization", key)
+ assert.Equal(t, "Bearer shared_token", value)
+}
