process subjects in batches (#67)

Signed-off-by: Brian DeHamer <bdehamer@github.com>

__tests__/main.test.ts

@@ -44,6 +44,9 @@ describe('action', () => {
   const originalEnv = process.env
   const originalContext = { ...github.context }
 
+  // Mock OIDC token endpoint
+  const tokenURL = 'https://token.url'
+
   // Fake an OIDC token
   const oidcSubject = 'foo@bar.com'
   const oidcPayload = { sub: oidcSubject, iss: '' }
@@ -62,9 +65,6 @@ describe('action', () => {
   beforeEach(() => {
     jest.clearAllMocks()
 
-    // Mock OIDC token endpoint
-    const tokenURL = 'https://token.url'
-
     nock(tokenURL)
       .get('/')
       .query({ audience: 'sigstore' })
@@ -289,10 +289,15 @@ describe('action', () => {
     })
   })
 
-  describe('when too many subjects are specified', () => {
+  describe('when the subject count exceeds the batch size', () => {
     let dir = ''
+    let scope: nock.Scope
 
     beforeEach(async () => {
+      // Start from scratch
+      nock.cleanAll()
+
+      const subjectCount = 5
       const filename = 'subject'
       const content = 'file content'
 
@@ -301,23 +306,51 @@ describe('action', () => {
       dir = await fs.mkdtemp(tmpDir + path.sep)
 
       // Add files for glob testing
-      for (let i = 0; i < 65; i++) {
+      for (let i = 0; i < subjectCount; i++) {
         await fs.writeFile(path.join(dir, `${filename}-${i}`), content)
+
+        // Set-up a Fulcio mock for each subject
+        await mockFulcio({
+          baseURL: 'https://fulcio.githubapp.com',
+          strict: false
+        })
+
+        // Set-up a TSA mock for each subject
+        await mockTSA({ baseURL: 'https://timestamp.githubapp.com' })
+
+        // Set-up a GH API mock for each subject
+        mockAgent
+          .get('https://api.github.com')
+          .intercept({
+            path: /^\/repos\/.*\/.*\/attestations$/,
+            method: 'post'
+          })
+          .reply(201, { id: attestationID })
       }
 
+      // Set-up a OIDC token mock for each subject
+      scope = nock(tokenURL)
+        .get('/')
+        .query({ audience: 'sigstore' })
+        .times(subjectCount)
+        .reply(200, { value: oidcToken })
+
       // Set the GH context with private repository visibility and a repo owner.
       setGHContext({
         payload: { repository: { visibility: 'private' } },
         repo: { owner: 'foo', repo: 'bar' }
       })
 
-      // Mock the action's inputs
+      const inputs = {
-      getInputMock.mockImplementation(
+        'subject-path': path.join(dir, `${filename}-*`),
-        mockInput({
+        'predicate-type': predicateType,
-          predicate: '{}',
+        predicate,
-          'subject-path': path.join(dir, `${filename}-*`)
+        'github-token': 'gh-token',
-        })
+        'batch-size': '2',
-      )
+        'batch-delay': '500'
+      }
+      getInputMock.mockImplementation(mockInput(inputs))
+      getBooleanInputMock.mockImplementation(() => false)
     })
 
     afterEach(async () => {
@@ -325,15 +358,24 @@ describe('action', () => {
       await fs.rm(dir, { recursive: true })
     })
 
-    it('sets a failed status', async () => {
+    it('invokes the action w/o error', async () => {
       await main.run()
 
       expect(runMock).toHaveReturned()
-      expect(setFailedMock).toHaveBeenCalledWith(
+      expect(setFailedMock).not.toHaveBeenCalled()
-        new Error(
+      expect(infoMock).toHaveBeenNthCalledWith(
-          'Too many subjects specified. The maximum number of subjects is 64.'
+        1,
-        )
+        expect.stringMatching('Processing subject batch 1/3')
       )
+      expect(infoMock).toHaveBeenNthCalledWith(
+        10,
+        expect.stringMatching('Processing subject batch 2/3')
+      )
+      expect(infoMock).toHaveBeenNthCalledWith(
+        19,
+        expect.stringMatching('Processing subject batch 3/3')
+      )
+      expect(scope.isDone()).toBe(true)
     })
   })
 })

dist/index.js

@@ -79954,7 +79954,8 @@ const COLOR_CYAN = '\x1B[36m';
 const COLOR_GRAY = '\x1B[38;5;244m';
 const COLOR_DEFAULT = '\x1B[39m';
 const ATTESTATION_FILE_NAME = 'attestation.jsonl';
-const MAX_SUBJECT_COUNT = 64;
+const DEFAULT_BATCH_SIZE = 50;
+const DEFAULT_BATCH_DELAY = 5000;
 const OCI_TIMEOUT = 2000;
 const OCI_RETRY = 3;
 /* istanbul ignore next */
@@ -79982,23 +79983,33 @@ async function run() {
         if (!process.env.ACTIONS_ID_TOKEN_REQUEST_URL) {
             throw new Error('missing "id-token" permission. Please add "permissions: id-token: write" to your workflow.');
         }
-        // Gather list of subjets
         const subjects = await (0, subject_1.subjectFromInputs)();
-        if (subjects.length > MAX_SUBJECT_COUNT) {
-            throw new Error(`Too many subjects specified. The maximum number of subjects is ${MAX_SUBJECT_COUNT}.`);
-        }
         const predicate = (0, predicate_1.predicateFromInputs)();
         const outputPath = path_1.default.join(tempDir(), ATTESTATION_FILE_NAME);
-        // Generate attestations for each subject serially
+        // Batch size and delay for rate limiting
-        for (const subject of subjects) {
+        const batchSize = parseInt(core.getInput('batch-size')) || DEFAULT_BATCH_SIZE;
-            const att = await createAttestation(subject, predicate, sigstoreInstance);
+        const batchDelay = parseInt(core.getInput('batch-delay')) || DEFAULT_BATCH_DELAY;
-            // Write attestation bundle to output file
+        const subjectChunks = chunkArray(subjects, batchSize);
-            fs_1.default.writeFileSync(outputPath, JSON.stringify(att.bundle) + os_1.default.EOL, {
+        let chunkCount = 0;
-                encoding: 'utf-8',
+        // Generate attestations for each subject serially, working in batches
-                flag: 'a'
+        for (const subjectChunk of subjectChunks) {
-            });
+            // Delay between batches (only when chunkCount > 0)
-            if (att.attestationID) {
+            if (chunkCount++) {
-                atts.push({ subject, attestationID: att.attestationID });
+                await new Promise(resolve => setTimeout(resolve, batchDelay));
+            }
+            if (subjectChunks.length > 1) {
+                core.info(`Processing subject batch ${chunkCount}/${subjectChunks.length}`);
+            }
+            for (const subject of subjectChunk) {
+                const att = await createAttestation(subject, predicate, sigstoreInstance);
+                // Write attestation bundle to output file
+                fs_1.default.writeFileSync(outputPath, JSON.stringify(att.bundle) + os_1.default.EOL, {
+                    encoding: 'utf-8',
+                    flag: 'a'
+                });
+                if (att.attestationID) {
+                    atts.push({ subject, attestationID: att.attestationID });
+                }
             }
         }
         if (atts.length > 0) {
@@ -80082,6 +80093,11 @@ const tempDir = () => {
     }
     return fs_1.default.mkdtempSync(path_1.default.join(basePath, path_1.default.sep));
 };
+// Transforms an array into an array of arrays, each containing at most
+// `chunkSize` elements.
+const chunkArray = (array, chunkSize) => {
+    return Array.from({ length: Math.ceil(array.length / chunkSize) }, (_, index) => array.slice(index * chunkSize, (index + 1) * chunkSize));
+};
 // Returns the subject's digest as a formatted string of the form
 // "<algorithm>:<digest>".
 const subjectDigest = (subject) => {

src/main.ts

@@ -17,7 +17,8 @@ const COLOR_GRAY = '\x1B[38;5;244m'
 const COLOR_DEFAULT = '\x1B[39m'
 const ATTESTATION_FILE_NAME = 'attestation.jsonl'
 
-const MAX_SUBJECT_COUNT = 64
+const DEFAULT_BATCH_SIZE = 50
+const DEFAULT_BATCH_DELAY = 5000
 
 const OCI_TIMEOUT = 2000
 const OCI_RETRY = 3
@@ -54,29 +55,48 @@ export async function run(): Promise<void> {
       )
     }
 
-    // Gather list of subjets
     const subjects = await subjectFromInputs()
-    if (subjects.length > MAX_SUBJECT_COUNT) {
-      throw new Error(
-        `Too many subjects specified. The maximum number of subjects is ${MAX_SUBJECT_COUNT}.`
-      )
-    }
-
     const predicate = predicateFromInputs()
     const outputPath = path.join(tempDir(), ATTESTATION_FILE_NAME)
 
-    // Generate attestations for each subject serially
+    // Batch size and delay for rate limiting
-    for (const subject of subjects) {
+    const batchSize =
-      const att = await createAttestation(subject, predicate, sigstoreInstance)
+      parseInt(core.getInput('batch-size')) || DEFAULT_BATCH_SIZE
+    const batchDelay =
+      parseInt(core.getInput('batch-delay')) || DEFAULT_BATCH_DELAY
 
-      // Write attestation bundle to output file
+    const subjectChunks = chunkArray(subjects, batchSize)
-      fs.writeFileSync(outputPath, JSON.stringify(att.bundle) + os.EOL, {
+    let chunkCount = 0
-        encoding: 'utf-8',
-        flag: 'a'
-      })
 
-      if (att.attestationID) {
+    // Generate attestations for each subject serially, working in batches
-        atts.push({ subject, attestationID: att.attestationID })
+    for (const subjectChunk of subjectChunks) {
+      // Delay between batches (only when chunkCount > 0)
+      if (chunkCount++) {
+        await new Promise(resolve => setTimeout(resolve, batchDelay))
+      }
+
+      if (subjectChunks.length > 1) {
+        core.info(
+          `Processing subject batch ${chunkCount}/${subjectChunks.length}`
+        )
+      }
+
+      for (const subject of subjectChunk) {
+        const att = await createAttestation(
+          subject,
+          predicate,
+          sigstoreInstance
+        )
+
+        // Write attestation bundle to output file
+        fs.writeFileSync(outputPath, JSON.stringify(att.bundle) + os.EOL, {
+          encoding: 'utf-8',
+          flag: 'a'
+        })
+
+        if (att.attestationID) {
+          atts.push({ subject, attestationID: att.attestationID })
+        }
       }
     }
 
@@ -194,6 +214,15 @@ const tempDir = (): string => {
   return fs.mkdtempSync(path.join(basePath, path.sep))
 }
 
+// Transforms an array into an array of arrays, each containing at most
+// `chunkSize` elements.
+const chunkArray = <T>(array: T[], chunkSize: number): T[][] => {
+  return Array.from(
+    { length: Math.ceil(array.length / chunkSize) },
+    (_, index) => array.slice(index * chunkSize, (index + 1) * chunkSize)
+  )
+}
+
 // Returns the subject's digest as a formatted string of the form
 // "<algorithm>:<digest>".
 const subjectDigest = (subject: Subject): string => {