package middleware import ( "context" "net/http/httptest" "sync" "testing" "time" "github.com/gin-gonic/gin" "fileshare/internal/audit" "fileshare/internal/model" ) // memSink 测试用内存落库实现。 type memSink struct { mu sync.Mutex logs []model.AuditLog notif chan struct{} } func newMemSink() *memSink { return &memSink{notif: make(chan struct{}, 16)} } func (m *memSink) Save(_ context.Context, logs []model.AuditLog) error { m.mu.Lock() m.logs = append(m.logs, logs...) m.mu.Unlock() m.notif <- struct{}{} return nil } func (m *memSink) snapshot() []model.AuditLog { m.mu.Lock() defer m.mu.Unlock() out := make([]model.AuditLog, len(m.logs)) copy(out, m.logs) return out } // waitFor 等待 sink 收到 n 条记录(带超时)。 func (m *memSink) waitFor(t *testing.T, n int) []model.AuditLog { t.Helper() deadline := time.After(2 * time.Second) for { logs := m.snapshot() if len(logs) >= n { return logs } select { case <-m.notif: case <-deadline: t.Fatalf("等待审计记录超时: 已收到 %d 条", len(logs)) } } } func auditRouter(svc *audit.Service) *gin.Engine { r := gin.New() r.Use(ClientIP(nil)) r.Use(Audit(svc, nil)) // 默认分类器 // 上传路由:命中默认分类器(POST /share/file) r.POST("/share/file", func(c *gin.Context) { // 模拟 handler 填充业务字段并显式落库 AuditSet(c, func(e *audit.Entry) { e.FileCode = "Ab3xY" e.FileName = "hello.zip" e.SizeBytes = 1024 }) AuditRecordRequest(c, svc, model.AuditResultSuccess, "") c.JSON(200, gin.H{"ok": true}) }) // 下载路由:命中默认分类器(GET /share/download) r.GET("/share/download", func(c *gin.Context) { AuditSet(c, func(e *audit.Entry) { e.FileCode = "Xy12Z" }) c.JSON(404, gin.H{"msg": "文件已过期删除"}) // 未显式落库 → 状态码兜底 }) // 普通路由:不命中,不应产生审计 r.GET("/plain", func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) return r } func TestAuditMiddlewareRecordsUpload(t *testing.T) { sink := newMemSink() svc := audit.NewService(sink) r := auditRouter(svc) w := httptest.NewRecorder() req := httptest.NewRequest("POST", "/share/file", nil) req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/120.0.0.0 Safari/537.36") r.ServeHTTP(w, req) if w.Code != 200 { t.Fatalf("上传应成功: %d", w.Code) } logs := sink.waitFor(t, 1) e := logs[0] if e.Action != audit.ActionUpload { t.Errorf("action = %s", e.Action) } if e.FileCode != "Ab3xY" || e.FileName != "hello.zip" { t.Errorf("file fields = %s/%s", e.FileCode, e.FileName) } if e.SizeBytes != 1024 { t.Errorf("size = %d", e.SizeBytes) } if e.Result != model.AuditResultSuccess { t.Errorf("result = %s", e.Result) } if e.DeviceOS != "Windows" || e.DeviceBrowser != "Chrome" || e.DeviceType != "desktop" { t.Errorf("device = %s/%s/%s", e.DeviceOS, e.DeviceBrowser, e.DeviceType) } if e.DurationMs < 0 { t.Errorf("duration = %d", e.DurationMs) } if e.Actor != audit.ActorGuest { t.Errorf("actor = %s", e.Actor) } } func TestAuditMiddlewareSkipsPlainRoutes(t *testing.T) { sink := newMemSink() svc := audit.NewService(sink) r := auditRouter(svc) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest("GET", "/plain", nil)) if w.Code != 200 { t.Fatalf("plain 路由应成功: %d", w.Code) } time.Sleep(150 * time.Millisecond) if logs := sink.snapshot(); len(logs) != 0 { t.Fatalf("普通路由不应产生审计记录: %v", logs) } } func TestAuditFailedDownloadFallback(t *testing.T) { sink := newMemSink() svc := audit.NewService(sink) r := auditRouter(svc) w := httptest.NewRecorder() req := httptest.NewRequest("GET", "/share/download?code=xyz", nil) req.Header.Set("User-Agent", "curl/8.4.0") r.ServeHTTP(w, req) logs := sink.waitFor(t, 1) e := logs[0] if e.Action != audit.ActionDownload { t.Errorf("action = %s", e.Action) } if e.Result != model.AuditResultFailed { t.Errorf("4xx 兜底 result = %s", e.Result) } if e.DeviceType != "bot" { t.Errorf("curl 应识别为 bot: %s", e.DeviceType) } if e.ErrorMsg == "" { t.Error("失败记录应包含错误信息") } } func TestAuditDeniedStatusMapping(t *testing.T) { sink := newMemSink() svc := audit.NewService(sink) r := gin.New() r.Use(ClientIP(nil)) r.Use(Audit(svc, nil)) r.GET("/share/select", func(c *gin.Context) { c.AbortWithStatus(429) }) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest("GET", "/share/select?code=abc", nil)) logs := sink.waitFor(t, 1) if logs[0].Result != model.AuditResultDenied { t.Errorf("429 应映射为 denied: %s", logs[0].Result) } } func TestAuditDownloadBytesCounted(t *testing.T) { sink := newMemSink() svc := audit.NewService(sink) r := gin.New() r.Use(ClientIP(nil)) r.Use(Audit(svc, nil)) r.GET("/share/download", func(c *gin.Context) { payload := []byte("0123456789abcdef") // 16 字节 c.Data(200, "application/octet-stream", payload) // 未显式落库 → 中间件兜底;TransferredBytes 应等于写出字节 }) w := httptest.NewRecorder() r.ServeHTTP(w, httptest.NewRequest("GET", "/share/download?code=bytes", nil)) logs := sink.waitFor(t, 1) if logs[0].TransferredBytes != 16 { t.Errorf("下载字节数 = %d, want 16", logs[0].TransferredBytes) } if logs[0].Result != model.AuditResultSuccess { t.Errorf("result = %s", logs[0].Result) } }