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.9KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  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. Count: 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. "google": {
  28. Nets: []string{"8.8.0.0/16"},
  29. MaxConcurrent: 128,
  30. MaxPerWindow: 256,
  31. },
  32. },
  33. },
  34. }
  35. func TestKeying(t *testing.T) {
  36. config := baseConfig
  37. config.postprocess()
  38. var limiter Limiter
  39. limiter.ApplyConfig(&config)
  40. key, maxConc, maxWin := limiter.addrToKey(easyParseIP("1.1.1.1"))
  41. assertEqual(key, "1.1.1.1/32", t)
  42. assertEqual(maxConc, 4, t)
  43. assertEqual(maxWin, 8, t)
  44. key, maxConc, maxWin = limiter.addrToKey(easyParseIP("2607:5301:201:3100::7426"))
  45. assertEqual(key, "2607:5301:201:3100::/64", t)
  46. assertEqual(maxConc, 4, t)
  47. assertEqual(maxWin, 8, t)
  48. key, maxConc, maxWin = limiter.addrToKey(easyParseIP("8.8.4.4"))
  49. assertEqual(key, "*google", t)
  50. assertEqual(maxConc, 128, t)
  51. assertEqual(maxWin, 256, t)
  52. }
  53. func TestLimits(t *testing.T) {
  54. regularIP := easyParseIP("2607:5301:201:3100::7426")
  55. config := baseConfig
  56. config.postprocess()
  57. var limiter Limiter
  58. limiter.ApplyConfig(&config)
  59. for i := 0; i < 4; i++ {
  60. err := limiter.AddClient(regularIP)
  61. if err != nil {
  62. t.Errorf("ip should not be blocked, but %v", err)
  63. }
  64. }
  65. err := limiter.AddClient(regularIP)
  66. if err != ErrLimitExceeded {
  67. t.Errorf("ip should be blocked, but %v", err)
  68. }
  69. limiter.RemoveClient(regularIP)
  70. err = limiter.AddClient(regularIP)
  71. if err != nil {
  72. t.Errorf("ip should not be blocked, but %v", err)
  73. }
  74. }