Attack wave samples (#360)

Author: PopoviciMarian

lib/agent/aikido_types/init_data.go

@@ -89,6 +89,8 @@ type AttackWaveState struct {
 	IpQueues map[string]*SlidingWindow
 	// Map of IP addresses to the last time an event was sent for that IP
 	LastSent map[string]int64
+	// Maximum number of samples to keep per IP, can not be higher than attackWaveThreshold
+	MaxSamplesPerIP int
 }
 
 type ListsConfigData struct {
@@ -238,11 +240,12 @@ func NewServerData() *ServerData {
 		Packages:                make(map[string]Package),
 		PollingData:             NewServerDataPolling(),
 		AttackWave: AttackWaveState{
-			Threshold:  15,             // Default: 15 requests
-			WindowSize: 1,              // Default: 1 minute
-			MinBetween: 20 * 60 * 1000, // Default: 20 minutes
-			IpQueues:   make(map[string]*SlidingWindow),
-			LastSent:   make(map[string]int64),
+			Threshold:       15,             // Default: 15 requests
+			WindowSize:      1,              // Default: 1 minute
+			MinBetween:      20 * 60 * 1000, // Default: 20 minutes
+			IpQueues:        make(map[string]*SlidingWindow),
+			LastSent:        make(map[string]int64),
+			MaxSamplesPerIP: 15,
 		},
 	}
 }

lib/agent/aikido_types/sliding_window.go

@@ -1,10 +1,16 @@
 package aikido_types
 
+type SuspiciousRequest struct {
+	Method string `json:"method"`
+	Url    string `json:"url"`
+}
+
 // SlidingWindow represents a time-based sliding window counter.
 // It maintains a queue of counts per time bucket and a running total.
 type SlidingWindow struct {
-	Total int        // Running total of all counts in the window
-	Queue Queue[int] // Queue of counts per time bucket
+	Total   int                 // Running total of all counts in the window
+	Queue   Queue[int]          // Queue of counts per time bucket
+	Samples []SuspiciousRequest // Sample requests collected for attack wave detection (max MaxSamplesPerIP)
 }
 
 // NewSlidingWindow creates a new sliding window with the specified size.
@@ -40,6 +46,25 @@ func (sw *SlidingWindow) Increment() {
 	sw.Total++
 }
 
+// AddSample adds a sample request to the sliding window for attack wave detection.
+// It maintains a maximum of MaxSamplesPerIP unique samples (based on method and URL).
+func (sw *SlidingWindow) AddSample(method, url string, maxSamplesPerIP int) {
+	// Check if this sample already exists
+	for _, sample := range sw.Samples {
+		if sample.Method == method && sample.Url == url {
+			return // Already exists, skip
+		}
+	}
+
+	// Add the sample if we haven't reached the limit
+	if len(sw.Samples) < maxSamplesPerIP {
+		sw.Samples = append(sw.Samples, SuspiciousRequest{
+			Method: method,
+			Url:    url,
+		})
+	}
+}
+
 // IsEmpty returns true if the total count is zero.
 func (sw *SlidingWindow) IsEmpty() bool {
 	return sw.Total == 0

lib/agent/aikido_types/sliding_window_test.go

@@ -266,3 +266,222 @@ func TestSlidingWindowIntegration(t *testing.T) {
 		assert.NotContains(t, windowMap, "endpoint3")
 	})
 }
+
+func TestAddSample(t *testing.T) {
+	t.Run("adds sample to empty sliding window", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("GET", "/api/users", 15)
+
+		assert.Equal(t, 1, len(sw.Samples))
+		assert.Equal(t, "GET", sw.Samples[0].Method)
+		assert.Equal(t, "/api/users", sw.Samples[0].Url)
+	})
+
+	t.Run("adds multiple different samples", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("GET", "/api/users", 15)
+		sw.AddSample("POST", "/api/login", 15)
+		sw.AddSample("DELETE", "/api/users/123", 15)
+
+		assert.Equal(t, 3, len(sw.Samples))
+		assert.Equal(t, "GET", sw.Samples[0].Method)
+		assert.Equal(t, "/api/users", sw.Samples[0].Url)
+		assert.Equal(t, "POST", sw.Samples[1].Method)
+		assert.Equal(t, "/api/login", sw.Samples[1].Url)
+		assert.Equal(t, "DELETE", sw.Samples[2].Method)
+		assert.Equal(t, "/api/users/123", sw.Samples[2].Url)
+	})
+
+	t.Run("prevents duplicate samples with same method and URL", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("GET", "/api/users", 15)
+		sw.AddSample("GET", "/api/users", 15) // duplicate
+		sw.AddSample("GET", "/api/users", 15) // duplicate
+
+		assert.Equal(t, 1, len(sw.Samples))
+		assert.Equal(t, "GET", sw.Samples[0].Method)
+		assert.Equal(t, "/api/users", sw.Samples[0].Url)
+	})
+
+	t.Run("allows same URL with different methods", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("GET", "/api/users", 15)
+		sw.AddSample("POST", "/api/users", 15)
+		sw.AddSample("DELETE", "/api/users", 15)
+
+		assert.Equal(t, 3, len(sw.Samples))
+		assert.Equal(t, "GET", sw.Samples[0].Method)
+		assert.Equal(t, "POST", sw.Samples[1].Method)
+		assert.Equal(t, "DELETE", sw.Samples[2].Method)
+	})
+
+	t.Run("allows same method with different URLs", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("GET", "/api/users", 15)
+		sw.AddSample("GET", "/api/posts", 15)
+		sw.AddSample("GET", "/api/comments", 15)
+
+		assert.Equal(t, 3, len(sw.Samples))
+		assert.Equal(t, "/api/users", sw.Samples[0].Url)
+		assert.Equal(t, "/api/posts", sw.Samples[1].Url)
+		assert.Equal(t, "/api/comments", sw.Samples[2].Url)
+	})
+
+	t.Run("enforces maximum of 10 samples", func(t *testing.T) {
+		sw := NewSlidingWindow()
+
+		// Add 12 unique samples
+		for i := 0; i < 12; i++ {
+			sw.AddSample("GET", "/api/endpoint"+string(rune('0'+i)), 10)
+		}
+
+		assert.Equal(t, 10, len(sw.Samples))
+	})
+
+	t.Run("does not add 11th sample even if unique", func(t *testing.T) {
+		sw := NewSlidingWindow()
+
+		// Add exactly 10 samples
+		for i := 0; i < 10; i++ {
+			sw.AddSample("GET", "/api/endpoint"+string(rune('0'+i)), 10)
+		}
+		assert.Equal(t, 10, len(sw.Samples))
+
+		// Try to add an 11th unique sample
+		sw.AddSample("POST", "/api/new-endpoint", 10)
+		assert.Equal(t, 10, len(sw.Samples))
+
+		// Verify the 11th sample was not added
+		found := false
+		for _, sample := range sw.Samples {
+			if sample.Method == "POST" && sample.Url == "/api/new-endpoint" {
+				found = true
+				break
+			}
+		}
+		assert.False(t, found)
+	})
+
+	t.Run("duplicates do not count toward 10 sample limit", func(t *testing.T) {
+		sw := NewSlidingWindow()
+
+		// Add 5 unique samples
+		for i := 0; i < 5; i++ {
+			sw.AddSample("GET", "/api/endpoint"+string(rune('0'+i)), 15)
+		}
+
+		// Try to add duplicates
+		sw.AddSample("GET", "/api/endpoint0", 15) // duplicate
+		sw.AddSample("GET", "/api/endpoint1", 15) // duplicate
+		sw.AddSample("GET", "/api/endpoint2", 15) // duplicate
+
+		assert.Equal(t, 5, len(sw.Samples))
+
+		// Add 5 more unique samples to reach 10
+		for i := 5; i < 10; i++ {
+			sw.AddSample("GET", "/api/endpoint"+string(rune('0'+i)), 15)
+		}
+		assert.Equal(t, 10, len(sw.Samples))
+
+		// Try to add more duplicates - should still be 10
+		sw.AddSample("GET", "/api/endpoint5", 15)
+		sw.AddSample("GET", "/api/endpoint9", 15)
+		assert.Equal(t, 10, len(sw.Samples))
+	})
+
+	t.Run("preserves samples during window operations", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("GET", "/api/users", 15)
+		sw.AddSample("POST", "/api/login", 15)
+		sw.Increment()
+
+		// Advance the window
+		sw.Advance(5)
+
+		// Samples should still be present
+		assert.Equal(t, 2, len(sw.Samples))
+		assert.Equal(t, "GET", sw.Samples[0].Method)
+		assert.Equal(t, "/api/users", sw.Samples[0].Url)
+		assert.Equal(t, "POST", sw.Samples[1].Method)
+		assert.Equal(t, "/api/login", sw.Samples[1].Url)
+	})
+
+	t.Run("empty method and URL are valid samples", func(t *testing.T) {
+		sw := NewSlidingWindow()
+		sw.AddSample("", "", 15)
+		sw.AddSample("GET", "", 15)
+		sw.AddSample("", "/api/users", 15)
+
+		assert.Equal(t, 3, len(sw.Samples))
+	})
+}
+
+func TestSlidingWindowSamplesIntegration(t *testing.T) {
+	t.Run("simulates attack wave detection with samples", func(t *testing.T) {
+		// Create a map simulating per-IP tracking
+		ipMap := map[string]*SlidingWindow{
+			"192.168.1.100": NewSlidingWindow(),
+		}
+
+		ip := "192.168.1.100"
+		sw := ipMap[ip]
+
+		// Simulate suspicious requests
+		requests := []struct {
+			method string
+			url    string
+		}{
+			{"GET", "/admin"},
+			{"GET", "/admin"}, // duplicate
+			{"POST", "/admin"},
+			{"GET", "/wp-admin"},
+			{"GET", "/.env"},
+			{"GET", "/config.php"},
+			{"POST", "/login"},
+			{"GET", "/admin"}, // duplicate
+		}
+
+		for _, req := range requests {
+			sw.Increment()
+			sw.AddSample(req.method, req.url, 15)
+		}
+
+		// Should have 6 unique samples (2 duplicates removed)
+		assert.Equal(t, 6, len(sw.Samples))
+		assert.Equal(t, 8, sw.Total) // But total count should be 8
+
+		// Verify samples are unique
+		uniqueCheck := make(map[string]bool)
+		for _, sample := range sw.Samples {
+			key := sample.Method + ":" + sample.Url
+			assert.False(t, uniqueCheck[key], "Found duplicate sample: "+key)
+			uniqueCheck[key] = true
+		}
+	})
+
+	t.Run("samples persist across window advances until removal", func(t *testing.T) {
+		windowMap := map[string]*SlidingWindow{
+			"10.0.0.1": NewSlidingWindow(),
+		}
+
+		sw := windowMap["10.0.0.1"]
+		sw.AddSample("GET", "/api/v1/users", 15)
+		sw.AddSample("POST", "/api/v1/login", 15)
+		sw.Increment()
+		sw.Increment()
+
+		// Advance window multiple times
+		AdvanceSlidingWindowMap(windowMap, 3)
+		AdvanceSlidingWindowMap(windowMap, 3)
+
+		// Samples should still be there
+		assert.Contains(t, windowMap, "10.0.0.1")
+		assert.Equal(t, 2, len(windowMap["10.0.0.1"].Samples))
+
+		// Advance until window is empty
+		AdvanceSlidingWindowMap(windowMap, 3)
+
+		// Window should be removed when total reaches 0
+		assert.NotContains(t, windowMap, "10.0.0.1")
+	})
+}

lib/agent/grpc/request.go

@@ -201,12 +201,12 @@ func isRateLimitingThresholdExceeded(config *RateLimitingConfig, countsMap map[s
 
 // updateAttackWaveCountsAndDetect implements the attack wave detection logic:
 //  1. Validates the request is from a web scanner and has a valid IP address
-//  2. Increments the sliding window counter for this IP
+//  2. Increments the sliding window counter for this IP and collects request samples
 //  3. Applies throttling: if an event was recently sent for this IP (within minBetween window),
 //     returns early without checking threshold or sending another event
 //  4. Checks if the total count within the sliding window exceeds the threshold
-//  5. If threshold exceeded: records the event time on the queue, logs the detection, and sends event to cloud
-func updateAttackWaveCountsAndDetect(server *ServerData, isWebScanner bool, ip string, userId string, userAgent string) bool {
+//  5. If threshold exceeded: records the event time on the queue, logs the detection, and sends event with samples to cloud
+func updateAttackWaveCountsAndDetect(server *ServerData, isWebScanner bool, ip string, userId string, userAgent string, method string, url string) bool {
 	if !isWebScanner || ip == "" {
 		return false
 	}
@@ -224,6 +224,11 @@ func updateAttackWaveCountsAndDetect(server *ServerData, isWebScanner bool, ip s
 		return false
 	}
 
+	// Add this request as a sample
+	if queue != nil {
+		queue.AddSample(method, url, server.AttackWave.MaxSamplesPerIP)
+	}
+
 	// check threshold within window
 	if queue == nil || queue.Total < server.AttackWave.Threshold {
 		return false // threshold not reached
@@ -234,11 +239,18 @@ func updateAttackWaveCountsAndDetect(server *ServerData, isWebScanner bool, ip s
 	if server.Logger != nil {
 		log.Infof(server.Logger, "Attack wave detected from IP: %s", ip)
 	}
+
 	// report event to cloud
 	cloud.SendAttackDetectedEvent(server, &protos.AttackDetected{
 		Token:   server.AikidoConfig.Token,
 		Request: &protos.Request{IpAddress: ip, UserAgent: userAgent},
-		Attack:  &protos.Attack{Metadata: []*protos.Metadata{}, UserId: userId},
+		Attack: &protos.Attack{
+			Metadata: []*protos.Metadata{{
+				Key:   "samples",
+				Value: utils.JsonMarshal(queue.Samples),
+			}},
+			UserId: userId,
+		},
 	}, "detected_attack_wave")
 	return true
 }

lib/agent/grpc/request_test.go

@@ -37,7 +37,7 @@ func TestAttackWaveThrottling(t *testing.T) {
 		server.AttackWave.IpQueues[ip] = sw
 
 		// Should return false (throttled) because last event was only 30 seconds ago (< 60s MinBetween)
-		assert.False(t, updateAttackWaveCountsAndDetect(server, true, ip, "", ""))
+		assert.False(t, updateAttackWaveCountsAndDetect(server, true, ip, "", "", "", ""))
 	})
 
 	t.Run("returns true and populates LastSent map when IP reaches threshold for first time", func(t *testing.T) {
@@ -66,7 +66,7 @@ func TestAttackWaveThrottling(t *testing.T) {
 		assert.False(t, exists, "IP should not be in LastSent map before threshold is reached")
 
 		// Should return true (event sent) because this is the first time reaching threshold
-		assert.True(t, updateAttackWaveCountsAndDetect(server, true, ip, "", ""))
+		assert.True(t, updateAttackWaveCountsAndDetect(server, true, ip, "", "", "", ""))
 
 		// Verify LastSent map was populated
 		assert.True(t, server.AttackWave.LastSent[ip] > 0, "LastSent should be set after event is sent")

lib/agent/grpc/server.go

@@ -82,7 +82,7 @@ func (s *GrpcServer) OnRequestShutdown(ctx context.Context, req *protos.RequestM
 		go storeRoute(server, req.GetMethod(), req.GetRouteParsed(), req.GetApiSpec(), req.GetRateLimited())
 		go updateRateLimitingCounts(server, req.GetMethod(), req.GetRoute(), req.GetRouteParsed(), req.GetUser(), req.GetIp(), req.GetRateLimitGroup())
 	}
-	go updateAttackWaveCountsAndDetect(server, req.GetIsWebScanner(), req.GetIp(), req.GetUser(), req.GetUserAgent())
+	go updateAttackWaveCountsAndDetect(server, req.GetIsWebScanner(), req.GetIp(), req.GetUser(), req.GetUserAgent(), req.GetMethod(), req.GetUrl())
 
 	atomic.StoreUint32(&server.GotTraffic, 1)
 	return &emptypb.Empty{}, nil

lib/agent/utils/utils.go

@@ -1,6 +1,7 @@
 package utils
 
 import (
+	"encoding/json"
 	"fmt"
 	. "main/aikido_types"
 	"main/config"
@@ -146,3 +147,11 @@ func AnonymizeToken(token string) string {
 	}
 	return token[len(token)-4:]
 }
+
+func JsonMarshal(v any) string {
+	bytes, err := json.Marshal(v)
+	if err != nil {
+		return ""
+	}
+	return string(bytes)
+}