FileShare/server/internal/middleware/audit_test.go

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)
	}
}