From 4874ebc2afa42531397bbce4cbeda43b21bdcff4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fredrik=20Adel=C3=B6w?= <freben@gmail.com>
Date: Wed, 4 May 2022 15:26:41 +0200
Subject: [PATCH] add caching as well
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Fredrik Adelöw <freben@gmail.com>
---
 .../src/tokens/ServerTokenManager.test.ts     | 29 +++++++++-----
 .../src/tokens/ServerTokenManager.ts          | 38 +++++++++++++++----
 2 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/packages/backend-common/src/tokens/ServerTokenManager.test.ts b/packages/backend-common/src/tokens/ServerTokenManager.test.ts
index 57435fe486..2827634aac 100644
--- a/packages/backend-common/src/tokens/ServerTokenManager.test.ts
+++ b/packages/backend-common/src/tokens/ServerTokenManager.test.ts
@@ -169,26 +169,37 @@ describe('ServerTokenManager', () => {
       );
     });
 
-    it('should throw for expired tokens', async () => {
-      jest.useFakeTimers('modern');
-      jest.setSystemTime(new Date('2020-02-02T02:00:00.0000000Z'));
+    it('should throw for expired tokens, and re-issue new ones', async () => {
+      jest.useFakeTimers();
 
       const tokenManager = ServerTokenManager.fromConfig(configWithSecret, {
         logger,
       });
 
-      const { token } = await tokenManager.getToken();
-      await expect(tokenManager.authenticate(token)).resolves.not.toThrow();
+      const { token: token1 } = await tokenManager.getToken();
+      await expect(tokenManager.authenticate(token1)).resolves.not.toThrow();
 
-      jest.setSystemTime(new Date('2020-02-02T02:59:00.0000000Z'));
-      await expect(tokenManager.authenticate(token)).resolves.not.toThrow();
+      // Less than ten minutes before expiry, it still returns the same token
+      jest.advanceTimersByTime(49 * 60 * 1000);
+      const { token: token1Again } = await tokenManager.getToken();
+      expect(token1).toEqual(token1Again);
+      await expect(tokenManager.authenticate(token1)).resolves.not.toThrow();
 
-      jest.setSystemTime(new Date('2020-02-02T03:00:01.0000000Z'));
+      // Right before the expiry, the old ones are still valid but returning a new token
+      jest.advanceTimersByTime(10 * 60 * 1000);
+      const { token: token2 } = await tokenManager.getToken();
+      expect(token1).not.toEqual(token2);
+      await expect(tokenManager.authenticate(token1)).resolves.not.toThrow();
+      await expect(tokenManager.authenticate(token2)).resolves.not.toThrow();
+
+      // After expiry, the newest one is still valid
+      jest.advanceTimersByTime(2 * 60 * 1000);
       await expect(
-        tokenManager.authenticate(token),
+        tokenManager.authenticate(token1),
       ).rejects.toThrowErrorMatchingInlineSnapshot(
         '"Invalid server token: JWTExpired: \\"exp\\" claim timestamp check failed"',
       );
+      await expect(tokenManager.authenticate(token2)).resolves.not.toThrow();
     });
   });
 
diff --git a/packages/backend-common/src/tokens/ServerTokenManager.ts b/packages/backend-common/src/tokens/ServerTokenManager.ts
index ba6a276421..65354f270f 100644
--- a/packages/backend-common/src/tokens/ServerTokenManager.ts
+++ b/packages/backend-common/src/tokens/ServerTokenManager.ts
@@ -17,11 +17,13 @@
 import { Config } from '@backstage/config';
 import { AuthenticationError, NotAllowedError } from '@backstage/errors';
 import { base64url, exportJWK, generateSecret, jwtVerify, SignJWT } from 'jose';
+import { DateTime, Duration } from 'luxon';
 import { Logger } from 'winston';
 import { TokenManager } from './types';
 
 const TOKEN_ALG = 'HS256';
 const TOKEN_SUB = 'backstage-server';
+const TOKEN_EXPIRY = Duration.fromObject({ hours: 1 });
 
 /**
  * A token manager that issues static dummy tokens and never fails
@@ -57,9 +59,10 @@ export interface ServerTokenManagerOptions {
  */
 export class ServerTokenManager implements TokenManager {
   private readonly options: ServerTokenManagerOptions;
-  private verificationKeys: Uint8Array[];
+  private readonly verificationKeys: Uint8Array[];
   private signingKey: Uint8Array;
-  private privateKeyPromise?: Promise<void>;
+  private privateKeyPromise: Promise<void> | undefined;
+  private currentTokenPromise: Promise<{ token: string }> | undefined;
 
   /**
    * Creates a token manager that issues static dummy tokens and never fails
@@ -140,13 +143,32 @@ export class ServerTokenManager implements TokenManager {
       await this.generateKeys();
     }
 
-    const jwt = await new SignJWT({})
-      .setProtectedHeader({ alg: TOKEN_ALG })
-      .setSubject(TOKEN_SUB)
-      .setExpirationTime('1h')
-      .sign(this.signingKey);
+    if (this.currentTokenPromise) {
+      return this.currentTokenPromise;
+    }
 
-    return { token: jwt };
+    const result = Promise.resolve().then(async () => {
+      const jwt = await new SignJWT({})
+        .setProtectedHeader({ alg: TOKEN_ALG })
+        .setSubject(TOKEN_SUB)
+        .setExpirationTime(DateTime.now().plus(TOKEN_EXPIRY).toUnixInteger())
+        .sign(this.signingKey);
+      return { token: jwt };
+    });
+
+    this.currentTokenPromise = result;
+
+    result
+      .then(() => {
+        setTimeout(() => {
+          this.currentTokenPromise = undefined;
+        }, TOKEN_EXPIRY.minus({ minutes: 5 }).toMillis());
+      })
+      .catch(() => {
+        this.currentTokenPromise = undefined;
+      });
+
+    return result;
   }
 
   async authenticate(token: string): Promise<void> {