NONE fix conflicts token installatin id (#2019)

atlassian · Mar 30, 2023 · 816842b · 816842b
1 parent 5d9ea63
commit 816842b
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 9 deletions.
diff --git a/src/github/client/github-installation-client.ts b/src/github/client/github-installation-client.ts
@@ -340,6 +340,7 @@ export class GitHubInstallationClient extends GitHubClient {
 	private async installationAuthenticationHeaders(): Promise<Partial<AxiosRequestConfig>> {
 		const installationToken = await this.installationTokenCache.getInstallationToken(
 			this.githubInstallationId.installationId,
+			this.gitHubServerAppId,
 			() => this.createInstallationToken(this.githubInstallationId.installationId));
 		return {
 			headers: {

diff --git a/src/github/client/installation-token-cache.test.ts b/src/github/client/installation-token-cache.test.ts
@@ -18,6 +18,96 @@ describe("InstallationTokenCache", () => {
 		jest.useRealTimers();
 	});
 
+	it("Reuse same token for cloud installation", async () => {
+
+		const GITHUB_INSTALLATION_ID = 1;
+		jest.setSystemTime(now);
+		const token1 = new AuthToken("token1", in10Minutes);
+		const token2 = new AuthToken("token2", in10Minutes);
+
+		const cache1 = InstallationTokenCache.getInstance();
+		const cache2 = InstallationTokenCache.getInstance();
+
+		const foundToken1 = await cache1.getInstallationToken(GITHUB_INSTALLATION_ID, undefined, () => Promise.resolve(token1));
+		const foundToken2 = await cache2.getInstallationToken(GITHUB_INSTALLATION_ID, undefined, () => Promise.resolve(token2));
+
+		expect(foundToken1).toEqual(foundToken2);
+
+	});
+
+	it("Reuse same token for GHE installation", async () => {
+
+		const GITHUB_INSTALLATION_ID = 1;
+		const GITHUB_APP_ID = 1;
+		jest.setSystemTime(now);
+		const token1 = new AuthToken("token1", in10Minutes);
+		const token2 = new AuthToken("token2", in10Minutes);
+
+		const cache1 = InstallationTokenCache.getInstance();
+		const cache2 = InstallationTokenCache.getInstance();
+
+		const foundToken1 = await cache1.getInstallationToken(GITHUB_INSTALLATION_ID, GITHUB_APP_ID, () => Promise.resolve(token1));
+		const foundToken2 = await cache2.getInstallationToken(GITHUB_INSTALLATION_ID, GITHUB_APP_ID, () => Promise.resolve(token2));
+
+		expect(foundToken1).toEqual(foundToken2);
+
+	});
+
+	it("won't have conflicts on the token for cloud installations", async () => {
+
+		const GITHUB_INSTALLATION_ID_1 = 21;
+		const GITHUB_INSTALLATION_ID_2 = 22;
+		jest.setSystemTime(now);
+		const token1 = new AuthToken("token1", in10Minutes);
+		const token2 = new AuthToken("token2", in10Minutes);
+
+		const cache1 = InstallationTokenCache.getInstance();
+		const cache2 = InstallationTokenCache.getInstance();
+
+		const foundToken1 = await cache1.getInstallationToken(GITHUB_INSTALLATION_ID_1, undefined, () => Promise.resolve(token1));
+		const foundToken2 = await cache2.getInstallationToken(GITHUB_INSTALLATION_ID_2, undefined, () => Promise.resolve(token2));
+
+		expect(foundToken1).not.toEqual(foundToken2);
+
+	});
+
+	it("won't have conflicts on the token for different GHE installations", async () => {
+
+		const CONFLICTIN_GITHUB_INSTALLATION_ID = 31;
+		const GITHUB_APP_ID_1 = 31;
+		const GITHUB_APP_ID_2 = 32;
+		jest.setSystemTime(now);
+		const token1 = new AuthToken("token1", in10Minutes);
+		const token2 = new AuthToken("token2", in10Minutes);
+
+		const cache1 = InstallationTokenCache.getInstance();
+		const cache2 = InstallationTokenCache.getInstance();
+
+		const foundToken1 = await cache1.getInstallationToken(CONFLICTIN_GITHUB_INSTALLATION_ID, GITHUB_APP_ID_1, () => Promise.resolve(token1));
+		const foundToken2 = await cache2.getInstallationToken(CONFLICTIN_GITHUB_INSTALLATION_ID, GITHUB_APP_ID_2, () => Promise.resolve(token2));
+
+		expect(foundToken1).not.toEqual(foundToken2);
+
+	});
+
+	it("won't have conflicts on the token for different cloud and GHE installations", async () => {
+
+		const CONFLICTIN_GITHUB_INSTALLATION_ID = 41;
+		const GITHUB_APP_ID = 41;
+		jest.setSystemTime(now);
+		const token1 = new AuthToken("token1", in10Minutes);
+		const token2 = new AuthToken("token2", in10Minutes);
+
+		const cache1 = InstallationTokenCache.getInstance();
+		const cache2 = InstallationTokenCache.getInstance();
+
+		const foundToken1 = await cache1.getInstallationToken(CONFLICTIN_GITHUB_INSTALLATION_ID, undefined, () => Promise.resolve(token1));
+		const foundToken2 = await cache2.getInstallationToken(CONFLICTIN_GITHUB_INSTALLATION_ID, GITHUB_APP_ID, () => Promise.resolve(token2));
+
+		expect(foundToken1).not.toEqual(foundToken2);
+
+	});
+
 	it("Re-generates expired tokens", async () => {
 		const initialInstallationToken = new AuthToken("initial installation token", in10Minutes);
 		const generateInitialInstallationToken = jest.fn().mockImplementation(() => Promise.resolve(initialInstallationToken));
@@ -29,21 +119,21 @@ describe("InstallationTokenCache", () => {
 		const installationTokenCache = new InstallationTokenCache();
 
 		jest.setSystemTime(now);
-		const token1 = await installationTokenCache.getInstallationToken(githubInstallationId, generateInitialInstallationToken);
+		const token1 = await installationTokenCache.getInstallationToken(githubInstallationId, undefined, generateInitialInstallationToken);
 		expect(token1).toEqual(initialInstallationToken);
 		expect(generateInitialInstallationToken).toHaveBeenCalledTimes(1);
 		expect(generateFreshInstallationToken).toHaveBeenCalledTimes(0);
 
 		// after 5 minutes we still expect the same token because it's still valid
 		jest.setSystemTime(in5Minutes);
-		const token2 = await installationTokenCache.getInstallationToken(githubInstallationId, generateFreshInstallationToken);
+		const token2 = await installationTokenCache.getInstallationToken(githubInstallationId, undefined, generateFreshInstallationToken);
 		expect(token2).toEqual(initialInstallationToken);
 		expect(generateInitialInstallationToken).toHaveBeenCalledTimes(1);
 		expect(generateFreshInstallationToken).toHaveBeenCalledTimes(0);
 
 		// after 10 minutes we expect a new token because the old one has expired
 		jest.setSystemTime(in10Minutes);
-		const token3 = await installationTokenCache.getInstallationToken(githubInstallationId, generateFreshInstallationToken);
+		const token3 = await installationTokenCache.getInstallationToken(githubInstallationId, undefined, generateFreshInstallationToken);
 		expect(token3).toEqual(freshInstallationToken);
 		expect(generateInitialInstallationToken).toHaveBeenCalledTimes(1);
 		expect(generateFreshInstallationToken).toHaveBeenCalledTimes(1);

diff --git a/src/github/client/installation-token-cache.ts b/src/github/client/installation-token-cache.ts
@@ -10,7 +10,7 @@ import { AuthToken } from "./auth-token";
  */
 export class InstallationTokenCache {
 	private static instance: InstallationTokenCache;
-	private readonly installationTokenCache: LRUCache<number, AuthToken>;
+	private readonly installationTokenCache: LRUCache<string, AuthToken>;
 
 	/**
 	 * Creates a new InstallationTokenCache. This cache should be shared between all GitHub clients so that the clients don't
@@ -19,7 +19,7 @@ export class InstallationTokenCache {
 	 * number, the least recently used tokens are evicted from the cache.
 	 */
 	constructor() {
-		this.installationTokenCache = new LRUCache<number, AuthToken>({ max: 1000 });
+		this.installationTokenCache = new LRUCache<string, AuthToken>({ max: 1000 });
 	}
 
 	public static getInstance(): InstallationTokenCache {
@@ -39,12 +39,13 @@ export class InstallationTokenCache {
 	 */
 	public async getInstallationToken(
 		githubInstallationId: number,
+		gitHubAppId: number | undefined,
 		generateNewInstallationToken: () => Promise<AuthToken>): Promise<AuthToken> {
-		let token = this.installationTokenCache.get(githubInstallationId);
+		let token = this.installationTokenCache.get(this.key(githubInstallationId, gitHubAppId));
 
 		if (!token || token.isAboutToExpire()) {
 			token = await generateNewInstallationToken();
-			this.installationTokenCache.set(githubInstallationId, token, token.millisUntilAboutToExpire());
+			this.installationTokenCache.set(this.key(githubInstallationId, gitHubAppId), token, token.millisUntilAboutToExpire());
 		}
 
 		return token;
@@ -53,4 +54,8 @@ export class InstallationTokenCache {
 	public clear(): void {
 		this.installationTokenCache.reset();
 	}
+
+	private key(githubInstallationId: number, gitHubAppId: number | undefined): string {
+		return `${githubInstallationId}_${gitHubAppId}`;
+	}
 }
diff --git a/src/github/client/token-cache.test.ts b/src/github/client/token-cache.test.ts
@@ -31,8 +31,8 @@ describe("InstallationTokenCache & AppTokenHolder", () => {
 		const generateInitialInstallationToken = jest.fn().mockImplementation(() => Promise.resolve(initialInstallationToken));
 
 		jest.setSystemTime(date);
-		const token1 = await installationTokenCache.getInstallationToken(githubInstallationId, generateInitialInstallationToken);
-		const token2 = await installationTokenCache.getInstallationToken(githubInstallationId, generateInitialInstallationToken);
+		const token1 = await installationTokenCache.getInstallationToken(githubInstallationId, undefined, generateInitialInstallationToken);
+		const token2 = await installationTokenCache.getInstallationToken(githubInstallationId, undefined, generateInitialInstallationToken);
 		expect(token1).toEqual(initialInstallationToken);
 		expect(token2).toEqual(initialInstallationToken);
 		expect(generateInitialInstallationToken).toHaveBeenCalledTimes(2);