refactor null check and race condition

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

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

@@ -15,8 +15,11 @@
  */
 
 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;
 
@@ -25,6 +28,10 @@ class StaticTokenCredential implements TokenCredential {
   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,
@@ -41,6 +48,7 @@ describe('AzureIdentityKubernetesAuthTranslator tests', () => {
 
   it('should decorate cluster with Azure token', async () => {
     const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      logger,
       new StaticTokenCredential(5 * 60 * 1000),
     );
 
@@ -50,6 +58,7 @@ describe('AzureIdentityKubernetesAuthTranslator tests', () => {
 
   it('should re-use token before expiry', async () => {
     const authTranslator = new AzureIdentityKubernetesAuthTranslator(
+      logger,
       new StaticTokenCredential(20 * 60 * 1000),
     );
 
@@ -62,15 +71,41 @@ describe('AzureIdentityKubernetesAuthTranslator tests', () => {
 
   it('should issue new token 15 minutes befory expiry', async () => {
     const authTranslator = new AzureIdentityKubernetesAuthTranslator(
-      new StaticTokenCredential(16 * 60 * 1000), // token expires in 11m
+      logger,
+      new StaticTokenCredential(16 * 60 * 1000), // token expires in 16m
     );
 
     const response = await authTranslator.decorateClusterDetailsWithAuth(cd);
     expect(response.serviceAccountToken).toEqual('MY_TOKEN_1');
 
-    jest.useFakeTimers().setSystemTime(Date.now() + 1 * 60 * 1000); // advance time by 1min
+    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 16m
+    );
+
+    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');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 2 * 60 * 1000); // advance time by 2mins
+
+    const response3 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response3.serviceAccountToken).toEqual('MY_TOKEN_2');
+
+    jest.useFakeTimers().setSystemTime(Date.now() + 2 * 60 * 1000); // advance time by 2mins
+
+    const response4 = await authTranslator.decorateClusterDetailsWithAuth(cd);
+    expect(response4.serviceAccountToken).toEqual('MY_TOKEN_4');
+  });
 });

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

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { Logger } from 'winston';
 import { KubernetesAuthTranslator } from './types';
 import { AzureClusterDetails } from '../types/types';
 import {
@@ -27,9 +28,11 @@ const aksScope = '6dae42f8-4368-4678-94ff-3960e28e3630/.default'; // This scope
 export class AzureIdentityKubernetesAuthTranslator
   implements KubernetesAuthTranslator
 {
-  private accessToken: AccessToken | null = null;
+  private accessToken: AccessToken = { token: '', expiresOnTimestamp: 0 };
+  private newToken: Promise<string> | undefined;
 
   constructor(
+    private readonly logger: Logger,
     private readonly tokenCredential: TokenCredential = new DefaultAzureCredential(),
   ) {}
 
@@ -41,23 +44,45 @@ export class AzureIdentityKubernetesAuthTranslator
       clusterDetails,
     );
 
-    if (this.tokenExpired()) {
-      this.accessToken = await this.tokenCredential.getToken(aksScope);
-
-      if (!this.accessToken) {
-        throw new Error('Unable to retrieve Azure token');
-      }
-    }
-
-    clusterDetailsWithAuthToken.serviceAccountToken = this.accessToken!.token;
+    clusterDetailsWithAuthToken.serviceAccountToken = await this.getToken();
     return clusterDetailsWithAuthToken;
   }
 
-  private tokenExpired(): boolean {
-    if (!this.accessToken) return true;
+  private async getToken(): Promise<string> {
+    if (this.isTokenValid()) {
+      return this.accessToken.token;
+    }
 
+    if (!this.newToken) {
+      this.newToken = this.fetchNewToken();
+    }
+
+    return this.newToken;
+  }
+
+  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);
+    }
+
+    this.newToken = undefined;
+    return this.accessToken.token;
+  }
+
+  private isTokenValid(): boolean {
     // Set tokens to expire 15 minutes before its actual expiry time
     const expiresOn = this.accessToken.expiresOnTimestamp - 15 * 60 * 1000;
-    return Date.now() >= expiresOn;
+    return expiresOn >= Date.now();
   }
 }

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(logger, 'google');
     expect(authTranslator instanceof GoogleKubernetesAuthTranslator).toBe(true);
   });
 
   it('can return an auth translator for aws auth', () => {
     const authTranslator: KubernetesAuthTranslator =
-      sut.getKubernetesAuthTranslatorInstance('aws');
+      sut.getKubernetesAuthTranslatorInstance(logger, 'aws');
     expect(authTranslator instanceof AwsIamKubernetesAuthTranslator).toBe(true);
   });
 
   it('can return an auth translator for serviceAccount auth', () => {
     const authTranslator: KubernetesAuthTranslator =
-      sut.getKubernetesAuthTranslatorInstance('serviceAccount');
+      sut.getKubernetesAuthTranslatorInstance(logger, 'serviceAccount');
     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(logger, 'oidc');
     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(logger, 'linode'),
+    ).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';
@@ -24,6 +25,7 @@ import { OidcKubernetesAuthTranslator } from './OidcKubernetesAuthTranslator';
 
 export class KubernetesAuthTranslatorGenerator {
   static getKubernetesAuthTranslatorInstance(
+    logger: Logger,
     authProvider: string,
   ): KubernetesAuthTranslator {
     switch (authProvider) {
@@ -34,7 +36,7 @@ export class KubernetesAuthTranslatorGenerator {
         return new AwsIamKubernetesAuthTranslator();
       }
       case 'azure': {
-        return new AzureIdentityKubernetesAuthTranslator();
+        return new AzureIdentityKubernetesAuthTranslator(logger);
       }
       case 'serviceAccount': {
         return new ServiceAccountKubernetesAuthTranslator();

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

@@ -293,6 +293,7 @@ export class KubernetesFanOutHandler {
 
     this.authTranslators[provider] =
       KubernetesAuthTranslatorGenerator.getKubernetesAuthTranslatorInstance(
+        this.logger,
         provider,
       );
     return this.authTranslators[provider];