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.

throttler_test.go 3.0KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124
  1. // Copyright (c) 2018 Shivaram Lingamneni
  2. // released under the MIT license
  3. package connection_limits
  4. import (
  5. "net"
  6. "reflect"
  7. "testing"
  8. "time"
  9. )
  10. func assertEqual(supplied, expected interface{}, t *testing.T) {
  11. if !reflect.DeepEqual(supplied, expected) {
  12. t.Errorf("expected %v but got %v", expected, supplied)
  13. }
  14. }
  15. func TestGenericThrottle(t *testing.T) {
  16. minute, _ := time.ParseDuration("1m")
  17. second, _ := time.ParseDuration("1s")
  18. zero, _ := time.ParseDuration("0s")
  19. throttler := GenericThrottle{
  20. Duration: minute,
  21. Limit: 2,
  22. }
  23. now := time.Now()
  24. throttled, remaining := throttler.touch(now)
  25. assertEqual(throttled, false, t)
  26. assertEqual(remaining, zero, t)
  27. now = now.Add(second)
  28. throttled, remaining = throttler.touch(now)
  29. assertEqual(throttled, false, t)
  30. assertEqual(remaining, zero, t)
  31. now = now.Add(second)
  32. throttled, remaining = throttler.touch(now)
  33. assertEqual(throttled, true, t)
  34. assertEqual(remaining, 58*second, t)
  35. now = now.Add(minute)
  36. throttled, remaining = throttler.touch(now)
  37. assertEqual(throttled, false, t)
  38. assertEqual(remaining, zero, t)
  39. }
  40. func TestGenericThrottleDisabled(t *testing.T) {
  41. minute, _ := time.ParseDuration("1m")
  42. throttler := GenericThrottle{
  43. Duration: minute,
  44. Limit: 0,
  45. }
  46. for i := 0; i < 1024; i += 1 {
  47. throttled, _ := throttler.Touch()
  48. if throttled {
  49. t.Error("disabled throttler should not throttle")
  50. }
  51. }
  52. }
  53. func makeTestThrottler(v4len, v6len int) *Limiter {
  54. minute, _ := time.ParseDuration("1m")
  55. maxConnections := 3
  56. config := LimiterConfig{
  57. RawLimiterConfig: RawLimiterConfig{
  58. Limit: false,
  59. Throttle: true,
  60. CidrLenIPv4: v4len,
  61. CidrLenIPv6: v6len,
  62. MaxPerWindow: maxConnections,
  63. Window: minute,
  64. },
  65. }
  66. config.postprocess()
  67. var limiter Limiter
  68. limiter.ApplyConfig(&config)
  69. return &limiter
  70. }
  71. func TestConnectionThrottle(t *testing.T) {
  72. throttler := makeTestThrottler(32, 64)
  73. addr := net.ParseIP("8.8.8.8")
  74. for i := 0; i < 3; i += 1 {
  75. err := throttler.AddClient(addr)
  76. assertEqual(err, nil, t)
  77. }
  78. err := throttler.AddClient(addr)
  79. assertEqual(err, ErrThrottleExceeded, t)
  80. }
  81. func TestConnectionThrottleIPv6(t *testing.T) {
  82. throttler := makeTestThrottler(32, 64)
  83. var err error
  84. err = throttler.AddClient(net.ParseIP("2001:0db8::1"))
  85. assertEqual(err, nil, t)
  86. err = throttler.AddClient(net.ParseIP("2001:0db8::2"))
  87. assertEqual(err, nil, t)
  88. err = throttler.AddClient(net.ParseIP("2001:0db8::3"))
  89. assertEqual(err, nil, t)
  90. err = throttler.AddClient(net.ParseIP("2001:0db8::4"))
  91. assertEqual(err, ErrThrottleExceeded, t)
  92. }
  93. func TestConnectionThrottleIPv4(t *testing.T) {
  94. throttler := makeTestThrottler(24, 64)
  95. var err error
  96. err = throttler.AddClient(net.ParseIP("192.168.1.101"))
  97. assertEqual(err, nil, t)
  98. err = throttler.AddClient(net.ParseIP("192.168.1.102"))
  99. assertEqual(err, nil, t)
  100. err = throttler.AddClient(net.ParseIP("192.168.1.103"))
  101. assertEqual(err, nil, t)
  102. err = throttler.AddClient(net.ParseIP("192.168.1.104"))
  103. assertEqual(err, ErrThrottleExceeded, t)
  104. }