Skip to content

Commit 0443a77

Browse files
mtrmacmtrmac
authored
Merge pull request #415 from mtrmac/shared-token
Only obtain a bearer token once at a time
2 parents bf349ae + de2a06c commit 0443a77

File tree

2 files changed

+128
-72
lines changed

2 files changed

+128
-72
lines changed

‎image/docker/docker_client.go‎

Lines changed: 122 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import (
3535
"go.podman.io/image/v5/types"
3636
"go.podman.io/storage/pkg/fileutils"
3737
"go.podman.io/storage/pkg/homedir"
38+
"golang.org/x/sync/semaphore"
3839
)
3940

4041
const (
@@ -86,8 +87,19 @@ type extensionSignatureList struct {
8687
Signatures []extensionSignature `json:"signatures"`
8788
}
8889

89-
// bearerToken records a cached token we can use to authenticate.
90+
// bearerToken records a cached token we can use to authenticate, or a pending process to obtain one.
91+
//
92+
// The goroutine obtaining the token holds lock to block concurrent token requests, and fills the structure (err and possibly the other fields)
93+
// before releasing the lock.
94+
// Other goroutines obtain lock to block on the token request, if any; and then inspect err to see if the token is usable.
95+
// If it is not, they try to get a new one.
9096
type bearerToken struct {
97+
// lock is held while obtaining the token. Potentially nested inside dockerClient.tokenCacheLock.
98+
// This is a counting semaphore only because we need a cancellable lock operation.
99+
lock *semaphore.Weighted
100+
101+
// The following fields can only be accessed with lock held.
102+
err error // nil if the token was successfully obtained (but may be expired); an error if the next lock holder _must_ obtain a new token.
91103
token string
92104
expirationTime time.Time
93105
}
@@ -116,8 +128,9 @@ type dockerClient struct {
116128
challenges []challenge
117129
supportsSignatures bool
118130

119-
// Private state for setupRequestAuth (key: string, value: bearerToken)
120-
tokenCache sync.Map
131+
// Private state for setupRequestAuth
132+
tokenCacheLock sync.Mutex // Protects tokenCache.
133+
tokenCache map[string]*bearerToken
121134
// Private state for detectProperties:
122135
detectPropertiesOnce sync.Once // detectPropertiesOnce is used to execute detectProperties() at most once.
123136
detectPropertiesError error // detectPropertiesError caches the initial error.
@@ -274,6 +287,7 @@ func newDockerClient(sys *types.SystemContext, registry, reference string) (*doc
274287
registry: registry,
275288
userAgent: userAgent,
276289
tlsClientConfig: tlsClientConfig,
290+
tokenCache: map[string]*bearerToken{},
277291
reportedWarnings: set.New[string](),
278292
}, nil
279293
}
@@ -716,50 +730,11 @@ func (c *dockerClient) setupRequestAuth(req *http.Request, extraScope *authScope
716730
req.SetBasicAuth(c.auth.Username, c.auth.Password)
717731
return nil
718732
case "bearer":
719-
registryToken := c.registryToken
720-
if registryToken == "" {
721-
cacheKey := ""
722-
scopes := []authScope{c.scope}
723-
if extraScope != nil {
724-
// Using ':' as a separator here is unambiguous because getBearerToken below
725-
// uses the same separator when formatting a remote request (and because
726-
// repository names that we create can't contain colons, and extraScope values
727-
// coming from a server come from `parseAuthScope`, which also splits on colons).
728-
cacheKey = fmt.Sprintf("%s:%s:%s", extraScope.resourceType, extraScope.remoteName, extraScope.actions)
729-
if colonCount := strings.Count(cacheKey, ":"); colonCount != 2 {
730-
return fmt.Errorf(
731-
"Internal error: there must be exactly 2 colons in the cacheKey ('%s') but got %d",
732-
cacheKey,
733-
colonCount,
734-
)
735-
}
736-
scopes = append(scopes, *extraScope)
737-
}
738-
var token bearerToken
739-
t, inCache := c.tokenCache.Load(cacheKey)
740-
if inCache {
741-
token = t.(bearerToken)
742-
}
743-
if !inCache || time.Now().After(token.expirationTime) {
744-
var (
745-
t *bearerToken
746-
err error
747-
)
748-
if c.auth.IdentityToken != "" {
749-
t, err = c.getBearerTokenOAuth2(req.Context(), challenge, scopes)
750-
} else {
751-
t, err = c.getBearerToken(req.Context(), challenge, scopes)
752-
}
753-
if err != nil {
754-
return err
755-
}
756-
757-
token = *t
758-
c.tokenCache.Store(cacheKey, token)
759-
}
760-
registryToken = token.token
733+
token, err := c.obtainBearerToken(req.Context(), challenge, extraScope)
734+
if err != nil {
735+
return err
761736
}
762-
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", registryToken))
737+
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
763738
return nil
764739
default:
765740
logrus.Debugf("no handler for %s authentication", challenge.Scheme)
@@ -769,16 +744,94 @@ func (c *dockerClient) setupRequestAuth(req *http.Request, extraScope *authScope
769744
return nil
770745
}
771746

772-
func (c *dockerClient) getBearerTokenOAuth2(ctx context.Context, challenge challenge,
773-
scopes []authScope) (*bearerToken, error) {
747+
// obtainBearerToken gets an "Authorization: Bearer" token if one is available, or obtains a fresh one.
748+
func (c *dockerClient) obtainBearerToken(ctx context.Context, challenge challenge, extraScope *authScope) (string, error) {
749+
if c.registryToken != "" {
750+
return c.registryToken, nil
751+
}
752+
753+
cacheKey := ""
754+
scopes := []authScope{c.scope}
755+
if extraScope != nil {
756+
// Using ':' as a separator here is unambiguous because getBearerToken below
757+
// uses the same separator when formatting a remote request (and because
758+
// repository names that we create can't contain colons, and extraScope values
759+
// coming from a server come from `parseAuthScope`, which also splits on colons).
760+
cacheKey = fmt.Sprintf("%s:%s:%s", extraScope.resourceType, extraScope.remoteName, extraScope.actions)
761+
if colonCount := strings.Count(cacheKey, ":"); colonCount != 2 {
762+
return "", fmt.Errorf(
763+
"Internal error: there must be exactly 2 colons in the cacheKey ('%s') but got %d",
764+
cacheKey,
765+
colonCount,
766+
)
767+
}
768+
scopes = append(scopes, *extraScope)
769+
}
770+
771+
token, newEntry, err := func() (*bearerToken, bool, error) { // A scope for defer
772+
c.tokenCacheLock.Lock()
773+
defer c.tokenCacheLock.Unlock()
774+
token, ok := c.tokenCache[cacheKey]
775+
if ok {
776+
return token, false, nil
777+
} else {
778+
token = &bearerToken{
779+
lock: semaphore.NewWeighted(1),
780+
}
781+
// If this is a new *bearerToken, lock the entry before adding it to the cache, so that any other goroutine that finds
782+
// this entry blocks until we obtain the token for the first time, and does not see an empty object
783+
// (and does not try to obtain the token itself when we are going to do so).
784+
if err := token.lock.Acquire(ctx, 1); err != nil {
785+
// We do not block on this Acquire, so we don’t really expect to fail here — but if ctx is canceled,
786+
// there is no point in trying to continue anyway.
787+
return nil, false, err
788+
}
789+
c.tokenCache[cacheKey] = token
790+
return token, true, nil
791+
}
792+
}()
793+
if err != nil {
794+
return "", err
795+
}
796+
if !newEntry {
797+
// If this is an existing *bearerToken, obtain the lock only after releasing c.tokenCacheLock,
798+
// so that users of other cacheKey values are not blocked for the whole duration of our HTTP roundtrip.
799+
if err := token.lock.Acquire(ctx, 1); err != nil {
800+
return "", err
801+
}
802+
}
803+
804+
defer token.lock.Release(1)
805+
806+
if !newEntry && token.err == nil && !time.Now().After(token.expirationTime) {
807+
return token.token, nil // We have a usable token already.
808+
}
809+
810+
if c.auth.IdentityToken != "" {
811+
err = c.getBearerTokenOAuth2(ctx, token, challenge, scopes)
812+
} else {
813+
err = c.getBearerToken(ctx, token, challenge, scopes)
814+
}
815+
token.err = err
816+
if token.err != nil {
817+
return "", token.err
818+
}
819+
return token.token, nil
820+
}
821+
822+
// getBearerTokenOAuth2 obtains an "Authorization: Bearer" token using a pre-existing identity token per
823+
// https://github.com/distribution/distribution/blob/main/docs/spec/auth/oauth.md for challenge and scopes,
824+
// and writes it into dest.
825+
func (c *dockerClient) getBearerTokenOAuth2(ctx context.Context, dest *bearerToken, challenge challenge,
826+
scopes []authScope) error {
774827
realm, ok := challenge.Parameters["realm"]
775828
if !ok {
776-
return nil, errors.New("missing realm in bearer auth challenge")
829+
return errors.New("missing realm in bearer auth challenge")
777830
}
778831

779832
authReq, err := http.NewRequestWithContext(ctx, http.MethodPost, realm, nil)
780833
if err != nil {
781-
return nil, err
834+
return err
782835
}
783836

784837
// Make the form data required against the oauth2 authentication
@@ -803,26 +856,29 @@ func (c *dockerClient) getBearerTokenOAuth2(ctx context.Context, challenge chall
803856
logrus.Debugf("%s %s", authReq.Method, authReq.URL.Redacted())
804857
res, err := c.client.Do(authReq)
805858
if err != nil {
806-
return nil, err
859+
return err
807860
}
808861
defer res.Body.Close()
809862
if err := httpResponseToError(res, "Trying to obtain access token"); err != nil {
810-
return nil, err
863+
return err
811864
}
812865

813-
return newBearerTokenFromHTTPResponseBody(res)
866+
return dest.readFromHTTPResponseBody(res)
814867
}
815868

816-
func (c *dockerClient) getBearerToken(ctx context.Context, challenge challenge,
817-
scopes []authScope) (*bearerToken, error) {
869+
// getBearerToken obtains an "Authorization: Bearer" token using a GET request, per
870+
// https://github.com/distribution/distribution/blob/main/docs/spec/auth/token.md for challenge and scopes,
871+
// and writes it into dest.
872+
func (c *dockerClient) getBearerToken(ctx context.Context, dest *bearerToken, challenge challenge,
873+
scopes []authScope) error {
818874
realm, ok := challenge.Parameters["realm"]
819875
if !ok {
820-
return nil, errors.New("missing realm in bearer auth challenge")
876+
return errors.New("missing realm in bearer auth challenge")
821877
}
822878

823879
authReq, err := http.NewRequestWithContext(ctx, http.MethodGet, realm, nil)
824880
if err != nil {
825-
return nil, err
881+
return err
826882
}
827883

828884
params := authReq.URL.Query()
@@ -850,22 +906,22 @@ func (c *dockerClient) getBearerToken(ctx context.Context, challenge challenge,
850906
logrus.Debugf("%s %s", authReq.Method, authReq.URL.Redacted())
851907
res, err := c.client.Do(authReq)
852908
if err != nil {
853-
return nil, err
909+
return err
854910
}
855911
defer res.Body.Close()
856912
if err := httpResponseToError(res, "Requesting bearer token"); err != nil {
857-
return nil, err
913+
return err
858914
}
859915

860-
return newBearerTokenFromHTTPResponseBody(res)
916+
return dest.readFromHTTPResponseBody(res)
861917
}
862918

863-
// newBearerTokenFromHTTPResponseBody parses a http.Response to obtain a bearerToken.
919+
// readFromHTTPResponseBody sets token data by parsing a http.Response.
864920
// The caller is still responsible for ensuring res.Body is closed.
865-
func newBearerTokenFromHTTPResponseBody(res *http.Response) (*bearerToken, error) {
921+
func (bt *bearerToken) readFromHTTPResponseBody(res *http.Response) error {
866922
blob, err := iolimits.ReadAtMost(res.Body, iolimits.MaxAuthTokenBodySize)
867923
if err != nil {
868-
return nil, err
924+
return err
869925
}
870926

871927
var token struct {
@@ -881,12 +937,10 @@ func newBearerTokenFromHTTPResponseBody(res *http.Response) (*bearerToken, error
881937
if len(bodySample) > bodySampleLength {
882938
bodySample = bodySample[:bodySampleLength]
883939
}
884-
return nil, fmt.Errorf("decoding bearer token (last URL %q, body start %q): %w", res.Request.URL.Redacted(), string(bodySample), err)
940+
return fmt.Errorf("decoding bearer token (last URL %q, body start %q): %w", res.Request.URL.Redacted(), string(bodySample), err)
885941
}
886942

887-
bt := &bearerToken{
888-
token: token.Token,
889-
}
943+
bt.token = token.Token
890944
if bt.token == "" {
891945
bt.token = token.AccessToken
892946
}
@@ -899,7 +953,7 @@ func newBearerTokenFromHTTPResponseBody(res *http.Response) (*bearerToken, error
899953
token.IssuedAt = time.Now().UTC()
900954
}
901955
bt.expirationTime = token.IssuedAt.Add(time.Duration(token.ExpiresIn) * time.Second)
902-
return bt, nil
956+
return nil
903957
}
904958

905959
// detectPropertiesHelper performs the work of detectProperties which executes

‎image/docker/docker_client_test.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ func testTokenHTTPResponse(t *testing.T, body string) *http.Response {
106106
}
107107
}
108108

109-
func TestNewBearerTokenFromHTTPResponseBody(t *testing.T) {
109+
func TestBearerTokenReadFromHTTPResponseBody(t *testing.T) {
110110
for _, c := range []struct {
111111
input string
112112
expected *bearerToken // or nil if a failure is expected
@@ -128,7 +128,8 @@ func TestNewBearerTokenFromHTTPResponseBody(t *testing.T) {
128128
expected: &bearerToken{token: "IAmAToken", expirationTime: time.Unix(1514800802+60, 0)},
129129
},
130130
} {
131-
token, err := newBearerTokenFromHTTPResponseBody(testTokenHTTPResponse(t, c.input))
131+
token := &bearerToken{}
132+
err := token.readFromHTTPResponseBody(testTokenHTTPResponse(t, c.input))
132133
if c.expected == nil {
133134
assert.Error(t, err, c.input)
134135
} else {
@@ -140,11 +141,12 @@ func TestNewBearerTokenFromHTTPResponseBody(t *testing.T) {
140141
}
141142
}
142143

143-
func TestNewBearerTokenFromHTTPResponseBodyIssuedAtZero(t *testing.T) {
144+
func TestBearerTokenReadFromHTTPResponseBodyIssuedAtZero(t *testing.T) {
144145
zeroTime := time.Time{}.Format(time.RFC3339)
145146
now := time.Now()
146147
tokenBlob := fmt.Sprintf(`{"token":"IAmAToken","expires_in":100,"issued_at":"%s"}`, zeroTime)
147-
token, err := newBearerTokenFromHTTPResponseBody(testTokenHTTPResponse(t, tokenBlob))
148+
token := &bearerToken{}
149+
err := token.readFromHTTPResponseBody(testTokenHTTPResponse(t, tokenBlob))
148150
require.NoError(t, err)
149151
expectedExpiration := now.Add(time.Duration(100) * time.Second)
150152
require.False(t, token.expirationTime.Before(expectedExpiration),

0 commit comments

Comments
 (0)