package store import ( "net" "path/filepath" "strings" "testing" "time" "mail_go/internal/db" "gorm.io/driver/sqlite" "gorm.io/gorm" ) func newTestStores(t *testing.T) *Stores { t.Helper() gdb, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "test.db")), &gorm.Config{}) if err != nil { t.Fatalf("open sqlite: %v", err) } if err := gdb.AutoMigrate(&db.User{}, &db.Domain{}, &db.Message{}, &db.Attachment{}, &db.BanEntry{}, &db.OutboundMessage{}, &db.ProtocolLog{}); err != nil { t.Fatalf("migrate: %v", err) } return NewStores(gdb) } // TestRecordAuthFailureFreeTriggers 验证前 3 次达到阈值只计数不封禁, // 第 4 次起封禁(第 1 次封禁 = 配置时长)。 func TestRecordAuthFailureFreeTriggers(t *testing.T) { s := newTestStores(t) const ip = "203.0.113.10" const maxFail = 2 failOnce := func() bool { banned, _ := s.RecordAuthFailure(ip, maxFail, 30, "登录失败次数过多") return banned } // 第 1 次触发需要 maxFail 次失败 for f := 0; f < maxFail; f++ { if failOnce() { t.Fatalf("trigger 1 (fail %d) should not ban yet", f+1) } } // 达到阈值后失败计数持续累计,之后每次失败都会再次触发 for i := 2; i <= 3; i++ { if failOnce() { t.Fatalf("trigger %d should not ban yet", i) } } entry, err := s.Bans.GetByIP(ip) if err != nil { t.Fatalf("get entry: %v", err) } if entry.BanCount != 3 { t.Fatalf("ban_count = %d, want 3", entry.BanCount) } if !entry.ExpiresAt.IsZero() { t.Fatal("observation record must not have expiry") } // 第 4 次触发封禁,时长 = firstBanMin(30 分钟) if !failOnce() { t.Fatal("4th trigger should ban the IP") } entry, err = s.Bans.GetByIP(ip) if err != nil { t.Fatalf("get entry: %v", err) } wantExpiry := time.Now().Add(30 * time.Minute) if entry.ExpiresAt.Before(wantExpiry.Add(-time.Minute)) || entry.ExpiresAt.After(wantExpiry.Add(time.Minute)) { t.Fatalf("ban expiry = %v, want ~%v", entry.ExpiresAt, wantExpiry) } if !strings.Contains(entry.Reason, "第1次封禁") { t.Fatalf("reason = %q, want 第1次封禁", entry.Reason) } if entry.BanCount != 4 { t.Fatalf("ban_count = %d, want 4", entry.BanCount) } // IP 现在处于封禁状态 if banned, _ := s.Bans.IsBanned(ip); !banned { t.Fatal("IP should be banned") } } // TestStagedBanEscalation 验证封禁档位递增:30分钟 → 3小时 → 3个月 → 半年(上限)。 func TestStagedBanEscalation(t *testing.T) { s := newTestStores(t) const ip = "203.0.113.11" const maxFail = 2 failOnce := func() bool { banned, _ := s.RecordAuthFailure(ip, maxFail, 30, "登录失败次数过多") return banned } // 第 1 次触发需要 maxFail 次失败;此后每次失败即触发下一轮 if failOnce() { t.Fatal("fail 1 must not trigger") } if failOnce() { // 触发 1 t.Fatal("trigger 1 must not ban") } if failOnce() { // 触发 2 t.Fatal("trigger 2 must not ban") } if failOnce() { // 触发 3 t.Fatal("trigger 3 must not ban") } // 第 4 次触发:30 分钟 if !failOnce() { t.Fatal("trigger 4 should ban") } expectBanDuration(t, s, ip, 4, 30*time.Minute) // 第 5 次:3 小时 expireBan(t, s, ip) if !failOnce() { t.Fatal("trigger 5 should ban") } expectBanDuration(t, s, ip, 5, 3*time.Hour) // 第 6 次:3 个月 expireBan(t, s, ip) if !failOnce() { t.Fatal("trigger 6 should ban") } expectBanDuration(t, s, ip, 6, 90*24*time.Hour) // 第 7 次:半年 expireBan(t, s, ip) if !failOnce() { t.Fatal("trigger 7 should ban") } expectBanDuration(t, s, ip, 7, 180*24*time.Hour) // 第 8 次:仍为半年(上限) expireBan(t, s, ip) if !failOnce() { t.Fatal("trigger 8 should ban") } expectBanDuration(t, s, ip, 8, 180*24*time.Hour) entry, _ := s.Bans.GetByIP(ip) if !strings.Contains(entry.Reason, "第5次封禁") { t.Fatalf("reason = %q, want 第5次封禁", entry.Reason) } } // expectBanDuration 断言该 IP 当前封禁时长约为 min(允许 2 分钟误差)。 func expectBanDuration(t *testing.T, s *Stores, ip string, trigger int, min time.Duration) { t.Helper() entry, err := s.Bans.GetByIP(ip) if err != nil { t.Fatalf("trigger %d: %v", trigger, err) } diff := entry.ExpiresAt.Sub(time.Now()) if diff < min-2*time.Minute || diff > min+2*time.Minute { t.Fatalf("trigger %d: ban duration = %v, want ~%v", trigger, diff, min) } } // expireBan 把该 IP 的封禁记录改成已过期(模拟时间流逝)。 func expireBan(t *testing.T, s *Stores, ip string) { t.Helper() entry, err := s.Bans.GetByIP(ip) if err != nil { t.Fatalf("get entry: %v", err) } entry.ExpiresAt = time.Now().Add(-time.Minute) if err := s.Bans.Update(entry); err != nil { t.Fatalf("update entry: %v", err) } } // TestBanListOnlyBannedOrExpired 验证列表只返回封禁(含已过期)记录, // 仅计数的观察记录不出现。 func TestBanListOnlyBannedOrExpired(t *testing.T) { s := newTestStores(t) // 观察记录:失败计数,未封禁(无到期时间) if _, err := s.Bans.IncrementFail("203.0.113.20"); err != nil { t.Fatalf("increment: %v", err) } // 当前生效的封禁 if err := s.Bans.Create(&db.BanEntry{ IPAddress: "198.51.100.21", Reason: "第1次封禁:登录失败次数过多(第4次触发,失败5次)", FailCount: 5, BanCount: 4, ExpiresAt: time.Now().Add(30 * time.Minute), }); err != nil { t.Fatalf("create ban: %v", err) } // 已过期的封禁(历史) if err := s.Bans.Create(&db.BanEntry{ IPAddress: "198.51.100.22", Reason: "第2次封禁:登录失败次数过多(第5次触发,失败6次)", FailCount: 6, BanCount: 5, ExpiresAt: time.Now().Add(-24 * time.Hour), }); err != nil { t.Fatalf("create expired ban: %v", err) } entries, total, err := s.Bans.List(1, 10) if err != nil { t.Fatalf("list: %v", err) } if total != 2 { t.Fatalf("total = %d, want 2", total) } if len(entries) != 2 { t.Fatalf("len = %d, want 2", len(entries)) } for _, e := range entries { if e.ExpiresAt.Before(time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)) { t.Fatalf("observation record leaked into list: %+v", e) } } } // TestRecordAuthFailureEmptyIPSafe 空 IP 不应产生副作用。 func TestRecordAuthFailureEmptyIPSafe(t *testing.T) { s := newTestStores(t) banned, count := s.RecordAuthFailure("", 3, 30, "登录失败次数过多") if banned || count != 0 { t.Fatalf("empty IP must be a no-op: banned=%v count=%d", banned, count) } if _, err := s.Bans.GetByIP(""); err == nil { t.Fatal("empty IP should not be recorded") } } // TestRecordAuthFailureWebAndProtocolShared 协议层与 Web 层共用封禁记录。 func TestRecordAuthFailureWebAndProtocolShared(t *testing.T) { s := newTestStores(t) const ip = "198.51.100.20" // Web 层已封禁(直接建记录模拟),协议层认证必须被拒绝 s.Bans.Create(&db.BanEntry{ IPAddress: ip, Reason: "web login failures", FailCount: 5, ExpiresAt: time.Now().Add(30 * time.Minute), }) if banned, _ := s.Bans.IsBanned(ip); !banned { t.Fatal("IP should be banned for both web and protocol auth") } } func TestClientIPFromAddr(t *testing.T) { cases := []struct { addr net.Addr want string }{ {nil, ""}, {addrMock("203.0.113.5:12345"), "203.0.113.5"}, {addrMock("[2001:db8::1]:993"), "2001:db8::1"}, {addrMock("bad-format"), "bad-format"}, } for _, tc := range cases { if got := ClientIPFromAddr(tc.addr); got != tc.want { t.Errorf("ClientIPFromAddr(%v) = %q, want %q", tc.addr, got, tc.want) } } } // addrMock 实现 net.Addr 的最小桩。 type addrMock string func (a addrMock) Network() string { return "tcp" } func (a addrMock) String() string { return string(a) } // P3 #13:配额原子预扣——并发/超额场景下不得绕过配额。 func TestTryReserveQuota(t *testing.T) { s := newTestStores(t) // 用户配额 1000 user := &db.User{ Username: "quota_user", PasswordHash: "x", DomainID: 0, QuotaBytes: 1000, IsActive: true, } if err := s.Users.Create(user); err != nil { t.Fatalf("create user: %v", err) } if user.ID == 0 { t.Fatal("user ID must be assigned") } // 预扣 600 成功 ok, err := s.Users.TryReserveQuota(user.ID, 600) if err != nil || !ok { t.Fatalf("reserve 600: ok=%v err=%v", ok, err) } // 再扣 400 正好用完 ok, err = s.Users.TryReserveQuota(user.ID, 400) if err != nil || !ok { t.Fatalf("reserve 400: ok=%v err=%v", ok, err) } // 超出配额被拒且不改变 used_bytes ok, err = s.Users.TryReserveQuota(user.ID, 1) if err != nil { t.Fatalf("reserve beyond quota: %v", err) } if ok { t.Fatal("reserve beyond quota must fail") } got, _ := s.Users.GetByID(user.ID) if got.UsedBytes != 1000 { t.Fatalf("used_bytes = %d, want 1000 (no partial charge)", got.UsedBytes) } // 释放后可以再次预扣 if err := s.Users.UpdateUsedBytes(user.ID, -500); err != nil { t.Fatalf("release: %v", err) } ok, err = s.Users.TryReserveQuota(user.ID, 500) if err != nil || !ok { t.Fatalf("reserve after release: ok=%v err=%v", ok, err) } } // P3 #13:非正 delta 不允许(防御)。 func TestTryReserveQuotaNonPositiveDelta(t *testing.T) { s := newTestStores(t) user := &db.User{Username: "u", PasswordHash: "x", QuotaBytes: 100} if err := s.Users.Create(user); err != nil { t.Fatal(err) } for _, delta := range []int64{0, -10} { ok, err := s.Users.TryReserveQuota(user.ID, delta) if err != nil { t.Fatalf("delta %d: %v", delta, err) } if ok { t.Fatalf("delta %d must not reserve", delta) } } got, _ := s.Users.GetByID(user.ID) if got.UsedBytes != 0 { t.Fatalf("used_bytes = %d, want 0", got.UsedBytes) } }