NextChapterSoftware · mahdi-torabi · Jan 17, 2024
diff --git a/.env.example b/.env.example
@@ -4,6 +4,7 @@
 # Use format "owner/repo".
 GITHUB_REPOSITORY=YOUR_ORG/YOUR_REPO
 GITHUB_RUN_ID=SomeExampleValue
+GITHUB_ACTOR=SomeUser
 
 INPUT_AWS_ACCESS_KEY_ID=
 INPUT_AWS_SECRET_ACCESS_KEY=

diff --git a/.github/workflows/aws-test-all.yml b/.github/workflows/aws-test-all.yml
@@ -4,14 +4,15 @@ on:
   push:
     branches:
       - main
+      - add_support_for_just_intime_runners
 jobs:
   package:
     permissions:
       actions: write
       contents: read
       pull-requests: read
     name: test action with all distros
-    if: github.ref == 'refs/heads/main'
+    #if: github.ref == 'refs/heads/main'
     uses: ./.github/workflows/aws-test-on-demand.yml
     strategy:
       matrix:

diff --git a/.github/workflows/aws-test-on-demand.yml b/.github/workflows/aws-test-on-demand.yml
@@ -1,12 +1,5 @@
 name: Test Distro
 
-#on:
-#  push:
-#    branches:
-#      - main
-#      - update_packages
- # pull_request:
- # workflow_dispatch:
 on:
   workflow_call:
     inputs:

diff --git a/README.md b/README.md
@@ -53,9 +53,10 @@ Users can provide their own custom AMI image pre-loaded with all the necessary t
 
 ### 1. Create GitHub Personal Access Token
 1. Create a [fine-grained personal access token](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens#creating-a-fine-grained-personal-access-token)
-2. Edit the token permissions and select `All repositories` for `Repository access`
-3. Grant `Read and Write access to administration` under Repository permissions
-4. Add the token to GitHub Action secrets and note the secret name
+2. Edit the token permissions and select `Only select repositories` for `Repository access`
+3. Select any repositories you wish to use with this action
+4. Grant `Read and Write access` for `Administration` access level under Repository permissions
+5. Add the token to GitHub Action secrets and note the secret name
 
 ### 2. Setup GitHub Secrets for IAM credentials
 
@@ -166,7 +167,7 @@ jobs:
 ## How it all works under the hood
 
 ### General instance launch flow
-- Your GitHub personal token is used to obtain a Runner Registration token
+- Your GitHub personal token is used to obtain a just-in-time runner configuration
 - If no explicit runner version has been provided, it will retrieve the latest version number
 - It then uses all the provided info to compile an EC2 user-data script which does the following:
   - Set a max TTL on the EC2 instance on startup 
@@ -178,7 +179,7 @@ jobs:
 - Once EC2 boot has completed, user-data script is executed
 - Runner binary registers itself with GitHub API using the current job ID
 - Once the Runner is registered, control is transferred to the next job (this is your build job)
-- Upon a job completion (failure/success), Shutdown script is triggered to kill the instance with a 1 minute delay
+- Upon a job completion (failure/success), Shutdown script is triggered to `terminate` the instance with a 1 minute delay
 
 ### Spot instance provisioning
 - Script looks up On-Demand price for the supplied instance type

diff --git a/dist/index.js b/dist/index.js
@@ -644,29 +644,6 @@ exports.Ec2Pricing = Ec2Pricing;
 
 "use strict";
 
-var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
-    if (k2 === undefined) k2 = k;
-    var desc = Object.getOwnPropertyDescriptor(m, k);
-    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
-      desc = { enumerable: true, get: function() { return m[k]; } };
-    }
-    Object.defineProperty(o, k2, desc);
-}) : (function(o, m, k, k2) {
-    if (k2 === undefined) k2 = k;
-    o[k2] = m[k];
-}));
-var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
-    Object.defineProperty(o, "default", { enumerable: true, value: v });
-}) : function(o, v) {
-    o["default"] = v;
-});
-var __importStar = (this && this.__importStar) || function (mod) {
-    if (mod && mod.__esModule) return mod;
-    var result = {};
-    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
-    __setModuleDefault(result, mod);
-    return result;
-};
 var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
     function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
     return new (P || (P = Promise))(function (resolve, reject) {
@@ -678,7 +655,6 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
 };
 Object.defineProperty(exports, "__esModule", ({ value: true }));
 exports.UserData = void 0;
-const github = __importStar(__nccwpck_require__(95438));
 const github_1 = __nccwpck_require__(51294);
 class UserData {
     constructor(config) {
@@ -688,7 +664,7 @@ class UserData {
         return __awaiter(this, void 0, void 0, function* () {
             const ghClient = new github_1.GithubClient(this.config);
             const githubActionRunnerVersion = yield ghClient.getRunnerVersion();
-            const runnerRegistrationToken = yield ghClient.getRunnerRegistrationToken();
+            const jitRunnerRegistrationConfig = yield ghClient.getJITRunnerRegistrationConfig();
             if (!this.config.githubActionRunnerLabel)
                 throw Error("failed to object job ID for label");
             const cmds = [
@@ -707,8 +683,7 @@ class UserData {
                 "export RUNNER_ALLOW_RUNASROOT=1",
                 `RUNNER_NAME=${this.config.githubJobId}-$(hostname)-ec2`,
                 "[ -n \"$(command -v yum)\" ] && yum install libicu -y",
-                `./config.sh --unattended  --ephemeral --url https://github.com/${github.context.repo.owner}/${github.context.repo.repo} --token ${runnerRegistrationToken.token} --labels ${this.config.githubActionRunnerLabel} --name $RUNNER_NAME`,
-                "./run.sh",
+                `./run.sh --jitconfig ${jitRunnerRegistrationConfig.encoded_jit_config}`,
             ];
             return Buffer.from(cmds.join("\n")).toString("base64");
         });
@@ -793,7 +768,6 @@ class GithubClient {
                             return { name: label };
                         }),
                     };
-                    //const mockRunnerData = {"total_count":2,"runners":[{"id":319,"name":"ip-192-168-0-139","os":"Linux","status":"online","busy":true,"labels":[{"id":297,"name":"self-hosted","type":"read-only"},{"id":298,"name":"Linux","type":"read-only"},{"id":299,"name":"X64","type":"read-only"},{"id":314,"name":"dow0w","type":"custom"}]},{"id":320,"name":"ip-192-168-11-102","os":"Linux","status":"online","busy":true,"labels":[{"id":297,"name":"self-hosted","type":"read-only"},{"id":298,"name":"Linux","type":"read-only"},{"id":299,"name":"X64","type":"read-only"},{"id":315,"name":"2pdrq","type":"custom"}]}]}
                     const matches = _.filter(runners.data.runners, searchLabels);
                     return matches.length > 0 ? matches[0] : null;
                 }
@@ -805,18 +779,30 @@ class GithubClient {
             return null;
         });
     }
-    getRunnerRegistrationToken() {
+    getJITRunnerRegistrationConfig() {
         return __awaiter(this, void 0, void 0, function* () {
             const octokit = github.getOctokit(this.config.githubToken);
             try {
-                const response = yield octokit.rest.actions.createRegistrationTokenForRepo({
+                this.config.githubActionRunnerLabel;
+                const response = yield octokit.rest.actions.generateRunnerJitconfigForRepo({
                     owner: github.context.repo.owner,
                     repo: github.context.repo.repo,
+                    name: `${this.config.githubRepo}-${this.config.githubJobId}-${github.context.actor}`,
+                    runner_group_id: 1,
+                    labels: [
+                        'self-hosted',
+                        this.config.githubActionRunnerLabel,
+                        github.context.actor
+                    ],
+                    work_folder: '_work',
+                    /*headers: {
+                      'X-GitHub-Api-Version': '2022-11-28'
+                    }*/
                 });
                 return response.data;
             }
             catch (error) {
-                core.error(`Failed to get Runner registration token: ${error}`);
+                core.error(`Failed to get just-in-time runner registration config: ${error}`);
                 throw error;
             }
         });
@@ -1018,14 +1004,29 @@ function stop() {
             const config = new config_1.ActionConfig();
             const ec2Client = new ec2_1.Ec2Instance(config);
             const ghClient = new github_1.GithubClient(config);
+            // Following cleanup tasks are best-effort!
+            // Unused JIT runners will be automatically removed by GitHub
+            // EC2 instances are configured with a hard TTL and will terminate once TTL expires
+            // Try to remove runner
+            const deleted = yield ghClient.removeRunnerWithLabels([config.githubActionRunnerLabel]);
+            if (deleted) {
+                core.info(`Removed runner with job id label ${config.githubActionRunnerLabel}`);
+            }
+            else {
+                core.error(`Failed to clean up runner with job id label ${config.githubActionRunnerLabel}`);
+                core.info(`GitHub will automatically cleanup unused runners after sometime`);
+            }
+            // Try to remove EC2 instance
             const instanceId = yield ec2Client.getInstancesForTags();
             if (instanceId === null || instanceId === void 0 ? void 0 : instanceId.InstanceId)
                 yield ec2Client.terminateInstances(instanceId === null || instanceId === void 0 ? void 0 : instanceId.InstanceId);
             const result = yield ghClient.removeRunnerWithLabels([config.githubJobId]);
-            if (result)
+            if (result) {
                 core.info("Finished instance cleanup");
-            else
-                throw Error("Failed to cleanup instance");
+            }
+            else {
+                core.error("Failed to terminate ec2 instance. It will be automatically terminated once its TTL expires");
+            }
         }
         catch (error) {
             core.info(error);

diff --git a/src/ec2/userdata.ts b/src/ec2/userdata.ts
@@ -12,7 +12,7 @@ export class UserData {
   async getUserData(): Promise<string> {
     const ghClient = new GithubClient(this.config);
     const githubActionRunnerVersion = await ghClient.getRunnerVersion();
-    const runnerRegistrationToken = await ghClient.getRunnerRegistrationToken();
+    const jitRunnerRegistrationConfig = await ghClient.getJITRunnerRegistrationConfig();
     if (!this.config.githubActionRunnerLabel)
       throw Error("failed to object job ID for label");
 
@@ -32,8 +32,7 @@ export class UserData {
       "export RUNNER_ALLOW_RUNASROOT=1",
       `RUNNER_NAME=${this.config.githubJobId}-$(hostname)-ec2`,
       "[ -n \"$(command -v yum)\" ] && yum install libicu -y",
-      `./config.sh --unattended  --ephemeral --url https://github.com/${github.context.repo.owner}/${github.context.repo.repo} --token ${runnerRegistrationToken.token} --labels ${this.config.githubActionRunnerLabel} --name $RUNNER_NAME`,
-      "./run.sh",
+      `./run.sh --jitconfig ${jitRunnerRegistrationConfig.encoded_jit_config}`,
     ];
 
     return Buffer.from(cmds.join("\n")).toString("base64");

diff --git a/src/github/github.ts b/src/github/github.ts
@@ -44,7 +44,6 @@ export class GithubClient {
           }),
         };
 
-        //const mockRunnerData = {"total_count":2,"runners":[{"id":319,"name":"ip-192-168-0-139","os":"Linux","status":"online","busy":true,"labels":[{"id":297,"name":"self-hosted","type":"read-only"},{"id":298,"name":"Linux","type":"read-only"},{"id":299,"name":"X64","type":"read-only"},{"id":314,"name":"dow0w","type":"custom"}]},{"id":320,"name":"ip-192-168-11-102","os":"Linux","status":"online","busy":true,"labels":[{"id":297,"name":"self-hosted","type":"read-only"},{"id":298,"name":"Linux","type":"read-only"},{"id":299,"name":"X64","type":"read-only"},{"id":315,"name":"2pdrq","type":"custom"}]}]}
         const matches = _.filter(runners.data.runners, searchLabels);
         return matches.length > 0 ? matches[0] : null;
       } catch (error) {
@@ -55,18 +54,30 @@ export class GithubClient {
     return null;
   }
 
-  async getRunnerRegistrationToken() {
+  async getJITRunnerRegistrationConfig() {
     const octokit = github.getOctokit(this.config.githubToken);
     try {
+      this.config.githubActionRunnerLabel
       const response =
-        await octokit.rest.actions.createRegistrationTokenForRepo({
-          owner: github.context.repo.owner,
-          repo: github.context.repo.repo,
-        });
+          await octokit.rest.actions.generateRunnerJitconfigForRepo({
+            owner: github.context.repo.owner,
+            repo: github.context.repo.repo,
+            name: `${this.config.githubRepo}-${this.config.githubJobId}-${github.context.actor}`,
+            runner_group_id: 1,
+            labels: [
+              'self-hosted',
+              this.config.githubActionRunnerLabel,
+              github.context.actor
+            ],
+            work_folder: '_work',
+            /*headers: {
+              'X-GitHub-Api-Version': '2022-11-28'
+            }*/
+          })
 
       return response.data;
     } catch (error) {
-      core.error(`Failed to get Runner registration token: ${error}`);
+      core.error(`Failed to get just-in-time runner registration config: ${error}`);
       throw error;
     }
   }

diff --git a/src/main.ts b/src/main.ts
@@ -72,14 +72,30 @@ async function stop() {
     const config = new ActionConfig();
     const ec2Client = new Ec2Instance(config);
     const ghClient = new GithubClient(config);
+
+    // Following cleanup tasks are best-effort!
+    // Unused JIT runners will be automatically removed by GitHub
+    // EC2 instances are configured with a hard TTL and will terminate once TTL expires
+
+    // Try to remove runner
+    const deleted = await ghClient.removeRunnerWithLabels([config.githubActionRunnerLabel])
+    if(deleted){
+      core.info(`Removed runner with job id label ${config.githubActionRunnerLabel}`);
+    } else {
+      core.error(`Failed to clean up runner with job id label ${config.githubActionRunnerLabel}`);
+      core.info(`GitHub will automatically cleanup unused runners after sometime`);
+    }
+
+    // Try to remove EC2 instance
     const instanceId = await ec2Client.getInstancesForTags();
     if (instanceId?.InstanceId)
       await ec2Client.terminateInstances(instanceId?.InstanceId);
     const result = await ghClient.removeRunnerWithLabels([config.githubJobId]);
-    if(result)
+    if(result) {
       core.info("Finished instance cleanup");
-    else
-      throw Error("Failed to cleanup instance")
+    } else {
+      core.error("Failed to terminate ec2 instance. It will be automatically terminated once its TTL expires")
+    }
   } catch(error){
     core.info(error)
   }

diff --git a/tests/github/github.ts b/tests/github/github.ts
@@ -17,17 +17,36 @@ describe('Github API tests', () => {
         expect(runners).not.throw
     });
 
-    it('get runner registration token for repo', async () => {
-        const token = await githubClient.getRunnerRegistrationToken()
-        expect(token.token).is.not.undefined;
-        expect(token.token.length).to.greaterThan(0);
-        expect(token.expires_at).is.not.undefined;
-        expect(token.expires_at.length).to.greaterThan(0);
+
+    it('get jit runner registration config for repo', async () => {
+        const jitConfig = await githubClient.getJITRunnerRegistrationConfig();
+        expect(jitConfig.encoded_jit_config).is.not.undefined;
+        expect(jitConfig.encoded_jit_config.length).to.greaterThan(0);
+        const runners = await githubClient.getRunnerWithLabels([config.githubActionRunnerLabel])
+        expect(runners).not.throw
+        expect(runners).is.not.empty
+        expect(runners).is.not.equal(null)
     });
 
-    it('list runners with labels for repo', async () => {
-        const runners = await githubClient.removeRunnerWithLabels(["foo", "bar"])
+    it('remove runners with labels for repo', async () => {
+        let runners = await githubClient.removeRunnerWithLabels(["foo", "bar"])
         expect(runners).is.true
+
+        // Check if runner exists before removing
+        runners = await githubClient.getRunnerWithLabels([config.githubActionRunnerLabel])
+        expect(runners).not.throw
+        expect(runners).is.not.empty
+        expect(runners).is.not.equal(null)
+
+        // Remove runner
+        const result = await githubClient.removeRunnerWithLabels([config.githubActionRunnerLabel])
+        expect(result).is.true
+
+        // Check if runner has been deleted
+        runners = await githubClient.getRunnerWithLabels([config.githubActionRunnerLabel])
+        expect(runners).not.throw
+        expect(runners).is.equal(null)
+
     });
 
 });