cache azure token for kubernetes

Signed-off-by: goenning <me@goenning.net>

.changeset/healthy-pets-mix.md

@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-kubernetes-backend': patch
+---
+
+cache and refresh Azure tokens to avoid excesive calls to Azure Identity

plugins/kubernetes-backend/src/kubernetes-auth-translator/AzureIdentityKubernetesAuthTranslator.test.ts

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2020 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { AccessToken, TokenCredential } from '@azure/identity';
+import { AzureIdentityKubernetesAuthTranslator } from './AzureIdentityKubernetesAuthTranslator';
+
+class StaticTokenCredential implements TokenCredential {
+  private count: number = 0;
+
+  constructor(private expiryInMs: number) {}
+
+  getToken(): Promise<AccessToken | null> {
+    this.count++;
+
+    return Promise.resolve({
+      token: `MY_TOKEN_${this.count}`,
+      expiresOnTimestamp: Date.now() + this.expiryInMs,
+    });
+  }
+}
+
+describe('AzureIdentityKubernetesAuthTranslator tests', () => {
+  const cd = {
+    authProvider: 'Azure',
+    name: 'My Cluster',
+    url: 'mycluster.privatelink.westeurope.azmk8s.io',
+  };
+
+  it('should decorate cluster with Azure token', async () => {
+    const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      new StaticTokenCredential(5 * 60 * 1000),
+    );
+
+    const response = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
+  });
+
+  it('should re-use token before expiry', async () => {
+    const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      new StaticTokenCredential(5 * 60 * 1000),
+    );
+
+    const response = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
+
+    const response2 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response2.serviceAccountToken).toEqual('MY_TOKEN_1');
+  });
+
+  it('should reissue new token 2 minutes befory expiry', async () => {
+    const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      new StaticTokenCredential(3 * 60 * 1000), // token expires in 3m
+    );
+
+    const response = await authTranslator.decorateClusterDetailsWithAuth({
+      authProvider: 'Azure',
+      name: 'My Cluster',
+      url: 'mycluster.privatelink.westeurope.azmk8s.io',
+    });
+    expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 1 * 60 * 1000); // advance time by 1min
+
+    const response2 = await authTranslator.decorateClusterDetailsWithAuth({
+      authProvider: 'Azure',
+      name: 'My Cluster',
+      url: 'mycluster.privatelink.westeurope.azmk8s.io',
+    });
+    expect(response2.serviceAccountToken).toEqual('MY_TOKEN_2');
+  });
+});

plugins/kubernetes-backend/src/kubernetes-auth-translator/AzureIdentityKubernetesAuthTranslator.ts

@@ -16,13 +16,24 @@
 
 import { KubernetesAuthTranslator } from './types';
 import { AzureClusterDetails } from '../types/types';
-import { DefaultAzureCredential } from '@azure/identity';
+import {
+  AccessToken,
+  DefaultAzureCredential,
+  TokenCredential,
+} from '@azure/identity';
 
 const aksScope = '6dae42f8-4368-4678-94ff-3960e28e3630/.default'; // This scope is the same for all Azure Managed Kubernetes
 
 export class AzureIdentityKubernetesAuthTranslator
   implements KubernetesAuthTranslator
 {
+  private tokenCredential: TokenCredential;
+  private accessToken: AccessToken | null = null;
+
+  constructor(tokenCredential?: TokenCredential) {
+    this.tokenCredential = tokenCredential || new DefaultAzureCredential();
+  }
+
   async decorateClusterDetailsWithAuth(
     clusterDetails: AzureClusterDetails,
   ): Promise<AzureClusterDetails> {
@@ -31,11 +42,23 @@ export class AzureIdentityKubernetesAuthTranslator
       clusterDetails,
     );
 
-    const credentials = new DefaultAzureCredential();
+    if (!this.accessToken || this.tokenExpired()) {
+      this.accessToken = await this.tokenCredential.getToken(aksScope);
 
-    // TODO: can we cache this? It's inneficiant to get a new token every time
-    const accessToken = await credentials.getToken(aksScope);
-    clusterDetailsWithAuthToken.serviceAccountToken = accessToken.token;
+      if (!this.accessToken) {
+        throw new Error('Unable to retrieve Azure token');
+      }
+    }
+
+    clusterDetailsWithAuthToken.serviceAccountToken = this.accessToken.token;
     return clusterDetailsWithAuthToken;
   }
+
+  private tokenExpired(): boolean {
+    if (!this.accessToken) return true;
+
+    // Set tokens to expire 2 minutes before its actual expiry time
+    const expiresOn = this.accessToken.expiresOnTimestamp - 2 * 60 * 1000;
+    return Date.now() >= expiresOn;
+  }
 }

plugins/kubernetes-backend/src/service/KubernetesFanOutHandler.ts

@@ -155,6 +155,7 @@ export class KubernetesFanOutHandler {
   private readonly serviceLocator: KubernetesServiceLocator;
   private readonly customResources: CustomResource[];
   private readonly objectTypesToFetch: Set<ObjectToFetch>;
+  private readonly authTranslators: Record<string, KubernetesAuthTranslator>;
 
   constructor({
     logger,
@@ -168,6 +169,7 @@ export class KubernetesFanOutHandler {
     this.serviceLocator = serviceLocator;
     this.customResources = customResources;
     this.objectTypesToFetch = new Set(objectTypesToFetch);
+    this.authTranslators = {};
   }
 
   async getKubernetesObjectsByEntity(
@@ -183,14 +185,9 @@ export class KubernetesFanOutHandler {
 
     // Execute all of these async actions simultaneously/without blocking sequentially as no common object is modified by them
     const promises: Promise<ClusterDetails>[] = clusterDetails.map(cd => {
-      const kubernetesAuthTranslator: KubernetesAuthTranslator =
-        KubernetesAuthTranslatorGenerator.getKubernetesAuthTranslatorInstance(
-          cd.authProvider,
-        );
-      return kubernetesAuthTranslator.decorateClusterDetailsWithAuth(
-        cd,
-        requestBody,
-      );
+      return this.getAuthTranslator(
+        cd.authProvider,
+      ).decorateClusterDetailsWithAuth(cd, requestBody);
     });
     const clusterDetailsDecoratedForAuth: ClusterDetails[] = await Promise.all(
       promises,
@@ -288,4 +285,16 @@ export class KubernetesFanOutHandler {
 
     return Promise.all([result, Promise.all(podMetrics)]);
   }
+
+  private getAuthTranslator(provider: string): KubernetesAuthTranslator {
+    if (this.authTranslators[provider]) {
+      return this.authTranslators[provider];
+    }
+
+    this.authTranslators[provider] =
+      KubernetesAuthTranslatorGenerator.getKubernetesAuthTranslatorInstance(
+        provider,
+      );
+    return this.authTranslators[provider];
+  }
 }