Merge pull request #11485 from goenning/go/cache-azure-tokens

[Kubernetes Plugin] Cache Azure tokens to avoid excessive calls to identity

.changeset/healthy-pets-mix.md

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

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

@@ -0,0 +1,130 @@
+/*
+ * 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 { getVoidLogger } from '@backstage/backend-common';
+import { AzureIdentityKubernetesAuthTranslator } from './AzureIdentityKubernetesAuthTranslator';
+
+const logger = getVoidLogger();
+
+class StaticTokenCredential implements TokenCredential {
+  private count: number = 0;
+
+  constructor(private expiryInMs: number) {}
+
+  getToken(): Promise<AccessToken | null> {
+    this.count++;
+
+    if (this.count === 3) {
+      return Promise.reject(new Error('Third time never works.'));
+    }
+
+    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(
+      logger,
+      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(
+      logger,
+      new StaticTokenCredential(20 * 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 issue new token 15 minutes befory expiry', async () => {
+    const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      logger,
+      new StaticTokenCredential(16 * 60 * 1000), // token expires in 16min
+    );
+
+    const response = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 2 * 60 * 1000); // advance time by 2mins
+
+    const response2 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response2.serviceAccountToken).toEqual('MY_TOKEN_2');
+  });
+
+  it('should re-use existing token if there is afailure', async () => {
+    const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      logger,
+      new StaticTokenCredential(16 * 60 * 1000), // new tokens expires in 16min
+    );
+
+    const response = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 2 * 60 * 1000); // advance time by 2min
+
+    const response2 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response2.serviceAccountToken).toEqual('MY_TOKEN_2');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 2 * 60 * 1000); // advance time by 2min
+
+    const response3 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response3.serviceAccountToken).toEqual('MY_TOKEN_2');
+
+    const response4 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response4.serviceAccountToken).toEqual('MY_TOKEN_4');
+  });
+
+  it('should throw if existing token expired and failed to fetch a new one', async () => {
+    const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      logger,
+      new StaticTokenCredential(16 * 60 * 1000), // new tokens expires in 16min
+    );
+
+    const response = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 2 * 60 * 1000); // advance time by 2min
+
+    const response2 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response2.serviceAccountToken).toEqual('MY_TOKEN_2');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 17 * 60 * 1000); // advance time by 17min
+
+    await expect(
+      authTranslator.decorateClusterDetailsWithAuth(cd),
+    ).rejects.toThrow();
+  });
+});

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

@@ -14,15 +14,28 @@
  * limitations under the License.
  */
 
+import { Logger } from 'winston';
 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 accessToken: AccessToken = { token: '', expiresOnTimestamp: 0 };
+  private newTokenPromise: Promise<string> | undefined;
+
+  constructor(
+    private readonly logger: Logger,
+    private readonly tokenCredential: TokenCredential = new DefaultAzureCredential(),
+  ) {}
+
   async decorateClusterDetailsWithAuth(
     clusterDetails: AzureClusterDetails,
   ): Promise<AzureClusterDetails> {
@@ -31,11 +44,54 @@ export class AzureIdentityKubernetesAuthTranslator
       clusterDetails,
     );
 
-    const credentials = new DefaultAzureCredential();
-
-    // 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;
+    clusterDetailsWithAuthToken.serviceAccountToken = await this.getToken();
     return clusterDetailsWithAuthToken;
   }
+
+  private async getToken(): Promise<string> {
+    if (!this.tokenRequiresRefresh()) {
+      return this.accessToken.token;
+    }
+
+    if (!this.newTokenPromise) {
+      this.newTokenPromise = this.fetchNewToken();
+    }
+
+    return this.newTokenPromise;
+  }
+
+  private async fetchNewToken(): Promise<string> {
+    try {
+      this.logger.info('Fetching new Azure token for AKS');
+
+      const newAccessToken = await this.tokenCredential.getToken(aksScope, {
+        requestOptions: { timeout: 10_000 }, // 10 seconds
+      });
+      if (!newAccessToken) {
+        throw new Error('AccessToken is null');
+      }
+
+      this.accessToken = newAccessToken;
+    } catch (err) {
+      this.logger.error('Unable to fetch Azure token', err);
+
+      // only throw the error if the token has already expired, otherwise re-use existing until we're able to fetch a new token
+      if (this.tokenExpired()) {
+        throw err;
+      }
+    }
+
+    this.newTokenPromise = undefined;
+    return this.accessToken.token;
+  }
+
+  private tokenRequiresRefresh(): boolean {
+    // Set tokens to expire 15 minutes before its actual expiry time
+    const expiresOn = this.accessToken.expiresOnTimestamp - 15 * 60 * 1000;
+    return Date.now() >= expiresOn;
+  }
+
+  private tokenExpired(): boolean {
+    return Date.now() >= this.accessToken.expiresOnTimestamp;
+  }
 }

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

@@ -20,25 +20,28 @@ import { KubernetesAuthTranslatorGenerator } from './KubernetesAuthTranslatorGen
 import { ServiceAccountKubernetesAuthTranslator } from './ServiceAccountKubernetesAuthTranslator';
 import { AwsIamKubernetesAuthTranslator } from './AwsIamKubernetesAuthTranslator';
 import { OidcKubernetesAuthTranslator } from './OidcKubernetesAuthTranslator';
+import { getVoidLogger } from '@backstage/backend-common';
+
+const logger = getVoidLogger();
 
 describe('getKubernetesAuthTranslatorInstance', () => {
   const sut = KubernetesAuthTranslatorGenerator;
 
   it('can return an auth translator for google auth', () => {
     const authTranslator: KubernetesAuthTranslator =
-      sut.getKubernetesAuthTranslatorInstance('google');
+      sut.getKubernetesAuthTranslatorInstance('google', { logger });
     expect(authTranslator instanceof GoogleKubernetesAuthTranslator).toBe(true);
   });
 
   it('can return an auth translator for aws auth', () => {
     const authTranslator: KubernetesAuthTranslator =
-      sut.getKubernetesAuthTranslatorInstance('aws');
+      sut.getKubernetesAuthTranslatorInstance('aws', { logger });
     expect(authTranslator instanceof AwsIamKubernetesAuthTranslator).toBe(true);
   });
 
   it('can return an auth translator for serviceAccount auth', () => {
     const authTranslator: KubernetesAuthTranslator =
-      sut.getKubernetesAuthTranslatorInstance('serviceAccount');
+      sut.getKubernetesAuthTranslatorInstance('serviceAccount', { logger });
     expect(
       authTranslator instanceof ServiceAccountKubernetesAuthTranslator,
     ).toBe(true);
@@ -46,12 +49,14 @@ describe('getKubernetesAuthTranslatorInstance', () => {
 
   it('can return an auth translator for oidc auth', () => {
     const authTranslator: KubernetesAuthTranslator =
-      sut.getKubernetesAuthTranslatorInstance('oidc');
+      sut.getKubernetesAuthTranslatorInstance('oidc', { logger });
     expect(authTranslator instanceof OidcKubernetesAuthTranslator).toBe(true);
   });
 
   it('throws an error when asked for an auth translator for an unsupported auth type', () => {
-    expect(() => sut.getKubernetesAuthTranslatorInstance('linode')).toThrow(
+    expect(() =>
+      sut.getKubernetesAuthTranslatorInstance('linode', { logger }),
+    ).toThrow(
       'authProvider "linode" has no KubernetesAuthTranslator associated with it',
     );
   });

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

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { Logger } from 'winston';
 import { KubernetesAuthTranslator } from './types';
 import { GoogleKubernetesAuthTranslator } from './GoogleKubernetesAuthTranslator';
 import { ServiceAccountKubernetesAuthTranslator } from './ServiceAccountKubernetesAuthTranslator';
@@ -25,6 +26,9 @@ import { OidcKubernetesAuthTranslator } from './OidcKubernetesAuthTranslator';
 export class KubernetesAuthTranslatorGenerator {
   static getKubernetesAuthTranslatorInstance(
     authProvider: string,
+    options: {
+      logger: Logger;
+    },
   ): KubernetesAuthTranslator {
     switch (authProvider) {
       case 'google': {
@@ -34,7 +38,7 @@ export class KubernetesAuthTranslatorGenerator {
         return new AwsIamKubernetesAuthTranslator();
       }
       case 'azure': {
-        return new AzureIdentityKubernetesAuthTranslator();
+        return new AzureIdentityKubernetesAuthTranslator(options.logger);
       }
       case 'serviceAccount': {
         return new ServiceAccountKubernetesAuthTranslator();

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

@@ -161,6 +161,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,
@@ -174,6 +175,7 @@ export class KubernetesFanOutHandler {
     this.serviceLocator = serviceLocator;
     this.customResources = customResources;
     this.objectTypesToFetch = new Set(objectTypesToFetch);
+    this.authTranslators = {};
   }
 
   async getKubernetesObjectsByEntity(
@@ -189,14 +191,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,
@@ -294,4 +291,19 @@ 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,
+        {
+          logger: this.logger,
+        },
+      );
+    return this.authTranslators[provider];
+  }
 }