Also synchronize pkce cache

Signed-off-by: Haytham Abuelfutuh <[email protected]>
flyteorg · Apr 27, 2024 · aa5e6c6 · aa5e6c6
1 parent e41c5d8
commit aa5e6c6
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 9 deletions.
diff --git a/cmd/core/cmd.go b/cmd/core/cmd.go
@@ -73,10 +73,10 @@ func generateCommandFunc(cmdEntry CommandEntry) func(cmd *cobra.Command, args []
 		cmdCtx := NewCommandContextNoClient(cmd.OutOrStdout())
 		if !cmdEntry.DisableFlyteClient {
 			clientSet, err := admin.ClientSetBuilder().WithConfig(admin.GetConfig(ctx)).
-				WithTokenCache(pkce.TokenCacheKeyringProvider{
-					ServiceUser: fmt.Sprintf("%s:%s", adminCfg.Endpoint.String(), pkce.KeyRingServiceUser),
-					ServiceName: pkce.KeyRingServiceName,
-				}).Build(ctx)
+				WithTokenCache(pkce.NewTokenCacheKeyringProvider(
+					pkce.KeyRingServiceName,
+					fmt.Sprintf("%s:%s", adminCfg.Endpoint.String(), pkce.KeyRingServiceUser),
+				)).Build(ctx)
 			if err != nil {
 				return err
 			}

diff --git a/pkg/pkce/token_cache_keyring.go b/pkg/pkce/token_cache_keyring.go
@@ -3,6 +3,7 @@ package pkce
 import (
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	"github.com/zalando/go-keyring"
 	"golang.org/x/oauth2"
@@ -17,19 +18,22 @@ const (
 type TokenCacheKeyringProvider struct {
 	ServiceName string
 	ServiceUser string
+	mu          *sync.Mutex
 }
 
-func (t TokenCacheKeyringProvider) Purge() {
+func (t *TokenCacheKeyringProvider) Purge() {
 	_ = keyring.Delete(t.ServiceName, t.ServiceUser)
 }
 
-func (t TokenCacheKeyringProvider) Lock() {
+func (t *TokenCacheKeyringProvider) Lock() {
+	t.mu.Lock()
 }
 
-func (t TokenCacheKeyringProvider) Unlock() {
+func (t *TokenCacheKeyringProvider) Unlock() {
+	t.mu.Unlock()
 }
 
-func (t TokenCacheKeyringProvider) SaveToken(token *oauth2.Token) error {
+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)
@@ -48,7 +52,7 @@ func (t TokenCacheKeyringProvider) SaveToken(token *oauth2.Token) error {
 	return nil
 }
 
-func (t TokenCacheKeyringProvider) GetToken() (*oauth2.Token, error) {
+func (t *TokenCacheKeyringProvider) GetToken() (*oauth2.Token, error) {
 	// get saved token
 	tokenJSON, err := keyring.Get(t.ServiceName, t.ServiceUser)
 	if len(tokenJSON) == 0 {
@@ -66,3 +70,11 @@ func (t TokenCacheKeyringProvider) GetToken() (*oauth2.Token, error) {
 
 	return &token, nil
 }
+
+func NewTokenCacheKeyringProvider(serviceName, serviceUser string) *TokenCacheKeyringProvider {
+	return &TokenCacheKeyringProvider{
+		mu:          &sync.Mutex{},
+		ServiceName: serviceName,
+		ServiceUser: serviceUser,
+	}
+}