Merge from vscode 2c306f762bf9c3db82dc06c7afaa56ef46d72f79 (#14050)

* Merge from vscode 2c306f762bf9c3db82dc06c7afaa56ef46d72f79

* Fix breaks

* Extension management fixes

* Fix breaks in windows bundling

* Fix/skip failing tests

* Update distro

* Add clear to nuget.config

* Add hygiene task

* Bump distro

* Fix hygiene issue

* Add build to hygiene exclusion

* Update distro

* Update hygiene

* Hygiene exclusions

* Update tsconfig

* Bump distro for server breaks

* Update build config

* Update darwin path

* Add done calls to notebook tests

* Skip failing tests

* Disable smoke tests

extensions/microsoft-authentication/src/AADHelper.ts

@@ -65,6 +65,7 @@ export interface ITokenResponse {
 	refresh_token: string;
 	scope: string;
 	token_type: string;
+	id_token?: string;
 }
 
 function parseQuery(uri: vscode.Uri) {
@@ -89,14 +90,20 @@ export class AzureActiveDirectoryService {
 	private _tokens: IToken[] = [];
 	private _refreshTimeouts: Map<string, NodeJS.Timeout> = new Map<string, NodeJS.Timeout>();
 	private _uriHandler: UriEventHandler;
+	private _disposables: vscode.Disposable[] = [];
+
+	// Used to keep track of current requests when not using the local server approach.
+	private _pendingStates = new Map<string, string[]>();
+	private _codeExchangePromises = new Map<string, Promise<vscode.AuthenticationSession>>();
+	private _codeVerfifiers = new Map<string, string>();
 
 	constructor() {
 		this._uriHandler = new UriEventHandler();
-		vscode.window.registerUriHandler(this._uriHandler);
+		this._disposables.push(vscode.window.registerUriHandler(this._uriHandler));
 	}
 
 	public async initialize(): Promise<void> {
-		const storedData = await keychain.getToken();
+		const storedData = await keychain.getToken() || await keychain.tryMigrate();
 		if (storedData) {
 			try {
 				const sessions = this.parseStoredData(storedData);
@@ -136,7 +143,7 @@ export class AzureActiveDirectoryService {
 			}
 		}
 
-		this.pollForChange();
+		this._disposables.push(vscode.authentication.onDidChangePassword(() => this.checkForUpdates));
 	}
 
 	private parseStoredData(data: string): IStoredSession[] {
@@ -156,67 +163,63 @@ export class AzureActiveDirectoryService {
 		await keychain.setToken(JSON.stringify(serializedData));
 	}
 
-	private pollForChange() {
-		setTimeout(async () => {
-			const addedIds: string[] = [];
-			let removedIds: string[] = [];
-			const storedData = await keychain.getToken();
-			if (storedData) {
-				try {
-					const sessions = this.parseStoredData(storedData);
-					let promises = sessions.map(async session => {
-						const matchesExisting = this._tokens.some(token => token.scope === session.scope && token.sessionId === session.id);
-						if (!matchesExisting && session.refreshToken) {
-							try {
-								await this.refreshToken(session.refreshToken, session.scope, session.id);
-								addedIds.push(session.id);
-							} catch (e) {
-								if (e.message === REFRESH_NETWORK_FAILURE) {
-									// Ignore, will automatically retry on next poll.
-								} else {
-									await this.logout(session.id);
-								}
+	private async checkForUpdates(): Promise<void> {
+		const addedIds: string[] = [];
+		let removedIds: string[] = [];
+		const storedData = await keychain.getToken();
+		if (storedData) {
+			try {
+				const sessions = this.parseStoredData(storedData);
+				let promises = sessions.map(async session => {
+					const matchesExisting = this._tokens.some(token => token.scope === session.scope && token.sessionId === session.id);
+					if (!matchesExisting && session.refreshToken) {
+						try {
+							await this.refreshToken(session.refreshToken, session.scope, session.id);
+							addedIds.push(session.id);
+						} catch (e) {
+							if (e.message === REFRESH_NETWORK_FAILURE) {
+								// Ignore, will automatically retry on next poll.
+							} else {
+								await this.logout(session.id);
 							}
 						}
-					});
+					}
+				});
 
-					promises = promises.concat(this._tokens.map(async token => {
-						const matchesExisting = sessions.some(session => token.scope === session.scope && token.sessionId === session.id);
-						if (!matchesExisting) {
-							await this.logout(token.sessionId);
-							removedIds.push(token.sessionId);
-						}
-					}));
+				promises = promises.concat(this._tokens.map(async token => {
+					const matchesExisting = sessions.some(session => token.scope === session.scope && token.sessionId === session.id);
+					if (!matchesExisting) {
+						await this.logout(token.sessionId);
+						removedIds.push(token.sessionId);
+					}
+				}));
 
-					await Promise.all(promises);
-				} catch (e) {
-					Logger.error(e.message);
-					// if data is improperly formatted, remove all of it and send change event
-					removedIds = this._tokens.map(token => token.sessionId);
-					this.clearSessions();
-				}
-			} else {
-				if (this._tokens.length) {
-					// Log out all, remove all local data
-					removedIds = this._tokens.map(token => token.sessionId);
-					Logger.info('No stored keychain data, clearing local data');
-
-					this._tokens = [];
-
-					this._refreshTimeouts.forEach(timeout => {
-						clearTimeout(timeout);
-					});
-
-					this._refreshTimeouts.clear();
-				}
+				await Promise.all(promises);
+			} catch (e) {
+				Logger.error(e.message);
+				// if data is improperly formatted, remove all of it and send change event
+				removedIds = this._tokens.map(token => token.sessionId);
+				this.clearSessions();
 			}
+		} else {
+			if (this._tokens.length) {
+				// Log out all, remove all local data
+				removedIds = this._tokens.map(token => token.sessionId);
+				Logger.info('No stored keychain data, clearing local data');
 
-			if (addedIds.length || removedIds.length) {
-				onDidChangeSessions.fire({ added: addedIds, removed: removedIds, changed: [] });
+				this._tokens = [];
+
+				this._refreshTimeouts.forEach(timeout => {
+					clearTimeout(timeout);
+				});
+
+				this._refreshTimeouts.clear();
 			}
+		}
 
-			this.pollForChange();
-		}, 1000 * 30);
+		if (addedIds.length || removedIds.length) {
+			onDidChangeSessions.fire({ added: addedIds, removed: removedIds, changed: [] });
+		}
 	}
 
 	private async convertToSession(token: IToken): Promise<vscode.AuthenticationSession> {
@@ -270,7 +273,7 @@ export class AzureActiveDirectoryService {
 		}
 
 		return new Promise(async (resolve, reject) => {
-			if (vscode.env.uiKind === vscode.UIKind.Web) {
+			if (vscode.env.remoteName !== undefined) {
 				resolve(this.loginWithoutLocalServer(scope));
 				return;
 			}
@@ -340,6 +343,11 @@ export class AzureActiveDirectoryService {
 		});
 	}
 
+	public dispose(): void {
+		this._disposables.forEach(disposable => disposable.dispose());
+		this._disposables = [];
+	}
+
 	private getCallbackEnvironment(callbackUri: vscode.Uri): string {
 		if (callbackUri.authority.endsWith('.workspaces.github.com') || callbackUri.authority.endsWith('.github.dev')) {
 			return `${callbackUri.authority},`;
@@ -353,7 +361,7 @@ export class AzureActiveDirectoryService {
 			case 'online.dev.core.vsengsaas.visualstudio.com':
 				return 'vsodev,';
 			default:
-				return '';
+				return `${callbackUri.scheme},`;
 		}
 	}
 
@@ -379,10 +387,28 @@ export class AzureActiveDirectoryService {
 			}, 1000 * 60 * 5);
 		});
 
-		return Promise.race([this.handleCodeResponse(state, codeVerifier, scope), timeoutPromise]);
+		const existingStates = this._pendingStates.get(scope) || [];
+		this._pendingStates.set(scope, [...existingStates, state]);
+
+		// Register a single listener for the URI callback, in case the user starts the login process multiple times
+		// before completing it.
+		let existingPromise = this._codeExchangePromises.get(scope);
+		if (!existingPromise) {
+			existingPromise = this.handleCodeResponse(scope);
+			this._codeExchangePromises.set(scope, existingPromise);
+		}
+
+		this._codeVerfifiers.set(state, codeVerifier);
+
+		return Promise.race([existingPromise, timeoutPromise])
+			.finally(() => {
+				this._pendingStates.delete(scope);
+				this._codeExchangePromises.delete(scope);
+				this._codeVerfifiers.delete(state);
+			});
 	}
 
-	private async handleCodeResponse(state: string, codeVerifier: string, scope: string): Promise<vscode.AuthenticationSession> {
+	private async handleCodeResponse(scope: string): Promise<vscode.AuthenticationSession> {
 		let uriEventListener: vscode.Disposable;
 		return new Promise((resolve: (value: vscode.AuthenticationSession) => void, reject) => {
 			uriEventListener = this._uriHandler.event(async (uri: vscode.Uri) => {
@@ -390,12 +416,18 @@ export class AzureActiveDirectoryService {
 					const query = parseQuery(uri);
 					const code = query.code;
 
+					const acceptedStates = this._pendingStates.get(scope) || [];
 					// Workaround double encoding issues of state in web
-					if (query.state !== state && decodeURIComponent(query.state) !== state) {
+					if (!acceptedStates.includes(query.state) && !acceptedStates.includes(decodeURIComponent(query.state))) {
 						throw new Error('State does not match.');
 					}
 
-					const token = await this.exchangeCodeForToken(code, codeVerifier, scope);
+					const verifier = this._codeVerfifiers.get(query.state) ?? this._codeVerfifiers.get(decodeURIComponent(query.state));
+					if (!verifier) {
+						throw new Error('No available code verifier');
+					}
+
+					const token = await this.exchangeCodeForToken(code, verifier, scope);
 					this.setToken(token, scope);
 
 					const session = await this.convertToSession(token);
@@ -446,7 +478,19 @@ export class AzureActiveDirectoryService {
 	}
 
 	private getTokenFromResponse(json: ITokenResponse, scope: string, existingId?: string): IToken {
-		const claims = this.getTokenClaims(json.access_token);
+		let claims = undefined;
+
+		try {
+			claims = this.getTokenClaims(json.access_token);
+		} catch (e) {
+			if (json.id_token) {
+				Logger.info('Failed to fetch token claims from access_token. Attempting to parse id_token instead');
+				claims = this.getTokenClaims(json.id_token);
+			} else {
+				throw e;
+			}
+		}
+
 		return {
 			expiresIn: json.expires_in,
 			expiresAt: json.expires_in ? Date.now() + json.expires_in * 1000 : undefined,

extensions/microsoft-authentication/src/extension.ts

@@ -14,6 +14,7 @@ export async function activate(context: vscode.ExtensionContext) {
 	const telemetryReporter = new TelemetryReporter(name, version, aiKey);
 
 	const loginService = new AzureActiveDirectoryService();
+	context.subscriptions.push(loginService);
 
 	await loginService.initialize();
 

extensions/microsoft-authentication/src/keychain.ts

@@ -28,7 +28,8 @@ export type Keytar = {
 	deletePassword: typeof keytarType['deletePassword'];
 };
 
-const SERVICE_ID = `${vscode.env.uriScheme}-microsoft.login`;
+const OLD_SERVICE_ID = `${vscode.env.uriScheme}-microsoft.login`;
+const SERVICE_ID = `microsoft.login`;
 const ACCOUNT_ID = 'account';
 
 export class Keychain {
@@ -46,7 +47,7 @@ export class Keychain {
 
 	async setToken(token: string): Promise<void> {
 		try {
-			return await this.keytar.setPassword(SERVICE_ID, ACCOUNT_ID, token);
+			return await vscode.authentication.setPassword(SERVICE_ID, token);
 		} catch (e) {
 			Logger.error(`Setting token failed: ${e}`);
 
@@ -68,7 +69,7 @@ export class Keychain {
 
 	async getToken(): Promise<string | null | undefined> {
 		try {
-			return await this.keytar.getPassword(SERVICE_ID, ACCOUNT_ID);
+			return await vscode.authentication.getPassword(SERVICE_ID);
 		} catch (e) {
 			// Ignore
 			Logger.error(`Getting token failed: ${e}`);
@@ -76,15 +77,30 @@ export class Keychain {
 		}
 	}
 
-	async deleteToken(): Promise<boolean | undefined> {
+	async deleteToken(): Promise<void> {
 		try {
-			return await this.keytar.deletePassword(SERVICE_ID, ACCOUNT_ID);
+			return await vscode.authentication.deletePassword(SERVICE_ID);
 		} catch (e) {
 			// Ignore
 			Logger.error(`Deleting token failed: ${e}`);
 			return Promise.resolve(undefined);
 		}
 	}
+
+	async tryMigrate(): Promise<string | null> {
+		try {
+			const oldValue = await this.keytar.getPassword(OLD_SERVICE_ID, ACCOUNT_ID);
+			if (oldValue) {
+				await this.setToken(oldValue);
+				await this.keytar.deletePassword(OLD_SERVICE_ID, ACCOUNT_ID);
+			}
+
+			return oldValue;
+		} catch (_) {
+			// Ignore
+			return Promise.resolve(null);
+		}
+	}
 }
 
 export const keychain = new Keychain();