diff --git a/pkg/workloadmanager/client_cache.go b/pkg/workloadmanager/client_cache.go index 4e242b81..66e5eefd 100644 --- a/pkg/workloadmanager/client_cache.go +++ b/pkg/workloadmanager/client_cache.go @@ -30,6 +30,7 @@ import ( type clientCacheEntry struct { key string client *UserK8sClient + token string tokenExpiry time.Time // Token expiration time parsed from JWT element *list.Element } @@ -96,8 +97,7 @@ func NewClientCache(maxSize int) *ClientCache { // Get retrieves a client from cache based on key (service account) // Returns the client if found and cached token is not expired, nil otherwise -// Different tokens for the same service account can share the same client -func (c *ClientCache) Get(key string) *UserK8sClient { +func (c *ClientCache) Get(key, token string) *UserK8sClient { c.mu.Lock() defer c.mu.Unlock() @@ -105,6 +105,9 @@ func (c *ClientCache) Get(key string) *UserK8sClient { if !exists { return nil } + if entry.token != token { + return nil + } // Check if cached entry's token is expired if !entry.tokenExpiry.IsZero() && time.Now().After(entry.tokenExpiry) { @@ -133,6 +136,7 @@ func (c *ClientCache) Set(key, token string, client *UserK8sClient) { if entry, exists := c.cache[key]; exists { // Update existing entry entry.client = client + entry.token = token entry.tokenExpiry = tokenExpiry c.lruList.MoveToFront(entry.element) return @@ -147,6 +151,7 @@ func (c *ClientCache) Set(key, token string, client *UserK8sClient) { entry := &clientCacheEntry{ key: key, client: client, + token: token, tokenExpiry: tokenExpiry, } entry.element = c.lruList.PushFront(entry) diff --git a/pkg/workloadmanager/client_cache_test.go b/pkg/workloadmanager/client_cache_test.go index c0678d9f..5ce7f52f 100644 --- a/pkg/workloadmanager/client_cache_test.go +++ b/pkg/workloadmanager/client_cache_test.go @@ -171,7 +171,7 @@ func TestClientCache_SetAndGet(t *testing.T) { cache := NewClientCache(10) // Test Get on non-existent key - client := cache.Get("nonexistent-key") + client := cache.Get("nonexistent-key", "token") assert.Nil(t, client) // Test Set and Get @@ -186,7 +186,7 @@ func TestClientCache_SetAndGet(t *testing.T) { cache.Set(key, token, newClient) - retrieved := cache.Get(key) + retrieved := cache.Get(key, token) assert.NotNil(t, retrieved) assert.Equal(t, newClient, retrieved) assert.Equal(t, 1, cache.Size()) @@ -208,7 +208,7 @@ func TestClientCache_Get_ExpiredToken(t *testing.T) { cache.Set(key, token, client) // Should return nil because token is expired - retrieved := cache.Get(key) + retrieved := cache.Get(key, token) assert.Nil(t, retrieved) assert.Equal(t, 0, cache.Size(), "Expired entry should be removed") } @@ -230,11 +230,30 @@ func TestClientCache_Get_TokenWithoutExpiry(t *testing.T) { cache.Set(key, token, client) // Should return client because tokenExpiry is zero (no expiry check) - retrieved := cache.Get(key) + retrieved := cache.Get(key, token) assert.NotNil(t, retrieved) assert.Equal(t, client, retrieved) } +func TestClientCache_Get_TokenMismatch(t *testing.T) { + cache := NewClientCache(10) + + key := testCacheKey + token := createTestJWT(time.Now().Add(1 * time.Hour).Unix()) + scheme := runtime.NewScheme() + dynamicClient := dynamicfake.NewSimpleDynamicClient(scheme) + client := &UserK8sClient{ + dynamicClient: dynamicClient, + namespace: "default", + } + + cache.Set(key, token, client) + + retrieved := cache.Get(key, token+"-new") + assert.Nil(t, retrieved) + assert.Equal(t, 1, cache.Size()) +} + func TestClientCache_UpdateExisting(t *testing.T) { cache := NewClientCache(10) @@ -259,7 +278,10 @@ func TestClientCache_UpdateExisting(t *testing.T) { cache.Set(key, token2, client2) assert.Equal(t, 1, cache.Size(), "Size should not increase when updating") - retrieved := cache.Get(key) + oldClient := cache.Get(key, token1) + assert.Nil(t, oldClient, "Old token should not match updated cache entry") + + retrieved := cache.Get(key, token2) assert.Equal(t, client2, retrieved, "Should return updated client") } @@ -268,11 +290,11 @@ func TestClientCache_Eviction(t *testing.T) { scheme := runtime.NewScheme() dynamicClient := dynamicfake.NewSimpleDynamicClient(scheme) + token := createTestJWT(time.Now().Add(1 * time.Hour).Unix()) // Fill cache to max size for i := 0; i < 3; i++ { key := "default:sa" + string(rune('0'+i)) - token := createTestJWT(time.Now().Add(1 * time.Hour).Unix()) client := &UserK8sClient{ dynamicClient: dynamicClient, namespace: "default", @@ -294,9 +316,9 @@ func TestClientCache_Eviction(t *testing.T) { assert.Equal(t, 3, cache.Size(), "Size should remain at max") // First key should be evicted - assert.Nil(t, cache.Get("default:sa0")) + assert.Nil(t, cache.Get("default:sa0", token)) // New key should be present - assert.NotNil(t, cache.Get(newKey)) + assert.NotNil(t, cache.Get(newKey, newToken)) } func TestClientCache_LRUBehavior(t *testing.T) { @@ -317,7 +339,7 @@ func TestClientCache_LRUBehavior(t *testing.T) { } // Access first entry (should move to front) - cache.Get("default:sa0") + cache.Get("default:sa0", token) // Add new entry - should evict sa1 (least recently used) newKey := "default:sa3" @@ -328,13 +350,13 @@ func TestClientCache_LRUBehavior(t *testing.T) { cache.Set(newKey, token, newClient) // sa0 should still be present (was accessed) - assert.NotNil(t, cache.Get("default:sa0")) + assert.NotNil(t, cache.Get("default:sa0", token)) // sa1 should be evicted - assert.Nil(t, cache.Get("default:sa1")) + assert.Nil(t, cache.Get("default:sa1", token)) // sa2 should be present - assert.NotNil(t, cache.Get("default:sa2")) + assert.NotNil(t, cache.Get("default:sa2", token)) // sa3 should be present - assert.NotNil(t, cache.Get(newKey)) + assert.NotNil(t, cache.Get(newKey, token)) } func TestClientCache_Remove(t *testing.T) { @@ -374,7 +396,7 @@ func TestClientCache_Remove(t *testing.T) { cache.Remove(tt.key) assert.Equal(t, 0, cache.Size()) - assert.Nil(t, cache.Get(tt.key)) + assert.Nil(t, cache.Get(tt.key, "token")) }) } } diff --git a/pkg/workloadmanager/k8s_client.go b/pkg/workloadmanager/k8s_client.go index 2c44d8e3..301c93ad 100644 --- a/pkg/workloadmanager/k8s_client.go +++ b/pkg/workloadmanager/k8s_client.go @@ -197,7 +197,7 @@ func (c *K8sClient) GetOrCreateUserK8sClient(userToken, namespace, serviceAccoun cacheKey := makeCacheKey(namespace, serviceAccountName) // Try to get from cache - if cachedClient := c.clientCache.Get(cacheKey); cachedClient != nil { + if cachedClient := c.clientCache.Get(cacheKey, userToken); cachedClient != nil { return cachedClient, nil } diff --git a/pkg/workloadmanager/k8s_client_test.go b/pkg/workloadmanager/k8s_client_test.go index b14526eb..b16d016f 100644 --- a/pkg/workloadmanager/k8s_client_test.go +++ b/pkg/workloadmanager/k8s_client_test.go @@ -25,6 +25,7 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/labels" listersv1 "k8s.io/client-go/listers/core/v1" + "k8s.io/client-go/rest" ) // Helper function to create a pod with owner reference @@ -210,3 +211,31 @@ func TestGetSandboxPodIP_InvalidPodStatus(t *testing.T) { }) } } + +func TestGetOrCreateUserK8sClient_CacheByToken(t *testing.T) { + client := &K8sClient{ + baseConfig: &rest.Config{Host: "https://example.com"}, + clientCache: NewClientCache(10), + } + + firstClient, err := client.GetOrCreateUserK8sClient("first-token", "default", "runner") + assert.NoError(t, err) + + sameTokenClient, err := client.GetOrCreateUserK8sClient("first-token", "default", "runner") + assert.NoError(t, err) + if firstClient != sameTokenClient { + t.Errorf("expected same token to reuse cached client") + } + + secondClient, err := client.GetOrCreateUserK8sClient("second-token", "default", "runner") + assert.NoError(t, err) + if firstClient == secondClient { + t.Errorf("expected changed token to create a new client") + } + + otherServiceAccountClient, err := client.GetOrCreateUserK8sClient("first-token", "default", "worker") + assert.NoError(t, err) + if secondClient == otherServiceAccountClient { + t.Errorf("expected different service account to create a new client") + } +}