mirror of
https://github.com/ergochat/ergo.git
synced 2024-11-22 11:59:40 +01:00
96 lines
2.2 KiB
Go
96 lines
2.2 KiB
Go
// Copyright (c) 2018 Shivaram Lingamneni
|
|
// released under the MIT license
|
|
|
|
package connection_limits
|
|
|
|
import (
|
|
"crypto/md5"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ergochat/ergo/irc/flatip"
|
|
)
|
|
|
|
func easyParseIP(ipstr string) (result flatip.IP) {
|
|
result, err := flatip.ParseIP(ipstr)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return
|
|
}
|
|
|
|
var baseConfig = LimiterConfig{
|
|
rawLimiterConfig: rawLimiterConfig{
|
|
Count: true,
|
|
MaxConcurrent: 4,
|
|
|
|
Throttle: true,
|
|
Window: time.Second * 600,
|
|
MaxPerWindow: 8,
|
|
|
|
CidrLenIPv4: 32,
|
|
CidrLenIPv6: 64,
|
|
|
|
Exempted: []string{"localhost"},
|
|
|
|
CustomLimits: map[string]CustomLimitConfig{
|
|
"google": {
|
|
Nets: []string{"8.8.0.0/16"},
|
|
MaxConcurrent: 128,
|
|
MaxPerWindow: 256,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
func TestKeying(t *testing.T) {
|
|
config := baseConfig
|
|
config.postprocess()
|
|
var limiter Limiter
|
|
limiter.ApplyConfig(&config)
|
|
|
|
// an ipv4 /32 looks like a /128 to us after applying the 4-in-6 mapping
|
|
key, _, maxConc, maxWin := limiter.addrToKey(easyParseIP("1.1.1.1"))
|
|
assertEqual(key.prefixLen, uint8(128), t)
|
|
assertEqual(key.maskedIP[12:], []byte{1, 1, 1, 1}, t)
|
|
assertEqual(maxConc, 4, t)
|
|
assertEqual(maxWin, 8, t)
|
|
|
|
testIPv6 := easyParseIP("2607:5301:201:3100::7426")
|
|
key, _, maxConc, maxWin = limiter.addrToKey(testIPv6)
|
|
assertEqual(key.prefixLen, uint8(64), t)
|
|
assertEqual(flatip.IP(key.maskedIP), easyParseIP("2607:5301:201:3100::"), t)
|
|
assertEqual(maxConc, 4, t)
|
|
assertEqual(maxWin, 8, t)
|
|
|
|
key, _, maxConc, maxWin = limiter.addrToKey(easyParseIP("8.8.4.4"))
|
|
assertEqual(key.prefixLen, uint8(0), t)
|
|
assertEqual([16]byte(key.maskedIP), md5.Sum([]byte("google")), t)
|
|
assertEqual(maxConc, 128, t)
|
|
assertEqual(maxWin, 256, t)
|
|
}
|
|
|
|
func TestLimits(t *testing.T) {
|
|
regularIP := easyParseIP("2607:5301:201:3100::7426")
|
|
config := baseConfig
|
|
config.postprocess()
|
|
var limiter Limiter
|
|
limiter.ApplyConfig(&config)
|
|
|
|
for i := 0; i < 4; i++ {
|
|
err := limiter.AddClient(regularIP)
|
|
if err != nil {
|
|
t.Errorf("ip should not be blocked, but %v", err)
|
|
}
|
|
}
|
|
err := limiter.AddClient(regularIP)
|
|
if err != ErrLimitExceeded {
|
|
t.Errorf("ip should be blocked, but %v", err)
|
|
}
|
|
limiter.RemoveClient(regularIP)
|
|
err = limiter.AddClient(regularIP)
|
|
if err != nil {
|
|
t.Errorf("ip should not be blocked, but %v", err)
|
|
}
|
|
}
|