From 847169ae36bbee2f4da4dc4a6c2f561910344f71 Mon Sep 17 00:00:00 2001 From: Haytham Abuelfutuh Date: Wed, 1 May 2024 12:22:42 -0700 Subject: [PATCH] Update interface implementation Signed-off-by: Haytham Abuelfutuh --- go.mod | 2 +- go.sum | 4 +- pkg/pkce/token_cache_keyring.go | 148 +++++++++++++++++--------------- 3 files changed, 84 insertions(+), 70 deletions(-) diff --git a/go.mod b/go.mod index 2fc09823..f55c996f 100644 --- a/go.mod +++ b/go.mod @@ -178,7 +178,7 @@ require ( ) replace ( - github.com/flyteorg/flyte/flyteidl => github.com/flyteorg/flyte/flyteidl v1.11.1-b1.0.20240429144225-f947ff324417 + github.com/flyteorg/flyte/flyteidl => github.com/flyteorg/flyte/flyteidl v1.11.1-b1.0.20240501191822-99436ee05b08 github.com/flyteorg/flyte/flyteplugins => github.com/flyteorg/flyte/flyteplugins v1.11.1-b1.0.20240427064045-c2a1e27e0a92 github.com/flyteorg/flyte/flytepropeller => github.com/flyteorg/flyte/flytepropeller v1.11.1-b1.0.20240427064045-c2a1e27e0a92 github.com/flyteorg/flyte/flytestdlib => github.com/flyteorg/flyte/flytestdlib v1.11.1-b1.0.20240427064045-c2a1e27e0a92 diff --git a/go.sum b/go.sum index 4b60c447..5c351b15 100644 --- a/go.sum +++ b/go.sum @@ -352,8 +352,8 @@ github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w= github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= -github.com/flyteorg/flyte/flyteidl v1.11.1-b1.0.20240429144225-f947ff324417 h1:hU/LX8HUmZ8N/155lZbTPdt2Wcm3/3kp2RjMo6JERSQ= -github.com/flyteorg/flyte/flyteidl v1.11.1-b1.0.20240429144225-f947ff324417/go.mod h1:4nlOzFdnc+9cQMsaXJEECek3hkvKmo1RxpDAGiELOCo= +github.com/flyteorg/flyte/flyteidl v1.11.1-b1.0.20240501191822-99436ee05b08 h1:dpv5Ln+2V2wdndVVeAw5r5V3/jQOhID4o4cyRcudTdQ= +github.com/flyteorg/flyte/flyteidl v1.11.1-b1.0.20240501191822-99436ee05b08/go.mod h1:4nlOzFdnc+9cQMsaXJEECek3hkvKmo1RxpDAGiELOCo= github.com/flyteorg/flyte/flyteplugins v1.11.1-b1.0.20240427064045-c2a1e27e0a92 h1:mtX9C2NSbrUUBoUFL2S5puQyHo4zuu6+b5VKam1vFUk= github.com/flyteorg/flyte/flyteplugins v1.11.1-b1.0.20240427064045-c2a1e27e0a92/go.mod h1:bnW+Jb8u60I7FlufsCu/nE7XZZOlY4m7ngWPA7YFnQc= github.com/flyteorg/flyte/flytepropeller v1.11.1-b1.0.20240427064045-c2a1e27e0a92 h1:kH0APNrAaNFoEj3RIXj2BKg1M86jX4HajjBRuHvbKqg= diff --git a/pkg/pkce/token_cache_keyring.go b/pkg/pkce/token_cache_keyring.go index c28af1d9..87efc040 100644 --- a/pkg/pkce/token_cache_keyring.go +++ b/pkg/pkce/token_cache_keyring.go @@ -1,97 +1,111 @@ package pkce import ( - "encoding/json" - "fmt" - "sync" + "encoding/json" + "fmt" + "sync" - "github.com/zalando/go-keyring" - "golang.org/x/oauth2" + "github.com/flyteorg/flyte/flyteidl/clients/go/admin/cache" + "github.com/zalando/go-keyring" + "golang.org/x/oauth2" ) const ( - KeyRingServiceUser = "flytectl-user" - KeyRingServiceName = "flytectl" -) - -var ( - ErrNotFound = fmt.Errorf("secret not found in keyring") + KeyRingServiceUser = "flytectl-user" + KeyRingServiceName = "flytectl" ) // TokenCacheKeyringProvider wraps the logic to save and retrieve tokens from the OS's keyring implementation. type TokenCacheKeyringProvider struct { - ServiceName string - ServiceUser string - mu *sync.Mutex + ServiceName string + ServiceUser string + mu *sync.Mutex + cond *sync.Cond } func (t *TokenCacheKeyringProvider) PurgeIfEquals(existing *oauth2.Token) (bool, error) { - if existingBytes, err := json.Marshal(existing); err != nil { - return false, fmt.Errorf("unable to marshal token to save in cache due to %w", err) - } else if tokenJSON, err := keyring.Get(t.ServiceName, t.ServiceUser); err != nil { - if err.Error() == "secret not found in keyring" { - return false, fmt.Errorf("unable to read token from cache. Error: %w", ErrNotFound) - } - - return false, fmt.Errorf("unable to read token from cache. Error: %w", err) - } else if tokenJSON != string(existingBytes) { - return false, nil - } - - _ = keyring.Delete(t.ServiceName, t.ServiceUser) - return true, nil + if existingBytes, err := json.Marshal(existing); err != nil { + return false, fmt.Errorf("unable to marshal token to save in cache due to %w", err) + } else if tokenJSON, err := keyring.Get(t.ServiceName, t.ServiceUser); err != nil { + if err.Error() == "secret not found in keyring" { + return false, fmt.Errorf("unable to read token from cache. Error: %w", cache.ErrNotFound) + } + + return false, fmt.Errorf("unable to read token from cache. Error: %w", err) + } else if tokenJSON != string(existingBytes) { + return false, nil + } + + _ = keyring.Delete(t.ServiceName, t.ServiceUser) + return true, nil } func (t *TokenCacheKeyringProvider) Lock() { - t.mu.Lock() + t.mu.Lock() } func (t *TokenCacheKeyringProvider) Unlock() { - t.mu.Unlock() + t.mu.Unlock() +} + +// TryLock the cache. +func (t *TokenCacheKeyringProvider) TryLock() bool { + return t.mu.TryLock() +} + +// CondWait waits for the condition to be true. +func (t *TokenCacheKeyringProvider) CondWait() { + t.cond.Wait() +} + +// CondBroadcast signals the condition. +func (t *TokenCacheKeyringProvider) CondBroadcast() { + t.cond.Broadcast() } func (t *TokenCacheKeyringProvider) SaveToken(token *oauth2.Token) error { - var tokenBytes []byte - if token.AccessToken == "" { - return fmt.Errorf("cannot save empty token with expiration %v", token.Expiry) - } - - var err error - if tokenBytes, err = json.Marshal(token); err != nil { - return fmt.Errorf("unable to marshal token to save in cache due to %w", err) - } - - // set token in keyring - if err = keyring.Set(t.ServiceName, t.ServiceUser, string(tokenBytes)); err != nil { - return fmt.Errorf("unable to save token. Error: %w", err) - } - - return nil + var tokenBytes []byte + if token.AccessToken == "" { + return fmt.Errorf("cannot save empty token with expiration %v", token.Expiry) + } + + var err error + if tokenBytes, err = json.Marshal(token); err != nil { + return fmt.Errorf("unable to marshal token to save in cache due to %w", err) + } + + // set token in keyring + if err = keyring.Set(t.ServiceName, t.ServiceUser, string(tokenBytes)); err != nil { + return fmt.Errorf("unable to save token. Error: %w", err) + } + + return nil } func (t *TokenCacheKeyringProvider) GetToken() (*oauth2.Token, error) { - // get saved token - tokenJSON, err := keyring.Get(t.ServiceName, t.ServiceUser) - if len(tokenJSON) == 0 { - return nil, fmt.Errorf("no token found in the cache") - } - - if err != nil { - return nil, err - } - - token := oauth2.Token{} - if err = json.Unmarshal([]byte(tokenJSON), &token); err != nil { - return nil, fmt.Errorf("unmarshalling error for saved token. Error: %w", err) - } - - return &token, nil + // get saved token + tokenJSON, err := keyring.Get(t.ServiceName, t.ServiceUser) + if len(tokenJSON) == 0 { + return nil, fmt.Errorf("no token found in the cache") + } + + if err != nil { + return nil, err + } + + token := oauth2.Token{} + if err = json.Unmarshal([]byte(tokenJSON), &token); err != nil { + return nil, fmt.Errorf("unmarshalling error for saved token. Error: %w", err) + } + + return &token, nil } func NewTokenCacheKeyringProvider(serviceName, serviceUser string) *TokenCacheKeyringProvider { - return &TokenCacheKeyringProvider{ - mu: &sync.Mutex{}, - ServiceName: serviceName, - ServiceUser: serviceUser, - } + return &TokenCacheKeyringProvider{ + mu: &sync.Mutex{}, + cond: sync.NewCond(&sync.Mutex{}), + ServiceName: serviceName, + ServiceUser: serviceUser, + } }