feat(remediation): support custom remediations

internal/remediation/root.go

@@ -1,37 +1,140 @@
 package remediation
 
-// The order matters since we use slices.Max to get the max value
+import (
+	"sync"
+)
+
+// Default weights for built-in remediations
+// Allow=0, Unknown=1, then expand others to allow custom remediations to slot in between
+const (
+	WeightAllow   = 0
+	WeightUnknown = 1
+	WeightCaptcha = 10
+	WeightBan     = 20
+)
+
+// Remediation represents a remediation type as a string
+type Remediation string
+
+// Built-in remediation constants
 const (
-	Allow   Remediation = iota // Allow remediation
-	Unknown                    // Unknown remediation (Unknown is used to have a value for remediation we don't support EG "MFA")
-	Captcha                    // Captcha remediation
-	Ban                        // Ban remediation
+	Allow   Remediation = "allow"
+	Unknown Remediation = "unknown"
+	Captcha Remediation = "captcha"
+	Ban     Remediation = "ban"
 )
 
-type Remediation uint8 // Remediation type is smallest uint to save space
+// registry manages remediation weights
+type registry struct {
+	mu      sync.RWMutex
+	weights map[string]int // Maps remediation name to its weight
+}
+
+var globalRegistry = &registry{
+	weights: make(map[string]int),
+}
+
+//nolint:gochecknoinits // init() is required to initialize default weights
+func init() {
+	// Initialize built-in remediations with default weights
+	globalRegistry.mu.Lock()
+	defer globalRegistry.mu.Unlock()
+
+	globalRegistry.weights["allow"] = WeightAllow
+	globalRegistry.weights["unknown"] = WeightUnknown
+	globalRegistry.weights["captcha"] = WeightCaptcha
+	globalRegistry.weights["ban"] = WeightBan
+}
+
+// SetWeight sets a custom weight for a remediation (for configuration)
+func SetWeight(name string, weight int) {
+	globalRegistry.mu.Lock()
+	defer globalRegistry.mu.Unlock()
+
+	globalRegistry.weights[name] = weight
+}
+
+// GetWeight returns the weight for a remediation name
+func GetWeight(name string) int {
+	globalRegistry.mu.RLock()
+	defer globalRegistry.mu.RUnlock()
+
+	if weight, exists := globalRegistry.weights[name]; exists {
+		return weight
+	}
+	// Default to Unknown weight for unknown remediations
+	return WeightUnknown
+}
+
+// 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
+	}
+}
+
+// New creates a new Remediation from a string.
+func New(name string) Remediation {
+	return Remediation(name)
+}
+
+// String returns the remediation name
 func (r Remediation) String() string {
-	switch r {
-	case Ban:
-		return "ban"
-	case Captcha:
-		return "captcha"
-	case Unknown:
-		return "unknown"
-	default:
-		return "allow"
+	if r == "" {
+		return "allow" // Default fallback
 	}
+	return string(r)
+}
+
+// Compare returns:
+// - negative if a < b
+// - zero if a == b
+// - positive if a > b
+func Compare(a, b Remediation) int {
+	weightA := GetWeight(a.String())
+	weightB := GetWeight(b.String())
+	return weightA - weightB
+}
+
+// IsHigher returns true if a has a higher weight than b
+func IsHigher(a, b Remediation) bool {
+	return Compare(a, b) > 0
+}
+
+// IsLower returns true if a has a lower weight than b
+func IsLower(a, b Remediation) bool {
+	return Compare(a, b) < 0
+}
+
+// IsEqual returns true if a represents the same remediation as b.
+// This compares the remediation names (strings).
+func IsEqual(a, b Remediation) bool {
+	return a == b
+}
+
+// HasSameWeight returns true if a has the same weight as b.
+// This is useful for checking if two different remediations have the same priority.
+// Note: Two remediations with the same weight will be compared by name (alphabetical)
+// as a tie-breaker when determining priority.
+func HasSameWeight(a, b Remediation) bool {
+	return Compare(a, b) == 0
 }
 
+// IsWeighted returns true if r is not Allow (has weight > Allow)
+// This is useful for checking if a remediation should be applied
+func IsWeighted(r Remediation) bool {
+	return GetWeight(r.String()) > WeightAllow
+}
+
+// FromString creates a Remediation from a string (alias for New for backward compatibility)
 func FromString(s string) Remediation {
-	switch s {
-	case "ban":
-		return Ban
-	case "captcha":
-		return Captcha
-	case "allow":
-		return Allow
-	default:
-		return Unknown
-	}
+	return New(s)
+}
+
+// IsZero returns true if the remediation is zero-valued
+func (r Remediation) IsZero() bool {
+	return r == ""
 }

pkg/cfg/config.go

@@ -7,6 +7,7 @@ import (
 	"gopkg.in/yaml.v2"
 
 	"github.com/crowdsecurity/crowdsec-spoa/internal/geo"
+	"github.com/crowdsecurity/crowdsec-spoa/internal/remediation"
 	"github.com/crowdsecurity/crowdsec-spoa/pkg/host"
 	cslogging "github.com/crowdsecurity/crowdsec-spoa/pkg/logging"
 	"github.com/crowdsecurity/go-cs-lib/csyaml"
@@ -36,6 +37,32 @@ type BouncerConfig struct {
 	ListenUnix       string                  `yaml:"listen_unix"`
 	PrometheusConfig PrometheusConfig        `yaml:"prometheus"`
 	PprofConfig      PprofConfig             `yaml:"pprof"`
+	// RemediationWeights allows users to configure custom weights for remediations.
+	//
+	// Format:
+	//   remediation_weights:
+	//     <remediation_name>: <weight>
+	//
+	// Example:
+	//   remediation_weights:
+	//     mfa: 15   # slots between captcha (10) and ban (20)
+	//
+	// Valid weight range: integer values >= 0. Lower values are less severe; higher values are more severe.
+	// Recommended: Use values between 0 and 100.
+	//
+	// Built-in defaults:
+	//   allow=0, unknown=1, captcha=10, ban=20
+	//
+	// Custom weights override or supplement built-in remediations. If a custom remediation is defined,
+	// its weight will be used for ordering and severity. Custom remediations can slot between built-in
+	// ones by choosing an appropriate weight value.
+	//
+	// Tie-breaking: If two remediations have the same weight, alphabetical order of the remediation
+	// name is used as a deterministic tie-breaker when determining priority.
+	//
+	// Note: Custom weights for built-in remediations (allow, unknown, captcha, ban) must be set
+	// before package initialization. After init(), package-level constants already have cached weights.
+	RemediationWeights map[string]int `yaml:"remediation_weights,omitempty"`
 }
 
 // MergedConfig() returns the byte content of the patched configuration file (with .yaml.local).
@@ -67,6 +94,11 @@ func NewConfig(reader io.Reader) (*BouncerConfig, error) {
 		return nil, fmt.Errorf("failed to setup logging: %w", err)
 	}
 
+	// Load custom remediation weights if configured (loads all weights at once on startup)
+	if config.RemediationWeights != nil {
+		remediation.LoadWeights(config.RemediationWeights)
+	}
+
 	if err := config.Validate(); err != nil {
 		return nil, err
 	}

pkg/dataset/bart_types.go

@@ -15,15 +15,15 @@ import (
 type BartAddOp struct {
 	Prefix netip.Prefix
 	Origin string
-	R      remediation.Remediation
+	R      string // Remediation name as string
 	IPType string
 	Scope  string
 }
 
 // BartRemoveOp represents a single prefix removal operation for batch processing
 type BartRemoveOp struct {
 	Prefix netip.Prefix
-	R      remediation.Remediation
+	R      string // Remediation name as string
 	Origin string
 	IPType string
 	Scope  string
@@ -91,7 +91,7 @@ func (s *BartRangeSet) initializeBatch(operations []BartAddOp) {
 		// Only build logging fields if trace level is enabled
 		var valueLog *log.Entry
 		if s.logger.Logger.IsLevelEnabled(log.TraceLevel) {
-			valueLog = s.logger.WithField("prefix", prefix.String()).WithField("remediation", op.R.String())
+			valueLog = s.logger.WithField("prefix", prefix.String()).WithField("remediation", op.R)
 			valueLog.Trace("initial load: collecting prefix operations")
 		}
 
@@ -101,7 +101,7 @@ func (s *BartRangeSet) initializeBatch(operations []BartAddOp) {
 			data = RemediationMap{}
 		}
 		// Add the remediation (this handles merging if prefix already seen)
-		data.Add(valueLog, op.R, op.Origin)
+		data.Add(valueLog, remediation.FromString(op.R), op.Origin)
 		prefixMap[prefix] = data
 	}
 
@@ -134,7 +134,7 @@ func (s *BartRangeSet) updateBatch(cur *bart.Table[RemediationMap], operations [
 		// Only build logging fields if trace level is enabled
 		var valueLog *log.Entry
 		if s.logger.Logger.IsLevelEnabled(log.TraceLevel) {
-			valueLog = s.logger.WithField("prefix", prefix.String()).WithField("remediation", op.R.String())
+			valueLog = s.logger.WithField("prefix", prefix.String()).WithField("remediation", op.R)
 			valueLog.Trace("adding to bart trie")
 		}
 
@@ -146,15 +146,15 @@ func (s *BartRangeSet) updateBatch(cur *bart.Table[RemediationMap], operations [
 					valueLog.Trace("exact prefix exists, merging remediations")
 				}
 				// bart already cloned via our Cloner interface, modify directly
-				existingData.Add(valueLog, op.R, op.Origin)
+				existingData.Add(valueLog, remediation.FromString(op.R), op.Origin)
 				return existingData, false // false = don't delete
 			}
 			if valueLog != nil {
 				valueLog.Trace("creating new entry")
 			}
 			// Create new data
 			newData := make(RemediationMap)
-			newData.Add(valueLog, op.R, op.Origin)
+			newData.Add(valueLog, remediation.FromString(op.R), op.Origin)
 			return newData, false // false = don't delete
 		})
 	}
@@ -193,7 +193,7 @@ func (s *BartRangeSet) RemoveBatch(operations []BartRemoveOp) []*BartRemoveOp {
 		// Only build logging fields if trace level is enabled
 		var valueLog *log.Entry
 		if s.logger.Logger.IsLevelEnabled(log.TraceLevel) {
-			valueLog = s.logger.WithField("prefix", prefix.String()).WithField("remediation", op.R.String())
+			valueLog = s.logger.WithField("prefix", prefix.String()).WithField("remediation", op.R)
 			valueLog.Trace("removing from bart trie")
 		}
 
@@ -210,11 +210,12 @@ func (s *BartRangeSet) RemoveBatch(operations []BartRemoveOp) []*BartRemoveOp {
 
 			// Check if the remediation exists with the matching origin before removing
 			// This prevents removing decisions when the origin has been overwritten (e.g., by CAPI)
-			if !existingData.HasRemediationWithOrigin(op.R, op.Origin) {
+			if !existingData.HasRemediationWithOrigin(remediation.FromString(op.R), op.Origin) {
 				// Origin doesn't match - this decision was likely overwritten by another origin
 				// Don't remove it, as it's not the decision we're trying to delete
 				if valueLog != nil {
-					storedOrigin, exists := existingData[op.R]
+					r := remediation.FromString(op.R)
+					storedOrigin, exists := existingData[r]
 					if exists {
 						valueLog.Tracef("remediation exists but origin mismatch (stored: %s, requested: %s), skipping removal", storedOrigin, op.Origin)
 					} else {
@@ -228,7 +229,7 @@ func (s *BartRangeSet) RemoveBatch(operations []BartRemoveOp) []*BartRemoveOp {
 			// bart already cloned via our Cloner interface, modify directly
 			// Remove returns an error if remediation doesn't exist (duplicate delete)
 			// We already checked origin above, so this should succeed
-			err := existingData.Remove(valueLog, op.R)
+			err := existingData.Remove(valueLog, remediation.FromString(op.R))
 			if errors.Is(err, ErrRemediationNotFound) {
 				// This shouldn't happen since we checked above, but handle it gracefully
 				if valueLog != nil {
@@ -288,11 +289,11 @@ func (s *BartRangeSet) Contains(ip netip.Addr) (remediation.Remediation, string)
 		return remediation.Allow, ""
 	}
 
-	remediationResult, origin := data.GetRemediationAndOrigin()
+	r, origin := data.GetRemediationAndOrigin()
 	if valueLog != nil {
-		valueLog.Tracef("bart result: %s (data: %+v)", remediationResult.String(), data)
+		valueLog.Tracef("bart result: %s (data: %+v)", r.String(), data)
 	}
-	return remediationResult, origin
+	return r, origin
 }
 
 // HasRemediation checks if an exact prefix has a specific remediation with a specific origin.