mirror of
https://github.com/andwati/telecom-platform.git
synced 2026-10-18 11:02:54 +00:00
feat: Add comprehensive Redis rate limiter tests with mock client, middleware validation, and multi-tier support
- Add MockRedisClient struct with Eval/Del/ZCount methods for Redis operation mocking - Add TestRedisRateLimiter_Allow with under/over limit scenarios and mock Redis response validation - Add TestRedisRateLimitMiddleware with IP-based rate limiting and empty key bypass tests - Add TestIPKeyExtractor/TestUserKeyExtractor/TestAPIKeyExtractor with valid/invalid input validation - Add TestEnd
This commit is contained in:
nutcas3
@@ -0,0 +1,339 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// MockRedisClient is a mock for Redis client
|
||||
type MockRedisClient struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockRedisClient) Eval(ctx context.Context, script string, keys []string, args ...interface{}) *redis.Cmd {
|
||||
args = append([]interface{}{script, keys}, args...)
|
||||
return m.Called(args...).Get(0).(*redis.Cmd)
|
||||
}
|
||||
|
||||
func (m *MockRedisClient) Del(ctx context.Context, keys ...string) *redis.IntCmd {
|
||||
return m.Called(keys...).Get(0).(*redis.IntCmd)
|
||||
}
|
||||
|
||||
func (m *MockRedisClient) ZCount(ctx context.Context, key, min, max string) *redis.IntCmd {
|
||||
return m.Called(key, min, max).Get(0).(*redis.IntCmd)
|
||||
}
|
||||
|
||||
func TestRedisRateLimiter_Allow(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
limit int64
|
||||
window time.Duration
|
||||
expected bool
|
||||
}{
|
||||
{"Under limit", 5, time.Minute, true},
|
||||
{"Over limit", 1, time.Minute, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create mock Redis client
|
||||
mockClient := new(MockRedisClient)
|
||||
|
||||
// Mock Redis response based on test case
|
||||
var mockResult []interface{}
|
||||
if tt.expected {
|
||||
mockResult = []interface{}{int64(1), int64(tt.limit - 1), int64(time.Now().Unix() + 60)}
|
||||
} else {
|
||||
mockResult = []interface{}{int64(0), int64(0), int64(time.Now().Unix() + 60), int64(30)}
|
||||
}
|
||||
|
||||
cmd := redis.NewCmd(context.Background())
|
||||
cmd.SetVal(mockResult)
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(cmd)
|
||||
|
||||
// Create rate limiter
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", tt.limit, tt.window)
|
||||
limiter.client = mockClient // Replace with mock
|
||||
|
||||
// Test rate limiting
|
||||
result, err := limiter.Allow(context.Background(), "test-key")
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, tt.expected, result.Allowed)
|
||||
assert.Equal(t, tt.limit, result.TotalLimit)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisRateLimitMiddleware(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
keyExtractor func(*gin.Context) string
|
||||
expectedStatus int
|
||||
shouldHaveHeaders bool
|
||||
}{
|
||||
{
|
||||
name: "IP-based rate limiting - allowed",
|
||||
keyExtractor: func(c *gin.Context) string {
|
||||
return "127.0.0.1"
|
||||
},
|
||||
expectedStatus: http.StatusOK,
|
||||
shouldHaveHeaders: true,
|
||||
},
|
||||
{
|
||||
name: "Empty key - bypass rate limiting",
|
||||
keyExtractor: func(c *gin.Context) string {
|
||||
return ""
|
||||
},
|
||||
expectedStatus: http.StatusOK,
|
||||
shouldHaveHeaders: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create mock Redis client
|
||||
mockClient := new(MockRedisClient)
|
||||
|
||||
// Mock successful rate limit check
|
||||
cmd := redis.NewCmd(context.Background())
|
||||
cmd.SetVal([]interface{}{int64(1), int64(4), int64(time.Now().Unix() + 60)})
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(cmd)
|
||||
|
||||
// Create rate limiter
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", 5, time.Minute)
|
||||
limiter.client = mockClient // Replace with mock
|
||||
|
||||
// Create Gin router with middleware
|
||||
router := gin.New()
|
||||
router.Use(RedisRateLimitMiddleware(limiter, tt.keyExtractor))
|
||||
router.GET("/test", func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"message": "success"})
|
||||
})
|
||||
|
||||
// Create request
|
||||
req := httptest.NewRequest("GET", "/test", nil)
|
||||
req.RemoteAddr = "127.0.0.1:12345"
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
// Perform request
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
// Check response
|
||||
assert.Equal(t, tt.expectedStatus, w.Code)
|
||||
|
||||
// Check rate limit headers
|
||||
if tt.shouldHaveHeaders {
|
||||
assert.Contains(t, w.Header(), "X-RateLimit-Limit")
|
||||
assert.Contains(t, w.Header(), "X-RateLimit-Remaining")
|
||||
assert.Contains(t, w.Header(), "X-RateLimit-Reset")
|
||||
assert.Equal(t, "5", w.Header().Get("X-RateLimit-Limit"))
|
||||
} else {
|
||||
assert.NotContains(t, w.Header(), "X-RateLimit-Limit")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIPKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
clientIP string
|
||||
expected string
|
||||
}{
|
||||
{"Valid IP", "192.168.1.1", "192.168.1.1"},
|
||||
{"IPv6", "::1", "::1"},
|
||||
{"Empty IP", "", ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.RemoteAddr = tt.clientIP + ":12345"
|
||||
|
||||
result := IPKeyExtractor(c)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
userID interface{}
|
||||
clientIP string
|
||||
expected string
|
||||
}{
|
||||
{"With user ID", "user123", "192.168.1.1", "user:user123"},
|
||||
{"Without user ID", nil, "192.168.1.1", "192.168.1.1"},
|
||||
{"Invalid user ID type", 123, "192.168.1.1", "192.168.1.1"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.RemoteAddr = tt.clientIP + ":12345"
|
||||
|
||||
if tt.userID != nil {
|
||||
c.Set("user_id", tt.userID)
|
||||
}
|
||||
|
||||
result := UserKeyExtractor(c)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
apiKey string
|
||||
userID string
|
||||
clientIP string
|
||||
expected string
|
||||
}{
|
||||
{"With API key", "api-key-123", "user123", "192.168.1.1", "api:api-key-123"},
|
||||
{"Without API key", "", "user123", "192.168.1.1", "user:user123"},
|
||||
{"No API key or user", "", "", "192.168.1.1", "192.168.1.1"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/", nil)
|
||||
c.Request.RemoteAddr = tt.clientIP + ":12345"
|
||||
|
||||
if tt.apiKey != "" {
|
||||
c.Request.Header.Set("X-API-Key", tt.apiKey)
|
||||
}
|
||||
if tt.userID != "" {
|
||||
c.Set("user_id", tt.userID)
|
||||
}
|
||||
|
||||
result := APIKeyExtractor(c)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEndpointKeyExtractor(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
baseExtractor := func(c *gin.Context) string {
|
||||
return "user123"
|
||||
}
|
||||
|
||||
extractor := EndpointKeyExtractor(baseExtractor)
|
||||
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest("GET", "/api/v1/users", nil)
|
||||
c.Request.RemoteAddr = "192.168.1.1:12345"
|
||||
|
||||
result := extractor(c)
|
||||
expected := "user123:GET:/api/v1/users"
|
||||
assert.Equal(t, expected, result)
|
||||
}
|
||||
|
||||
func TestMultiTierRateLimiter(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
// Create mock Redis client
|
||||
mockClient := new(MockRedisClient)
|
||||
|
||||
// Mock different tier responses
|
||||
freeCmd := redis.NewCmd(context.Background())
|
||||
freeCmd.SetVal([]interface{}{int64(1), int64(9), int64(time.Now().Unix() + 60)})
|
||||
|
||||
premiumCmd := redis.NewCmd(context.Background())
|
||||
premiumCmd.SetVal([]interface{}{int64(1), int64(99), int64(time.Now().Unix() + 60)})
|
||||
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(freeCmd).Once()
|
||||
mockClient.On("Eval", mock.AnythingOfType("[]interface {}")).Return(premiumCmd).Once()
|
||||
|
||||
// Create multi-tier rate limiter
|
||||
tiers := map[string]struct {
|
||||
Limit int64
|
||||
Window time.Duration
|
||||
Prefix string
|
||||
}{
|
||||
"free": {Limit: 10, Window: time.Minute, Prefix: "free"},
|
||||
"premium": {Limit: 100, Window: time.Minute, Prefix: "premium"},
|
||||
}
|
||||
|
||||
multiLimiter := NewMultiTierRateLimiter(&redis.Client{}, tiers)
|
||||
|
||||
// Replace client with mock for testing
|
||||
for tier, limiter := range multiLimiter.limiters {
|
||||
if tier == "free" {
|
||||
limiter.client = mockClient
|
||||
} else if tier == "premium" {
|
||||
limiter.client = mockClient
|
||||
}
|
||||
}
|
||||
|
||||
// Test free tier
|
||||
result, err := multiLimiter.CheckTier(context.Background(), "free", "user123")
|
||||
assert.NoError(t, err)
|
||||
assert.True(t, result.Allowed)
|
||||
assert.Equal(t, int64(10), result.TotalLimit)
|
||||
|
||||
// Test premium tier
|
||||
result, err = multiLimiter.CheckTier(context.Background(), "premium", "user456")
|
||||
assert.NoError(t, err)
|
||||
assert.True(t, result.Allowed)
|
||||
assert.Equal(t, int64(100), result.TotalLimit)
|
||||
}
|
||||
|
||||
func TestRedisRateLimiter_Reset(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
mockClient := new(MockRedisClient)
|
||||
cmd := redis.NewIntCmd(context.Background())
|
||||
cmd.SetVal(1)
|
||||
mockClient.On("Del", "test:user123").Return(cmd)
|
||||
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", 10, time.Minute)
|
||||
limiter.client = mockClient
|
||||
|
||||
err := limiter.Reset(context.Background(), "user123")
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestRedisRateLimiter_GetStats(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
mockClient := new(MockRedisClient)
|
||||
cmd := redis.NewIntCmd(context.Background())
|
||||
cmd.SetVal(3)
|
||||
mockClient.On("ZCount", "test:user123", mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(cmd)
|
||||
|
||||
limiter := NewRedisRateLimiter(&redis.Client{}, "test", 10, time.Minute)
|
||||
limiter.client = mockClient
|
||||
|
||||
result, err := limiter.GetStats(context.Background(), "user123")
|
||||
assert.NoError(t, err)
|
||||
assert.NotNil(t, result)
|
||||
assert.Equal(t, int64(10), result.TotalLimit)
|
||||
assert.Equal(t, int64(7), result.Remaining) // 10 - 3 = 7 remaining
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
# Packet Gateway with eBPF
|
||||
|
||||
A high-performance packet filtering gateway using eBPF XDP programs for real-time traffic control and credit management.
|
||||
|
||||
## Architecture
|
||||
|
||||
The packet gateway consists of two main components:
|
||||
|
||||
1. **eBPF Kernel Program** (`src/ebpf/packet_filter.c`): Runs in kernel space, filters packets at line rate
|
||||
2. **Userspace Control Plane** (`src/main.rs`): Manages eBPF maps, syncs with Redis, handles configuration
|
||||
|
||||
## Features
|
||||
|
||||
- **XDP Packet Filtering**: Kernel-space packet filtering for maximum performance
|
||||
- **Credit Management**: Real-time credit deduction and enforcement
|
||||
- **Dynamic Blocking**: Block/unblock users via eBPF maps
|
||||
- **Redis Integration**: Sync stats and configuration with Redis
|
||||
- **Graceful Shutdown**: Clean up resources and sync final state
|
||||
|
||||
## eBPF Maps
|
||||
|
||||
### packet_stats
|
||||
- **Key**: `u32` (IP address)
|
||||
- **Value**: `u64` (byte count)
|
||||
- **Purpose**: Track traffic usage per IP
|
||||
|
||||
### user_credits
|
||||
- **Key**: `u32` (IP address)
|
||||
- **Value**: `i64` (credit balance)
|
||||
- **Purpose**: Enforce credit limits, block when depleted
|
||||
|
||||
### block_list
|
||||
- **Key**: `u32` (IP address)
|
||||
- **Value**: `u8` (block flag: 0=allow, 1=block)
|
||||
- **Purpose**: Manual blocking of specific IPs
|
||||
|
||||
## Build Requirements
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
# Install eBPF development tools
|
||||
sudo apt-get install -y linux-headers-$(uname -r) clang llvm
|
||||
|
||||
# Install Rust eBPF tools
|
||||
cargo install aya-cli
|
||||
```
|
||||
|
||||
### Build Process
|
||||
```bash
|
||||
# Build eBPF program
|
||||
cargo xtask build-ebpf
|
||||
|
||||
# Build userspace application
|
||||
cargo build --release
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Usage
|
||||
```bash
|
||||
# Run with default settings
|
||||
sudo ./target/release/packet-gateway
|
||||
|
||||
# Specify interface and Redis
|
||||
sudo ./target/release/packet-gateway \
|
||||
--interface eth0 \
|
||||
--redis-url redis://127.0.0.1/ \
|
||||
--sync-interval 1
|
||||
```
|
||||
|
||||
### Redis Integration
|
||||
|
||||
The gateway syncs with Redis using these key patterns:
|
||||
|
||||
**Packet Stats:**
|
||||
```
|
||||
packet_stats:192.168.1.100 -> 1048576 # bytes transferred
|
||||
```
|
||||
|
||||
**User Credits:**
|
||||
```
|
||||
user_credit:192.168.1.100 -> 10485760 # remaining credit in bytes
|
||||
```
|
||||
|
||||
**Blocked Users:**
|
||||
```
|
||||
blocked_user:192.168.1.100 -> 1 # blocked flag
|
||||
```
|
||||
|
||||
### Managing Credits
|
||||
|
||||
```bash
|
||||
# Set user credit (Redis)
|
||||
redis-cli set user_credit:192.168.1.100 10485760
|
||||
|
||||
# Block a user
|
||||
redis-cli set blocked_user:192.168.1.100 1
|
||||
|
||||
# Unblock a user
|
||||
redis-cli del blocked_user:192.168.1.100
|
||||
```
|
||||
|
||||
## eBPF Program Details
|
||||
|
||||
### Packet Processing Flow
|
||||
|
||||
1. **Parse Packet**: Extract source/destination IPs
|
||||
2. **Check Block List**: Block if IP is in block list
|
||||
3. **Check Credit**: Block if insufficient credit
|
||||
4. **Update Stats**: Increment packet statistics
|
||||
5. **Deduct Credit**: Decrease user credit balance
|
||||
6. **Return Action**: `XDP_PASS` or `XDP_DROP`
|
||||
|
||||
### Key Functions
|
||||
|
||||
- `packet_filter()`: Main XDP program entry point
|
||||
- `should_block_packet()`: Determine if packet should be blocked
|
||||
- `update_packet_stats()`: Update traffic statistics
|
||||
- `update_credit_usage()`: Deduct credit from user balance
|
||||
|
||||
## Performance
|
||||
|
||||
- **Line Rate Processing**: XDP enables kernel-space packet filtering
|
||||
- **Map Operations**: O(1) hash map lookups for credit/block checks
|
||||
- **Atomic Updates**: Lock-free map updates for high concurrency
|
||||
|
||||
## Monitoring
|
||||
|
||||
### Metrics Available
|
||||
- Packets processed per second
|
||||
- Bytes transferred per IP
|
||||
- Active user count
|
||||
- Low credit warnings
|
||||
- Block list size
|
||||
|
||||
### Log Levels
|
||||
```bash
|
||||
# Debug logging
|
||||
RUST_LOG=debug ./packet-gateway
|
||||
|
||||
# Production logging
|
||||
RUST_LOG=info ./packet-gateway
|
||||
```
|
||||
|
||||
## Security Considerations
|
||||
|
||||
- **Root Required**: eBPF XDP programs require root privileges
|
||||
- **Network Isolation**: Test in isolated network environments
|
||||
- **Memory Safety**: eBPF verifier ensures kernel safety
|
||||
- **Resource Limits**: Map sizes prevent memory exhaustion
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
**Permission Denied:**
|
||||
```bash
|
||||
# Run with sudo
|
||||
sudo ./packet-gateway
|
||||
```
|
||||
|
||||
**Interface Not Found:**
|
||||
```bash
|
||||
# Check available interfaces
|
||||
ip link show
|
||||
```
|
||||
|
||||
**Redis Connection Failed:**
|
||||
```bash
|
||||
# Verify Redis is running
|
||||
redis-cli ping
|
||||
```
|
||||
|
||||
**eBPF Program Failed to Load:**
|
||||
```bash
|
||||
# Check kernel logs
|
||||
dmesg | grep -i ebpf
|
||||
```
|
||||
|
||||
### Debug Mode
|
||||
|
||||
Enable detailed logging to troubleshoot issues:
|
||||
```bash
|
||||
RUST_LOG=debug ./packet-gateway --interface eth0
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
### Adding New Features
|
||||
|
||||
1. **eBPF Side**: Add new maps or modify packet_filter.c
|
||||
2. **Userspace Side**: Add corresponding methods in EbpfManager
|
||||
3. **Redis Integration**: Update sync functions
|
||||
4. **Testing**: Add unit tests and integration tests
|
||||
|
||||
### Testing
|
||||
|
||||
```bash
|
||||
# Run unit tests
|
||||
cargo test
|
||||
|
||||
# Run integration tests (requires test environment)
|
||||
cargo test --test integration
|
||||
|
||||
# Build and run in debug mode
|
||||
cargo build && sudo ./target/debug/packet-gateway
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
This project is licensed under the GPL v2 license, compatible with eBPF kernel requirements.
|
||||
Reference in New Issue
Block a user