Feature/support for jit (machulav#261)

* first version of jit compability
* modified readme to meet the new implementations

Author: Preen

.env.example

@@ -22,3 +22,5 @@ GITHUB_REPOSITORY=
 INPUT_EC2-VOLUME-SIZE=
 INPUT_EC2-DEVICE-NAME=
 INPUT_EC2-VOLUME-TYPE=
+INPUT_USE-JIT=
+INPUT_RUNNER-GROUP-ID=

CHANGELOG.md

@@ -0,0 +1,37 @@
+# Changelog
+
+## [Unreleased]
+
+### Added
+- JIT (Just-In-Time) runner support via new `use-jit` input (default: `false`).
+  JIT runners use GitHub's `generate-jitconfig` API, skip `config.sh`,
+  and auto-deregister after completing one job.
+- New `runner-group-id` input for specifying the runner group when using JIT mode (default: `1`).
+- Validation: `use-jit` and `run-runner-as-service` cannot be used together (JIT is single-use).
+- New `runner-debug` input (default: `false`) for verbose debug logging. When enabled,
+  injects detailed echo statements into the setup script and polls EC2 serial console
+  output during runner registration. Requires `ec2:GetConsoleOutput` IAM permission.
+- New `availability-zones-config` input for multi-AZ failover. The action tries each
+  configuration in sequence until an instance is successfully launched.
+- New `metadata-options` input for configuring EC2 instance metadata (e.g. IMDSv2).
+- New `packages` input for installing packages via cloud-init during boot.
+- New `region` output for tracking which AWS region the instance was launched in.
+- EC2 console output polling via `GetConsoleOutputCommand` for remote debugging.
+- Test suite expanded to 25 tests covering JIT, debug mode, cloud-boothook,
+  runuser, tolerant chown, stale config cleanup, and package installation.
+
+### Changed
+- Switched user-data format from `#cloud-config` with `runcmd` to `#cloud-boothook`.
+  This fixes compatibility with Amazon Linux 2023 and other AMIs where
+  `cloud_final_modules` may be empty or misconfigured.
+- Replaced `su <user> -c` with `runuser -u <user> --` to avoid password prompts
+  in non-interactive cloud-init contexts.
+- Made `chown` tolerant of permission errors (`|| true`) to prevent `set -e`
+  from killing the script when `_diag/` files are owned by root.
+- Setup script now removes stale runner config files (`.runner`, `.credentials`,
+  `.credentials_rsaparams`) before `config.sh` to handle AMIs created from
+  previously configured runner instances.
+- Setup script logs are written to `/tmp/runner-setup.log` instead of
+  `/var/log/user-data.log` and `/dev/console` (which may not be accessible).
+- Updated README with full documentation for all new inputs, IAM requirements
+  for debug mode, and advanced usage sections (JIT, Multi-AZ, Debug).

README.md

@@ -135,6 +135,21 @@ inputs:
       Example: '["git", "docker.io", "nodejs"]'
     required: false
     default: '[]'
+  use-jit:
+    description: >-
+      Enable JIT (Just-In-Time) runner configuration. Uses GitHub's
+      generate-jitconfig API instead of the traditional registration-token approach.
+      JIT runners are single-use and auto-deregister after completing one job.
+      Incompatible with 'run-runner-as-service: true'.
+    required: false
+    default: 'false'
+  runner-group-id:
+    description: >-
+      The ID of the runner group to register the JIT runner in.
+      Defaults to 1, which is the "Default" runner group for repository-level runners.
+      Only used when 'use-jit' is true.
+    required: false
+    default: '1'
   runner-debug:
     description: >-
       Enable verbose debug logging for the runner setup.

action.yml

@@ -145105,10 +145105,62 @@ function buildRunCommands(githubRegistrationToken, label) {
   return userData.filter(Boolean);
 }
 
+// Build the commands to run on the instance for JIT mode.
+// JIT runners skip config.sh entirely and pass the encoded config directly to run.sh.
+function buildJitRunCommands(encodedJitConfig) {
+  const debug = config.input.runnerDebug;
+  const dbg = (cmd) => debug ? cmd : null;
+
+  // Common preamble: fail-fast and log capture
+  const preamble = [
+    '#!/bin/bash',
+    'LOGFILE=/tmp/runner-setup.log',
+    'exec > >(tee -a "$LOGFILE") 2>&1',
+    'set -e',
+    dbg('echo "[RUNNER] =========================================="'),
+    dbg('echo "[RUNNER] JIT Setup script started at $(date -u)"'),
+    dbg('echo "[RUNNER] =========================================="'),
+  ].filter(Boolean);
+
+  let userData;
+  if (config.input.runnerHomeDir) {
+    userData = [
+      ...preamble,
+      `cd "${config.input.runnerHomeDir}"`,
+      'source /tmp/pre-runner-script.sh',
+      'export RUNNER_ALLOW_RUNASROOT=1',
+      // Remove stale runner config from AMI so run.sh doesn't get confused
+      'rm -f .runner .credentials .credentials_rsaparams',
+    ];
+  } else {
+    userData = [
+      ...preamble,
+      'mkdir actions-runner && cd actions-runner',
+      'source /tmp/pre-runner-script.sh',
+      'case $(uname -m) in aarch64) ARCH="arm64" ;; amd64|x86_64) ARCH="x64" ;; esac && export RUNNER_ARCH=${ARCH}',
+      `RUNNER_VERSION=$(curl -s "https://api.github.com/repos/actions/runner/releases/latest" | grep -o '"tag_name"[[:space:]]*:[[:space:]]*"[^"]*"' | sed 's/.*"tag_name"[[:space:]]*:[[:space:]]*"\\([^"]*\\)".*/\\1/' | tr -d "v")`,
+      'curl -O -L https://github.com/actions/runner/releases/download/v${RUNNER_VERSION}/actions-runner-linux-${RUNNER_ARCH}-${RUNNER_VERSION}.tar.gz',
+      'tar xzf ./actions-runner-linux-${RUNNER_ARCH}-${RUNNER_VERSION}.tar.gz',
+      'export RUNNER_ALLOW_RUNASROOT=1',
+    ];
+  }
+
+  if (config.input.runAsUser) {
+    userData.push(`chown -R ${config.input.runAsUser} . 2>&1 || true`);
+    userData.push(`runuser -u ${config.input.runAsUser} -- ./run.sh --jitconfig ${encodedJitConfig}`);
+  } else {
+    userData.push(`./run.sh --jitconfig ${encodedJitConfig}`);
+  }
+
+  return userData;
+}
+
 // Build user data as a cloud-boothook (runs during cloud-init init stage,
 // bypassing cloud_final_modules which may be empty on some AMIs)
-function buildUserDataScript(githubRegistrationToken, label) {
-  const runCommands = buildRunCommands(githubRegistrationToken, label);
+function buildUserDataScript(githubRegistrationToken, label, encodedJitConfig) {
+  const runCommands = encodedJitConfig
+    ? buildJitRunCommands(encodedJitConfig)
+    : buildRunCommands(githubRegistrationToken, label);
 
   const lines = [];
 
@@ -145173,12 +145225,12 @@ function buildMarketOptions() {
   };
 }
 
-async function createEc2InstanceWithParams(imageId, subnetId, securityGroupId, label, githubRegistrationToken, region) {
+async function createEc2InstanceWithParams(imageId, subnetId, securityGroupId, label, githubRegistrationToken, region, encodedJitConfig) {
   // Region is always specified now, so we can directly use it
   const ec2ClientOptions = { region };
   const ec2 = new EC2Client(ec2ClientOptions);
 
-  const userData = buildUserDataScript(githubRegistrationToken, label);
+  const userData = buildUserDataScript(githubRegistrationToken, label, encodedJitConfig);
 
   const params = {
     ImageId: imageId,
@@ -145215,7 +145267,7 @@ async function createEc2InstanceWithParams(imageId, subnetId, securityGroupId, l
   return ec2InstanceId;
 }
 
-async function startEc2Instance(label, githubRegistrationToken) {
+async function startEc2Instance(label, githubRegistrationToken, encodedJitConfig) {
   core.info(`Attempting to start EC2 instance using ${config.availabilityZones.length} availability zone configuration(s)`);
 
   const errors = [];
@@ -145235,7 +145287,8 @@ async function startEc2Instance(label, githubRegistrationToken) {
         azConfig.securityGroupId,
         label,
         githubRegistrationToken,
-        region
+        region,
+        encodedJitConfig
       );
 
       core.info(`Successfully started AWS EC2 instance ${ec2InstanceId} using availability zone configuration ${i + 1} in region ${region}`);
@@ -145341,6 +145394,8 @@ module.exports = {
   terminateEc2Instance,
   waitForInstanceRunning,
   getInstanceConsoleOutput,
+  // Exposed for testing only
+  _buildUserDataScriptForTest: buildUserDataScript,
 };
 
 
@@ -145379,12 +145434,14 @@ class Config {
       availabilityZonesConfig: core.getInput('availability-zones-config'),
       metadataOptions: JSON.parse(core.getInput('metadata-options') || '{}'),
       packages: JSON.parse(core.getInput('packages') || '[]'),
+      useJit: core.getInput('use-jit') === 'true',
+      runnerGroupId: parseInt(core.getInput('runner-group-id') || '1', 10),
       runnerDebug: core.getInput('runner-debug') === 'true',
     };
 
     // Get the AWS_REGION environment variable
     this.defaultRegion = process.env.AWS_REGION;
-    
+
     const tags = JSON.parse(core.getInput('aws-resource-tags'));
     this.tagSpecifications = null;
     if (tags.length > 0) {
@@ -145422,12 +145479,12 @@ class Config {
       if (this.input.availabilityZonesConfig) {
         try {
           this.availabilityZones = JSON.parse(this.input.availabilityZonesConfig);
-          
+
           // Validate each availability zone configuration
           if (!Array.isArray(this.availabilityZones)) {
             throw new Error('availability-zones-config must be a JSON array');
           }
-          
+
           this.availabilityZones.forEach((az, index) => {
             if (!az.imageId) {
               throw new Error(`Missing imageId in availability-zones-config at index ${index}`);
@@ -145460,7 +145517,7 @@ class Config {
             `Either provide 'availability-zones-config' or all of the following: 'ec2-image-id', 'subnet-id', 'security-group-id'`
           );
         }
-        
+
         // Convert individual parameters to a single availability zone config
         this.availabilityZones.push({
           imageId: this.input.ec2ImageId,
@@ -145469,10 +145526,17 @@ class Config {
           // Add default region when using legacy configuration
           region: this.defaultRegion
         });
-        
+
         core.info('Using individual parameters as a single availability zone configuration');
       }
 
+      if (this.input.useJit && this.input.runAsService) {
+        throw new Error(
+          "The 'use-jit' and 'run-runner-as-service' inputs are incompatible. " +
+          'JIT runners are single-use and cannot run as a service.'
+        );
+      }
+
       if (this.marketType?.length > 0 && this.input.marketType !== 'spot') {
         throw new Error('Invalid `market-type` input. Allowed values: spot.');
       }
@@ -145500,6 +145564,8 @@ try {
   core.setFailed(error.message);
 }
 
+module.exports.Config = Config;
+
 
 /***/ }),
 
@@ -145539,6 +145605,33 @@ async function getRegistrationToken() {
   }
 }
 
+// generate a JIT (Just-In-Time) runner configuration via the GitHub API
+async function getJitRunnerConfig(label) {
+  const octokit = github.getOctokit(config.input.githubToken);
+
+  try {
+    const response = await octokit.request(
+      'POST /repos/{owner}/{repo}/actions/runners/generate-jitconfig',
+      {
+        ...config.githubContext,
+        name: `ec2-${label}`,
+        runner_group_id: config.input.runnerGroupId,
+        labels: [label],
+        work_folder: '_work',
+      }
+    );
+
+    core.info('GitHub JIT runner configuration is received');
+    return {
+      runnerId: response.data.runner.id,
+      encodedJitConfig: response.data.encoded_jit_config,
+    };
+  } catch (error) {
+    core.error('GitHub JIT runner configuration generation error');
+    throw error;
+  }
+}
+
 async function removeRunner() {
   const runner = await getRunner(config.input.label);
   const octokit = github.getOctokit(config.input.githubToken);
@@ -145604,6 +145697,7 @@ async function waitForRunnerRegistered(label, onPollCallback) {
 
 module.exports = {
   getRegistrationToken,
+  getJitRunnerConfig,
   removeRunner,
   waitForRunnerRegistered,
 };
@@ -147590,8 +147684,19 @@ function setOutput(label, ec2InstanceId, region) {
 
 async function start() {
   const label = config.input.label ? config.input.label : config.generateUniqueLabel();
-  const githubRegistrationToken = await gh.getRegistrationToken();
-  const result = await aws.startEc2Instance(label, githubRegistrationToken);
+
+  let githubRegistrationToken = null;
+  let encodedJitConfig = null;
+
+  if (config.input.useJit) {
+    const jitConfig = await gh.getJitRunnerConfig(label);
+    encodedJitConfig = jitConfig.encodedJitConfig;
+    core.info(`JIT runner created with runner ID: ${jitConfig.runnerId}`);
+  } else {
+    githubRegistrationToken = await gh.getRegistrationToken();
+  }
+
+  const result = await aws.startEc2Instance(label, githubRegistrationToken, encodedJitConfig);
   const ec2InstanceId = result.ec2InstanceId;
   const region = result.region;
 
@@ -147623,7 +147728,12 @@ async function start() {
 
 async function stop() {
   await aws.terminateEc2Instance();
-  await gh.removeRunner();
+
+  if (config.input.useJit) {
+    core.info('JIT runner auto-deregisters after job completion. Skipping runner removal.');
+  } else {
+    await gh.removeRunner();
+  }
 }
 
 (async function () {