feat(remediation): support custom remediations #139
Conversation
- Replace iota-based Remediation enum with string-based system - Add weight system: Allow=0, Unknown=1, Captcha=10, Ban=20 - Implement string deduplication using pointers to reduce allocations - Add remediation_weights configuration option for custom remediations - Update RemediationMap to use string keys instead of enum - Update GetRemediationAndOrigin to compare weights for priority - Update all dataset and SPOA code to use string-based remediations - Fix metrics and AppSec checks to use > WeightAllow instead of > WeightUnknown to properly include custom remediations that default to Unknown weight Fixes #136
…lication Change RemediationMap from map[string]string to map[remediation.Remediation]string. This leverages the deduplicated string pointers in Remediation structs to automatically reduce memory allocations when storing remediations in maps. Benefits: - Automatic string deduplication via shared *string pointers - Reduced memory allocations in RemediationMap - Type-safe map keys - Efficient pointer-based comparisons Updated all methods (Add, Remove, HasRemediationWithOrigin, GetRemediationAndOrigin) and Contains methods to use remediation.Remediation types directly.
…n comparisons Tests were comparing remediation.Remediation with strings after Contains methods were updated to return remediation.Remediation directly. Updated assertions to use IsEqual() method for proper type-safe comparisons.
Replace all r.String() == "..." comparisons with r.IsEqual(remediation.X) for type-safe remediation comparisons. Updated switch statement to use if/else chain with IsEqual() checks.
Add nolint comment to suppress gocritic ifElseChain warning. We use if-else with IsEqual() for type-safe remediation comparisons instead of switch on String() to maintain type safety.
Replace if-else chain with switch statement on remediation.Remediation struct. Since Remediation is comparable (all fields are comparable), we can switch on it directly using the known remediation constants (Allow, Ban, Captcha).
Replace if-else chain with switch statement on remediation.Remediation struct. Since Remediation is comparable (all fields are comparable), we can switch on it directly using the known remediation constants (Allow, Ban, Captcha). This satisfies the linter's ifElseChain warning while maintaining type safety.
There was a problem hiding this comment.
Pull request overview
This PR refactors the remediation system from a simple uint8 enum to a struct-based approach with configurable weights to support custom remediations. The implementation aims to enable users to define custom remediation types (e.g., "MFA") with weights that slot between built-in types (Allow=0, Unknown=1, Captcha=10, Ban=20). However, the implementation has several critical bugs that will prevent it from working correctly.
Key Issues:
- Critical: Struct-based Remediations used as map keys will fail equality checks due to pointer comparison
- Critical: Direct
!=and switch statement comparisons won't work with struct types - Critical: SetWeight() configuration happens too late - built-in constants already have cached weights
- Missing: No test coverage for the new remediation implementation
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| internal/remediation/root.go | Complete rewrite of Remediation from uint8 to struct with weight-based comparison and string pointer deduplication. Introduces registry for managing weights and deduplicated strings. Critical bugs in map key usage and weight caching. |
| pkg/cfg/config.go | Adds RemediationWeights configuration field and applies custom weights via SetWeight(). Documentation missing; weight configuration timing issue prevents built-in weight overrides. |
| pkg/spoa/root.go | Updates remediation comparisons to use new IsEqual() and IsWeighted() methods. Direct != comparison on line 317 will fail with struct types. |
| pkg/dataset/types.go | Updates RemediationMap to use Remediation struct as keys with weight-based priority via IsHigher(). Map key comparison will fail due to pointer inequality. |
| pkg/dataset/root.go | Converts between Remediation and string representations in Add/Remove operations. Inconsistent approach between IP/Range (string) and Country (struct) operations. |
| pkg/dataset/ipmap.go | Changes operation structs to store remediation as string, converting back to Remediation when needed for map operations. |
| pkg/dataset/bart_types.go | Changes operation structs to store remediation as string, converting back to Remediation when needed for map operations. |
| pkg/dataset/root_test.go | Updates test assertions to use IsEqual() method for remediation comparison. |
| pkg/dataset/metrics_test.go | Updates test assertions to use IsEqual() method for remediation comparison. |
| pkg/dataset/benchmark_test.go | Updates benchmark code to use .String() method when creating test operations, and Compare() method for assertions. |
Comments suppressed due to low confidence (1)
pkg/dataset/root.go:134
- Inconsistency in how Remediation is passed to batch operations. IP and Range operations convert to string (
remediationName), but Country operations pass the Remediation struct directly (r). This makes the codebase harder to maintain and understand. Consider using a consistent approach across all operation types.
remediationName := r.String() // Convert to string for operations structs
switch scope {
case "ip":
ip, err := netip.ParseAddr(*decision.Value)
if err != nil {
log.Errorf("Error parsing IP address %s: %s", *decision.Value, err.Error())
continue
}
// Check for no-op: same IP, same remediation, same origin already exists
if d.IPMap.HasRemediation(ip, r, origin) {
// Exact duplicate - skip processing (no-op)
continue
}
ipType := "ipv4"
if ip.Is6() {
ipType = "ipv6"
}
// Check if we're overwriting an existing decision with different origin
if existingR, existingOrigin, found := d.IPMap.Contains(ip); found && existingR.IsEqual(r) && existingOrigin != origin {
// Decrement old origin's metric before incrementing new one
// Label order: origin, ip_type, scope (as defined in metrics.go)
metrics.TotalActiveDecisions.WithLabelValues(existingOrigin, ipType, "ip").Dec()
}
// Individual IPs go to IPMap for memory efficiency
ipOps = append(ipOps, IPAddOp{IP: ip, Origin: origin, R: remediationName, IPType: ipType})
// Label order: origin, ip_type, scope (as defined in metrics.go)
metrics.TotalActiveDecisions.WithLabelValues(origin, ipType, "ip").Inc()
case "range":
prefix, err := netip.ParsePrefix(*decision.Value)
if err != nil {
log.Errorf("Error parsing prefix %s: %s", *decision.Value, err.Error())
continue
}
// Check for no-op: same prefix, same remediation, same origin already exists
if d.RangeSet.HasRemediation(prefix, r, origin) {
// Exact duplicate - skip processing (no-op)
continue
}
ipType := "ipv4"
if prefix.Addr().Is6() {
ipType = "ipv6"
}
// Check if we're overwriting an existing decision with different origin
if existingOrigin, found := d.RangeSet.GetOriginForRemediation(prefix, r); found && existingOrigin != origin {
// Decrement old origin's metric before incrementing new one
// Label order: origin, ip_type, scope (as defined in metrics.go)
metrics.TotalActiveDecisions.WithLabelValues(existingOrigin, ipType, "range").Dec()
}
// Ranges go to BART for LPM support
rangeOps = append(rangeOps, BartAddOp{Prefix: prefix, Origin: origin, R: remediationName, IPType: ipType, Scope: "range"})
// Label order: origin, ip_type, scope (as defined in metrics.go)
metrics.TotalActiveDecisions.WithLabelValues(origin, ipType, "range").Inc()
case "country":
// Clone country code to break reference to Decision struct memory
cn := strings.Clone(*decision.Value)
// Check for no-op: same country, same remediation, same origin already exists
if d.CNSet.HasRemediation(cn, r, origin) {
// Exact duplicate - skip processing (no-op)
continue
}
// Check if we're overwriting an existing decision with different origin
if existingR, existingOrigin := d.CNSet.Contains(cn); existingR.IsEqual(r) && existingOrigin != "" && existingOrigin != origin {
// Decrement old origin's metric before incrementing new one
// Label order: origin, ip_type, scope (as defined in metrics.go)
metrics.TotalActiveDecisions.WithLabelValues(existingOrigin, "", "country").Dec()
}
cnOps = append(cnOps, cnOp{cn: cn, origin: origin, r: r})
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Fix IsEqual() to compare both name pointer and weight for true equality - Add comprehensive documentation for RemediationWeights configuration - Add clarifying comments about struct comparison and custom remediation handling - Note: Direct struct comparison (r != remediation.X) works because Remediation is comparable (all fields are comparable: *string and int)
- Add HasSameWeight() method as suggested by Copilot for checking weight equality - Update GetRemediationAndOrigin() to handle remediations with same weight using alphabetical order as a deterministic tie-breaker - Add documentation explaining tie-breaking behavior when remediations have same weight - This ensures deterministic behavior when users configure multiple remediations with the same weight value
Add documentation explaining that all Remediation instances must be created via New() or FromString() to ensure proper deduplication for map key equality. Direct struct initialization will create different string pointers, causing map lookups to fail. This addresses Copilot's concern about map key comparison - the deduplication mechanism ensures all Remediation instances with the same name share the same *string pointer, making map key equality work correctly.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 13 comments.
Comments suppressed due to low confidence (1)
pkg/dataset/types.go:97
- The
GetRemediationAndOrigin()function returns a zero-valuedRemediationwhen the map is empty. This zero-value hasname == nil, which could cause nil pointer dereferences when calling.String()or other methods.
Issue: If RemediationMap is empty, line 96 returns a zero-valued Remediation and an empty string. Calling .String() on this will hit the nil check on line 144 of internal/remediation/root.go and return "allow", but this is misleading - an empty map should not imply "allow".
Recommendation: Return an explicit value to indicate no remediation found, or check for empty map before calling this function. Consider:
func (rM RemediationMap) GetRemediationAndOrigin() (remediation.Remediation, string) {
if len(rM) == 0 {
return remediation.Allow, ""
}
// ... rest of implementation
}func (rM RemediationMap) GetRemediationAndOrigin() (remediation.Remediation, string) {
var maxRemediation remediation.Remediation
var maxOrigin string
first := true
for r, origin := range rM {
if first {
maxRemediation = r
maxOrigin = origin
first = false
continue
}
// Compare by weight first
if r.IsHigher(maxRemediation) {
maxRemediation = r
maxOrigin = origin
} else if r.HasSameWeight(maxRemediation) {
// Tie-breaker: use alphabetical order of name for deterministic behavior
if r.String() < maxRemediation.String() {
maxRemediation = r
maxOrigin = origin
}
}
}
return maxRemediation, maxOrigin
}
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…ment - Change Remediation from struct to string type - Move weight storage to registry level only - Add LoadWeights() function for startup initialization - Convert all comparison methods to registry-level functions - Update all call sites to use new function signatures - Simplify config loading to use LoadWeights() - All tests passing
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 19 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| // If no IP-specific remediation, check country-based remediation | ||
| if r < remediation.Unknown { | ||
| // If no IP-specific remediation (Allow), check country-based remediation | ||
| if remediation.IsEqual(r, remediation.Allow) { |
There was a problem hiding this comment.
Using remediation.IsEqual for comparison is more verbose than the direct equality check. Since Remediation is a comparable type (string), the direct comparison "r == remediation.Allow" is clearer and more idiomatic. The IsEqual function doesn't provide additional value here.
| if remediation.IsEqual(r, remediation.Allow) { | |
| if r == remediation.Allow { |
| // New creates a new Remediation from a string. | ||
| func New(name string) Remediation { | ||
| return Remediation(name) | ||
| } |
There was a problem hiding this comment.
The New function does not normalize the input string (e.g., to lowercase). This creates a critical inconsistency where "Ban", "ban", and "BAN" are treated as different remediations with potentially different weights. Since weight lookups are case-sensitive, this could lead to incorrect priority comparisons and unexpected behavior. The function should normalize the input, similar to how scope is normalized with strings.ToLower in pkg/dataset/root.go line 65.
| // Count blocked requests (not Allow) | ||
| // This includes Unknown, Captcha, Ban, and any custom remediations | ||
| if remediation.IsWeighted(r) { |
There was a problem hiding this comment.
The comment states "Count blocked requests (not Allow)" and "This includes Unknown, Captcha, Ban, and any custom remediations", but the condition uses IsWeighted which returns true for any remediation with weight > 0. This means Unknown (weight=1) is counted as a blocked request, which is semantically incorrect since Unknown represents remediations that are not supported/recognized and shouldn't be treated as active blocks.
| // Count blocked requests (not Allow) | |
| // This includes Unknown, Captcha, Ban, and any custom remediations | |
| if remediation.IsWeighted(r) { | |
| // Count blocked requests (not Allow or Unknown) | |
| // This includes Captcha, Ban, and any custom remediations | |
| if r != remediation.Allow && r != remediation.Unknown { |
| // LoadWeights loads weights for multiple remediations at once (for startup initialization) | ||
| func LoadWeights(weights map[string]int) { | ||
| globalRegistry.mu.Lock() | ||
| defer globalRegistry.mu.Unlock() | ||
|
|
||
| for name, weight := range weights { | ||
| globalRegistry.weights[name] = weight | ||
| } | ||
| } |
There was a problem hiding this comment.
The LoadWeights function does not normalize remediation names to lowercase. This creates inconsistency with the built-in remediations (initialized as lowercase in init) and can cause weight lookups to fail if the case doesn't match. The name keys should be normalized using strings.ToLower.
| // Remediation represents a remediation type as a string | ||
| type Remediation string |
There was a problem hiding this comment.
The PR description claims "Automatic deduplication: Shared *string pointers deduplicate map keys" and "Reduced allocations: Eliminates duplicate string headers in maps", but this implementation uses string values (not pointers) as the Remediation type. The claimed benefits about pointer-based deduplication are not realized by this design, making the documentation misleading.
| // LoadWeights loads weights for multiple remediations at once (for startup initialization) | ||
| func LoadWeights(weights map[string]int) { | ||
| globalRegistry.mu.Lock() | ||
| defer globalRegistry.mu.Unlock() | ||
|
|
||
| for name, weight := range weights { | ||
| globalRegistry.weights[name] = weight | ||
| } | ||
| } |
There was a problem hiding this comment.
The LoadWeights function does not validate inputs. Missing validation includes: (1) No check for negative weights, which violates the documented requirement that weights must be >= 0; (2) No normalization of remediation names to lowercase, creating case-sensitivity issues; (3) No protection against extremely large weights that could cause integer overflow in Compare. These issues could lead to unexpected behavior or security problems from malicious configuration.
| // Unknown or custom remediation: currently, only HTTP data is parsed for AppSec processing. | ||
| // If a custom remediation requires special handling (like Ban), this must be implemented explicitly. |
There was a problem hiding this comment.
The comment states "Unknown or custom remediation" but doesn't clarify the intended behavior. The code parses HTTP data for AppSec processing, but it's unclear whether custom remediations should always behave this way or if they need explicit configuration. The comment should clarify the design decision and expected behavior for custom remediations.
| // Unknown or custom remediation: currently, only HTTP data is parsed for AppSec processing. | |
| // If a custom remediation requires special handling (like Ban), this must be implemented explicitly. | |
| // Unknown or custom remediation: | |
| // By default, only HTTP data is parsed for AppSec processing for unknown or custom remediations. | |
| // If you introduce a custom remediation that requires special handling (e.g., setting cookies, injecting key-values, etc.), | |
| // you must implement that logic explicitly here. Otherwise, only HTTP data will be available for AppSec. |
| // Load custom remediation weights if configured (loads all weights at once on startup) | ||
| if config.RemediationWeights != nil { | ||
| remediation.LoadWeights(config.RemediationWeights) | ||
| } |
There was a problem hiding this comment.
The RemediationWeights configuration is loaded without validation. Malicious or misconfigured values (e.g., negative weights, extremely large weights that could cause integer overflow in comparisons) are not checked before being passed to remediation.LoadWeights. This could lead to unexpected behavior or security issues.
| Prefix netip.Prefix | ||
| Origin string | ||
| R remediation.Remediation | ||
| R string // Remediation name as string |
There was a problem hiding this comment.
The operation struct now stores remediation as a string instead of a Remediation type, then converts it back using FromString on every operation (lines 104, 149, 157, 213, 232). This introduces unnecessary string-to-Remediation conversions. Consider storing Remediation type directly in the operation struct to avoid repeated conversions.
| // defer a function that always add the remediation to the request at end of processing | ||
| defer func() { | ||
| if matchedHost == nil && r == remediation.Captcha { | ||
| if matchedHost == nil && remediation.IsEqual(r, remediation.Captcha) { |
There was a problem hiding this comment.
Using remediation.IsEqual for comparison is more verbose than the direct equality check. Since Remediation is a comparable type (string), the direct comparison "r == remediation.Captcha" is clearer and more idiomatic. The IsEqual function doesn't provide additional value here.
| if matchedHost == nil && remediation.IsEqual(r, remediation.Captcha) { | |
| if matchedHost == nil && r == remediation.Captcha { |
Allow custom
remediation.RemediationChanges
Benefits
*stringpointers deduplicate map keysuntested purely draft for now