From 3dea7630a584051b7e2eb152f71f8423e55fc827 Mon Sep 17 00:00:00 2001 From: Tim Nguyen Date: Wed, 9 Mar 2022 17:04:58 -0500 Subject: [PATCH] fix: Throw HTTP Status 429 error when there are too many get Sagemaker Presigned URL requests (#942) --- .../environment-sc-connection-service.test.js | 12 +++ .../environment-sc-connection-service.js | 80 ++++++++++--------- .../packages/services-container/lib/boom.js | 1 + 3 files changed, 57 insertions(+), 36 deletions(-) diff --git a/addons/addon-base-raas/packages/base-raas-services/lib/environment/service-catalog/__tests__/environment-sc-connection-service.test.js b/addons/addon-base-raas/packages/base-raas-services/lib/environment/service-catalog/__tests__/environment-sc-connection-service.test.js index 28be8569bf..903615c252 100644 --- a/addons/addon-base-raas/packages/base-raas-services/lib/environment/service-catalog/__tests__/environment-sc-connection-service.test.js +++ b/addons/addon-base-raas/packages/base-raas-services/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/addons/addon-base-raas/packages/base-raas-services/lib/environment/service-catalog/environment-sc-connection-service.js b/addons/addon-base-raas/packages/base-raas-services/lib/environment/service-catalog/environment-sc-connection-service.js index 0b32c8cb44..0f7f650773 100644 --- a/addons/addon-base-raas/packages/base-raas-services/lib/environment/service-catalog/environment-sc-connection-service.js +++ b/addons/addon-base-raas/packages/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 index f66187e67e..e0e33a3a00 100644 --- 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], );