From 3dd678d2b020146ffaf1e7de62324fbfc29058c1 Mon Sep 17 00:00:00 2001
From: Ladislau Szomoru <3372902+lszomoru@users.noreply.github.com>
Date: Tue, 9 May 2023 20:10:00 +0200
Subject: [PATCH] GitHub - update authentication flow for branch protection
 (#181924)

* GitHub - tweak authentication request

* Add tracing messages
---
 extensions/github/src/auth.ts             | 39 +++++++++++++----------
 extensions/github/src/branchProtection.ts | 20 ++++++++++--
 2 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/extensions/github/src/auth.ts b/extensions/github/src/auth.ts
index c96d21187d0..4fdf9bcd517 100644
--- a/extensions/github/src/auth.ts
+++ b/extensions/github/src/auth.ts
@@ -10,6 +10,8 @@ import { Octokit } from '@octokit/rest';
 import { httpsOverHttp } from 'tunnel';
 import { URL } from 'url';
 
+export class AuthenticationError extends Error { }
+
 function getAgent(url: string | undefined = process.env.HTTPS_PROXY): Agent {
 	if (!url) {
 		return globalAgent;
@@ -57,25 +59,30 @@ export function getOctokit(): Promise<Octokit> {
 
 let _octokitGraphql: Promise<graphql> | undefined;
 
-export function getOctokitGraphql(): Promise<graphql> {
+export async function getOctokitGraphql(): Promise<graphql> {
 	if (!_octokitGraphql) {
-		_octokitGraphql = getSession()
-			.then(async session => {
-				const token = session.accessToken;
-				const { graphql } = await import('@octokit/graphql');
+		try {
+			const session = await authentication.getSession('github', scopes, { createIfNone: false });
 
-				return graphql.defaults({
-					headers: {
-						authorization: `token ${token}`
-					},
-					request: {
-						agent: getAgent()
-					}
-				});
-			}).then(null, async err => {
-				_octokitGraphql = undefined;
-				throw err;
+			if (!session) {
+				throw new AuthenticationError('No GitHub authentication session available.');
+			}
+
+			const token = session.accessToken;
+			const { graphql } = await import('@octokit/graphql');
+
+			return graphql.defaults({
+				headers: {
+					authorization: `token ${token}`
+				},
+				request: {
+					agent: getAgent()
+				}
 			});
+		} catch (err) {
+			_octokitGraphql = undefined;
+			throw err;
+		}
 	}
 
 	return _octokitGraphql;
diff --git a/extensions/github/src/branchProtection.ts b/extensions/github/src/branchProtection.ts
index 0ca239d2054..f2f93aa73f9 100644
--- a/extensions/github/src/branchProtection.ts
+++ b/extensions/github/src/branchProtection.ts
@@ -3,9 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { EventEmitter, LogOutputChannel, Memento, Uri, workspace } from 'vscode';
+import { authentication, EventEmitter, LogOutputChannel, Memento, Uri, workspace } from 'vscode';
 import { Repository as GitHubRepository, RepositoryRuleset } from '@octokit/graphql-schema';
-import { getOctokitGraphql } from './auth';
+import { AuthenticationError, getOctokitGraphql } from './auth';
 import { API, BranchProtection, BranchProtectionProvider, BranchProtectionRule, Repository } from './typings/git';
 import { DisposableStore, getRepositoryFromUrl } from './util';
 
@@ -163,10 +163,12 @@ export class GithubBranchProtectionProvider implements BranchProtectionProvider
 				}
 
 				// Repository details
+				this.logger.trace(`Fetching repository details for "${repository.owner}/${repository.repo}".`);
 				const repositoryDetails = await this.getRepositoryDetails(repository.owner, repository.repo);
 
 				// Check repository write permission
 				if (repositoryDetails.viewerPermission !== 'ADMIN' && repositoryDetails.viewerPermission !== 'MAINTAIN' && repositoryDetails.viewerPermission !== 'WRITE') {
+					this.logger.trace(`Skipping branch protection for "${repository.owner}/${repository.repo}" due to missing repository write permission.`);
 					continue;
 				}
 
@@ -189,8 +191,20 @@ export class GithubBranchProtectionProvider implements BranchProtectionProvider
 
 			// Save branch protection to global state
 			await this.globalState.update(this.globalStateKey, branchProtection);
+			this.logger.trace(`Branch protection for "${this.repository.rootUri.toString()}": ${JSON.stringify(branchProtection)}.`);
 		} catch (err) {
-			// noop
+			if (err instanceof AuthenticationError) {
+				// Since there is no GitHub authentication session available we need to wait
+				// until the user signs in with their GitHub account so that we can query the
+				// repository rulesets
+				const disposable = authentication.onDidChangeSessions(e => {
+					if (e.provider.id === 'github') {
+						disposable.dispose();
+						this.updateRepositoryBranchProtection();
+					}
+				});
+			}
+
 			this.logger.warn(`Failed to update repository branch protection: ${err.message}`);
 		}
 	}