From c41964f89bddf696fe7fd495b976466fb82f686a Mon Sep 17 00:00:00 2001
From: Tejas Kumar <tejask@spotify.com>
Date: Tue, 18 May 2021 11:41:31 +0200
Subject: [PATCH] Scaffolder: Enable branch protection on new repo

Signed-off-by: Tejas Kumar <tejask@spotify.com>
---
 .../src/scaffolder/stages/publish/github.ts   | 23 ++++++-
 .../src/scaffolder/stages/publish/helpers.ts  | 61 +++++++++++++++++++
 2 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/plugins/scaffolder-backend/src/scaffolder/stages/publish/github.ts b/plugins/scaffolder-backend/src/scaffolder/stages/publish/github.ts
index 829856ec0c..599cd32337 100644
--- a/plugins/scaffolder-backend/src/scaffolder/stages/publish/github.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/stages/publish/github.ts
@@ -15,7 +15,10 @@
  */
 
 import { PublisherBase, PublisherOptions, PublisherResult } from './types';
-import { initRepoAndPush } from './helpers';
+import {
+  enableBranchProtectionOnDefaultRepoBranch,
+  initRepoAndPush,
+} from './helpers';
 import {
   GitHubIntegrationConfig,
   GithubCredentialsProvider,
@@ -99,6 +102,19 @@ export class GithubPublisher implements PublisherBase {
       /\.git$/,
       '/blob/master/catalog-info.yaml',
     );
+
+    try {
+      await enableBranchProtectionOnDefaultRepoBranch({
+        owner,
+        client,
+        repoName: name,
+      });
+    } catch (e) {
+      throw new Error(
+        `Failed to add branch protection to '${name}': ${e.message}`,
+      );
+    }
+
     return { remoteUrl, catalogInfoUrl };
   }
 
@@ -130,7 +146,7 @@ export class GithubPublisher implements PublisherBase {
             description,
           });
 
-    const { data } = await repoCreationPromise;
+    const { data: newRepo } = await repoCreationPromise;
 
     try {
       if (access?.startsWith(`${owner}/`)) {
@@ -156,6 +172,7 @@ export class GithubPublisher implements PublisherBase {
         `Failed to add access to '${access}'. Status ${e.status} ${e.message}`,
       );
     }
-    return data?.clone_url;
+
+    return newRepo.clone_url;
   }
 }
diff --git a/plugins/scaffolder-backend/src/scaffolder/stages/publish/helpers.ts b/plugins/scaffolder-backend/src/scaffolder/stages/publish/helpers.ts
index ffefb88cf0..6527a30fe5 100644
--- a/plugins/scaffolder-backend/src/scaffolder/stages/publish/helpers.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/stages/publish/helpers.ts
@@ -17,6 +17,7 @@
 import globby from 'globby';
 import { Logger } from 'winston';
 import { Git } from '@backstage/backend-common';
+import { Octokit, RestEndpointMethodTypes } from '@octokit/rest';
 
 export async function initRepoAndPush({
   dir,
@@ -67,3 +68,63 @@ export async function initRepoAndPush({
     remote: 'origin',
   });
 }
+
+type BranchProtectionOptions = {
+  client: Octokit;
+  owner: string;
+  repoName: string;
+  isRetry?: boolean;
+};
+
+export const enableBranchProtectionOnDefaultRepoBranch = async ({
+  repoName,
+  client,
+  owner,
+  isRetry = false,
+}: BranchProtectionOptions): Promise<
+  RestEndpointMethodTypes['repos']['updateBranchProtection']['response']
+> => {
+  const { data: repo } = await client.repos.get({
+    owner,
+    repo: repoName,
+  });
+
+  try {
+    const response = await client.repos.updateBranchProtection({
+      headers: {
+        Accept:
+          /**
+           * 👇 we need this header because allowing a custom
+           * reviewer count on branch protection is a preview
+           * feature.
+           *
+           * More here: https://docs.github.com/en/rest/overview/api-previews#require-multiple-approving-reviews
+           */
+          'application/vnd.github.luke-cage-preview+json',
+      },
+      owner,
+      repo: repoName,
+      branch: repo.default_branch,
+      required_status_checks: { strict: true, contexts: [] },
+      restrictions: null,
+      enforce_admins: true,
+      required_pull_request_reviews: { required_approving_review_count: 1 },
+    });
+
+    return response;
+  } catch (e) {
+    if (!isRetry && e.message.includes('Branch not found')) {
+      // GitHub has eventual consistency. Fail silently, wait, and try again.
+      await new Promise(resolve => setTimeout(resolve, 600));
+
+      return await enableBranchProtectionOnDefaultRepoBranch({
+        repoName,
+        client,
+        owner,
+        isRetry: true,
+      });
+    }
+
+    throw e;
+  }
+};