fix: Add maxDownloadSize for artifact directory download (#3232)

screwdriver-cd · Nov 5, 2024 · 03303f6 · 03303f6
1 parent ffdfc1d
commit 03303f6
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 2 deletions.
diff --git a/config/custom-environment-variables.yaml b/config/custom-environment-variables.yaml
@@ -489,6 +489,9 @@ build:
   environment:
     __name: CLUSTER_ENVIRONMENT_VARIABLES
     __format: json
+  artifacts:
+    # max artifact download size (in GB)
+    maxDownloadSize: MAX_DOWNLOAD_SIZE
 
 rateLimit:
   __name: RATE_LIMIT_VARIABLES

diff --git a/config/default.yaml b/config/default.yaml
@@ -390,6 +390,9 @@ log:
 build:
   environment:
     SD_VERSION: 4
+  artifacts:
+    # max artifact download size (in GB)
+    maxDownloadSize: 2
 
 rateLimit:
   # set true to enable rate limiting on auth token

diff --git a/lib/server.js b/lib/server.js
@@ -211,6 +211,8 @@ module.exports = async config => {
                 expiresIn
             );
         server.app.buildFactory.executor.tokenGen = server.app.buildFactory.tokenGen;
+        server.app.buildFactory.maxDownloadSize =
+            parseInt(config.build.artifacts.maxDownloadSize, 10) * 1024 * 1024 * 1024;
 
         server.app.jobFactory.apiUri = server.info.uri;
         server.app.jobFactory.tokenGen = (username, metadata, scmContext, scope = ['user']) =>

diff --git a/plugins/builds/artifacts/get.js b/plugins/builds/artifacts/get.js
@@ -31,6 +31,7 @@ module.exports = config => ({
             const { credentials } = req.auth;
             const { canAccessPipeline } = req.server.plugins.pipelines;
             const { buildFactory, eventFactory } = req.server.app;
+            const { maxDownloadSize } = buildFactory;
 
             return buildFactory.get(buildId)
                 .then(build => {
@@ -70,6 +71,23 @@ module.exports = config => ({
                             }).text();
                             const manifestArray = manifest.trim().split('\n');
                             const directoryArray = manifestArray.filter(f => f.startsWith(`./${artifact}/`));
+                            let totalSize = 0;
+
+                            // Check file sizes by fetching metadata
+                            for (const file of directoryArray) {
+                                if (file) {
+                                    const fileMetaResponse = await request.head(`${baseUrl}/${file}?token=${token}&type=download`);
+                                    const fileSize = parseInt(fileMetaResponse.headers['content-length'], 10);
+
+                                    // Accumulate total size
+                                    totalSize += fileSize;
+
+                                    // If total size exceeds allowed limit, stop further processing
+                                    if (totalSize > maxDownloadSize) {
+                                        throw new Error(`Total size of files exceeds the allowed limit of ${maxDownloadSize/1024/1024/1024}GB.`);
+                                    }
+                                }
+                            }
 
                             // Create a stream and set up archiver
                             const archive = archiver('zip', { zlib: { level: 9 } });

diff --git a/test/lib/server.test.js b/test/lib/server.test.js
@@ -35,6 +35,11 @@ describe('server case', () => {
             cookieName: 'test_cookie',
             cookieValue: 'test_value',
             cookieTimeout: 2
+        },
+        build: {
+            artifacts: {
+                maxDownloadSize: 2
+            }
         }
     };
 

diff --git a/test/plugins/builds.test.js b/test/plugins/builds.test.js
@@ -147,7 +147,8 @@ describe('build plugin test', () => {
                 getPrInfo: sinon.stub()
             },
             getLatestBuilds: sinon.stub(),
-            getBuildStatuses: sinon.stub()
+            getBuildStatuses: sinon.stub(),
+            maxDownloadSize: 3000
         };
         stepFactoryMock = {
             get: sinon.stub(),
@@ -7247,6 +7248,13 @@ describe('build plugin test', () => {
             nock.disableNetConnect();
             nock.cleanAll();
             nock.enableNetConnect();
+            nock(logBaseUrl)
+                .persist()
+                .defaultReplyHeaders({
+                    'content-length': 10
+                })
+                .head(/\/v1\/builds\/12345\/ARTIFACTS\/.+?token=sign&type=download/)
+                .reply(200);
             nock(logBaseUrl)
                 .defaultReplyHeaders(expectedHeaders)
                 .get('/v1/builds/12345/ARTIFACTS/manifest.txt?token=sign')
@@ -7257,7 +7265,7 @@ describe('build plugin test', () => {
                 .get(/\/v1\/builds\/12345\/ARTIFACTS\/.+?token=sign&type=download/)
                 .reply(200, testManifest);
 
-            options.url = `/builds/${id}/artifacts/./artifacts/?type=download`;
+            options.url = `/builds/${id}/artifacts/./artifacts?type=download&dir=true`;
 
             return server.inject(options).then(reply => {
                 assert.equal(reply.statusCode, 200);
@@ -7276,6 +7284,13 @@ describe('build plugin test', () => {
             nock.disableNetConnect();
             nock.cleanAll();
             nock.enableNetConnect();
+            nock(logBaseUrl)
+                .persist()
+                .defaultReplyHeaders({
+                    'content-length': 2000
+                })
+                .head(/\/v1\/builds\/12345\/ARTIFACTS\/.+?token=sign&type=download/)
+                .reply(200);
             nock(logBaseUrl)
                 .defaultReplyHeaders(expectedHeaders)
                 .get('/v1/builds/12345/ARTIFACTS/manifest.txt?token=sign')
@@ -7297,13 +7312,58 @@ describe('build plugin test', () => {
             });
         });
 
+        it('throws error for a too large artifact download request for directory', async () => {
+            const expectedHeaders = {
+                'content-type': 'application/zip',
+                'content-disposition': 'attachment; filename="artifacts_dir.zip"',
+                'cache-control': 'no-cache'
+            };
+
+            nock.disableNetConnect();
+            nock.cleanAll();
+            nock.enableNetConnect();
+            nock(logBaseUrl)
+                .persist()
+                .defaultReplyHeaders({
+                    'content-length': 2000
+                })
+                .head(/\/v1\/builds\/12345\/ARTIFACTS\/.+?token=sign&type=download/)
+                .reply(200);
+            nock(logBaseUrl)
+                .defaultReplyHeaders(expectedHeaders)
+                .get('/v1/builds/12345/ARTIFACTS/manifest.txt?token=sign')
+                .reply(200, testManifest);
+
+            options.url = `/builds/${id}/artifacts/./artifacts?type=download&dir=true`;
+
+            try {
+                await server.inject(options);
+            } catch (err) {
+                assert.isOk(err);
+                assert.calledThrice(loggerMock.error);
+                assert.calledWithMatch(
+                    loggerMock.error.getCall(0),
+                    'Error downloading file: ./artifacts/sample-mp4-file.mp4'
+                );
+                assert.calledWithMatch(loggerMock.error.getCall(1), 'Archiver error:');
+                assert.calledWithMatch(loggerMock.error.getCall(2), 'PassThrough stream error:');
+            }
+        });
+
         it('emits error for an error downloading artifact directory', async () => {
             const expectedHeaders = {
                 'content-type': 'application/zip',
                 'content-disposition': 'attachment; filename="artifacts_dir.zip"',
                 'cache-control': 'no-cache'
             };
 
+            nock(logBaseUrl)
+                .persist()
+                .defaultReplyHeaders({
+                    'content-length': 2000
+                })
+                .head(/\/v1\/builds\/12345\/ARTIFACTS\/.+?token=sign&type=download/)
+                .reply(200);
             nock(logBaseUrl)
                 .defaultReplyHeaders(expectedHeaders)
                 .get('/v1/builds/12345/ARTIFACTS/manifest.txt?token=sign')