From 223aa3fe0e6392f6b13a6086a5e085af598c1fe8 Mon Sep 17 00:00:00 2001
From: Amir Omidi <amomidi@microsoft.com>
Date: Thu, 2 Apr 2020 09:02:04 -0700
Subject: [PATCH] We can't authenticate to every subscription (#9826)

* We can't authenticate to every subscription

* Refresh tenants and subs when the application starts
---
 .../src/account-provider/auths/azureAuth.ts   | 39 +++++++++++++------
 .../src/azureResource/tree/accountTreeNode.ts | 10 +++++
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/extensions/azurecore/src/account-provider/auths/azureAuth.ts b/extensions/azurecore/src/account-provider/auths/azureAuth.ts
index acd6576dd0..cc5e6b5157 100644
--- a/extensions/azurecore/src/account-provider/auths/azureAuth.ts
+++ b/extensions/azurecore/src/account-provider/auths/azureAuth.ts
@@ -138,29 +138,39 @@ export abstract class AzureAuth implements vscode.Disposable {
 
 	public dispose() { }
 
-	public async refreshAccess(account: azdata.Account): Promise<azdata.Account> {
-		const response = await this.getCachedToken(account.key);
+	public async refreshAccess(oldAccount: azdata.Account): Promise<azdata.Account> {
+		const response = await this.getCachedToken(oldAccount.key);
 		if (!response) {
-			account.isStale = true;
-			return account;
+			oldAccount.isStale = true;
+			return oldAccount;
 		}
 
 		const refreshToken = response.refreshToken;
 		if (!refreshToken || !refreshToken.key) {
-			account.isStale = true;
-			return account;
+			oldAccount.isStale = true;
+			return oldAccount;
 		}
 
 		try {
-			await this.refreshAccessToken(account.key, refreshToken);
+			// Refresh the access token
+			const tokenResponse = await this.refreshAccessToken(oldAccount.key, refreshToken);
+			const tenants = await this.getTenants(tokenResponse.accessToken);
+
+			// Recreate account object
+			const newAccount = this.createAccount(tokenResponse.tokenClaims, tokenResponse.accessToken.key, tenants);
+
+			const subscriptions = await this.getSubscriptions(newAccount);
+			newAccount.properties.subscriptions = subscriptions;
+
+			return newAccount;
 		} catch (ex) {
-			account.isStale = true;
+			oldAccount.isStale = true;
 			if (ex.message) {
 				await vscode.window.showErrorMessage(ex.message);
 			}
 			console.log(ex);
 		}
-		return account;
+		return oldAccount;
 	}
 
 
@@ -375,7 +385,7 @@ export abstract class AzureAuth implements vscode.Disposable {
 		}
 	}
 
-	private async refreshAccessToken(account: azdata.AccountKey, rt: RefreshToken, tenant?: Tenant, resource?: Resource): Promise<void> {
+	private async refreshAccessToken(account: azdata.AccountKey, rt: RefreshToken, tenant?: Tenant, resource?: Resource): Promise<TokenRefreshResponse> {
 		const postData: { [key: string]: string } = {
 			grant_type: 'refresh_token',
 			refresh_token: rt.token,
@@ -387,7 +397,10 @@ export abstract class AzureAuth implements vscode.Disposable {
 			postData.resource = resource.endpoint;
 		}
 
-		const { accessToken, refreshToken } = await this.getToken(postData, tenant?.id, resource?.id);
+		const getTokenResponse = await this.getToken(postData, tenant?.id, resource?.id);
+
+		const accessToken = getTokenResponse?.accessToken;
+		const refreshToken = getTokenResponse?.refreshToken;
 
 		if (!accessToken || !refreshToken) {
 			console.log('Access or refresh token were undefined');
@@ -395,7 +408,9 @@ export abstract class AzureAuth implements vscode.Disposable {
 			throw new Error(msg);
 		}
 
-		return this.setCachedToken(account, accessToken, refreshToken, resource?.id, tenant?.id);
+		await this.setCachedToken(account, accessToken, refreshToken, resource?.id, tenant?.id);
+
+		return getTokenResponse;
 	}
 
 
diff --git a/extensions/azurecore/src/azureResource/tree/accountTreeNode.ts b/extensions/azurecore/src/azureResource/tree/accountTreeNode.ts
index d411178455..fbc32fb663 100644
--- a/extensions/azurecore/src/azureResource/tree/accountTreeNode.ts
+++ b/extensions/azurecore/src/azureResource/tree/accountTreeNode.ts
@@ -79,6 +79,16 @@ export class AzureResourceAccountTreeNode extends AzureResourceContainerTreeNode
 			if (subscriptions.length === 0) {
 				return [AzureResourceMessageTreeNode.create(AzureResourceAccountTreeNode.noSubscriptionsLabel, this)];
 			} else {
+				// Filter out everything that we can't authenticate to.
+				subscriptions = subscriptions.filter(s => {
+					const token = tokens[s.id];
+					if (!token) {
+						console.info(`Account does not have permissions to view subscription ${JSON.stringify(s)}.`);
+						return false;
+					}
+					return true;
+				});
+
 				let subTreeNodes = await Promise.all(subscriptions.map(async (subscription) => {
 					const token = tokens[subscription.id];
 					const tenantId = await this._tenantService.getTenantId(subscription, this.account, new TokenCredentials(token.token, token.tokenType));