feat(QueryLimits): Support bypassing query limits when prefer:wait=-1 header is provided BED-6991 #2264
Conversation
Useinovski
added
the
api
WalkthroughThis PR adds bypass-limit handling for the Prefer header in API middleware. It introduces validation logic for requested wait durations, a helper function to normalize timeouts when they equal the bypass limit, and corresponding test coverage for negative values and bypass behavior. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@cmd/api/src/api/middleware/middleware_internal_test.go`:
- Around line 126-135: In the test cases where only the error is being asserted,
avoid ineffectual assignments by replacing the left-hand variable with the blank
identifier; specifically, change assignments like `duration, err =
parsePreferHeaderWait("5")` and `duration, err = parsePreferHeaderWait("five")`
to `_ , err = parsePreferHeaderWait("5")` and `_ , err =
parsePreferHeaderWait("five")` respectively so only `err` is used; the other
assertions (e.g., the valid parse using `duration, err =
parsePreferHeaderWait("wait=-1")`) should remain unchanged and still reference
`duration` and `err`.
🧹 Nitpick comments (1)
cmd/api/src/config/config.go (1)
173-175: GuardloadedConfigaccess if configuration can reload after startup.
GetLoadedConfig()is read per request; ifGetConfiguration()is ever called concurrently, this introduces a data race. Please confirm initialization order or add synchronization.🛠️ Example with atomic.Value
+import "sync/atomic" ... -var ( - loadedConfig Configuration -) +var loadedConfig atomic.Value // stores Configuration ... -func GetLoadedConfig() Configuration { - return loadedConfig -} +func GetLoadedConfig() Configuration { + if cfg, ok := loadedConfig.Load().(Configuration); ok { + return cfg + } + return Configuration{} +} ... - loadedConfig = cfg + loadedConfig.Store(cfg)Also applies to: 283-295
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
cmd/api/src/api/middleware/logging.gocmd/api/src/api/middleware/middleware.gocmd/api/src/api/middleware/middleware_internal_test.gocmd/api/src/config/config.go
💤 Files with no reviewable changes (1)
- cmd/api/src/api/middleware/logging.go
🧰 Additional context used
🧬 Code graph analysis (2)
cmd/api/src/api/middleware/middleware.go (2)
cmd/api/src/config/config.go (1)
GetLoadedConfig(283-285)cmd/api/src/ctx/ctx.go (3)
Get(75-85)Set(88-90)RequestID(93-95)
cmd/api/src/api/middleware/middleware_internal_test.go (1)
cmd/api/src/api/middleware/middleware.go (1)
RequestWaitDuration(86-104)
🪛 golangci-lint (2.5.0)
cmd/api/src/api/middleware/middleware_internal_test.go
[major] 130-130: ineffectual assignment to duration
(ineffassign)
[major] 133-133: ineffectual assignment to duration
(ineffassign)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Build BloodHound Container Image / Build and Package Container
- GitHub Check: build-ui
- GitHub Check: run-tests
- GitHub Check: run-analysis
🔇 Additional comments (5)
cmd/api/src/api/middleware/middleware_internal_test.go (2)
66-77: Good negative-wait coverage.
Covers invalid negative wait values and asserts zero duration on error.
137-143: Nice coverage for bypass timeout normalization.
The test clearly verifies the bypass and normal paths forsetUserTimeout.cmd/api/src/api/middleware/middleware.go (3)
86-103: Bypass validation for Prefer wait looks solid.
The guardrails for< -1and config-gated-1behavior are clear.
110-152: Preference-Applied header + timeout normalization are consistent.
The bypass signaling and normalized timeout in context align with the new semantics.
167-173: Helper is clear and focused.
setUserTimeoutcleanly normalizes the bypass value.
✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.
| } else if requestedWaitDuration < bypassLimit { | ||
| return 0, errors.New("incorrect bypass limit value") | ||
| } else if requestedWaitDuration == bypassLimit && !canBypassLimits { | ||
| return 0, errors.New("failed to bypass limits") |
There was a problem hiding this comment.
Let's add some detail to the return statement here: failed to bypass limits: feature disabled
| StartTime: startTime, | ||
| RequestID: requestID, | ||
| Timeout: requestedWaitDuration, | ||
| Timeout: setUserTimeout(requestedWaitDuration, bypassLimit), |
There was a problem hiding this comment.
Setting a zero timeout here will mean the context is already expired. When you want to bypass the timeout, you will instead need to create a new context an empty (NOT zero) Timeout field. Here, check out how these dynamics play out:
https://go.dev/play/p/pWTyGNo0R-U
You'll need to use the third block for the case where the limit is being bypassed.
There was a problem hiding this comment.
I might be wrong but I think we can just remove the setUserTimeout call and this will work as expected. On line 141 we're setting context.WithTimeout(Go context) to whatever the requestedWaitDuration is. If we're bypassing timeouts, we have our sentinel value of -1, or if no timeout is specified it will be 0. Since we're just setting the Timeout field in bhCtx(not Go context), we can store the value directly with no conversions.
There was a problem hiding this comment.
I believe I agree with Stephanie here. I believe contextWithTimeout() is still responsible for handling the timeout and not the Timeout: field in our struct.
We have a "Context" struct in ctx.go that stores a time.duration but doesn't seem to do anything with it beyond that. I believe for logging purposes? But I don't see anything indicating that the Timeout: field is responsible for handling the timeout.
There was a problem hiding this comment.
My logic behind the setUserTimeouts() function was to not pass around negative values in Context for the Timeout: field. I don't believe a negative, 0, or positive time.duration() value will directly affect the functionality. I'm more than happy to revert it though if we think this is overly excessive.
There was a problem hiding this comment.
I did indeed miss the block above this, where we're setting timeout only if requestedWaitDuration>0. I'm good with the changes here. As for setUserTimeout you can use go's builtin functions to achieve this without having to declare a separate function.
time.Duration(max(requestedWaitDuration.Seconds(), 0)) should give you 0 seconds
| requestCtx, cancel = context.WithTimeout(request.Context(), requestedWaitDuration) | ||
| defer cancel() | ||
| } else if requestedWaitDuration == bypassLimit && canBypassLimits { | ||
| response.Header().Set(headers.PreferenceApplied.String(), fmt.Sprintf("wait=%.2f; bypass=enabled", requestedWaitDuration.Seconds())) |
There was a problem hiding this comment.
Also, since we aren't using time, this can just be something like:
`response.Header().Set(headers.PreferenceApplied.String(), "wait=-1; bypass=enabled")
| requestID string | ||
| canBypassLimits = config.GetLoadedConfig().DisableTimeoutLimit | ||
| ) | ||
| const bypassLimit = time.Second * time.Duration(-1) |
| err error | ||
| canBypassLimits = config.GetLoadedConfig().DisableTimeoutLimit | ||
| ) | ||
| const bypassLimit = time.Second * time.Duration(-1) |
There was a problem hiding this comment.
I like this to avoid magic numbers! But I think since this is just a sentinel value, we dont need to convert to seconds. It can just be const bypassLimit = -1
There was a problem hiding this comment.
We set requestedWaitDuration, err = parsePreferHeaderWait(preferValue) and parsePreferHeaderWait will return a time.Duration.
Therefore the if check of requrestedWaitDuration < bypassLimit we get incorrect behavior always resulting in true if we pass in a -1. This should only be true for values less than -1
-1000000000 < -1 is the comparison
I also tried to change the return value we get from parsePreferHeaderWait() to return only a time.Duration and not a time.Second * time.duration() but resulted in getting a 401 "Token Authorization failed" so I thought this was the best approach to not breaking anything else in the application.
There was a problem hiding this comment.
Might be worthwhile adding another constant to make that clear.
bypassLimitFlag=-1
bypassLimit = time.Second * time.Duration(bypassLimitFlag)
)
or if you have a better name for it feel free to use something else
There was a problem hiding this comment.
So in parsePreferHeaderWait, we can add a check if we are bypassing the limit and just return -1. Something like this:
func parsePreferHeaderWait(value string) (time.Duration, error) {
if waitStr, hasWaitSpecified := ParseHeaderValues(value)["wait"]; hasWaitSpecified {
if parsedNumSeconds, err := strconv.Atoi(waitStr); err != nil {
return 0, err
} else {
if parsedNumSeconds == -1 {
return -1, nil
}
return time.Second * time.Duration(parsedNumSeconds), nil
}
} else {
return 0, errors.New("leave field empty or specify with : 'wait=x'")
}
}
Just add a comment that the check is for bypassing query limits.
irshadaj
left a comment
irshadaj
left a comment
There was a problem hiding this comment.
Approved pending 2 suggestions added in comments
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@cmd/api/src/api/middleware/middleware.go`:
- Around line 145-146: The branch checking requestedWaitDuration == bypassLimit
is currently unreachable because RequestWaitDuration rejects the special wait=-1
value when the bypass feature flag is commented out; to fix this, restore or
implement the feature flag check inside RequestWaitDuration so that
RequestWaitDuration accepts bypassLimit (e.g., wait=-1) when the bypass feature
is enabled, and ensure downstream code (the branch that sets
headers.PreferenceApplied via response.Header().Set("Preference-Applied",
fmt.Sprintf("wait=-1; bypass=enabled"))) is only hit when RequestWaitDuration
returns without error for that value; alternatively, if the feature is
intentionally disabled, remove the unreachable branch to avoid dead code.
- Around line 97-100: RequestWaitDuration currently rejects the bypass value
unconditionally because the canBypassLimits feature-check was commented out;
restore and use the canBypassLimits check (or call the existing
GetTimeoutLimitParameter/api.timeout_limit feature flag helper) inside
RequestWaitDuration so that when requestedWaitDuration == bypassLimit (e.g., -1)
you only return the "failed to bypass limits" error if canBypassLimits is false,
otherwise accept the bypass and return the special wait value; also make
ContextMiddleware and RequestWaitDuration consistent by ensuring
ContextMiddleware only sets wait=-1 and bypass=true when the same
canBypassLimits flag (via GetTimeoutLimitParameter/api.timeout_limit) yields
true.
In `@cmd/api/src/model/appcfg/parameter.go`:
- Around line 529-539: The TimeoutLimitParameter.Enabled field name is ambiguous
about allowing timeout bypass; update its semantic clarity by either renaming
the field (e.g., TimeoutLimitParameter.BypassAllowed or
TimeoutLimitParameter.TimeoutBypassEnabled) and adjusting all references
(notably GetTimeoutLimitParameter and any mapping code that expects Enabled), or
add a clear comment on the Enabled field stating that Enabled:true means clients
may bypass timeout limits via wait=-1; ensure the chosen name/comment is
reflected in GetTimeoutLimitParameter's default (currently Enabled: true) and
any middleware checks that consume this parameter.
🧹 Nitpick comments (3)
cmd/api/src/api/middleware/middleware.go (3)
90-91: Remove commented-out debug code before merging.These commented-out lines (
//canBypassLimits = ...,//fmt.Println...) appear to be work-in-progress artifacts. They should be removed to keep the codebase clean, or if the feature flag integration is pending, consider adding a TODO comment with the ticket reference instead.Suggested cleanup
- //canBypassLimits = appcfg.GetTimeoutLimitParameter(request.Context(), ___) + // TODO(BED-6991): Wire up feature flag check + // canBypassLimits := appcfg.GetTimeoutLimitParameter(request.Context(), service)Also applies to: 116-117
92-92: Extract duplicatedbypassLimitconstant to package level.The
bypassLimitconstant is defined identically in bothRequestWaitDurationandContextMiddleware. Extract it to a package-level constant to avoid duplication and ensure consistency.Suggested refactor
+const bypassLimit = time.Second * time.Duration(-1) + func RequestWaitDuration(request *http.Request) (time.Duration, error) { var ( requestedWaitDuration time.Duration err error ) - const bypassLimit = time.Second * time.Duration(-1)Then remove the duplicate declaration from
ContextMiddlewareas well.Also applies to: 114-114
169-175: SimplifysetUserTimeoutusing built-inmaxfunction.Per past review discussion, this can be simplified. However, note that the current logic returns
0whenduration == limit(i.e.,-1s), which is the intended behavior for bypass. The function is correct but could be more concise.Alternative using max (if Go 1.21+)
-func setUserTimeout(duration time.Duration, limit time.Duration) time.Duration { - if duration == limit { - return time.Second * time.Duration(0) - } - return duration -} +func setUserTimeout(duration time.Duration, limit time.Duration) time.Duration { + if duration == limit { + return 0 + } + return duration +}Note:
time.Second * time.Duration(0)is equivalent to just0fortime.Duration.
| } else if requestedWaitDuration < bypassLimit { | ||
| return 0, errors.New("incorrect bypass limit value") | ||
| } else if requestedWaitDuration == bypassLimit { // revert back to check canBypassLimits | ||
| return 0, errors.New("failed to bypass limits: feature disabled") |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, let's check the file and see lines around 97-100
echo "=== Checking middleware.go around lines 97-100 ==="
cat -n cmd/api/src/api/middleware/middleware.go | sed -n '80,115p'Repository: SpecterOps/BloodHound
Length of output: 186
🏁 Script executed:
#!/bin/bash
# Check for GetTimeoutLimitParameter usages
rg -n "GetTimeoutLimitParameter" --type goRepository: SpecterOps/BloodHound
Length of output: 436
🏁 Script executed:
#!/bin/bash
# Check for canBypassLimits references
rg -n "canBypassLimits" --type goRepository: SpecterOps/BloodHound
Length of output: 682
🏁 Script executed:
#!/bin/bash
# Look for feature flag related to timeout/bypass in the middleware file
rg -n "disable_timeout_limit|bypass|feature" cmd/api/src/api/middleware/middleware.go -A 2 -B 2Repository: SpecterOps/BloodHound
Length of output: 2245
🏁 Script executed:
#!/bin/bash
# Check the GetTimeoutLimitParameter function definition
cat -n cmd/api/src/model/appcfg/parameter.go | sed -n '520,550p'Repository: SpecterOps/BloodHound
Length of output: 1331
🏁 Script executed:
#!/bin/bash
# Search for disable_timeout_limit references
rg -n "disable_timeout_limit" --type goRepository: SpecterOps/BloodHound
Length of output: 47
🏁 Script executed:
#!/bin/bash
# Search for any references to timeout limit configuration
rg -n "TimeoutLimit|timeout.*limit|disable.*timeout" --type go -i | head -30Repository: SpecterOps/BloodHound
Length of output: 2415
🏁 Script executed:
#!/bin/bash
# Check the full context of the bypass logic in middleware - lines 80-160
cat -n cmd/api/src/api/middleware/middleware.go | sed -n '80,160p'Repository: SpecterOps/BloodHound
Length of output: 3964
Feature flag check is commented out - bypass is incompletely implemented.
The canBypassLimits check is commented out at lines 90 and 112, causing wait=-1 to always fail with "failed to bypass limits: feature disabled" in RequestWaitDuration. While the underlying configuration (api.timeout_limit) and GetTimeoutLimitParameter function exist, the feature flag integration is incomplete.
Additionally, there's a logic inconsistency: RequestWaitDuration (line 99) rejects the bypass request unconditionally, but ContextMiddleware (line 145) accepts it and sets wait=-1; bypass=enabled. This needs to be resolved alongside integrating the deferred feature flag check.
🤖 Prompt for AI Agents
In `@cmd/api/src/api/middleware/middleware.go` around lines 97 - 100,
RequestWaitDuration currently rejects the bypass value unconditionally because
the canBypassLimits feature-check was commented out; restore and use the
canBypassLimits check (or call the existing
GetTimeoutLimitParameter/api.timeout_limit feature flag helper) inside
RequestWaitDuration so that when requestedWaitDuration == bypassLimit (e.g., -1)
you only return the "failed to bypass limits" error if canBypassLimits is false,
otherwise accept the bypass and return the special wait value; also make
ContextMiddleware and RequestWaitDuration consistent by ensuring
ContextMiddleware only sets wait=-1 and bypass=true when the same
canBypassLimits flag (via GetTimeoutLimitParameter/api.timeout_limit) yields
true.
| func GetTimeoutLimitParameter(ctx context.Context, service ParameterService) bool { | ||
| result := TimeoutLimitParameter{Enabled: true} | ||
|
|
||
| if cfg, err := service.GetConfigurationParameter(ctx, TimeoutLimit); err != nil { | ||
| slog.WarnContext(ctx, "Failed to fetch timeout limit configuration; returning default values") | ||
| } else if err := cfg.Map(&result); err != nil { | ||
| slog.WarnContext(ctx, fmt.Sprintf("Invalid timeout limit configuration supplied, %v. returning default values.", err)) | ||
| } | ||
|
|
||
| return result.Enabled | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check how this parameter is intended to be used in middleware
rg -n "GetTimeoutLimitParameter|TimeoutLimit" --type go -C 3Repository: SpecterOps/BloodHound
Length of output: 5015
🏁 Script executed:
#!/bin/bash
# Find the setUserTimeout function to understand bypass logic
rg -n "setUserTimeout\|canBypassLimits\|disable_timeout_limit" --type go -B 3 -A 5Repository: SpecterOps/BloodHound
Length of output: 47
🏁 Script executed:
#!/bin/bash
# Search for setUserTimeout function definition
rg -n "func setUserTimeout" --type go -A 10Repository: SpecterOps/BloodHound
Length of output: 870
🏁 Script executed:
#!/bin/bash
# Look at the middleware.go file around the commented code
cat -n cmd/api/src/api/middleware/middleware.go | sed -n '85,120p'Repository: SpecterOps/BloodHound
Length of output: 1840
Clarify the semantic meaning of TimeoutLimitParameter.Enabled.
The field name is ambiguous. Based on the middleware code (lines 99-100, 112), GetTimeoutLimitParameter is intended to control whether timeout bypass (via wait=-1) is allowed. The default Enabled: true means bypass is allowed by default, but the field name doesn't clearly convey this.
Consider renaming Enabled to BypassAllowed or TimeoutBypassEnabled, or add a clarifying comment explaining that Enabled: true means "users can bypass timeout limits."
Note: The function is currently unused in active code (calls are commented out at lines 90 and 112 of middleware.go pending integration).
🤖 Prompt for AI Agents
In `@cmd/api/src/model/appcfg/parameter.go` around lines 529 - 539, The
TimeoutLimitParameter.Enabled field name is ambiguous about allowing timeout
bypass; update its semantic clarity by either renaming the field (e.g.,
TimeoutLimitParameter.BypassAllowed or
TimeoutLimitParameter.TimeoutBypassEnabled) and adjusting all references
(notably GetTimeoutLimitParameter and any mapping code that expects Enabled), or
add a clear comment on the Enabled field stating that Enabled:true means clients
may bypass timeout limits via wait=-1; ensure the chosen name/comment is
reflected in GetTimeoutLimitParameter's default (currently Enabled: true) and
any middleware checks that consume this parameter.
Description
If the feature flag for query-limit bypass is enabled and the incoming request includes the header wait: -1, we should skip all query-limit checks and allow the request to execute without restriction.
We set in the configuration parameter in build.config.json '"disable_timeout_limit" : true' which will allow the user to bypass query limits on all end points.
Motivation and Context
Resolves BED-6991
To be able to bypass query limits for longer running queries that would be timed out with the default timeout set from the UI.
How Has This Been Tested?
Tested in Bruno, with all cases of the Prefer header value. When using the UI, we have a default prefer header set by client.ts which will be handled in future PR to be removed (BED-6985). Unit tests have also been created and or modified.
This default value is only applied in the search endpoint via the explore page.
"disable_timeout_limit" : true
Values tested: Corresponding Status code:
Prefer : "wait=40" 200/201 - Sets timeout of 40 seconds
Prefer : "wait=0" 200/201 - Bypasses query Limits (current behavior)
Prefer : "wait=-1" 200/201 - Bypasses query limits
Prefer : "wait=-30" 400 - Do not allow for negative values < -1
"disable_timeout_limit" : false
Values tested: Corresponding Status code:
Prefer : "wait=40" 200/201 - Sets timeout of 40 seconds
Prefer : "wait=0" 200/201 - Bypasses query Limits (current behavior)
Prefer : "wait=-1" 400 - Failed to bypass limits
Prefer : "wait=-30" 400 - Do not allow for negative values < -1
Screenshots (optional):
"disable_timeout_limit" : true
"disable_timeout_limit" : false
From UI
Types of changes
Checklist:
Summary by CodeRabbit
New Features
Bug Fixes
✏️ Tip: You can customize this high-level summary in your review settings.