diff --git a/src/main/java/burp/RequestSender.java b/src/main/java/burp/RequestSender.java index a36c412..7dddb5d 100644 --- a/src/main/java/burp/RequestSender.java +++ b/src/main/java/burp/RequestSender.java @@ -53,7 +53,8 @@ class RequestSender { // Rate limiting and circuit breaker state per host private static final Map REQUEST_COUNTS = new ConcurrentHashMap<>(); - private static final Map LAST_REQUEST_TIME = new ConcurrentHashMap<>(); + private static final Map RATE_LIMIT_TIMESTAMPS = new ConcurrentHashMap<>(); + private static final Map LAST_FAILURE_TIME = new ConcurrentHashMap<>(); private static final Map FAILURE_COUNTS = new ConcurrentHashMap<>(); private static final int MAX_REQUESTS_PER_SECOND = 10; private static final int CIRCUIT_BREAKER_THRESHOLD = 5; @@ -405,6 +406,15 @@ private static URL createNewURLWithSegment(URL orgURL, String additional, String * Uses Caffeine cache for high-performance caching with TTL. */ protected static Map retrieveResponseDetails(IHttpService service, byte[] request) { + String hostKey = service.getHost(); + String cacheKey = service.toString() + Arrays.hashCode(request); + + for (int attempt = 0; attempt < MAX_RETRIES; attempt++) { + try { + if (isCircuitOpen(hostKey)) { + BurpExtender.logDebug("Circuit breaker open for " + hostKey); + return null; + } return retrieveResponseDetails(service, request, 0); } @@ -430,53 +440,56 @@ private static Map retrieveResponseDetails(IHttpService service, return cached; } - // Adaptive delay based on host response times - long delay = calculateAdaptiveDelay(hostKey); - try { Thread.sleep(delay); } catch (InterruptedException ignored) {} + waitForRateLimit(hostKey); - long startTime = System.currentTimeMillis(); - IHttpRequestResponse response = BurpExtender.getCallbacks().makeHttpRequest(service, request); - long responseTime = System.currentTimeMillis() - startTime; - - if (response == null) { - recordFailure(hostKey); - if (retryCount < MAX_RETRIES) { - int delayMs = calculateRetryDelay(retryCount); - try { Thread.sleep(delayMs); } catch (InterruptedException ignored) {} - return retrieveResponseDetails(service, request, retryCount + 1); + Map cached = RESPONSE_CACHE.getIfPresent(cacheKey); + if (cached != null) { + return cached; } - return null; - } - IResponseInfo responseInfo = BurpExtender.getHelpers().analyzeResponse(response.getResponse()); - Map details = new HashMap<>(); - details.put("statusCode", (int) responseInfo.getStatusCode()); - details.put("headers", responseInfo.getHeaders()); - details.put("responseTime", responseTime); + long delay = calculateAdaptiveDelay(hostKey); + sleepRespectingInterrupts(delay); - byte[] responseBody = java.util.Arrays.copyOfRange(response.getResponse(), - responseInfo.getBodyOffset(), response.getResponse().length); - details.put("body", responseBody); + long startTime = System.currentTimeMillis(); + IHttpRequestResponse response = BurpExtender.getCallbacks().makeHttpRequest(service, request); + long responseTime = System.currentTimeMillis() - startTime; - // Cache successful responses - if (responseInfo.getStatusCode() >= 200 && responseInfo.getStatusCode() < 500) { - RESPONSE_CACHE.put(cacheKey, details); - recordSuccess(hostKey, responseTime); - } else { + if (response == null) { + recordFailure(hostKey); + if (attempt < MAX_RETRIES - 1) { + sleepRespectingInterrupts(calculateRetryDelay(attempt)); + continue; + } + return null; + } + + IResponseInfo responseInfo = BurpExtender.getHelpers().analyzeResponse(response.getResponse()); + Map details = new HashMap<>(); + details.put("statusCode", (int) responseInfo.getStatusCode()); + details.put("headers", responseInfo.getHeaders()); + details.put("responseTime", responseTime); + + byte[] responseBody = java.util.Arrays.copyOfRange(response.getResponse(), + responseInfo.getBodyOffset(), response.getResponse().length); + details.put("body", responseBody); + + if (responseInfo.getStatusCode() >= 200 && responseInfo.getStatusCode() < 500) { + RESPONSE_CACHE.put(cacheKey, details); + recordSuccess(hostKey, responseTime); + } else { + recordFailure(hostKey); + } + + return details; + } catch (Exception e) { recordFailure(hostKey); + if (attempt < MAX_RETRIES - 1) { + sleepRespectingInterrupts(calculateRetryDelay(attempt)); + } } - - return details; - } catch (Exception e) { - String hostKey = service.getHost(); - recordFailure(hostKey); - if (retryCount < MAX_RETRIES) { - int delayMs = calculateRetryDelay(retryCount); - try { Thread.sleep(delayMs); } catch (InterruptedException ignored) {} - return retrieveResponseDetails(service, request, retryCount + 1); - } - return null; } + + return null; } /** @@ -504,16 +517,14 @@ private static String buildServiceCacheKey(IHttpService service, byte[] request) private static boolean isCircuitOpen(String hostKey) { AtomicInteger failures = FAILURE_COUNTS.computeIfAbsent(hostKey, k -> new AtomicInteger(0)); if (failures.get() >= CIRCUIT_BREAKER_THRESHOLD) { - AtomicLong lastFailure = LAST_REQUEST_TIME.get(hostKey); + AtomicLong lastFailure = LAST_FAILURE_TIME.get(hostKey); if (lastFailure == null) { - // No failure time recorded yet, circuit should not be open return false; } long timeSinceLastFailure = System.currentTimeMillis() - lastFailure.get(); if (timeSinceLastFailure < CIRCUIT_BREAKER_RESET_MS) { return true; } else { - // Reset circuit breaker after reset period has elapsed failures.set(0); } } @@ -522,28 +533,27 @@ private static boolean isCircuitOpen(String hostKey) { private static boolean checkRateLimit(String hostKey) { AtomicInteger count = REQUEST_COUNTS.computeIfAbsent(hostKey, k -> new AtomicInteger(0)); - AtomicLong lastTime = LAST_REQUEST_TIME.computeIfAbsent(hostKey, k -> new AtomicLong(System.currentTimeMillis())); - + AtomicLong lastTime = RATE_LIMIT_TIMESTAMPS.computeIfAbsent(hostKey, k -> new AtomicLong(System.currentTimeMillis())); + long currentTime = System.currentTimeMillis(); long timeDiff = currentTime - lastTime.get(); - + if (timeDiff >= 1000) { - // Reset counter every second count.set(0); - lastTime.set(currentTime); } - + if (count.get() >= MAX_REQUESTS_PER_SECOND) { return false; } - + count.incrementAndGet(); + lastTime.set(currentTime); return true; } - + private static long calculateAdaptiveDelay(String hostKey) { // Start with base delay, adjust based on response times - AtomicLong lastTime = LAST_REQUEST_TIME.get(hostKey); + AtomicLong lastTime = RATE_LIMIT_TIMESTAMPS.get(hostKey); if (lastTime == null) { return 50; // Base delay } @@ -563,7 +573,24 @@ private static void recordSuccess(String hostKey, long responseTime) { private static void recordFailure(String hostKey) { FAILURE_COUNTS.computeIfAbsent(hostKey, k -> new AtomicInteger(0)).incrementAndGet(); - LAST_REQUEST_TIME.computeIfAbsent(hostKey, k -> new AtomicLong(System.currentTimeMillis())).set(System.currentTimeMillis()); + LAST_FAILURE_TIME.computeIfAbsent(hostKey, k -> new AtomicLong(System.currentTimeMillis())).set(System.currentTimeMillis()); + } + + private static void waitForRateLimit(String hostKey) { + while (!checkRateLimit(hostKey)) { + sleepRespectingInterrupts(100); + } + } + + private static void sleepRespectingInterrupts(long delayMs) { + if (delayMs <= 0) { + return; + } + try { + Thread.sleep(delayMs); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } } /**