feat(QueryLimits): Support bypassing query limits when prefer:wait=-1 header is provided BED-6991#2264
feat(QueryLimits): Support bypassing query limits when prefer:wait=-1 header is provided BED-6991#2264Useinovski merged 16 commits intomainfrom
Conversation
WalkthroughThreads a startup-read bypass-limits flag into Context and Logging middleware, extends Prefer: wait parsing to accept/validate negative (bypass) values, conditionally applies per-request timeouts or bypass behavior, updates middleware signatures and registration call sites, and adjusts tests for the new behavior. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Entrypoint
participant AppCfg as appcfg
participant DB
participant Router
participant ContextMW as ContextMiddleware
participant LoggingMW as LoggingMiddleware
participant Handler
Entrypoint->>AppCfg: GetTimeoutLimitParameter(context.Background(), db)
AppCfg->>DB: read timeout limit
AppCfg-->>Entrypoint: bypassLimitsParam
Entrypoint->>Router: register ContextMiddleware(bypassLimitsParam) and LoggingMiddleware(idResolver, bypassLimitsParam)
Client->>Router: HTTP Request
Router->>ContextMW: invoke middleware
ContextMW->>ContextMW: parse Prefer header (wait), validate against bypassLimitsParam
alt timeout applied
ContextMW->>Handler: set cancellable context with timeout and continue
else bypass enabled + wait=-1
ContextMW->>Handler: set bypass (no timeout) and continue
else invalid Prefer header
ContextMW-->>Client: respond 400 Bad Request
end
Handler->>LoggingMW: request flows through logging wrapper
LoggingMW->>LoggingMW: compute RequestWaitDuration(req, bypassLimitsParam)
LoggingMW-->>Client: response (with Request-ID and Preference-Applied / bypass header)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 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 |
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.
There was a problem hiding this comment.
I like this because it kicks out of the function sooner but we only return a -1 which resolves to -1 nano second. We would still need to multiply by time.Second() to set the requestedWaitDuration in func RequestWaitDuration() so that it is still being handled properly in the rest of the ContextMiddleware Func as well as logging.go
Please let me know if I missed something in my breakdown
There was a problem hiding this comment.
So basically we are not using the -1 as actual seconds. It is just a sentinel value. So here would be the flow:
- parsePreferHeaderWait returns -1 (the integer, as a sentinel)
- RequestWaitDuration returns -1
- In the middleware, we check:
if requestedWaitDuration == -1 - When it's -1, we don't call context.WithTimeout at all (no timeout enforcement)
- We store -1 in the bhCtx for logging/tracking
-1 is never being used as an actual time.Duration. Does this make sense?
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.
| 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.
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`:
- Line 63: Tests break because ContextMiddleware now has signature
ContextMiddleware(bool) which returns a middleware; update tests to call it with
the bool first and then pass the handler to the returned middleware. Replace
direct uses like middleware.ContextMiddleware(http.HandlerFunc(...)) with
wrappedHandler := middleware.ContextMiddleware(false)(http.HandlerFunc(func(rw
http.ResponseWriter, r *http.Request) {})) and then call
wrappedHandler.ServeHTTP(rr, req) (choose true/false per test intent); update
both test sites that construct the middleware (references to ContextMiddleware
and ServeHTTP in the test).
…or bypass check in ContextMiddleware BED-6991
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@cmd/api/src/api/middleware/middleware.go`:
- Around line 131-134: The error path calls api.BuildErrorResponse /
api.WriteErrorResponse while the BloodHound context (used by
ctx.FromRequest(request).RequestID) isn't initialized, so the response gets an
empty RequestID; fix by ensuring a minimal context containing the generated
requestID is attached to the request before invoking RequestWaitDuration error
handling (or move the context initialization earlier). Concretely, set up the
lightweight context with the requestID (the same value created at requestID
generation) on the request context used by RequestWaitDuration error handling so
api.BuildErrorResponse / api.WriteErrorResponse read the correct RequestID.
🧹 Nitpick comments (1)
cmd/api/src/api/middleware/middleware.go (1)
94-96: Consider extracting duplicate sentinel constants to package level.The bypass sentinel values (
bypassLimitValue/bypassLimitFlag = -1andbypassLimit = -1 second) are defined identically in bothRequestWaitDurationandContextMiddleware. Extracting these to package-level constants would improve maintainability and ensure consistency.♻️ Proposed refactor
Add at package level (e.g., after the imports):
const ( // bypassLimitSeconds is the sentinel value indicating bypass of query limits bypassLimitSeconds = -1 bypassLimitDuration = time.Second * time.Duration(bypassLimitSeconds) )Then reference these constants in both
RequestWaitDurationandContextMiddlewareinstead of the local definitions.Also applies to: 120-122
| if requestedWaitDuration, err := RequestWaitDuration(request, canBypassLimits); err != nil { | ||
| // If there is a failure or other expectation mismatch with the client, respond right away with the relevant | ||
| // error information | ||
| api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, fmt.Sprintf("Prefer header has an invalid value: %v", err), request), response) |
There was a problem hiding this comment.
Error response will have empty RequestID since context is not yet enriched.
When RequestWaitDuration returns an error, api.BuildErrorResponse is called, which internally invokes ctx.FromRequest(request).RequestID (see cmd/api/src/api/error.go lines 133-144). However, at this point the BloodHound context hasn't been set up yet (that happens in lines 156-167), so the error response will contain an empty or default RequestID rather than the generated requestID from line 128.
Consider setting a minimal context with the RequestID before handling the error, or restructuring to ensure the RequestID is available in error responses.
🐛 Proposed fix
if requestedWaitDuration, err := RequestWaitDuration(request, canBypassLimits); err != nil {
// If there is a failure or other expectation mismatch with the client, respond right away with the relevant
// error information
+ // Set minimal context so error response has the RequestID
+ errorCtx := ctx.Set(request.Context(), &ctx.Context{
+ RequestID: requestID,
+ })
+ request = request.WithContext(errorCtx)
api.WriteErrorResponse(request.Context(), api.BuildErrorResponse(http.StatusBadRequest, fmt.Sprintf("Prefer header has an invalid value: %v", err), request), response)
} else {🤖 Prompt for AI Agents
In `@cmd/api/src/api/middleware/middleware.go` around lines 131 - 134, The error
path calls api.BuildErrorResponse / api.WriteErrorResponse while the BloodHound
context (used by ctx.FromRequest(request).RequestID) isn't initialized, so the
response gets an empty RequestID; fix by ensuring a minimal context containing
the generated requestID is attached to the request before invoking
RequestWaitDuration error handling (or move the context initialization earlier).
Concretely, set up the lightweight context with the requestID (the same value
created at requestID generation) on the request context used by
RequestWaitDuration error handling so api.BuildErrorResponse /
api.WriteErrorResponse read the correct RequestID.
stephanieslamb
left a comment
There was a problem hiding this comment.
Looks good! Tested locally. Nicely done
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
Tests
✏️ Tip: You can customize this high-level summary in your review settings.