From 77bbbb6d3d0c66b708c008d16538a35d91a0b766 Mon Sep 17 00:00:00 2001 From: Chirag-CG Date: Mon, 22 Dec 2025 18:57:02 +0530 Subject: [PATCH] adding rate limit logic --- collectors/azure/collector.js | 156 +++++++++++++++++++++++++++++----- 1 file changed, 136 insertions(+), 20 deletions(-) diff --git a/collectors/azure/collector.js b/collectors/azure/collector.js index fabd6f2c20..c61d6e93f2 100644 --- a/collectors/azure/collector.js +++ b/collectors/azure/collector.js @@ -10,8 +10,8 @@ - api_calls: (Optional) If provided, will only query these APIs. - Example: { - "skip_locations": ["eastus", "westus"], - "api_calls": ["storageAccounts:list", "resourceGroups:list"] + "skip_locations": ["eastus", "westus"], + "api_calls": ["storageAccounts:list", "resourceGroups:list"] } - callback: Function to call when the collection is complete *********************/ @@ -44,8 +44,9 @@ function parseCollection(path, obj) { } } -let collect = function(AzureConfig, settings, callback) { +let collect = function(AzureConfig, settings, callback, fargateFlag) { // Used to gather info only + fargateFlag = true if (settings.gather) { return callback(null, calls, postcalls, tertiarycalls, specialcalls); } @@ -62,14 +63,59 @@ let collect = function(AzureConfig, settings, callback) { var collection = {}; let makeCall = function(localUrl, obj, cb, localData) { - helpers.call({ - url: localUrl, - post: obj.post, - token: obj.graph ? loginData.graphToken : (obj.vault ? loginData.vaultToken : loginData.token), - govcloud : AzureConfig.Govcloud - }, function(err, data) { - if (err) return cb(err); + const makeApiCall = async(retryAttempt = 0) => { + try { + const response = await new Promise((resolve, reject) => { + helpers.call({ + url: localUrl, + post: obj.post, + token: obj.graph ? loginData.graphToken : (obj.vault ? loginData.vaultToken : loginData.token), + govcloud: AzureConfig.Govcloud + }, (err, data, apiResponse) => { + if (err) { + reject({ error: err, response: apiResponse }); + } else { + resolve({ data, response: apiResponse }); + } + }); + }); + + handleResponse(response.data); + } catch (error) { + if (error.error && error.error.includes('ECONNRESET') && retryAttempt < 3) { + // Refresh token using callback patterns + helpers.login(AzureConfig, function(refreshErr, refreshedLoginData) { + if (refreshErr) { + return cb(`Failed to refresh token: ${refreshErr}`, null, error.response); + } + + // Update loginData with refreshed tokens + loginData = refreshedLoginData; + console.log('Token refreshed successfully. New token data:', JSON.stringify({ + environment: loginData.environment, + token: loginData.token ? '***' + loginData.token.slice(-8) : null, + graphToken: loginData.graphToken ? '***' + loginData.graphToken.slice(-8) : null, + vaultToken: loginData.vaultToken ? '***' + loginData.vaultToken.slice(-8) : null + }, null, 2)); + + // Retry with new token + makeApiCall(retryAttempt + 1) + .then(() => cb(null, null, null)) + .catch(err => cb(err, null, null)); + }); + } else { + return cb(error.error, null, error.response); + } + } + }; + + // Start the async process + makeApiCall().catch(error => { + console.log(`Unexpected error in makeApiCall: ${error}`); + cb(error, null, null); + }); + function handleResponse(data) { // If a new nextLink is provided, this will be updated. There shouldn't // be a need to hold on to the previous value if (data && obj.hasListResponse && data.length) data.value = data; @@ -85,25 +131,95 @@ let collect = function(AzureConfig, settings, callback) { return cb(null, localData); } - let resData = localData || data; + const resData = localData || data; if (data && ((obj.paginate && data[obj.paginate]) || data['nextLink']) && (!obj.limit || (obj.limit && resData && resData.value && resData.value.length < obj.limit))) { obj.nextUrl = data['nextLink'] || data[obj.paginate]; processCall(obj, cb, localData || data); } else { return cb(null, localData || data || []); } - }); + } }; let processCall = function(obj, cb, localData) { - let localUrl = obj.nextUrl || obj.url.replace(/\{subscriptionId\}/g, AzureConfig.SubscriptionID); - if (obj.rateLimit) { - setTimeout(function() { - console.log(`url: ${localUrl}`); - makeCall(localUrl, obj, cb, localData); - }, obj.rateLimit); + let callbackCalled = false; + const wrappedCallback = (err, data, response) => { + if (callbackCalled) return; + callbackCalled = true; + cb(err, data, response); + }; + + if (fargateFlag) { + const maxApiRetryAttempts = 15; + let initialResponse = null; + async.retry({ + times: maxApiRetryAttempts, + interval: function() { + let retryAfter = initialResponse && initialResponse.headers['retry-after'] ? parseInt(initialResponse.headers['retry-after']) * 1000 : 5000; + let localUrl = obj.nextUrl || obj.url.replace(/\{subscriptionId\}/g, AzureConfig.SubscriptionID); + console.log(`Retrying in ${retryAfter / 1000} seconds for api call : ${localUrl}`); + return retryAfter; + }, + errorFilter: function(err) { + const errorMessage = typeof err === 'string' ? err : err.message || err.toString(); + + // Azure throttling patterns + const throttlingPatterns = [ + 'TooManyRequests', + 'RateLimitExceeded', + 'Throttling', + 'Throttled', + 'RequestThrottled', + 'RequestLimitExceeded', + 'ServerBusy', + 'ServiceBusy', + 'toomanyrequests', + 'ratelimitexceeded', + 'throttling', + 'throttled', + 'requestthrottled', + 'requestlimitexceeded', + 'serverbusy', + 'servicebusy', + 'too many requests', + 'rate limit', + 'retry after', + 'the request is being throttled', + 'request rate is large', + 'rate exceeded' + ]; + + const errorMatch = throttlingPatterns.some(pattern => errorMessage.includes(pattern)); + return errorMatch; + } + }, function(retryCallback) { + let localUrl = obj.nextUrl || obj.url.replace(/\{subscriptionId\}/g, AzureConfig.SubscriptionID); + var rateLimit = obj.rateLimit && obj.rateLimit == 3000? 6000: obj.rateLimit; + if (rateLimit) { + setTimeout(function() { + console.log(`Fargate collector rate limited: url: ${localUrl}`); + makeCall(localUrl, obj, function(err, data, response) { + initialResponse = response; + retryCallback(err, data, response); + }, localData); + }, rateLimit); + } else { + makeCall(localUrl, obj, function(err, data, response) { + initialResponse = response; + retryCallback(err, data, response); + }, localData); + } + }, wrappedCallback); } else { - makeCall(localUrl, obj, cb, localData); + let localUrl = obj.nextUrl || obj.url.replace(/\{subscriptionId\}/g, AzureConfig.SubscriptionID); + if (obj.rateLimit) { + setTimeout(function() { + console.log(`url: ${localUrl}`); + makeCall(localUrl, obj, wrappedCallback, localData); + }, obj.rateLimit); + } else { + makeCall(localUrl, obj, wrappedCallback, localData); + } } }; @@ -494,4 +610,4 @@ let collect = function(AzureConfig, settings, callback) { }); }; -module.exports = collect; +module.exports = collect; \ No newline at end of file