You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

limiter_test.go 1.8KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. // Copyright (c) 2018 Shivaram Lingamneni
  2. // released under the MIT license
  3. package connection_limits
  4. import (
  5. "net"
  6. "testing"
  7. "time"
  8. )
  9. func easyParseIP(ipstr string) (result net.IP) {
  10. result = net.ParseIP(ipstr)
  11. if result == nil {
  12. panic(ipstr)
  13. }
  14. return
  15. }
  16. var baseConfig = LimiterConfig{
  17. RawLimiterConfig: RawLimiterConfig{
  18. Limit: true,
  19. MaxConcurrent: 4,
  20. Throttle: true,
  21. Window: time.Second * 600,
  22. MaxPerWindow: 8,
  23. CidrLenIPv4: 32,
  24. CidrLenIPv6: 64,
  25. Exempted: []string{"localhost"},
  26. CustomLimits: map[string]CustomLimitConfig{
  27. "8.8.0.0/16": {
  28. MaxConcurrent: 128,
  29. MaxPerWindow: 256,
  30. },
  31. },
  32. },
  33. }
  34. func TestKeying(t *testing.T) {
  35. config := baseConfig
  36. config.postprocess()
  37. var limiter Limiter
  38. limiter.ApplyConfig(&config)
  39. key, maxConc, maxWin := limiter.addrToKey(easyParseIP("1.1.1.1"))
  40. assertEqual(key, "1.1.1.1/32", t)
  41. assertEqual(maxConc, 4, t)
  42. assertEqual(maxWin, 8, t)
  43. key, maxConc, maxWin = limiter.addrToKey(easyParseIP("2607:5301:201:3100::7426"))
  44. assertEqual(key, "2607:5301:201:3100::/64", t)
  45. assertEqual(maxConc, 4, t)
  46. assertEqual(maxWin, 8, t)
  47. key, maxConc, maxWin = limiter.addrToKey(easyParseIP("8.8.4.4"))
  48. assertEqual(key, "8.8.0.0/16", t)
  49. assertEqual(maxConc, 128, t)
  50. assertEqual(maxWin, 256, t)
  51. }
  52. func TestLimits(t *testing.T) {
  53. regularIP := easyParseIP("2607:5301:201:3100::7426")
  54. config := baseConfig
  55. config.postprocess()
  56. var limiter Limiter
  57. limiter.ApplyConfig(&config)
  58. for i := 0; i < 4; i++ {
  59. err := limiter.AddClient(regularIP)
  60. if err != nil {
  61. t.Errorf("ip should not be blocked, but %v", err)
  62. }
  63. }
  64. err := limiter.AddClient(regularIP)
  65. if err != ErrLimitExceeded {
  66. t.Errorf("ip should be blocked, but %v", err)
  67. }
  68. limiter.RemoveClient(regularIP)
  69. err = limiter.AddClient(regularIP)
  70. if err != nil {
  71. t.Errorf("ip should not be blocked, but %v", err)
  72. }
  73. }