diff --git a/config/config.yaml b/config/config.yaml index 5d5072e..13b700f 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -3,6 +3,20 @@ ports: udp: 5001 ssh: 2222 +delays: + global: + min: 50 # milliseconds + max: 200 # milliseconds + per_port: + - port: 5000 + min: 100 + max: 300 + protocols: ["tcp"] + - port: 5001 + min: 75 + max: 150 + protocols: ["udp"] + rules_path: config/rules.yaml addresses: ["1.2.3.4", "5.4.3.2"] diff --git a/delay/delay.go b/delay/delay.go new file mode 100644 index 0000000..5c44950 --- /dev/null +++ b/delay/delay.go @@ -0,0 +1,47 @@ +package delay + +import ( + "math/rand" + "time" +) + +type DelayRange struct { + Min int `yaml:"min"` + Max int `yaml:"max"` +} + +type PortDelay struct { + Port int `yaml:"port"` + Min int `yaml:"min"` + Max int `yaml:"max"` + Protocols []string `yaml:"protocols"` +} + +type DelayConfig struct { + Global DelayRange `yaml:"global"` + PerPort []PortDelay `yaml:"per_port"` +} + +func NewDelayHandler(config DelayConfig) *DelayHandler { + return &DelayHandler{config: config} +} + +type DelayHandler struct { + config DelayConfig +} + +func (h *DelayHandler) GetDelay(port int, protocol string) time.Duration { + // Check port specific delay + for _, pd := range h.config.PerPort { + if pd.Port == port { + for _, p := range pd.Protocols { + if p == protocol { + return time.Duration(rand.Intn(pd.Max-pd.Min+1)+pd.Min) * time.Millisecond + } + } + } + } + + // Fall back to global delay + return time.Duration(rand.Intn(h.config.Global.Max-h.config.Global.Min+1)+h.config.Global.Min) * time.Millisecond +} diff --git a/delay/delay_test.go b/delay/delay_test.go new file mode 100644 index 0000000..322606c --- /dev/null +++ b/delay/delay_test.go @@ -0,0 +1,55 @@ +package delay + +import ( + "testing" + "time" +) + +func TestDelayHandler(t *testing.T) { + config := DelayConfig{ + Global: DelayRange{Min: 50, Max: 200}, + PerPort: []PortDelay{ + { + Port: 5000, + Min: 100, + Max: 300, + Protocols: []string{"tcp"}, + }, + }, + } + + handler := NewDelayHandler(config) + + tests := []struct { + name string + port int + protocol string + minDelay time.Duration + maxDelay time.Duration + }{ + { + name: "TCP port 5000", + port: 5000, + protocol: "tcp", + minDelay: 100 * time.Millisecond, + maxDelay: 300 * time.Millisecond, + }, + { + name: "Global fallback", + port: 5001, + protocol: "udp", + minDelay: 50 * time.Millisecond, + maxDelay: 200 * time.Millisecond, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + delay := handler.GetDelay(tt.port, tt.protocol) + if delay < tt.minDelay || delay > tt.maxDelay { + t.Errorf("Delay %v outside expected range [%v, %v]", + delay, tt.minDelay, tt.maxDelay) + } + }) + } +}