feat: Add Enterprise Runners support

README.md

@@ -127,6 +127,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 | <a name="input_disable_runner_autoupdate"></a> [disable\_runner\_autoupdate](#input\_disable\_runner\_autoupdate) | Disable the auto update of the github runner agent. Be aware there is a grace period of 30 days, see also the [GitHub article](https://github.blog/changelog/2022-02-01-github-actions-self-hosted-runners-can-now-disable-automatic-updates/) | `bool` | `false` | no |
 | <a name="input_enable_ami_housekeeper"></a> [enable\_ami\_housekeeper](#input\_enable\_ami\_housekeeper) | Option to disable the lambda to clean up old AMIs. | `bool` | `false` | no |
 | <a name="input_enable_cloudwatch_agent"></a> [enable\_cloudwatch\_agent](#input\_enable\_cloudwatch\_agent) | Enables the cloudwatch agent on the ec2 runner instances. The runner uses a default config that can be overridden via `cloudwatch_config`. | `bool` | `true` | no |
+| <a name="input_enable_enterprise_runners"></a> [enable\_enterprise\_runners](#input\_enable\_enterprise\_runners) | Enable enterprise-level runners; when true, authentication must use enterprise\_pat. | `bool` | `false` | no |
 | <a name="input_enable_ephemeral_runners"></a> [enable\_ephemeral\_runners](#input\_enable\_ephemeral\_runners) | Enable ephemeral runners, runners will only be used once. | `bool` | `false` | no |
 | <a name="input_enable_jit_config"></a> [enable\_jit\_config](#input\_enable\_jit\_config) | Overwrite the default behavior for JIT configuration. By default JIT configuration is enabled for ephemeral runners and disabled for non-ephemeral runners. In case of GHES check first if the JIT config API is avaialbe. In case you upgradeing from 3.x to 4.x you can set `enable_jit_config` to `false` to avoid a breaking change when having your own AMI. | `bool` | `null` | no |
 | <a name="input_enable_job_queued_check"></a> [enable\_job\_queued\_check](#input\_enable\_job\_queued\_check) | Only scale if the job event received by the scale up lambda is in the queued state. By default enabled for non ephemeral runners and disabled for ephemeral. Set this variable to overwrite the default behavior. | `bool` | `null` | no |
@@ -139,10 +140,12 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 | <a name="input_enable_ssm_on_runners"></a> [enable\_ssm\_on\_runners](#input\_enable\_ssm\_on\_runners) | Enable to allow access to the runner instances for debugging purposes via SSM. Note that this adds additional permissions to the runner instances. | `bool` | `false` | no |
 | <a name="input_enable_user_data_debug_logging_runner"></a> [enable\_user\_data\_debug\_logging\_runner](#input\_enable\_user\_data\_debug\_logging\_runner) | Option to enable debug logging for user-data, this logs all secrets as well. | `bool` | `false` | no |
 | <a name="input_enable_userdata"></a> [enable\_userdata](#input\_enable\_userdata) | Should the userdata script be enabled for the runner. Set this to false if you are using your own prebuilt AMI. | `bool` | `true` | no |
+| <a name="input_enterprise_pat"></a> [enterprise\_pat](#input\_enterprise\_pat) | GitHub enterprise PAT. Used only when enable\_enterprise\_runners is true. | `string` | `null` | no |
+| <a name="input_enterprise_slug"></a> [enterprise\_slug](#input\_enterprise\_slug) | Enterprise slug | `string` | `""` | no |
 | <a name="input_eventbridge"></a> [eventbridge](#input\_eventbridge) | Enable the use of EventBridge by the module. By enabling this feature events will be put on the EventBridge by the webhook instead of directly dispatching to queues for scaling.<br/><br/>    `enable`: Enable the EventBridge feature.<br/>    `accept_events`: List can be used to only allow specific events to be putted on the EventBridge. By default all events, empty list will be be interpreted as all events. | <pre>object({<br/>    enable        = optional(bool, true)<br/>    accept_events = optional(list(string), null)<br/>  })</pre> | `{}` | no |
 | <a name="input_ghes_ssl_verify"></a> [ghes\_ssl\_verify](#input\_ghes\_ssl\_verify) | GitHub Enterprise SSL verification. Set to 'false' when custom certificate (chains) is used for GitHub Enterprise Server (insecure). | `bool` | `true` | no |
 | <a name="input_ghes_url"></a> [ghes\_url](#input\_ghes\_url) | GitHub Enterprise Server URL. Example: https://github.internal.co - DO NOT SET IF USING PUBLIC GITHUB. However if you are using Github Enterprise Cloud with data-residency (ghe.com), set the endpoint here. Example - https://companyname.ghe.com | `string` | `null` | no |
-| <a name="input_github_app"></a> [github\_app](#input\_github\_app) | GitHub app parameters, see your github app.<br/>  You can optionally create the SSM parameters yourself and provide the ARN and name here, through the `*_ssm` attributes.<br/>  If you chose to provide the configuration values directly here,<br/>  please ensure the key is the base64-encoded `.pem` file (the output of `base64 app.private-key.pem`, not the content of `private-key.pem`).<br/>  Note: the provided SSM parameters arn and name have a precedence over the actual value (i.e `key_base64_ssm` has a precedence over `key_base64` etc). | <pre>object({<br/>    key_base64 = optional(string)<br/>    key_base64_ssm = optional(object({<br/>      arn  = string<br/>      name = string<br/>    }))<br/>    id = optional(string)<br/>    id_ssm = optional(object({<br/>      arn  = string<br/>      name = string<br/>    }))<br/>    webhook_secret = optional(string)<br/>    webhook_secret_ssm = optional(object({<br/>      arn  = string<br/>      name = string<br/>    }))<br/>  })</pre> | n/a | yes |
+| <a name="input_github_app"></a> [github\_app](#input\_github\_app) | GitHub app parameters. | <pre>object({<br/>    key_base64         = optional(string, null)<br/>    key_base64_ssm     = optional(object({ arn = string, name = string }))<br/>    id                 = optional(string, null)<br/>    id_ssm             = optional(object({ arn = string, name = string }))<br/>    webhook_secret     = optional(string)<br/>    webhook_secret_ssm = optional(object({ arn = string, name = string }))<br/>  })</pre> | n/a | yes |
 | <a name="input_idle_config"></a> [idle\_config](#input\_idle\_config) | List of time periods, defined as a cron expression, to keep a minimum amount of runners active instead of scaling down to 0. By defining this list you can ensure that in time periods that match the cron expression within 5 seconds a runner is kept idle. | <pre>list(object({<br/>    cron             = string<br/>    timeZone         = string<br/>    idleCount        = number<br/>    evictionStrategy = optional(string, "oldest_first")<br/>  }))</pre> | `[]` | no |
 | <a name="input_instance_allocation_strategy"></a> [instance\_allocation\_strategy](#input\_instance\_allocation\_strategy) | The allocation strategy for spot instances. AWS recommends using `price-capacity-optimized` however the AWS default is `lowest-price`. | `string` | `"lowest-price"` | no |
 | <a name="input_instance_max_spot_price"></a> [instance\_max\_spot\_price](#input\_instance\_max\_spot\_price) | Max price price for spot instances per hour. This variable will be passed to the create fleet as max spot price for the fleet. | `string` | `null` | no |

lambdas/functions/control-plane/src/aws/runners.d.ts

@@ -1,6 +1,6 @@
 import { DefaultTargetCapacityType, SpotAllocationStrategy } from '@aws-sdk/client-ec2';
 
-export type RunnerType = 'Org' | 'Repo';
+export type RunnerType = 'Enterprise' | 'Org' | 'Repo';
 
 export interface RunnerList {
   instanceId: string;

lambdas/functions/control-plane/src/aws/runners.test.ts

@@ -661,7 +661,7 @@ function createRunnerConfig(runnerConfig: RunnerConfig): RunnerInputParameters {
 }
 
 interface ExpectedFleetRequestValues {
-  type: 'Repo' | 'Org';
+  type: 'Enterprise' | 'Repo' | 'Org';
   capacityType: DefaultTargetCapacityType;
   allocationStrategy: SpotAllocationStrategy;
   maxSpotPrice?: string;

lambdas/functions/control-plane/src/github/octokit.test.ts

@@ -10,6 +10,10 @@ const mockOctokit = {
   },
 };
 
+function setDefaults() {
+  process.env.PARAMETER_ENTERPRISE_PAT_NAME = 'github-pat-id';
+}
+
 vi.mock('../github/auth', async () => ({
   createGithubInstallationAuth: vi.fn().mockImplementation(async (installationId) => {
     return { token: 'token', type: 'installation', installationId: installationId };
@@ -22,29 +26,46 @@ vi.mock('@octokit/rest', async () => ({
   Octokit: vi.fn().mockImplementation(() => mockOctokit),
 }));
 
+vi.mock('@aws-github-runner/aws-ssm-util', async () => {
+  const actual = (await vi.importActual(
+    '@aws-github-runner/aws-ssm-util',
+  )) as typeof import('@aws-github-runner/aws-ssm-util');
+
+  return {
+    ...actual,
+    getParameter: vi.fn(),
+  };
+});
+
 // We've already mocked '../github/auth' above
 
 describe('Test getOctokit', () => {
   const data = [
     {
       description: 'Should look-up org installation if installationId is 0.',
-      input: { orgLevelRunner: false, installationId: 0 },
-      output: { callReposInstallation: true, callOrgInstallation: false },
+      input: { enableEnterpriseLevel: false, orgLevelRunner: false, installationId: 0 },
+      output: { callEnterpriseToken: false, callReposInstallation: true, callOrgInstallation: false },
     },
     {
       description: 'Should look-up org installation if installationId is 0.',
-      input: { orgLevelRunner: true, installationId: 0 },
-      output: { callReposInstallation: false, callOrgInstallation: true },
+      input: { enableEnterpriseLevel: false, orgLevelRunner: true, installationId: 0 },
+      output: { callEnterpriseToken: false, callReposInstallation: false, callOrgInstallation: true },
     },
     {
       description: 'Should not look-up org installation if provided in payload.',
-      input: { orgLevelRunner: true, installationId: 1 },
-      output: { callReposInstallation: false, callOrgInstallation: false },
+      input: { enableEnterpriseLevel: false, orgLevelRunner: true, installationId: 1 },
+      output: { callEnterpriseToken: false, callReposInstallation: false, callOrgInstallation: false },
+    },
+    {
+      description: 'Should not look-up org installation if enterprise is enabled.',
+      input: { enableEnterpriseLevel: true, orgLevelRunner: false, installationId: 1 },
+      output: { callEnterpriseToken: true, callReposInstallation: false, callOrgInstallation: false },
     },
   ];
 
   beforeEach(() => {
     vi.clearAllMocks();
+    setDefaults();
   });
 
   it.each(data)(`$description`, async ({ input, output }) => {
@@ -64,9 +85,12 @@ describe('Test getOctokit', () => {
       mockOctokit.apps.getOrgInstallation.mockRejectedValue(new Error('Error'));
     }
 
-    await expect(getOctokit('', input.orgLevelRunner, payload)).resolves.toBeDefined();
+    await expect(getOctokit('', input.enableEnterpriseLevel, input.orgLevelRunner, payload)).resolves.toBeDefined();
 
-    if (output.callOrgInstallation) {
+    if (output.callEnterpriseToken) {
+      expect(mockOctokit.apps.getOrgInstallation).not.toHaveBeenCalled();
+      expect(mockOctokit.apps.getRepoInstallation).not.toHaveBeenCalled();
+    } else if (output.callOrgInstallation) {
       expect(mockOctokit.apps.getOrgInstallation).toHaveBeenCalled();
       expect(mockOctokit.apps.getRepoInstallation).not.toHaveBeenCalled();
     } else if (output.callReposInstallation) {

lambdas/functions/control-plane/src/github/octokit.ts

@@ -1,6 +1,7 @@
 import { Octokit } from '@octokit/rest';
 import { ActionRequestMessage } from '../scale-runners/scale-up';
 import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from './auth';
+import { getParameter } from '@aws-github-runner/aws-ssm-util';
 
 export async function getInstallationId(
   ghesApiUrl: string,
@@ -37,10 +38,17 @@ export async function getInstallationId(
  */
 export async function getOctokit(
   ghesApiUrl: string,
+  enableEnterpriseLevel: boolean,
   enableOrgLevel: boolean,
   payload: ActionRequestMessage,
 ): Promise<Octokit> {
-  const installationId = await getInstallationId(ghesApiUrl, enableOrgLevel, payload);
-  const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
-  return await createOctokitClient(ghAuth.token, ghesApiUrl);
+  let ghToken;
+  if (enableEnterpriseLevel) {
+    ghToken = await getParameter(process.env.PARAMETER_ENTERPRISE_PAT_NAME);
+  } else {
+    const installationId = await getInstallationId(ghesApiUrl, enableOrgLevel, payload);
+    ghToken = (await createGithubInstallationAuth(installationId, ghesApiUrl)).token;
+  }
+
+  return await createOctokitClient(ghToken, ghesApiUrl);
 }

lambdas/functions/control-plane/src/modules.d.ts

@@ -14,6 +14,7 @@ declare namespace NodeJS {
     PARAMETER_GITHUB_APP_CLIENT_SECRET_NAME: string;
     PARAMETER_GITHUB_APP_ID_NAME: string;
     PARAMETER_GITHUB_APP_KEY_BASE64_NAME: string;
+    PARAMETER_ENTERPRISE_PAT_NAME: string;
     RUNNER_OWNER: string;
     SCALE_DOWN_CONFIG: string;
     SSM_TOKEN_PATH: string;

lambdas/functions/control-plane/src/pool/pool.test.ts

@@ -1,5 +1,6 @@
 import { Octokit } from '@octokit/rest';
 import moment from 'moment-timezone';
+import * as aws_ssm from '@aws-github-runner/aws-ssm-util';
 import * as nock from 'nock';
 
 import { listEC2Runners } from '../aws/runners';
@@ -34,6 +35,10 @@ vi.mock('./../github/auth', async () => ({
   createOctokitClient: vi.fn(),
 }));
 
+vi.mock('@aws-github-runner/aws-ssm-util', () => ({
+  getParameter: vi.fn(),
+}));
+
 vi.mock('../scale-runners/scale-up', async () => ({
   scaleUp: vi.fn(),
   createRunners: vi.fn(),
@@ -47,6 +52,7 @@ vi.mock('../scale-runners/scale-up', async () => ({
 const mocktokit = Octokit as MockedClass<typeof Octokit>;
 const mockedAppAuth = vi.mocked(ghAuth.createGithubAppAuth);
 const mockedInstallationAuth = vi.mocked(ghAuth.createGithubInstallationAuth);
+const mockedGetParameter = vi.mocked(aws_ssm.getParameter);
 const mockCreateClient = vi.mocked(ghAuth.createOctokitClient);
 const mockListRunners = vi.mocked(listEC2Runners);
 
@@ -174,6 +180,7 @@ beforeEach(() => {
     installationId: 0,
   });
 
+  mockedGetParameter.mockResolvedValue('dummy-enterprise-pat');
   mockCreateClient.mockResolvedValue(new mocktokit());
 });
 
@@ -274,6 +281,24 @@ describe('Test simple pool.', () => {
         expect.anything(),
       );
     });
+
+    it('Uses Enterprise PAT from SSM when ENABLE_ENTERPRISE_RUNNERS=true', async () => {
+      process.env.ENABLE_ENTERPRISE_RUNNERS = 'true';
+      process.env.PARAMETER_ENTERPRISE_PAT_NAME = '/ssm/enterprise/pat';
+      mockedGetParameter.mockResolvedValue('enterprise-pat-123');
+
+      await adjust({ poolSize: 5 });
+
+      // createGithubInstallationAuth must NOT be called in enterprise mode
+      expect(ghAuth.createGithubInstallationAuth).not.toHaveBeenCalled();
+
+      // client is created with PAT
+      expect(ghAuth.createOctokitClient).toHaveBeenCalledWith(
+        'enterprise-pat-123',
+        'https://api.github.enterprise.something',
+      );
+      expect(createRunners).toHaveBeenCalled(); // still reaches scale-up if needed
+    });
   });
 
   describe('With Github Data Residency', () => {

lambdas/functions/control-plane/src/pool/pool.ts

@@ -1,11 +1,13 @@
 import { Octokit } from '@octokit/rest';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
+import { getParameter } from '@aws-github-runner/aws-ssm-util';
 import yn from 'yn';
 
 import { bootTimeExceeded, listEC2Runners } from '../aws/runners';
 import { RunnerList } from '../aws/runners.d';
 import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from '../github/auth';
 import { createRunners, getGitHubEnterpriseApiUrl } from '../scale-runners/scale-up';
+import { GhRunners } from '../scale-runners/cache';
 
 const logger = createChildLogger('pool');
 
@@ -20,6 +22,7 @@ interface RunnerStatus {
 
 export async function adjust(event: PoolEvent): Promise<void> {
   logger.info(`Checking current pool size against pool of size: ${event.poolSize}`);
+  const enableEnterpriseLevel = yn(process.env.ENABLE_ENTERPRISE_RUNNERS, { default: false });
   const runnerLabels = process.env.RUNNER_LABELS || '';
   const runnerGroup = process.env.RUNNER_GROUP_NAME || '';
   const runnerNamePrefix = process.env.RUNNER_NAME_PREFIX || '';
@@ -44,22 +47,27 @@ export async function adjust(event: PoolEvent): Promise<void> {
 
   const { ghesApiUrl, ghesBaseUrl } = getGitHubEnterpriseApiUrl();
 
-  const installationId = await getInstallationId(ghesApiUrl, runnerOwner);
-  const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl);
-  const githubInstallationClient = await createOctokitClient(ghAuth.token, ghesApiUrl);
+  const installationId = enableEnterpriseLevel ? undefined : await getInstallationId(ghesApiUrl, runnerOwner);
+  const ghToken = enableEnterpriseLevel
+    ? await getParameter(process.env.PARAMETER_ENTERPRISE_PAT_NAME)
+    : (await createGithubInstallationAuth(installationId, ghesApiUrl)).token;
+  const githubInstallationClient = await createOctokitClient(ghToken, ghesApiUrl);
 
   // Get statusses of runners registed in GitHub
   const runnerStatusses = await getGitHubRegisteredRunnnerStatusses(
+    enableEnterpriseLevel,
     githubInstallationClient,
     runnerOwner,
     runnerNamePrefix,
   );
 
+  const ec2RunnerType = enableEnterpriseLevel ? 'Enterprise' : 'Org';
+
   // Look up the managed ec2 runners in AWS, but running does not mean idle
   const ec2runners = await listEC2Runners({
     environment,
     runnerOwner,
-    runnerType: 'Org',
+    runnerType: ec2RunnerType,
     statuses: ['running'],
   });
 
@@ -77,7 +85,7 @@ export async function adjust(event: PoolEvent): Promise<void> {
         runnerGroup,
         runnerOwner,
         runnerNamePrefix,
-        runnerType: 'Org',
+        runnerType: ec2RunnerType,
         disableAutoUpdate: disableAutoUpdate,
         ssmTokenPath,
         ssmConfigPath,
@@ -140,14 +148,20 @@ function calculatePooSize(ec2runners: RunnerList[], runnerStatus: Map<string, Ru
 }
 
 async function getGitHubRegisteredRunnnerStatusses(
+  enableEnterpriseLevel: boolean,
   ghClient: Octokit,
   runnerOwner: string,
   runnerNamePrefix: string,
 ): Promise<Map<string, RunnerStatus>> {
-  const runners = await ghClient.paginate(ghClient.actions.listSelfHostedRunnersForOrg, {
-    org: runnerOwner,
-    per_page: 100,
-  });
+  const runners: GhRunners = enableEnterpriseLevel
+    ? await ghClient.paginate('GET /enterprises/{enterprise}/actions/runners', {
+        enterprise: runnerOwner,
+        per_page: 100,
+      })
+    : await ghClient.paginate(ghClient.actions.listSelfHostedRunnersForOrg, {
+        org: runnerOwner,
+        per_page: 100,
+      });
   const runnerStatus = new Map<string, RunnerStatus>();
   for (const runner of runners) {
     runner.name = runnerNamePrefix ? runner.name.replace(runnerNamePrefix, '') : runner.name;

lambdas/functions/control-plane/src/scale-runners/job-retry.test.ts

@@ -170,6 +170,42 @@ describe(`Test job retry check`, () => {
     expect(createSingleMetric).not.toHaveBeenCalled();
   });
 
+  it(`should publish a message for retry if retry is enabled and enterprise level is enabled.`, async () => {
+    // setup
+    mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({
+      data: {
+        status: 'queued',
+      },
+    }));
+
+    const message: ActionRequestMessageRetry = {
+      eventType: 'workflow_job',
+      id: 0,
+      installationId: 0,
+      repositoryName: 'test',
+      repositoryOwner: 'github-aws-runners',
+      repoOwnerType: 'Enterprise',
+      retryCounter: 0,
+    };
+    process.env.ENABLE_ENTERPRISE_RUNNERS = 'true';
+    process.env.RUNNER_NAME_PREFIX = 'test';
+    process.env.JOB_QUEUE_SCALE_UP_URL =
+      'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue';
+
+    // act
+    await checkAndRetryJob(message);
+
+    // assert
+    expect(publishMessage).toHaveBeenCalledWith(
+      JSON.stringify({
+        ...message,
+      }),
+      'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue',
+    );
+    expect(createSingleMetric).not.toHaveBeenCalled();
+  });
+
+
   it(`should publish a message for retry if retry is enabled and counter is below max attempts.`, async () => {
     // setup
     mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({