fix: Throw HTTP Status 429 error when there are too many get Sagemake…

…r Presigned URL requests (#942)
awslabs · Mar 9, 2022 · 3dea763 · 3dea763
1 parent 15eb4d3
commit 3dea763
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 36 deletions.
diff --git a/...vices/lib/environment/service-catalog/__tests__/environment-sc-connection-service.test.js b/...vices/lib/environment/service-catalog/__tests__/environment-sc-connection-service.test.js
@@ -17,6 +17,7 @@ const ServicesContainer = require('@aws-ee/base-services-container/lib/services-
 const JsonSchemaValidationService = require('@aws-ee/base-services/lib/json-schema-validation-service');
 const Logger = require('@aws-ee/base-services/lib/logger/logger-service');
 const crypto = require('crypto');
+const Boom = require('@aws-ee/base-services-container/lib/boom');
 
 // Mocked dependencies
 const AwsService = require('@aws-ee/base-services/lib/aws/aws-service');
@@ -337,6 +338,17 @@ describe('EnvironmentScConnectionService', () => {
       expect(envScService.getClientSdkWithEnvMgmtRole).toHaveBeenCalledTimes(2);
     });
 
+    it('should return too many requests error if request cannot obtain a lock', async () => {
+      // BUILD
+      lockService.tryWriteLockAndRun = jest.fn(() => {
+        throw service.boom.internalError('Could not obtain a lock', true);
+      });
+
+      // OPERATE, CHECK
+      await expect(service.createPrivateSageMakerUrl({}, 'envId1', {})).rejects.toEqual(
+        new Boom().tooManyRequests('Please wait 30 seconds before requesting Sagemaker URL', true),
+      );
+    });
     it('should return private SageMaker URL', async () => {
       // BUILD
       lockService.tryWriteLockAndRun = jest.fn((id, func) => {

diff --git a/...s/base-raas-services/lib/environment/service-catalog/environment-sc-connection-service.js b/...s/base-raas-services/lib/environment/service-catalog/environment-sc-connection-service.js
@@ -383,46 +383,54 @@ class EnvironmentScConnectionService extends Service {
 
   async createPrivateSageMakerUrl(requestContext, envId, connection, presign_retries = 10) {
     const lockService = await this.service('lockService');
-    const signedURL = await lockService.tryWriteLockAndRun({ id: `${envId}presign` }, async () => {
-      if (!(_.toLower(_.get(connection, 'type', '')) === 'sagemaker')) {
-        throw this.boom.badRequest(
-          `Cannot generate presigned URL for non-sagemaker connection ${connection.type}`,
-          true,
-        );
-      }
-      const environmentScService = await this.service('environmentScService');
-      const iam = await environmentScService.getClientSdkWithEnvMgmtRole(
-        requestContext,
-        { id: envId },
-        { clientName: 'IAM', options: { apiVersion: '2017-07-24' } },
-      );
-      const currentPolicyResponse = await this.getCurrentRolePolicy(iam, connection);
-      await this.updateRoleToIncludeCurrentIP(iam, connection, currentPolicyResponse);
-      const createPresignedURLFn = async () => {
-        const stsEnvMgmt = await environmentScService.getClientSdkWithEnvMgmtRole(
+    let signedURL = '';
+    try {
+      signedURL = await lockService.tryWriteLockAndRun({ id: `${envId}presign` }, async () => {
+        if (!(_.toLower(_.get(connection, 'type', '')) === 'sagemaker')) {
+          throw this.boom.badRequest(
+            `Cannot generate presigned URL for non-sagemaker connection ${connection.type}`,
+            true,
+          );
+        }
+        const environmentScService = await this.service('environmentScService');
+        const iam = await environmentScService.getClientSdkWithEnvMgmtRole(
           requestContext,
           { id: envId },
-          { clientName: 'STS', options: { apiVersion: '2017-07-24' } },
+          { clientName: 'IAM', options: { apiVersion: '2017-07-24' } },
         );
-        const sageMakerResponse = await this.createPresignedURL(stsEnvMgmt, connection);
-        return sageMakerResponse;
-      };
-      try {
-        // Give sufficient number of retries to create presigned URL.
-        // This is needed because IAM role takes a while to propagate
-        // call with a linear strategy where we wait 2 seconds between retries
-        // This makes it 20 seconds with default of 10 retries
-        const sageMakerResponse = await retry(createPresignedURLFn, presign_retries, () => linearInterval(1, 2000));
-        return _.get(sageMakerResponse, 'AuthorizedUrl');
-      } catch (error) {
-        throw this.boom.internalError(`Could not generate presigned URL`, true).cause(error);
-      } finally {
-        // restore the original policy document. This ensures that caller IP address which was responsible for
-        // creating the presigned URL doesn't have access
-        const oldPolicyDocument = decodeURIComponent(currentPolicyResponse.PolicyDocument);
-        await this.putRolePolicy(iam, connection, oldPolicyDocument);
+        const currentPolicyResponse = await this.getCurrentRolePolicy(iam, connection);
+        await this.updateRoleToIncludeCurrentIP(iam, connection, currentPolicyResponse);
+        const createPresignedURLFn = async () => {
+          const stsEnvMgmt = await environmentScService.getClientSdkWithEnvMgmtRole(
+            requestContext,
+            { id: envId },
+            { clientName: 'STS', options: { apiVersion: '2017-07-24' } },
+          );
+          const sageMakerResponse = await this.createPresignedURL(stsEnvMgmt, connection);
+          return sageMakerResponse;
+        };
+        try {
+          // Give sufficient number of retries to create presigned URL.
+          // This is needed because IAM role takes a while to propagate
+          // call with a linear strategy where we wait 2 seconds between retries
+          // This makes it 20 seconds with default of 10 retries
+          const sageMakerResponse = await retry(createPresignedURLFn, presign_retries, () => linearInterval(1, 2000));
+          return _.get(sageMakerResponse, 'AuthorizedUrl');
+        } catch (error) {
+          throw this.boom.internalError(`Could not generate presigned URL`, true).cause(error);
+        } finally {
+          // restore the original policy document. This ensures that caller IP address which was responsible for
+          // creating the presigned URL doesn't have access
+          const oldPolicyDocument = decodeURIComponent(currentPolicyResponse.PolicyDocument);
+          await this.putRolePolicy(iam, connection, oldPolicyDocument);
+        }
+      });
+    } catch (e) {
+      if (this.boom.is(e, 'internalError') && e.message === 'Could not obtain a lock') {
+        throw this.boom.tooManyRequests('Please wait 30 seconds before requesting Sagemaker URL', true);
       }
-    });
+      throw e;
+    }
     return signedURL;
   }
 

diff --git a/addons/addon-base/packages/services-container/lib/boom.js b/addons/addon-base/packages/services-container/lib/boom.js
@@ -33,6 +33,7 @@ class Boom {
       // revision of the same is updated by someone else before that)
       ['outdatedUpdateAttempt', 409],
       ['timeout', 408],
+      ['tooManyRequests', 429],
       ['badImplementation', 500],
       ['internalError', 500],
     );