From a140a4194e2a521ecac5bffb03573e6a120c4ad6 Mon Sep 17 00:00:00 2001
From: Mzack9999 <mzack9999@protonmail.com>
Date: Wed, 3 Apr 2024 19:40:09 +0200
Subject: [PATCH] boh - placing resize in wrapped method

---
 pkg/protocols/dns/request.go       |  2 +-
 pkg/protocols/http/request.go      |  4 ++--
 pkg/protocols/http/request_fuzz.go |  2 +-
 pkg/protocols/protocols.go         | 11 +++++++++++
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/pkg/protocols/dns/request.go b/pkg/protocols/dns/request.go
index ba9dd666..8aa8c4a5 100644
--- a/pkg/protocols/dns/request.go
+++ b/pkg/protocols/dns/request.go
@@ -143,7 +143,7 @@ func (request *Request) execute(input *contextargs.Context, domain string, metad
 		}
 	}
 
-	request.options.RateLimiter.Take()
+	request.options.RateLimitTake()
 
 	// Send the request to the target servers
 	response, err := dnsClient.Do(compiledRequest)
diff --git a/pkg/protocols/http/request.go b/pkg/protocols/http/request.go
index 2763c5c6..1b27aa7b 100644
--- a/pkg/protocols/http/request.go
+++ b/pkg/protocols/http/request.go
@@ -222,7 +222,7 @@ func (request *Request) executeParallelHTTP(input *contextargs.Context, dynamicV
 				return
 			case spmHandler.ResultChan <- func() error {
 				// putting ratelimiter here prevents any unnecessary waiting if any
-				request.options.RateLimiter.Take()
+				request.options.RateLimitTake()
 				previous := make(map[string]interface{})
 				return request.executeRequest(input, httpRequest, previous, false, wrappedCallback, 0)
 			}():
@@ -366,7 +366,7 @@ func (request *Request) ExecuteWithResults(input *contextargs.Context, dynamicVa
 		executeFunc := func(data string, payloads, dynamicValue map[string]interface{}) (bool, error) {
 			hasInteractMatchers := interactsh.HasMatchers(request.CompiledOperators)
 
-			request.options.RateLimiter.Take()
+			request.options.RateLimitTake()
 
 			ctx := request.newContext(input)
 			ctxWithTimeout, cancel := context.WithTimeout(ctx, time.Duration(request.options.Options.Timeout)*time.Second)
diff --git a/pkg/protocols/http/request_fuzz.go b/pkg/protocols/http/request_fuzz.go
index d6a29c18..aed17f54 100644
--- a/pkg/protocols/http/request_fuzz.go
+++ b/pkg/protocols/http/request_fuzz.go
@@ -145,7 +145,7 @@ func (request *Request) executeGeneratedFuzzingRequest(gr fuzz.GeneratedRequest,
 	if request.options.HostErrorsCache != nil && request.options.HostErrorsCache.Check(input.MetaInput.Input) {
 		return false
 	}
-	request.options.RateLimiter.Take()
+	request.options.RateLimitTake()
 	req := &generatedRequest{
 		request:        gr.Request,
 		dynamicValues:  gr.DynamicValues,
diff --git a/pkg/protocols/protocols.go b/pkg/protocols/protocols.go
index f9d90642..5ac0a2f0 100644
--- a/pkg/protocols/protocols.go
+++ b/pkg/protocols/protocols.go
@@ -128,6 +128,17 @@ type ExecutorOptions struct {
 	ExportReqURLPattern bool
 }
 
+// todo: centralizing components is not feasible with current clogged architecture
+// a possible approach could be an internal event bus with pub-subs? This would be less invasive than
+// reworking dep injection from scratch
+func (eo *ExecutorOptions) RateLimitTake() {
+	if eo.RateLimiter.GetLimit() != uint(eo.Options.RateLimit) {
+		eo.RateLimiter.SetLimit(uint(eo.Options.RateLimit))
+		eo.RateLimiter.SetDuration(eo.Options.RateLimitDuration)
+	}
+	eo.RateLimiter.Take()
+}
+
 // GetThreadsForPayloadRequests returns the number of threads to use as default for
 // given max-request of payloads
 func (e *ExecutorOptions) GetThreadsForNPayloadRequests(totalRequests int, currentThreads int) int {