runners; Add accountId to findAmiID

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.test.ts

@@ -49,9 +49,19 @@ const mockSSM = {
   describeParameters: jest.fn().mockReturnValue({ promise: mockSSMdescribeParametersRet }),
   putParameter: jest.fn().mockReturnValue({ promise: jest.fn() }),
 };
+
+const mockSTS = {
+  getCallerIdentity: jest.fn().mockReturnValue({
+    promise: jest.fn().mockImplementation(async () => {
+      return { Account: '123456789012' };
+    }),
+  }),
+};
+
 jest.mock('aws-sdk', () => ({
   EC2: jest.fn().mockImplementation(() => mockEC2),
   SSM: jest.fn().mockImplementation(() => mockSSM),
+  STS: jest.fn().mockImplementation(() => mockSTS),
   CloudWatch: jest.requireActual('aws-sdk').CloudWatch,
 }));
 jest.mock('./utils', () => ({
@@ -415,7 +425,7 @@ describe('findAmiID', () => {
     const result1 = await findAmiID(metrics, 'REGION', 'FILTER');
     expect(mockEC2.describeImages).toBeCalledTimes(1);
     expect(mockEC2.describeImages).toBeCalledWith({
-      Owners: ['amazon'],
+      Owners: ['123456789012', '391835788720', 'amazon'],
       Filters: [
         {
           Name: 'name',
@@ -439,6 +449,51 @@ describe('findAmiID', () => {
     await findAmiID(metrics, 'REGION', 'FILTER2');
     expect(mockEC2.describeImages).toBeCalledTimes(3);
   });
+
+  it('finds custom AMIs using account ID as owner', async () => {
+    const customAmiMock = jest.fn().mockImplementation(async () => {
+      return {
+        Images: [{ CreationDate: '2024-07-09T12:32:23+0000', ImageId: 'ami-custom123' }],
+      };
+    });
+
+    mockEC2.describeImages.mockClear();
+    mockEC2.describeImages.mockReturnValue({
+      promise: customAmiMock,
+    });
+
+    const result = await findAmiID(metrics, 'us-east-1', 'custom-runner-ami-*');
+
+    expect(mockEC2.describeImages).toBeCalledTimes(1);
+    expect(mockEC2.describeImages).toBeCalledWith({
+      Owners: ['123456789012', '391835788720', 'amazon'],
+      Filters: [
+        {
+          Name: 'name',
+          Values: ['custom-runner-ami-*'],
+        },
+        {
+          Name: 'state',
+          Values: ['available'],
+        },
+      ],
+    });
+    expect(result).toBe('ami-custom123');
+
+    // Reset mock to default behavior for other tests
+    mockEC2.describeImages.mockReturnValue({
+      promise: jest.fn().mockImplementation(async () => {
+        return {
+          Images: [
+            { CreationDate: '2024-07-09T12:32:23+0000', ImageId: 'ami-234' },
+            { CreationDate: '2024-07-09T13:55:23+0000', ImageId: 'ami-AGDGADU113' },
+            { CreationDate: '2024-07-09T13:32:23+0000', ImageId: 'ami-113' },
+            { CreationDate: '2024-07-09T13:32:23+0000', ImageId: 'ami-1113' },
+          ],
+        };
+      }),
+    });
+  });
 });
 
 describe('tryReuseRunner', () => {
@@ -1290,7 +1345,7 @@ describe('createRunner', () => {
       await createRunner(runnerParameters, metrics);
       expect(mockEC2.describeImages).toBeCalledTimes(1);
       expect(mockEC2.describeImages).toBeCalledWith({
-        Owners: ['amazon'],
+        Owners: ['123456789012', '391835788720', 'amazon'],
         Filters: [
           {
             Name: 'name',

terraform-aws-github-runner/modules/runners/lambdas/runners/src/scale-runners/runners.ts

@@ -1,4 +1,4 @@
-import { EC2, SSM } from 'aws-sdk';
+import { EC2, SSM, STS } from 'aws-sdk';
 import { PromiseResult } from 'aws-sdk/lib/request';
 import { RunnerInfo, expBackOff, getRepo, shuffleArrayInPlace } from './utils';
 
@@ -70,27 +70,56 @@
   ssmParametersCache.reset();
 }
 
-export async function findAmiID(metrics: Metrics, region: string, filter: string, owners = 'amazon'): Promise<string> {
+// Helper function to get the current AWS account ID
+async function getCurrentAccountId(): Promise<string | null> {
+  try {
+    const sts = new STS();
+    const identity = await sts.getCallerIdentity().promise();
+    if (!identity.Account) {
+      console.warn('Unable to retrieve AWS account ID, falling back to Amazon-only images');
+      return null;
+    }
+    return identity.Account;
+  } catch (error) {
+    console.warn('Error retrieving AWS account ID:', error);
+    return null;
+  }
+}
+
+export async function findAmiID(
+  metrics: Metrics, 
+  region: string, 
+  filter: string, 
+  owners: string[] = ['amazon']
+): Promise<string> {
   const ec2 = new EC2({ region: region });
+  const accountId = await getCurrentAccountId();
+  const lfAccountId = '391835788720';
+
+  const allOwners = accountId 
+    ? [accountId, lfAccountId, ...owners]
+    : [lfAccountId, ...owners];
+
   const filters = [
     { Name: 'name', Values: [filter] },
     { Name: 'state', Values: ['available'] },
   ];
-  return redisCached('awsEC2', `findAmiID-${region}-${filter}-${owners}`, 10 * 60, 0.5, () => {
+
+  return redisCached('awsEC2', `findAmiID-${region}-${filter}-${allOwners.join(',')}`, 10 * 60, 0.5, () => {
     return expBackOff(() => {
       return metrics.trackRequestRegion(
         region,
         metrics.ec2DescribeImagesSuccess,
         metrics.ec2DescribeImagesFailure,
         () => {
           return ec2
-            .describeImages({ Owners: [owners], Filters: filters })
+            .describeImages({ Owners: allOwners, Filters: filters })
             .promise()
             .then((data: EC2.DescribeImagesResult) => {
               /* istanbul ignore next */
               if (data.Images?.length === 0) {
-                console.error(`No availabe images found for filter '${filter}'`);
-                throw new Error(`No availabe images found for filter '${filter}'`);
+                console.error(`No available images found for filter '${filter}' in owners ${allOwners.join(', ')}`);
+                throw new Error(`No available images found for filter '${filter}'`);
               }
               sortByCreationDate(data);
               return data.Images?.shift()?.ImageId as string;