Handle SAML errors when fetching repository context

Author: joshmgross

src/api/handleSamlError.ts

@@ -0,0 +1,23 @@
+import {Octokit} from "@octokit/rest";
+import {getSession} from "../auth/auth";
+import {logDebug} from "../log";
+import {getClient} from "./api";
+
+export async function handleSamlError<T>(request: (client: Octokit) => Promise<T>): Promise<T> {
+  try {
+    const session = await getSession();
+    const client = getClient(session.accessToken);
+    return await request(client);
+  } catch (error) {
+    if ((error as Error).message.includes("Resource protected by organization SAML enforcement.")) {
+      logDebug("SAML error, re-authenticating");
+      const session = await getSession(
+        "Your organization is protected by SAML enforcement. Please sign-in again to continue."
+      );
+      const client = getClient(session.accessToken);
+      return await request(client);
+    } else {
+      throw error;
+    }
+  }
+}

src/auth/auth.ts

@@ -3,9 +3,15 @@ import * as vscode from "vscode";
 const AUTH_PROVIDER_ID = "github";
 const DEFAULT_SCOPES = ["repo", "workflow"];
 
-export async function getSession(): Promise<vscode.AuthenticationSession> {
+/**
+ * Retrieves a session from the GitHub authentication provider
+ * @param forceMessage Force a new session with a prompt to the user
+ * @returns A {@link vscode.AuthenticationSession}
+ */
+export async function getSession(forceMessage?: string): Promise<vscode.AuthenticationSession> {
   const existingSession = await vscode.authentication.getSession(AUTH_PROVIDER_ID, getScopes(), {
-    createIfNone: true
+    createIfNone: true,
+    forceNewSession: forceMessage ? {detail: forceMessage} : false
   });
 
   if (!existingSession) {

src/git/repository.ts

@@ -1,13 +1,11 @@
 import * as vscode from "vscode";
-
-import {logDebug, logError} from "../log";
-import {API, GitExtension, RefType, RepositoryState} from "../typings/git";
-
 import {Octokit} from "@octokit/rest";
-import {getClient} from "../api/api";
-import {getSession} from "../auth/auth";
+
+import {handleSamlError} from "../api/handleSamlError";
 import {getRemoteName} from "../configuration/configuration";
 import {Protocol} from "../external/protocol";
+import {logDebug, logError} from "../log";
+import {API, GitExtension, RefType, RepositoryState} from "../typings/git";
 
 interface GitHubUrls {
   workspaceUri: vscode.Uri;
@@ -146,8 +144,6 @@ export async function getGitHubContext(): Promise<GitHubContext | undefined> {
 
   try {
     const git = await getGitExtension();
-    const session = await getSession();
-    const client = getClient(session.accessToken);
 
     const protocolInfos = await getGitHubUrls();
     if (!protocolInfos) {
@@ -157,29 +153,31 @@ export async function getGitHubContext(): Promise<GitHubContext | undefined> {
 
     logDebug("Found protocol infos", protocolInfos.length.toString());
 
-    const repos = await Promise.all(
-      protocolInfos.map(async (protocolInfo): Promise<GitHubRepoContext> => {
-        logDebug("Getting infos for repository", protocolInfo.url);
-
-        const repoInfo = await client.repos.get({
-          repo: protocolInfo.protocol.repositoryName,
-          owner: protocolInfo.protocol.owner
-        });
-
-        const repo = git && git.getRepository(protocolInfo.workspaceUri);
-
-        return {
-          workspaceUri: protocolInfo.workspaceUri,
-          client,
-          repositoryState: repo?.state,
-          name: protocolInfo.protocol.repositoryName,
-          owner: protocolInfo.protocol.owner,
-          id: repoInfo.data.id,
-          defaultBranch: `refs/heads/${repoInfo.data.default_branch}`,
-          organizationOwned: repoInfo.data.owner.type === "Organization"
-        };
-      })
-    );
+    const repos = await handleSamlError(async (client: Octokit) => {
+      return await Promise.all(
+        protocolInfos.map(async (protocolInfo): Promise<GitHubRepoContext> => {
+          logDebug("Getting infos for repository", protocolInfo.url);
+
+          const repoInfo = await client.repos.get({
+            repo: protocolInfo.protocol.repositoryName,
+            owner: protocolInfo.protocol.owner
+          });
+
+          const repo = git && git.getRepository(protocolInfo.workspaceUri);
+
+          return {
+            workspaceUri: protocolInfo.workspaceUri,
+            client,
+            repositoryState: repo?.state,
+            name: protocolInfo.protocol.repositoryName,
+            owner: protocolInfo.protocol.owner,
+            id: repoInfo.data.id,
+            defaultBranch: `refs/heads/${repoInfo.data.default_branch}`,
+            organizationOwned: repoInfo.data.owner.type === "Organization"
+          };
+        })
+      );
+    });
 
     gitHubContext = Promise.resolve({
       repos,