From aa5e6c6b7179d18cc9ddab87de3e8bbb17b61597 Mon Sep 17 00:00:00 2001
From: Haytham Abuelfutuh <haytham@afutuh.com>
Date: Sat, 27 Apr 2024 16:56:40 -0700
Subject: [PATCH] Also synchronize pkce cache

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>
---
 cmd/core/cmd.go                 |  8 ++++----
 pkg/pkce/token_cache_keyring.go | 22 +++++++++++++++++-----
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/cmd/core/cmd.go b/cmd/core/cmd.go
index 2d00adc0..a64c9917 100644
--- 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
index e0b9c77f..20343436 100644
--- 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,
+	}
+}