check_client_ip_job_test.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  1. package job
  2. import (
  3. "reflect"
  4. "testing"
  5. )
  6. func TestMergeClientIps_EvictsStaleOldEntries(t *testing.T) {
  7. // #4077: after a ban expires, a single IP that reconnects used to get
  8. // banned again immediately because a long-disconnected IP stayed in the
  9. // DB with an ancient timestamp and kept "protecting" itself against
  10. // eviction. Guard against that regression here.
  11. old := []IPWithTimestamp{
  12. {IP: "1.1.1.1", Timestamp: 100}, // stale — client disconnected long ago
  13. {IP: "2.2.2.2", Timestamp: 1900}, // fresh — still connecting
  14. }
  15. new := []IPWithTimestamp{
  16. {IP: "2.2.2.2", Timestamp: 2000}, // same IP, newer log line
  17. }
  18. got := mergeClientIps(old, new, 1000)
  19. want := map[string]int64{"2.2.2.2": 2000}
  20. if !reflect.DeepEqual(got, want) {
  21. t.Fatalf("stale 1.1.1.1 should have been dropped\ngot: %v\nwant: %v", got, want)
  22. }
  23. }
  24. func TestMergeClientIps_KeepsFreshOldEntriesUnchanged(t *testing.T) {
  25. // Backwards-compat: entries that aren't stale are still carried forward,
  26. // so enforcement survives access-log rotation.
  27. old := []IPWithTimestamp{
  28. {IP: "1.1.1.1", Timestamp: 1500},
  29. }
  30. got := mergeClientIps(old, nil, 1000)
  31. want := map[string]int64{"1.1.1.1": 1500}
  32. if !reflect.DeepEqual(got, want) {
  33. t.Fatalf("fresh old IP should have been retained\ngot: %v\nwant: %v", got, want)
  34. }
  35. }
  36. func TestMergeClientIps_PrefersLaterTimestampForSameIp(t *testing.T) {
  37. old := []IPWithTimestamp{{IP: "1.1.1.1", Timestamp: 1500}}
  38. new := []IPWithTimestamp{{IP: "1.1.1.1", Timestamp: 1700}}
  39. got := mergeClientIps(old, new, 1000)
  40. if got["1.1.1.1"] != 1700 {
  41. t.Fatalf("expected latest timestamp 1700, got %d", got["1.1.1.1"])
  42. }
  43. }
  44. func TestMergeClientIps_DropsStaleNewEntries(t *testing.T) {
  45. // A log line with a clock-skewed old timestamp must not resurrect a
  46. // stale IP past the cutoff.
  47. new := []IPWithTimestamp{{IP: "1.1.1.1", Timestamp: 500}}
  48. got := mergeClientIps(nil, new, 1000)
  49. if len(got) != 0 {
  50. t.Fatalf("stale new IP should have been dropped, got %v", got)
  51. }
  52. }
  53. func TestMergeClientIps_NoStaleCutoffStillWorks(t *testing.T) {
  54. // Defensive: a zero cutoff (e.g. during very first run on a fresh
  55. // install) must not over-evict.
  56. old := []IPWithTimestamp{{IP: "1.1.1.1", Timestamp: 100}}
  57. new := []IPWithTimestamp{{IP: "2.2.2.2", Timestamp: 200}}
  58. got := mergeClientIps(old, new, 0)
  59. want := map[string]int64{"1.1.1.1": 100, "2.2.2.2": 200}
  60. if !reflect.DeepEqual(got, want) {
  61. t.Fatalf("zero cutoff should keep everything\ngot: %v\nwant: %v", got, want)
  62. }
  63. }
  64. func collectIps(entries []IPWithTimestamp) []string {
  65. out := make([]string, 0, len(entries))
  66. for _, e := range entries {
  67. out = append(out, e.IP)
  68. }
  69. return out
  70. }
  71. func TestPartitionLiveIps_SingleLiveNotStarvedByStillFreshHistoricals(t *testing.T) {
  72. // #4091: db holds A, B, C from minutes ago (still in the 30min
  73. // window) but they're not connecting anymore. only D is. old code
  74. // merged all four, sorted ascending, kept [A,B,C] and banned D
  75. // every tick. pin the new rule: only live ips count toward the limit.
  76. ipMap := map[string]int64{
  77. "A": 1000,
  78. "B": 1100,
  79. "C": 1200,
  80. "D": 2000,
  81. }
  82. observed := map[string]bool{"D": true}
  83. live, historical := partitionLiveIps(ipMap, observed)
  84. if got := collectIps(live); !reflect.DeepEqual(got, []string{"D"}) {
  85. t.Fatalf("live set should only contain the ip observed this scan\ngot: %v\nwant: [D]", got)
  86. }
  87. if got := collectIps(historical); !reflect.DeepEqual(got, []string{"A", "B", "C"}) {
  88. t.Fatalf("historical set should contain db-only ips in ascending order\ngot: %v\nwant: [A B C]", got)
  89. }
  90. }
  91. func TestPartitionLiveIps_ConcurrentLiveIpsStillBanNewcomers(t *testing.T) {
  92. // keep the "protect original, ban newcomer" policy when several ips
  93. // are really live. with limit=1, A must stay and B must be banned.
  94. ipMap := map[string]int64{
  95. "A": 5000,
  96. "B": 5500,
  97. }
  98. observed := map[string]bool{"A": true, "B": true}
  99. live, historical := partitionLiveIps(ipMap, observed)
  100. if got := collectIps(live); !reflect.DeepEqual(got, []string{"A", "B"}) {
  101. t.Fatalf("both live ips should be in the live set, ascending\ngot: %v\nwant: [A B]", got)
  102. }
  103. if len(historical) != 0 {
  104. t.Fatalf("no historical ips expected, got %v", historical)
  105. }
  106. }
  107. func TestPartitionLiveIps_EmptyScanLeavesDbIntact(t *testing.T) {
  108. // quiet tick: nothing observed => nothing live. everything merged
  109. // is historical. keeps the panel from wiping recent-but-idle ips.
  110. ipMap := map[string]int64{
  111. "A": 1000,
  112. "B": 1100,
  113. }
  114. observed := map[string]bool{}
  115. live, historical := partitionLiveIps(ipMap, observed)
  116. if len(live) != 0 {
  117. t.Fatalf("no live ips expected, got %v", live)
  118. }
  119. if got := collectIps(historical); !reflect.DeepEqual(got, []string{"A", "B"}) {
  120. t.Fatalf("all merged entries should flow to historical\ngot: %v\nwant: [A B]", got)
  121. }
  122. }