From 0051835458ebe085b7a9edde4d4a54c51c0fdd89 Mon Sep 17 00:00:00 2001 From: Alireza Sheikholmolouki Date: Fri, 20 Sep 2024 17:18:41 +0200 Subject: [PATCH] refactor: switched to a more advanced rate-limiter --- src/aws-app/utils/createRateLimiterFactory.ts | 3 +- .../common/AWSRateLimitExhaustionStrategy.ts | 28 +++ src/scanners/common/RateLimiter.ts | 171 ++++++++---------- src/scanners/common/RetryStrategy.ts | 28 +++ src/scanners/common/createGlobalScanner.ts | 10 +- src/scanners/common/createRegionalScanner.ts | 2 +- src/scanners/index.ts | 2 +- .../scan-functions/aws/_boilerplate.ts | 3 +- .../aws/athena-named-queries.ts | 3 +- .../scan-functions/aws/autoscaling-groups.ts | 3 +- .../aws/cloudfront-distributions.ts | 3 +- .../scan-functions/aws/cloudtrail-trails.ts | 3 +- .../aws/cloudwatch-metric-alarms.ts | 3 +- .../aws/cloudwatch-metric-streams.ts | 3 +- .../scan-functions/aws/dynamodb-tables.ts | 3 +- .../scan-functions/aws/ec2-instances.ts | 3 +- .../aws/ec2-internet-gateways.ts | 3 +- .../scan-functions/aws/ec2-nat-gateways.ts | 3 +- .../scan-functions/aws/ec2-network-acls.ts | 3 +- .../aws/ec2-network-interfaces.ts | 3 +- .../scan-functions/aws/ec2-route-tables.ts | 3 +- .../scan-functions/aws/ec2-subnets.ts | 3 +- .../aws/ec2-transit-gateways.ts | 3 +- .../scan-functions/aws/ec2-volumes.ts | 3 +- .../scan-functions/aws/ec2-vpc-endpoints.ts | 3 +- src/scanners/scan-functions/aws/ec2-vpcs.ts | 3 +- .../scan-functions/aws/ec2-vpn-gateways.ts | 3 +- src/scanners/scan-functions/aws/ecs.ts | 3 +- .../scan-functions/aws/efs-file-systems.ts | 3 +- .../scan-functions/aws/eks-clusters.ts | 3 +- .../aws/elasticache-clusters.ts | 3 +- .../aws/elbv1-load-balancers.ts | 3 +- .../aws/elbv2-load-balancers.ts | 3 +- .../scan-functions/aws/elbv2-target-groups.ts | 3 +- .../scan-functions/aws/lambda-functions.ts | 3 +- .../scan-functions/aws/rds-clusters.ts | 3 +- .../scan-functions/aws/rds-instances.ts | 3 +- .../scan-functions/aws/rds-proxies.ts | 3 +- .../scan-functions/aws/redshift-clusters.ts | 3 +- .../aws/route53-hosted-zones.ts | 3 +- src/scanners/scan-functions/aws/s3-buckets.ts | 5 +- src/scanners/scan-functions/aws/sns-topics.ts | 3 +- src/scanners/scan-functions/aws/sqs-queues.ts | 10 +- src/types.ts | 11 +- tests/mocks/RateLimiterMock.ts | 44 ++++- 45 files changed, 235 insertions(+), 181 deletions(-) create mode 100644 src/scanners/common/AWSRateLimitExhaustionStrategy.ts create mode 100644 src/scanners/common/RetryStrategy.ts diff --git a/src/aws-app/utils/createRateLimiterFactory.ts b/src/aws-app/utils/createRateLimiterFactory.ts index 37ee2e2..ee2135b 100644 --- a/src/aws-app/utils/createRateLimiterFactory.ts +++ b/src/aws-app/utils/createRateLimiterFactory.ts @@ -1,5 +1,6 @@ -import {RateLimiter, createRateLimiter} from '@/scanners' +import {createRateLimiter} from '@/scanners' import {GetRateLimiterFunction} from '@/scanners/types' +import {RateLimiter} from '@/types' /** * Returns the getRateLimiter function for a given service and region diff --git a/src/scanners/common/AWSRateLimitExhaustionStrategy.ts b/src/scanners/common/AWSRateLimitExhaustionStrategy.ts new file mode 100644 index 0000000..2189a12 --- /dev/null +++ b/src/scanners/common/AWSRateLimitExhaustionStrategy.ts @@ -0,0 +1,28 @@ +import {RetryOptions, RetryStrategy} from './RetryStrategy' + +export class AWSRateLimitExhaustionStrategy extends RetryStrategy { + constructor(options: Partial = {}) { + super({ + maxRetries: 5, + initialDelay: 1000, + maxDelay: 20000, + backoffFactor: 2, + ...options, + }) + } + + shouldRetry(error: any): boolean { + // AWS-specific error codes for rate limit exhaustion + const rateLimitErrorCodes = [ + 'ThrottlingException', + 'TooManyRequestsException', + 'RequestLimitExceeded', + 'Throttling', + 'RequestThrottled', + 'RequestThrottledException', + 'SlowDown', + ] + + return rateLimitErrorCodes.includes(error.code) || error.statusCode === 429 + } +} diff --git a/src/scanners/common/RateLimiter.ts b/src/scanners/common/RateLimiter.ts index 9dfe433..32ad11d 100644 --- a/src/scanners/common/RateLimiter.ts +++ b/src/scanners/common/RateLimiter.ts @@ -1,120 +1,103 @@ -const BASE_RETRY_DELAY = 1000 // 1 second -const MAX_RETRIES = 6 // last attempt will take 64 seconds (2^6 = 64 * base retry delay) and overall will take 127 seconds for all retries to complete, before giving up -const CAPACITY_USAGE_PERCENTAGE = 0.6 +import {RateLimiter} from '@/types' +import {RetryStrategy} from './RetryStrategy' -const ONE_SECOND = 1000 +export class RateLimiterImpl implements RateLimiter { + private executionTimesMs: number[] = [] + private pendingQueue: (() => void)[] = [] + private timeoutId: NodeJS.Timeout | null = null + private _isPaused = false + private readonly intervalMs: number + private retryStrategy: RetryStrategy | null -export const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) - -export interface RateLimiter { - throttle(fn: () => Promise): Promise -} - -class RateLimiterImpl implements RateLimiter { - private allowance: number - private lastCheck: number - private maxUsage: number - private queue: (() => void)[] - - constructor(rate: number) { - this.maxUsage = rate * CAPACITY_USAGE_PERCENTAGE - this.allowance = this.maxUsage - this.lastCheck = Date.now() - this.queue = [] + constructor( + private _rate: number, + retryStrategy: RetryStrategy | null = null, + ) { + if (!Number.isFinite(_rate) || _rate <= 0) { + throw new TypeError('Expected `rate` to be a positive finite number') + } + this.intervalMs = 1000 / _rate + this.retryStrategy = retryStrategy } - // Acquire permission to proceed with a function call - private async acquire() { - const current = Date.now() - const timePassed = current - this.lastCheck - this.lastCheck = current - this.allowance += (timePassed / ONE_SECOND) * this.maxUsage + private scheduleNextExecution(): void { + if (this.timeoutId) { + clearTimeout(this.timeoutId) + } - if (this.allowance > this.maxUsage) { - this.allowance = this.maxUsage + if (this._isPaused || this.pendingQueue.length === 0) { + return } - // If allowance is insufficient, wait for more capacity - if (this.allowance < 1) { - const waitTime = (1 - this.allowance) * (ONE_SECOND / this.maxUsage) - await sleep(waitTime) - this.allowance = 0 - } else { - this.allowance -= 1 + const now = Date.now() + let delay = 0 + + if (this.executionTimesMs.length > 0) { + const nextExecutionTime = this.executionTimesMs[this.executionTimesMs.length - 1] + this.intervalMs + delay = Math.max(0, nextExecutionTime - now) } + + this.timeoutId = setTimeout(() => { + const fn = this.pendingQueue.shift() + if (fn) { + this.executionTimesMs.push(Date.now()) + if (this.executionTimesMs.length > 10) { + this.executionTimesMs.shift() + } + fn() + } + this.scheduleNextExecution() + }, delay) } - // Throttle function execution - async throttle(fn: () => Promise): Promise { + throttle(fn: () => Promise): Promise { return new Promise((resolve, reject) => { - const execute = async () => { - try { - await this.acquire() - resolve(await fn()) - } catch (error: any) { - if (this.isRequestLimitError(error)) { - this.retry(fn, resolve, reject, 1) - } else { - reject(error) - } + const wrappedFn = () => { + if (this.retryStrategy) { + this.retryStrategy.retry(fn).then(resolve).catch(reject) + } else { + fn().then(resolve).catch(reject) } } - this.queue.push(execute) - - // If it's the only function in the queue, it means that we just added it. So we should start the process. - // But, if it's more in the queue, it means that we are already processing the queue, and the function will be processed when it's its turn. - if (this.queue.length === 1) { - this.dequeue() - } + this.pendingQueue.push(wrappedFn) + this.scheduleNextExecution() }) } - // Retry function with exponential backoff - private async retry( - fn: () => Promise, - resolve: (value: U | PromiseLike) => void, - reject: (reason?: any) => void, - attempt: number, - ) { - if (attempt < MAX_RETRIES) { - // Exponential backoff - // https://en.wikipedia.org/wiki/Exponential_backoff - await sleep(BASE_RETRY_DELAY * Math.pow(2, attempt)) - try { - // acquire permission to proceed and add enough sleep time - await this.acquire() - resolve(await fn()) - } catch (error: any) { - if (this.isRequestLimitError(error)) { - this.retry(fn, resolve, reject, attempt + 1) - } else { - reject(error) - } - } - } else { - reject(new Error('Max retries reached')) + pause(): void { + this._isPaused = true + if (this.timeoutId) { + clearTimeout(this.timeoutId) + this.timeoutId = null } } - // Check if error is a request limit error - private isRequestLimitError(error: Error) { - return [ - 'TooManyRequestsException', - 'ThrottlingException', - 'ProvisionedThroughputExceededException', - 'RequestLimitExceeded', - ].includes(error.name) + resume(): void { + this._isPaused = false + this.scheduleNextExecution() } - // Process the queue of function calls - // @idea Maybe we can add some lifecycle hooks here in the future? - private async dequeue() { - while (this.queue.length > 0) { - const fn = this.queue.shift() - if (fn) await fn() + abort(): void { + this.pendingQueue = [] + this.executionTimesMs = [] + if (this.timeoutId) { + clearTimeout(this.timeoutId) + this.timeoutId = null } } + + get queueSize(): number { + return this.pendingQueue.length + } + + get isPaused(): boolean { + return this._isPaused + } + + get rate(): number { + return this._rate + } } -export const createRateLimiter = (rate: number): RateLimiter => new RateLimiterImpl(rate) +export const createRateLimiter = (rate: number) => new RateLimiterImpl(rate) diff --git a/src/scanners/common/RetryStrategy.ts b/src/scanners/common/RetryStrategy.ts new file mode 100644 index 0000000..cdb6e46 --- /dev/null +++ b/src/scanners/common/RetryStrategy.ts @@ -0,0 +1,28 @@ +export interface RetryOptions { + maxRetries: number + initialDelay: number + maxDelay: number + backoffFactor: number +} + +export abstract class RetryStrategy { + constructor(protected options: RetryOptions) {} + + abstract shouldRetry(error: any): boolean + + async retry(fn: () => Promise, retryCount: number = 0): Promise { + try { + return await fn() + } catch (error) { + if (retryCount < this.options.maxRetries && this.shouldRetry(error)) { + const delay = Math.min( + this.options.initialDelay * Math.pow(this.options.backoffFactor, retryCount), + this.options.maxDelay, + ) + await new Promise((resolve) => setTimeout(resolve, delay)) + return this.retry(fn, retryCount + 1) + } + throw error + } + } +} diff --git a/src/scanners/common/createGlobalScanner.ts b/src/scanners/common/createGlobalScanner.ts index 2f3d402..ada26da 100644 --- a/src/scanners/common/createGlobalScanner.ts +++ b/src/scanners/common/createGlobalScanner.ts @@ -1,6 +1,12 @@ -import {Resources, ResourceDescription, GlobalScanFunction, Credentials, ScannerLifecycleHook} from '@/types' +import { + Resources, + ResourceDescription, + GlobalScanFunction, + Credentials, + ScannerLifecycleHook, + RateLimiter, +} from '@/types' import {CreateGlobalScannerFunction, GetRateLimiterFunction} from '@/scanners/types' -import {RateLimiter} from './RateLimiter' type GlobalScanResult = { resources: Resources diff --git a/src/scanners/common/createRegionalScanner.ts b/src/scanners/common/createRegionalScanner.ts index 113e7d5..103a935 100644 --- a/src/scanners/common/createRegionalScanner.ts +++ b/src/scanners/common/createRegionalScanner.ts @@ -5,9 +5,9 @@ import { RegionalScanFunction, Credentials, ScannerLifecycleHook, + RateLimiter, } from '@/types' import {CreateRegionalScannerFunction, GetRateLimiterFunction} from '@/scanners/types' -import {RateLimiter} from './RateLimiter' type RegionScanResult = { region: string diff --git a/src/scanners/index.ts b/src/scanners/index.ts index 7571ad4..c49fab3 100644 --- a/src/scanners/index.ts +++ b/src/scanners/index.ts @@ -1,4 +1,4 @@ -export {RateLimiter, createRateLimiter} from './common/RateLimiter' +export {createRateLimiter} from './common/RateLimiter' export {createGlobalScanner} from './common/createGlobalScanner' export {createRegionalScanner} from './common/createRegionalScanner' export {getAwsScanners} from './getAwsScanners' diff --git a/src/scanners/scan-functions/aws/_boilerplate.ts b/src/scanners/scan-functions/aws/_boilerplate.ts index 5198de3..1c920a6 100644 --- a/src/scanners/scan-functions/aws/_boilerplate.ts +++ b/src/scanners/scan-functions/aws/_boilerplate.ts @@ -1,6 +1,5 @@ import {DBInstance} from '@aws-sdk/client-rds' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' /** * 0️⃣ I put DBInstance just as an example of a real aws type. diff --git a/src/scanners/scan-functions/aws/athena-named-queries.ts b/src/scanners/scan-functions/aws/athena-named-queries.ts index 4c3fde3..13ddacb 100644 --- a/src/scanners/scan-functions/aws/athena-named-queries.ts +++ b/src/scanners/scan-functions/aws/athena-named-queries.ts @@ -1,6 +1,5 @@ import {AthenaClient, ListNamedQueriesCommand, GetNamedQueryCommand, NamedQuery} from '@aws-sdk/client-athena' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' import {buildARN} from './common/buildArn' import {getAwsAccountId} from './common/getAwsAccountId' diff --git a/src/scanners/scan-functions/aws/autoscaling-groups.ts b/src/scanners/scan-functions/aws/autoscaling-groups.ts index 1df4280..bba9aa0 100644 --- a/src/scanners/scan-functions/aws/autoscaling-groups.ts +++ b/src/scanners/scan-functions/aws/autoscaling-groups.ts @@ -1,6 +1,5 @@ import {AutoScalingClient, DescribeAutoScalingGroupsCommand, AutoScalingGroup} from '@aws-sdk/client-auto-scaling' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' async function getAutoScalingGroups( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/cloudfront-distributions.ts b/src/scanners/scan-functions/aws/cloudfront-distributions.ts index b8e1aea..80447a5 100644 --- a/src/scanners/scan-functions/aws/cloudfront-distributions.ts +++ b/src/scanners/scan-functions/aws/cloudfront-distributions.ts @@ -1,6 +1,5 @@ import {CloudFrontClient, ListDistributionsCommand, DistributionSummary} from '@aws-sdk/client-cloudfront' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getCloudFrontDistributions( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/cloudtrail-trails.ts b/src/scanners/scan-functions/aws/cloudtrail-trails.ts index 7ab6504..7c37d77 100644 --- a/src/scanners/scan-functions/aws/cloudtrail-trails.ts +++ b/src/scanners/scan-functions/aws/cloudtrail-trails.ts @@ -1,6 +1,5 @@ import {CloudTrailClient, ListTrailsCommand, ListTrailsCommandOutput, TrailInfo} from '@aws-sdk/client-cloudtrail' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getCloudTrailTrails( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/cloudwatch-metric-alarms.ts b/src/scanners/scan-functions/aws/cloudwatch-metric-alarms.ts index a09e63b..8aced9e 100644 --- a/src/scanners/scan-functions/aws/cloudwatch-metric-alarms.ts +++ b/src/scanners/scan-functions/aws/cloudwatch-metric-alarms.ts @@ -1,6 +1,5 @@ import {CloudWatchClient, DescribeAlarmsCommand, MetricAlarm} from '@aws-sdk/client-cloudwatch' -import {RateLimiter} from '@/scanners/common/RateLimiter' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getCloudWatchMetricAlarms( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/cloudwatch-metric-streams.ts b/src/scanners/scan-functions/aws/cloudwatch-metric-streams.ts index 54fd71a..6193ea6 100644 --- a/src/scanners/scan-functions/aws/cloudwatch-metric-streams.ts +++ b/src/scanners/scan-functions/aws/cloudwatch-metric-streams.ts @@ -1,6 +1,5 @@ import {CloudWatchClient, ListMetricStreamsCommand, MetricStreamEntry} from '@aws-sdk/client-cloudwatch' -import {RateLimiter} from '@/scanners/common/RateLimiter' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getCloudWatchMetricStreams( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/dynamodb-tables.ts b/src/scanners/scan-functions/aws/dynamodb-tables.ts index 5a94fb3..cb92bed 100644 --- a/src/scanners/scan-functions/aws/dynamodb-tables.ts +++ b/src/scanners/scan-functions/aws/dynamodb-tables.ts @@ -1,6 +1,5 @@ import {DynamoDBClient, ListTablesCommand, DescribeTableCommand, TableDescription} from '@aws-sdk/client-dynamodb' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getDynamoDBTables( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/ec2-instances.ts b/src/scanners/scan-functions/aws/ec2-instances.ts index 7f6993c..6cc745b 100644 --- a/src/scanners/scan-functions/aws/ec2-instances.ts +++ b/src/scanners/scan-functions/aws/ec2-instances.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeInstancesCommand, Instance} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2Instances( diff --git a/src/scanners/scan-functions/aws/ec2-internet-gateways.ts b/src/scanners/scan-functions/aws/ec2-internet-gateways.ts index 9e846a1..7770e1f 100644 --- a/src/scanners/scan-functions/aws/ec2-internet-gateways.ts +++ b/src/scanners/scan-functions/aws/ec2-internet-gateways.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeInternetGatewaysCommand, InternetGateway} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2InternetGateways( diff --git a/src/scanners/scan-functions/aws/ec2-nat-gateways.ts b/src/scanners/scan-functions/aws/ec2-nat-gateways.ts index c87ebc7..94c928e 100644 --- a/src/scanners/scan-functions/aws/ec2-nat-gateways.ts +++ b/src/scanners/scan-functions/aws/ec2-nat-gateways.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeNatGatewaysCommand, NatGateway} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2NatGateways( diff --git a/src/scanners/scan-functions/aws/ec2-network-acls.ts b/src/scanners/scan-functions/aws/ec2-network-acls.ts index 18588fb..7f1061a 100644 --- a/src/scanners/scan-functions/aws/ec2-network-acls.ts +++ b/src/scanners/scan-functions/aws/ec2-network-acls.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeNetworkAclsCommand, NetworkAcl} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2NetworkAcls( diff --git a/src/scanners/scan-functions/aws/ec2-network-interfaces.ts b/src/scanners/scan-functions/aws/ec2-network-interfaces.ts index 7f66fd7..f215aed 100644 --- a/src/scanners/scan-functions/aws/ec2-network-interfaces.ts +++ b/src/scanners/scan-functions/aws/ec2-network-interfaces.ts @@ -1,6 +1,5 @@ import {EC2Client, DescribeNetworkInterfacesCommand, NetworkInterface} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {buildARN} from './common/buildArn' import {getAwsAccountId} from './common/getAwsAccountId' diff --git a/src/scanners/scan-functions/aws/ec2-route-tables.ts b/src/scanners/scan-functions/aws/ec2-route-tables.ts index 5b43c2d..5842f0c 100644 --- a/src/scanners/scan-functions/aws/ec2-route-tables.ts +++ b/src/scanners/scan-functions/aws/ec2-route-tables.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeRouteTablesCommand, RouteTable} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2RouteTables( diff --git a/src/scanners/scan-functions/aws/ec2-subnets.ts b/src/scanners/scan-functions/aws/ec2-subnets.ts index cdfd0a4..0454dbb 100644 --- a/src/scanners/scan-functions/aws/ec2-subnets.ts +++ b/src/scanners/scan-functions/aws/ec2-subnets.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeSubnetsCommand, Subnet} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2Subnets( diff --git a/src/scanners/scan-functions/aws/ec2-transit-gateways.ts b/src/scanners/scan-functions/aws/ec2-transit-gateways.ts index 62faa08..9e6ae69 100644 --- a/src/scanners/scan-functions/aws/ec2-transit-gateways.ts +++ b/src/scanners/scan-functions/aws/ec2-transit-gateways.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeTransitGatewaysCommand, TransitGateway} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2TransitGateways( diff --git a/src/scanners/scan-functions/aws/ec2-volumes.ts b/src/scanners/scan-functions/aws/ec2-volumes.ts index c2b66f7..52aa486 100644 --- a/src/scanners/scan-functions/aws/ec2-volumes.ts +++ b/src/scanners/scan-functions/aws/ec2-volumes.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeVolumesCommand, Volume} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2Volumes( diff --git a/src/scanners/scan-functions/aws/ec2-vpc-endpoints.ts b/src/scanners/scan-functions/aws/ec2-vpc-endpoints.ts index 7f51523..939adc1 100644 --- a/src/scanners/scan-functions/aws/ec2-vpc-endpoints.ts +++ b/src/scanners/scan-functions/aws/ec2-vpc-endpoints.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeVpcEndpointsCommand, VpcEndpoint} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2VpcEndpoints( diff --git a/src/scanners/scan-functions/aws/ec2-vpcs.ts b/src/scanners/scan-functions/aws/ec2-vpcs.ts index 294c7de..90a3d01 100644 --- a/src/scanners/scan-functions/aws/ec2-vpcs.ts +++ b/src/scanners/scan-functions/aws/ec2-vpcs.ts @@ -1,7 +1,6 @@ import {EC2Client, DescribeVpcsCommand, Vpc} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getEC2Vpcs( diff --git a/src/scanners/scan-functions/aws/ec2-vpn-gateways.ts b/src/scanners/scan-functions/aws/ec2-vpn-gateways.ts index 99a3977..419ba7b 100644 --- a/src/scanners/scan-functions/aws/ec2-vpn-gateways.ts +++ b/src/scanners/scan-functions/aws/ec2-vpn-gateways.ts @@ -1,6 +1,5 @@ import {EC2Client, DescribeVpnGatewaysCommand, VpnGateway} from '@aws-sdk/client-ec2' -import {RateLimiter} from '@/scanners/common/RateLimiter' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' import {buildARN} from './common/buildArn' diff --git a/src/scanners/scan-functions/aws/ecs.ts b/src/scanners/scan-functions/aws/ecs.ts index 983de34..a0d524c 100644 --- a/src/scanners/scan-functions/aws/ecs.ts +++ b/src/scanners/scan-functions/aws/ecs.ts @@ -10,8 +10,7 @@ import { Service, Task, } from '@aws-sdk/client-ecs' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' async function getECSClusters(client: ECSClient, rateLimiter: RateLimiter): Promise { const listClustersCommand = new ListClustersCommand({}) diff --git a/src/scanners/scan-functions/aws/efs-file-systems.ts b/src/scanners/scan-functions/aws/efs-file-systems.ts index 669e435..80445e3 100644 --- a/src/scanners/scan-functions/aws/efs-file-systems.ts +++ b/src/scanners/scan-functions/aws/efs-file-systems.ts @@ -1,6 +1,5 @@ import {EFSClient, DescribeFileSystemsCommand, FileSystemDescription} from '@aws-sdk/client-efs' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getEFSFileSystems( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/eks-clusters.ts b/src/scanners/scan-functions/aws/eks-clusters.ts index 6c3f14e..e2e7a2c 100644 --- a/src/scanners/scan-functions/aws/eks-clusters.ts +++ b/src/scanners/scan-functions/aws/eks-clusters.ts @@ -1,6 +1,5 @@ import {EKSClient, ListClustersCommand, DescribeClusterCommand, Cluster} from '@aws-sdk/client-eks' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getEKSClusters( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/elasticache-clusters.ts b/src/scanners/scan-functions/aws/elasticache-clusters.ts index 96e9094..53427fc 100644 --- a/src/scanners/scan-functions/aws/elasticache-clusters.ts +++ b/src/scanners/scan-functions/aws/elasticache-clusters.ts @@ -1,6 +1,5 @@ import {ElastiCacheClient, DescribeCacheClustersCommand, CacheCluster} from '@aws-sdk/client-elasticache' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getElastiCacheClusters( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/elbv1-load-balancers.ts b/src/scanners/scan-functions/aws/elbv1-load-balancers.ts index 44bd334..88ba1a3 100644 --- a/src/scanners/scan-functions/aws/elbv1-load-balancers.ts +++ b/src/scanners/scan-functions/aws/elbv1-load-balancers.ts @@ -3,8 +3,7 @@ import { DescribeLoadBalancersCommand, LoadBalancerDescription, } from '@aws-sdk/client-elastic-load-balancing' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getELBv1LoadBalancers( diff --git a/src/scanners/scan-functions/aws/elbv2-load-balancers.ts b/src/scanners/scan-functions/aws/elbv2-load-balancers.ts index 9d0ea35..2df8b3d 100644 --- a/src/scanners/scan-functions/aws/elbv2-load-balancers.ts +++ b/src/scanners/scan-functions/aws/elbv2-load-balancers.ts @@ -3,8 +3,7 @@ import { DescribeLoadBalancersCommand, LoadBalancer, } from '@aws-sdk/client-elastic-load-balancing-v2' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getELBv2LoadBalancers( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/elbv2-target-groups.ts b/src/scanners/scan-functions/aws/elbv2-target-groups.ts index d32f06e..685230f 100644 --- a/src/scanners/scan-functions/aws/elbv2-target-groups.ts +++ b/src/scanners/scan-functions/aws/elbv2-target-groups.ts @@ -3,8 +3,7 @@ import { DescribeTargetGroupsCommand, TargetGroup, } from '@aws-sdk/client-elastic-load-balancing-v2' -import {RateLimiter} from '@/scanners/common/RateLimiter' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getELBv2TargetGroups( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/lambda-functions.ts b/src/scanners/scan-functions/aws/lambda-functions.ts index 721c6f9..756ae5f 100644 --- a/src/scanners/scan-functions/aws/lambda-functions.ts +++ b/src/scanners/scan-functions/aws/lambda-functions.ts @@ -1,6 +1,5 @@ import {LambdaClient, ListFunctionsCommand, GetFunctionCommand, FunctionConfiguration} from '@aws-sdk/client-lambda' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getLambdaFunctions( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/rds-clusters.ts b/src/scanners/scan-functions/aws/rds-clusters.ts index 693848d..e2a60ab 100644 --- a/src/scanners/scan-functions/aws/rds-clusters.ts +++ b/src/scanners/scan-functions/aws/rds-clusters.ts @@ -1,6 +1,5 @@ import {RDSClient, DescribeDBClustersCommand, DBCluster} from '@aws-sdk/client-rds' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getRDSClusters( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/rds-instances.ts b/src/scanners/scan-functions/aws/rds-instances.ts index 9561837..097ec16 100644 --- a/src/scanners/scan-functions/aws/rds-instances.ts +++ b/src/scanners/scan-functions/aws/rds-instances.ts @@ -1,6 +1,5 @@ import {RDSClient, DescribeDBInstancesCommand, DBInstance} from '@aws-sdk/client-rds' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getRDSInstances( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/rds-proxies.ts b/src/scanners/scan-functions/aws/rds-proxies.ts index 7ecab40..21b5911 100644 --- a/src/scanners/scan-functions/aws/rds-proxies.ts +++ b/src/scanners/scan-functions/aws/rds-proxies.ts @@ -1,6 +1,5 @@ import {RDSClient, DescribeDBProxiesCommand, DBProxy} from '@aws-sdk/client-rds' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getRDSProxies( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/redshift-clusters.ts b/src/scanners/scan-functions/aws/redshift-clusters.ts index 5b8b30e..a7a7270 100644 --- a/src/scanners/scan-functions/aws/redshift-clusters.ts +++ b/src/scanners/scan-functions/aws/redshift-clusters.ts @@ -1,7 +1,6 @@ import {RedshiftClient, DescribeClustersCommand, Cluster} from '@aws-sdk/client-redshift' -import {RateLimiter} from '@/scanners/common/RateLimiter' import {buildARN} from './common/buildArn' -import {Credentials, Resources} from '@/types' +import {Credentials, Resources, RateLimiter} from '@/types' import {getAwsAccountId} from './common/getAwsAccountId' export async function getRedshiftClusters( diff --git a/src/scanners/scan-functions/aws/route53-hosted-zones.ts b/src/scanners/scan-functions/aws/route53-hosted-zones.ts index 6cc6365..dd0c65c 100644 --- a/src/scanners/scan-functions/aws/route53-hosted-zones.ts +++ b/src/scanners/scan-functions/aws/route53-hosted-zones.ts @@ -1,6 +1,5 @@ import {Route53Client, ListHostedZonesCommand, GetHostedZoneCommand, HostedZone} from '@aws-sdk/client-route-53' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' import {buildARN} from './common/buildArn' import {getAwsAccountId} from './common/getAwsAccountId' diff --git a/src/scanners/scan-functions/aws/s3-buckets.ts b/src/scanners/scan-functions/aws/s3-buckets.ts index ae4e1bc..bf162a8 100644 --- a/src/scanners/scan-functions/aws/s3-buckets.ts +++ b/src/scanners/scan-functions/aws/s3-buckets.ts @@ -1,6 +1,5 @@ -import {S3Client, ListBucketsCommand, GetBucketLocationCommand, Bucket} from '@aws-sdk/client-s3' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {S3Client, ListBucketsCommand, Bucket} from '@aws-sdk/client-s3' +import {Credentials, Resources, RateLimiter} from '@/types' import {buildARN} from './common/buildArn' import {getAwsAccountId} from './common/getAwsAccountId' diff --git a/src/scanners/scan-functions/aws/sns-topics.ts b/src/scanners/scan-functions/aws/sns-topics.ts index 24c9ef9..33a3b6f 100644 --- a/src/scanners/scan-functions/aws/sns-topics.ts +++ b/src/scanners/scan-functions/aws/sns-topics.ts @@ -1,6 +1,5 @@ import {ListTopicsCommand, SNSClient, Topic} from '@aws-sdk/client-sns' -import {Credentials, Resources} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {Credentials, Resources, RateLimiter} from '@/types' export async function getSNSTopics( credentials: Credentials, diff --git a/src/scanners/scan-functions/aws/sqs-queues.ts b/src/scanners/scan-functions/aws/sqs-queues.ts index 9d2e5ee..f5e4eff 100644 --- a/src/scanners/scan-functions/aws/sqs-queues.ts +++ b/src/scanners/scan-functions/aws/sqs-queues.ts @@ -1,11 +1,5 @@ -import { - SQSClient, - ListQueuesCommand, - GetQueueAttributesCommand, - GetQueueAttributesCommandOutput, -} from '@aws-sdk/client-sqs' -import {Credentials, Resources, SQSQueue} from '@/types' -import {RateLimiter} from '@/scanners/common/RateLimiter' +import {SQSClient, ListQueuesCommand, GetQueueAttributesCommand} from '@aws-sdk/client-sqs' +import {Credentials, Resources, SQSQueue, RateLimiter} from '@/types' export async function getSQSQueues( credentials: Credentials, diff --git a/src/types.ts b/src/types.ts index b6a5ad1..fe12688 100644 --- a/src/types.ts +++ b/src/types.ts @@ -19,7 +19,6 @@ import type * as Redshift from '@aws-sdk/client-redshift' import type * as CloudWatch from '@aws-sdk/client-cloudwatch' import type * as Athena from '@aws-sdk/client-athena' -import type {RateLimiter} from './scanners/common/RateLimiter' import type {AwsCredentialIdentity} from '@aws-sdk/types' import {awsRegionIds} from './constants' @@ -72,6 +71,16 @@ export type Resources = { export type Credentials = AwsCredentialIdentity | undefined +export interface RateLimiter { + throttle(fn: () => Promise): Promise + pause(): void + resume(): void + abort(): void + readonly queueSize: number + readonly isPaused: boolean + readonly rate: number +} + export type RegionalScanFunction = ( credentials: Credentials, rateLimiter: RateLimiter, diff --git a/tests/mocks/RateLimiterMock.ts b/tests/mocks/RateLimiterMock.ts index 89808b6..dec8681 100644 --- a/tests/mocks/RateLimiterMock.ts +++ b/tests/mocks/RateLimiterMock.ts @@ -1,7 +1,47 @@ -import {RateLimiter} from '@/scanners' +import {RateLimiter} from '@/types' export class RateLimiterMockImpl implements RateLimiter { + private _isPaused: boolean = false + private _queueSize: number = 0 + private _rate: number = 0 + + constructor(rate: number = 0) { + this._rate = rate + } + async throttle(callback: () => Promise): Promise { - return callback() + if (this._isPaused) { + throw new Error('RateLimiter is paused') + } + this._queueSize++ + try { + return await callback() + } finally { + this._queueSize-- + } + } + + pause(): void { + this._isPaused = true + } + + resume(): void { + this._isPaused = false + } + + abort(): void { + this._queueSize = 0 + } + + get queueSize(): number { + return this._queueSize + } + + get isPaused(): boolean { + return this._isPaused + } + + get rate(): number { + return this._rate } }