From bb6e4e5ff1bae97eaafcc0c92c25c54d8adeacac Mon Sep 17 00:00:00 2001
From: Amir Omidi <amomidi@microsoft.com>
Date: Thu, 26 Mar 2020 17:51:51 -0700
Subject: [PATCH] Forces device code auth on SAW and fixes a small axios issue
 (#9756)

* Force device code and fix a device code issue

* Trailing comma
---
 .../azurecore/src/account-provider/auths/azureAuth.ts  | 10 +++++++---
 .../src/account-provider/auths/azureDeviceCode.ts      |  2 +-
 .../src/account-provider/azureAccountProvider.ts       |  9 +++++----
 .../account-provider/azureAccountProviderService.ts    |  5 ++++-
 4 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/extensions/azurecore/src/account-provider/auths/azureAuth.ts b/extensions/azurecore/src/account-provider/auths/azureAuth.ts
index 73ef29beb7..6c9c46a541 100644
--- a/extensions/azurecore/src/account-provider/auths/azureAuth.ts
+++ b/extensions/azurecore/src/account-provider/auths/azureAuth.ts
@@ -6,7 +6,7 @@
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import * as nls from 'vscode-nls';
-import axios, { AxiosResponse } from 'axios';
+import axios, { AxiosResponse, AxiosRequestConfig } from 'axios';
 import * as qs from 'qs';
 import * as url from 'url';
 
@@ -234,13 +234,17 @@ export abstract class AzureAuth {
 		return base64string.replace(/=/g, '').replace(/\+/g, '-').replace(/\//g, '_'); // Need to use base64url encoding
 	}
 
-	protected async makePostRequest(uri: string, postData: { [key: string]: string }) {
-		const config = {
+	protected async makePostRequest(uri: string, postData: { [key: string]: string }, validateStatus = false) {
+		const config: AxiosRequestConfig = {
 			headers: {
 				'Content-Type': 'application/x-www-form-urlencoded'
 			}
 		};
 
+		if (validateStatus) {
+			config.validateStatus = () => true;
+		}
+
 		return axios.post(uri, qs.stringify(postData), config);
 	}
 
diff --git a/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts b/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts
index f563db446b..d42ac1b225 100644
--- a/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts
+++ b/extensions/azurecore/src/account-provider/auths/azureDeviceCode.ts
@@ -140,7 +140,7 @@ export class AzureDeviceCode extends AzureAuth {
 				code: info.device_code
 			};
 
-			const postResult = await this.makePostRequest(uri, postData);
+			const postResult = await this.makePostRequest(uri, postData, true);
 
 			const result: DeviceCodeLoginResult = postResult.data;
 
diff --git a/extensions/azurecore/src/account-provider/azureAccountProvider.ts b/extensions/azurecore/src/account-provider/azureAccountProvider.ts
index 5eaa9f0249..99db70ab72 100644
--- a/extensions/azurecore/src/account-provider/azureAccountProvider.ts
+++ b/extensions/azurecore/src/account-provider/azureAccountProvider.ts
@@ -29,7 +29,8 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 	constructor(
 		metadata: AzureAccountProviderMetadata,
 		tokenCache: SimpleTokenCache,
-		context: vscode.ExtensionContext
+		context: vscode.ExtensionContext,
+		private readonly forceDeviceCode: boolean = false
 	) {
 		vscode.workspace.onDidChangeConfiguration((changeEvent) => {
 			const impact = changeEvent.affectsConfiguration(AzureAccountProvider.CONFIGURATION_SECTION);
@@ -52,10 +53,10 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		const codeGrantMethod: boolean = configuration.get('codeGrant');
 		const deviceCodeMethod: boolean = configuration.get('deviceCode');
 
-		if (codeGrantMethod === true) {
+		if (codeGrantMethod === true && !this.forceDeviceCode) {
 			this.authMappings.set(AzureAuthType.AuthCodeGrant, new AzureAuthCodeGrant(metadata, tokenCache, context));
 		}
-		if (deviceCodeMethod === true) {
+		if (deviceCodeMethod === true || this.forceDeviceCode) {
 			this.authMappings.set(AzureAuthType.DeviceCode, new AzureDeviceCode(metadata, tokenCache, context));
 		}
 	}
@@ -69,7 +70,7 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		if (authType) {
 			return this.authMappings.get(authType);
 		} else {
-			return this.authMappings.get(AzureAuthType.AuthCodeGrant);
+			return this.authMappings.values().next().value;
 		}
 	}
 
diff --git a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts
index 22a0f80041..7b56e7d510 100644
--- a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts
+++ b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts
@@ -131,7 +131,10 @@ export class AzureAccountProviderService implements vscode.Disposable {
 			let tokenCacheKey = `azureTokenCache-${provider.metadata.id}`;
 			let simpleTokenCache = new SimpleTokenCache(tokenCacheKey, this._userStoragePath, noSystemKeychain, this._credentialProvider);
 			await simpleTokenCache.init();
-			let accountProvider = new AzureAccountProvider(provider.metadata as AzureAccountProviderMetadata, simpleTokenCache, this._context);
+
+			const isSaw: boolean = vscode.env.appName.toLowerCase().indexOf('saw') > 0;
+			let accountProvider = new AzureAccountProvider(provider.metadata as AzureAccountProviderMetadata, simpleTokenCache, this._context, isSaw);
+
 			this._accountProviders[provider.metadata.id] = accountProvider;
 			this._accountDisposals[provider.metadata.id] = azdata.accounts.registerAccountProvider(provider.metadata, accountProvider);
 		} catch (e) {