diff --git a/internal/proxy/tempfile.go b/internal/proxy/tempfile.go new file mode 100644 index 0000000..fa31af7 --- /dev/null +++ b/internal/proxy/tempfile.go @@ -0,0 +1,491 @@ +// Package proxy 提供反向代理功能的临时文件处理。 +// +// 该文件包含代理响应的临时文件存储功能,用于保护内存: +// - 响应大小检测 +// - 超过阈值时写入临时文件 +// - 超过最大值时返回 502 +// - 响应完成后自动清理 +// +// 主要用途: +// +// 大响应场景下避免内存溢出。 +// +// 注意事项: +// - 使用 bodylimit.ParseSize 解析大小字符串 +// - 临时文件在响应完成后删除 +// - 无 Content-Length 时动态检测 +// +// 作者:xfy +package proxy + +import ( + "fmt" + "io" + "os" + "sync" + + "github.com/valyala/fasthttp" + "rua.plus/lolly/internal/middleware/bodylimit" +) + +// TempFileManager 临时文件管理器。 +// +// 管理代理响应的临时文件存储,保护内存免受大响应影响。 +// +// 注意事项: +// - 阈值和最大值的解析在初始化时完成 +// - 临时文件在响应后自动删除 +// - 并发安全 +type TempFileManager struct { + // tempPath 临时文件存储目录 + tempPath string + + // threshold 触发临时文件的阈值(字节) + threshold int64 + + // maxSize 最大允许的响应大小(字节) + maxSize int64 + + // activeFiles 正在使用的临时文件映射 + activeFiles map[string]*TempFile + + // mu 保护 activeFiles 的互斥锁 + mu sync.RWMutex +} + +// TempFile 临时文件包装器。 +// +// 包装临时文件,提供便捷的读写和自动清理功能。 +type TempFile struct { + // file 底层文件句柄 + file *os.File + + // path 文件路径 + path string + + // size 当前写入大小 + size int64 + + // maxSize 最大允许大小 + maxSize int64 + + // exceeded 是否超过最大大小 + exceeded bool +} + +// NewTempFileManager 创建临时文件管理器。 +// +// 参数: +// - tempPath: 临时文件存储目录 +// - threshold: 触发阈值字符串(如 "1m") +// - maxSize: 最大大小字符串(如 "1024m") +// +// 返回值: +// - *TempFileManager: 临时文件管理器实例 +// - error: 解析大小失败时的错误 +func NewTempFileManager(tempPath, threshold, maxSize string) (*TempFileManager, error) { + // 解析阈值 + thresholdBytes, err := bodylimit.ParseSize(threshold) + if err != nil { + return nil, fmt.Errorf("解析 temp_file_threshold 失败: %w", err) + } + + // 解析最大大小 + maxSizeBytes, err := bodylimit.ParseSize(maxSize) + if err != nil { + return nil, fmt.Errorf("解析 max_temp_file_size 失败: %w", err) + } + + // 使用默认临时目录 + if tempPath == "" { + tempPath = os.TempDir() + } + + // 确保临时目录存在 + if err := os.MkdirAll(tempPath, 0755); err != nil { + return nil, fmt.Errorf("创建临时目录失败: %w", err) + } + + return &TempFileManager{ + tempPath: tempPath, + threshold: thresholdBytes, + maxSize: maxSizeBytes, + activeFiles: make(map[string]*TempFile), + }, nil +} + +// ShouldUseTempFile 检查是否应该使用临时文件。 +// +// 参数: +// - contentLength: Content-Length 值(-1 表示未知) +// +// 返回值: +// - bool: 是否应使用临时文件 +func (m *TempFileManager) ShouldUseTempFile(contentLength int64) bool { + // 如果已知大小且超过阈值,使用临时文件 + if contentLength >= 0 && contentLength >= m.threshold { + return true + } + // 未知大小时由调用方决定(动态检测) + return false +} + +// CreateTempFile 创建临时文件。 +// +// 返回值: +// - *TempFile: 临时文件实例 +// - error: 创建失败时的错误 +func (m *TempFileManager) CreateTempFile() (*TempFile, error) { + // 创建临时文件 + file, err := os.CreateTemp(m.tempPath, "lolly-proxy-*") + if err != nil { + return nil, fmt.Errorf("创建临时文件失败: %w", err) + } + + tf := &TempFile{ + file: file, + path: file.Name(), + maxSize: m.maxSize, + } + + // 记录活动文件 + m.mu.Lock() + m.activeFiles[tf.path] = tf + m.mu.Unlock() + + return tf, nil +} + +// GetThreshold 获取触发阈值。 +func (m *TempFileManager) GetThreshold() int64 { + return m.threshold +} + +// GetMaxSize 获取最大大小限制。 +func (m *TempFileManager) GetMaxSize() int64 { + return m.maxSize +} + +// RemoveTempFile 移除临时文件记录。 +// +// 参数: +// - path: 临时文件路径 +func (m *TempFileManager) RemoveTempFile(path string) { + m.mu.Lock() + delete(m.activeFiles, path) + m.mu.Unlock() +} + +// GetActiveCount 获取活动临时文件数量。 +func (m *TempFileManager) GetActiveCount() int { + m.mu.RLock() + count := len(m.activeFiles) + m.mu.RUnlock() + return count +} + +// Write 写入数据到临时文件。 +// +// 参数: +// - data: 要写入的数据 +// +// 返回值: +// - int: 写入的字节数 +// - error: 写入失败或超过最大大小时的错误 +func (tf *TempFile) Write(data []byte) (int, error) { + if tf.exceeded { + return 0, fmt.Errorf("response exceeds max_temp_file_size") + } + + // 检查是否超过最大大小 + if tf.size+int64(len(data)) > tf.maxSize { + tf.exceeded = true + return 0, fmt.Errorf("response exceeds max_temp_file_size") + } + + n, err := tf.file.Write(data) + tf.size += int64(n) + return n, err +} + +// WriteTo 将临时文件内容写入响应。 +// +// 参数: +// - ctx: fasthttp 请求上下文 +// - statusCode: HTTP 状态码 +// +// 返回值: +// - error: 写入失败时的错误 +func (tf *TempFile) WriteTo(ctx *fasthttp.RequestCtx, statusCode int) error { + // 关闭文件以便读取 + if err := tf.file.Close(); err != nil { + return fmt.Errorf("关闭临时文件失败: %w", err) + } + + // 重新打开文件用于读取 + file, err := os.Open(tf.path) + if err != nil { + return fmt.Errorf("打开临时文件失败: %w", err) + } + defer func() { _ = file.Close() }() + + // 设置状态码 + ctx.Response.SetStatusCode(statusCode) + + // 流式传输文件内容 + buf := make([]byte, 64*1024) // 64KB 缓冲区 + for { + n, err := file.Read(buf) + if n > 0 { + ctx.Response.AppendBody(buf[:n]) + } + if err == io.EOF { + break + } + if err != nil { + return fmt.Errorf("读取临时文件失败: %w", err) + } + } + + return nil +} + +// Close 关闭并删除临时文件。 +func (tf *TempFile) Close() error { + if tf.file != nil { + _ = tf.file.Close() + } + if tf.path != "" { + _ = os.Remove(tf.path) + } + return nil +} + +// IsExceeded 检查是否超过最大大小。 +func (tf *TempFile) IsExceeded() bool { + return tf.exceeded +} + +// GetSize 获取当前写入大小。 +func (tf *TempFile) GetSize() int64 { + return tf.size +} + +// GetPath 获取临时文件路径。 +func (tf *TempFile) GetPath() string { + return tf.path +} + +// DynamicTempFileWriter 动态检测临时文件写入器。 +// +// 用于未知 Content-Length 时的动态阈值检测。 +type DynamicTempFileWriter struct { + // manager 临时文件管理器 + manager *TempFileManager + + // tempFile 当前使用的临时文件(如果已切换) + tempFile *TempFile + + // buffer 缓冲区(在未超过阈值前使用) + buffer []byte + + // totalSize 已接收的总大小 + totalSize int64 + + // threshold 触发阈值 + threshold int64 + + // maxSize 最大大小限制 + maxSize int64 + + // switched 是否已切换到临时文件 + switched bool + + // exceeded 是否超过最大大小 + exceeded bool +} + +// NewDynamicTempFileWriter 创建动态临时文件写入器。 +// +// 参数: +// - manager: 临时文件管理器 +// +// 返回值: +// - *DynamicTempFileWriter: 动态写入器实例 +func NewDynamicTempFileWriter(manager *TempFileManager) *DynamicTempFileWriter { + return &DynamicTempFileWriter{ + manager: manager, + buffer: make([]byte, 0, manager.GetThreshold()), + threshold: manager.GetThreshold(), + maxSize: manager.GetMaxSize(), + } +} + +// Write 写入数据。 +// +// 参数: +// - data: 要写入的数据 +// +// 返回值: +// - error: 写入失败或超过最大大小时的错误 +func (w *DynamicTempFileWriter) Write(data []byte) error { + if w.exceeded { + return fmt.Errorf("response exceeds max_temp_file_size") + } + + dataLen := int64(len(data)) + + // 检查是否超过最大大小 + if w.totalSize+dataLen > w.maxSize { + w.exceeded = true + return fmt.Errorf("response exceeds max_temp_file_size") + } + + // 如果已经切换到临时文件,直接写入 + if w.switched { + _, err := w.tempFile.Write(data) + return err + } + + // 检查是否需要切换到临时文件 + if w.totalSize+dataLen >= w.threshold { + // 创建临时文件 + tf, err := w.manager.CreateTempFile() + if err != nil { + return err + } + w.tempFile = tf + w.switched = true + + // 将缓冲区数据写入临时文件 + if len(w.buffer) > 0 { + _, err := w.tempFile.Write(w.buffer) + if err != nil { + return err + } + w.buffer = nil // 释放缓冲区 + } + + // 写入新数据 + _, err = w.tempFile.Write(data) + if err != nil { + return err + } + } else { + // 继续累积到缓冲区 + w.buffer = append(w.buffer, data...) + } + + w.totalSize += dataLen + return nil +} + +// Finalize 完成写入并返回结果。 +// +// 参数: +// - ctx: fasthttp 请求上下文 +// - statusCode: HTTP 状态码 +// +// 返回值: +// - error: 处理失败时的错误 +func (w *DynamicTempFileWriter) Finalize(ctx *fasthttp.RequestCtx, statusCode int) error { + if w.exceeded { + return fmt.Errorf("response exceeds max_temp_file_size") + } + + // 设置状态码 + ctx.Response.SetStatusCode(statusCode) + + // 如果使用了临时文件,流式传输 + if w.switched && w.tempFile != nil { + // 关闭文件以便读取 + if err := w.tempFile.file.Close(); err != nil { + return fmt.Errorf("关闭临时文件失败: %w", err) + } + + // 重新打开文件用于读取 + file, err := os.Open(w.tempFile.path) + if err != nil { + return fmt.Errorf("打开临时文件失败: %w", err) + } + defer func() { _ = file.Close() }() + + // 流式传输文件内容 + buf := make([]byte, 64*1024) // 64KB 缓冲区 + for { + n, err := file.Read(buf) + if n > 0 { + ctx.Response.AppendBody(buf[:n]) + } + if err == io.EOF { + break + } + if err != nil { + _ = file.Close() + return fmt.Errorf("读取临时文件失败: %w", err) + } + } + _ = file.Close() + + // 删除临时文件 + _ = os.Remove(w.tempFile.path) + w.manager.RemoveTempFile(w.tempFile.path) + return nil + } + + // 未使用临时文件,直接返回缓冲区内容 + ctx.Response.SetBody(w.buffer) + return nil +} + +// IsExceeded 检查是否超过最大大小。 +func (w *DynamicTempFileWriter) IsExceeded() bool { + return w.exceeded +} + +// GetTotalSize 获取总大小。 +func (w *DynamicTempFileWriter) GetTotalSize() int64 { + return w.totalSize +} + +// Cleanup 清理资源。 +func (w *DynamicTempFileWriter) Cleanup() { + if w.tempFile != nil { + _ = w.tempFile.Close() + if w.tempFile.path != "" { + w.manager.RemoveTempFile(w.tempFile.path) + } + w.tempFile = nil + } + w.buffer = nil +} + +// defaultTempFileManager 默认临时文件管理器(未配置时使用)。 +var defaultTempFileManager *TempFileManager +var defaultTempFileManagerOnce sync.Once + +// GetDefaultTempFileManager 获取默认临时文件管理器。 +// +// 默认配置: +// - temp_path: 系统临时目录 +// - temp_file_threshold: 1m +// - max_temp_file_size: 1024m +// +// 返回值: +// - *TempFileManager: 默认临时文件管理器 +func GetDefaultTempFileManager() *TempFileManager { + defaultTempFileManagerOnce.Do(func() { + manager, err := NewTempFileManager("", "1m", "1024m") + if err != nil { + // 默认配置不应该出错,但如果出错则创建一个最小配置 + manager = &TempFileManager{ + tempPath: os.TempDir(), + threshold: 1 << 20, // 1MB + maxSize: 1 << 30, // 1GB + activeFiles: make(map[string]*TempFile), + } + } + defaultTempFileManager = manager + }) + return defaultTempFileManager +} diff --git a/internal/proxy/tempfile_cleaner.go b/internal/proxy/tempfile_cleaner.go new file mode 100644 index 0000000..1dd5ed0 --- /dev/null +++ b/internal/proxy/tempfile_cleaner.go @@ -0,0 +1,279 @@ +// Package proxy 提供反向代理功能的临时文件清理。 +// +// 该文件包含孤儿临时文件的清理功能: +// - 定期扫描临时目录 +// - 删除过期的临时文件 +// - 后台 goroutine 运行 +// +// 主要用途: +// +// 清理异常退出时遗留的临时文件。 +// +// 注意事项: +// - 只清理以 "lolly-proxy-" 前缀开头的文件 +// - 清理超过 1 小时的文件 +// - 可通过 stopCh 停止清理器 +// +// 作者:xfy +package proxy + +import ( + "os" + "path/filepath" + "strings" + "sync" + "time" +) + +// TempFileCleaner 临时文件清理器。 +// +// 定期清理临时目录中的孤儿文件。 +// +// 注意事项: +// - 清理器在后台运行 +// - 可通过 Stop 方法停止 +// - 只清理以 "lolly-proxy-" 前缀开头的文件 +type TempFileCleaner struct { + // tempPath 临时文件目录 + tempPath string + + // interval 清理间隔 + interval time.Duration + + // maxAge 文件最大存活时间 + maxAge time.Duration + + // prefix 临时文件前缀 + prefix string + + // stopCh 停止信号 + stopCh chan struct{} + + // stopped 是否已停止 + stopped bool + + // mu 保护 stopped + mu sync.RWMutex +} + +// DefaultCleanupInterval 默认清理间隔(5 分钟)。 +const DefaultCleanupInterval = 5 * time.Minute + +// DefaultMaxFileAge 默认文件最大存活时间(1 小时)。 +const DefaultMaxFileAge = time.Hour + +// TempFilePrefix 临时文件前缀。 +const TempFilePrefix = "lolly-proxy-" + +// NewTempFileCleaner 创建临时文件清理器。 +// +// 参数: +// - tempPath: 临时文件目录 +// - interval: 清理间隔(0 使用默认值 5 分钟) +// - maxAge: 文件最大存活时间(0 使用默认值 1 小时) +// +// 返回值: +// - *TempFileCleaner: 清理器实例 +func NewTempFileCleaner(tempPath string, interval, maxAge time.Duration) *TempFileCleaner { + if interval <= 0 { + interval = DefaultCleanupInterval + } + if maxAge <= 0 { + maxAge = DefaultMaxFileAge + } + if tempPath == "" { + tempPath = os.TempDir() + } + + return &TempFileCleaner{ + tempPath: tempPath, + interval: interval, + maxAge: maxAge, + prefix: TempFilePrefix, + stopCh: make(chan struct{}), + } +} + +// Start 启动清理器。 +// +// 在后台启动一个 goroutine 定期清理临时文件。 +func (c *TempFileCleaner) Start() { + go c.run() +} + +// Stop 停止清理器。 +// +// 发送停止信号并等待清理器退出。 +func (c *TempFileCleaner) Stop() { + c.mu.Lock() + if c.stopped { + c.mu.Unlock() + return + } + c.stopped = true + c.mu.Unlock() + + close(c.stopCh) +} + +// IsStopped 检查清理器是否已停止。 +func (c *TempFileCleaner) IsStopped() bool { + c.mu.RLock() + stopped := c.stopped + c.mu.RUnlock() + return stopped +} + +// run 清理循环。 +func (c *TempFileCleaner) run() { + ticker := time.NewTicker(c.interval) + defer ticker.Stop() + + // 立即执行一次清理 + c.cleanup() + + for { + select { + case <-ticker.C: + c.cleanup() + case <-c.stopCh: + return + } + } +} + +// cleanup 执行一次清理。 +func (c *TempFileCleaner) cleanup() { + // 读取目录 + entries, err := os.ReadDir(c.tempPath) + if err != nil { + // 目录读取失败,跳过本次清理 + return + } + + cutoff := time.Now().Add(-c.maxAge) + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + + // 检查文件名前缀 + if !strings.HasPrefix(name, c.prefix) { + continue + } + + // 获取文件信息 + info, err := entry.Info() + if err != nil { + continue + } + + // 检查文件年龄 + if info.ModTime().After(cutoff) { + continue + } + + // 删除过期文件 + fullPath := filepath.Join(c.tempPath, name) + _ = os.Remove(fullPath) + } +} + +// GetTempPath 获取临时目录路径。 +func (c *TempFileCleaner) GetTempPath() string { + return c.tempPath +} + +// GetInterval 获取清理间隔。 +func (c *TempFileCleaner) GetInterval() time.Duration { + return c.interval +} + +// GetMaxAge 获取文件最大存活时间。 +func (c *TempFileCleaner) GetMaxAge() time.Duration { + return c.maxAge +} + +// CleanupNow 立即执行一次清理(用于测试)。 +func (c *TempFileCleaner) CleanupNow() { + c.cleanup() +} + +// CountOrphanFiles 统计孤儿临时文件数量。 +// +// 返回值: +// - int: 孤儿文件数量 +func (c *TempFileCleaner) CountOrphanFiles() int { + entries, err := os.ReadDir(c.tempPath) + if err != nil { + return 0 + } + + cutoff := time.Now().Add(-c.maxAge) + count := 0 + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + if !strings.HasPrefix(entry.Name(), c.prefix) { + continue + } + + info, err := entry.Info() + if err != nil { + continue + } + + if info.ModTime().Before(cutoff) { + count++ + } + } + + return count +} + +// globalCleaner 全局清理器实例。 +var globalCleaner *TempFileCleaner +var globalCleanerMu sync.RWMutex + +// StartGlobalTempFileCleaner 启动全局临时文件清理器。 +// +// 参数: +// - tempPath: 临时文件目录 +func StartGlobalTempFileCleaner(tempPath string) { + globalCleanerMu.Lock() + defer globalCleanerMu.Unlock() + + if globalCleaner != nil { + globalCleaner.Stop() + } + + globalCleaner = NewTempFileCleaner(tempPath, 0, 0) + globalCleaner.Start() +} + +// StopGlobalTempFileCleaner 停止全局临时文件清理器。 +func StopGlobalTempFileCleaner() { + globalCleanerMu.Lock() + defer globalCleanerMu.Unlock() + + if globalCleaner != nil { + globalCleaner.Stop() + globalCleaner = nil + } +} + +// GetGlobalTempFileCleaner 获取全局临时文件清理器。 +// +// 返回值: +// - *TempFileCleaner: 全局清理器实例(可能为 nil) +func GetGlobalTempFileCleaner() *TempFileCleaner { + globalCleanerMu.RLock() + defer globalCleanerMu.RUnlock() + return globalCleaner +} diff --git a/internal/proxy/tempfile_test.go b/internal/proxy/tempfile_test.go new file mode 100644 index 0000000..7e1f0fb --- /dev/null +++ b/internal/proxy/tempfile_test.go @@ -0,0 +1,542 @@ +// Package proxy 提供临时文件处理功能的测试。 +// +// 该文件测试临时文件模块的各项功能,包括: +// - 临时文件管理器创建 +// - 阈值判定逻辑 +// - 临时文件写入 +// - 动态检测切换 +// - 超过最大大小处理 +// - 清理功能 +// +// 作者:xfy +package proxy + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/valyala/fasthttp" +) + +// TestNewTempFileManager 测试临时文件管理器创建 +func TestNewTempFileManager(t *testing.T) { + tests := []struct { + name string + tempPath string + threshold string + maxSize string + wantErr bool + errContains string + }{ + { + name: "正常创建", + tempPath: t.TempDir(), + threshold: "1mb", + maxSize: "1024mb", + wantErr: false, + }, + { + name: "使用默认临时目录", + tempPath: "", + threshold: "1mb", + maxSize: "1024mb", + wantErr: false, + }, + { + name: "无效阈值格式", + tempPath: t.TempDir(), + threshold: "invalid", + maxSize: "1024mb", + wantErr: true, + errContains: "temp_file_threshold", + }, + { + name: "无效最大大小格式", + tempPath: t.TempDir(), + threshold: "1mb", + maxSize: "invalid", + wantErr: true, + errContains: "max_temp_file_size", + }, + { + name: "使用字节单位", + tempPath: t.TempDir(), + threshold: "1048576", + maxSize: "1073741824", + wantErr: false, + }, + { + name: "使用 kb 单位", + tempPath: t.TempDir(), + threshold: "1024kb", + maxSize: "1048576kb", + wantErr: false, + }, + { + name: "使用 gb 单位", + tempPath: t.TempDir(), + threshold: "0.001gb", + maxSize: "1gb", + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + manager, err := NewTempFileManager(tt.tempPath, tt.threshold, tt.maxSize) + + if tt.wantErr { + if err == nil { + t.Errorf("NewTempFileManager() 期望错误但未返回") + return + } + if tt.errContains != "" && !strContains(err.Error(), tt.errContains) { + t.Errorf("NewTempFileManager() 错误 = %v, 应包含 %q", err, tt.errContains) + } + return + } + + if err != nil { + t.Errorf("NewTempFileManager() 意外错误 = %v", err) + return + } + + if manager == nil { + t.Error("NewTempFileManager() 返回 nil") + return + } + + // 验证阈值和最大值 + if manager.GetThreshold() <= 0 { + t.Error("GetThreshold() 应返回正数") + } + if manager.GetMaxSize() <= 0 { + t.Error("GetMaxSize() 应返回正数") + } + }) + } +} + +// TestTempFileManager_ShouldUseTempFile 测试阈值判定 +func TestTempFileManager_ShouldUseTempFile(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "1mb", "1024mb") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + tests := []struct { + name string + contentLength int64 + want bool + }{ + { + name: "正好等于阈值", + contentLength: 1 << 20, // 1MB + want: true, + }, + { + name: "超过阈值", + contentLength: 2 << 20, // 2MB + want: true, + }, + { + name: "低于阈值", + contentLength: 512 << 10, // 512KB + want: false, + }, + { + name: "未知大小", + contentLength: -1, + want: false, + }, + { + name: "零大小", + contentLength: 0, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := manager.ShouldUseTempFile(tt.contentLength) + if got != tt.want { + t.Errorf("ShouldUseTempFile(%d) = %v, want %v", tt.contentLength, got, tt.want) + } + }) + } +} + +// TestTempFile_Write 测试临时文件写入 +func TestTempFile_Write(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "1mb", "10mb") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + t.Run("正常写入", func(t *testing.T) { + tf, err := manager.CreateTempFile() + if err != nil { + t.Fatalf("创建临时文件失败: %v", err) + } + defer func() { _ = tf.Close() }() + + data := []byte("test data") + n, err := tf.Write(data) + if err != nil { + t.Errorf("Write() 错误 = %v", err) + return + } + if n != len(data) { + t.Errorf("Write() 写入字节数 = %d, want %d", n, len(data)) + } + if tf.GetSize() != int64(len(data)) { + t.Errorf("GetSize() = %d, want %d", tf.GetSize(), len(data)) + } + }) + + t.Run("超过最大大小", func(t *testing.T) { + // 创建小阈值管理器 + smallManager, err := NewTempFileManager(t.TempDir(), "1mb", "100b") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + tf, err := smallManager.CreateTempFile() + if err != nil { + t.Fatalf("创建临时文件失败: %v", err) + } + defer func() { _ = tf.Close() }() + + data := make([]byte, 200) // 超过 100b + _, err = tf.Write(data) + if err == nil { + t.Error("Write() 应返回错误(超过最大大小)") + } + if !tf.IsExceeded() { + t.Error("IsExceeded() 应为 true") + } + }) +} + +// TestTempFile_WriteTo 测试临时文件写入响应 +func TestTempFile_WriteTo(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "1mb", "10mb") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + tf, err := manager.CreateTempFile() + if err != nil { + t.Fatalf("创建临时文件失败: %v", err) + } + defer func() { _ = tf.Close() }() + + // 写入测试数据 + data := []byte("response body content") + _, err = tf.Write(data) + if err != nil { + t.Fatalf("写入数据失败: %v", err) + } + + // 写入响应 + ctx := &fasthttp.RequestCtx{} + err = tf.WriteTo(ctx, 200) + if err != nil { + t.Errorf("WriteTo() 错误 = %v", err) + return + } + + // 验证状态码 + if ctx.Response.StatusCode() != 200 { + t.Errorf("StatusCode() = %d, want 200", ctx.Response.StatusCode()) + } + + // 验证内容 + body := string(ctx.Response.Body()) + if body != "response body content" { + t.Errorf("Body() = %q, want %q", body, "response body content") + } +} + +// TestTempFile_Close 测试临时文件关闭和清理 +func TestTempFile_Close(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "1mb", "10mb") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + tf, err := manager.CreateTempFile() + if err != nil { + t.Fatalf("创建临时文件失败: %v", err) + } + + path := tf.GetPath() + + // 关闭文件 + err = tf.Close() + if err != nil { + t.Errorf("Close() 错误 = %v", err) + } + + // 验证文件已删除 + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Error("Close() 后文件应被删除") + } +} + +// TestDynamicTempFileWriter 测试动态临时文件写入器 +func TestDynamicTempFileWriter(t *testing.T) { + t.Run("小响应使用缓冲区", func(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "1mb", "10mb") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + writer := NewDynamicTempFileWriter(manager) + defer writer.Cleanup() + + // 写入小数据(低于阈值) + data := []byte("small data") + err = writer.Write(data) + if err != nil { + t.Errorf("Write() 错误 = %v", err) + return + } + + // 验证最终化 + ctx := &fasthttp.RequestCtx{} + err = writer.Finalize(ctx, 200) + if err != nil { + t.Errorf("Finalize() 错误 = %v", err) + return + } + + // 验证内容 + body := string(ctx.Response.Body()) + if body != "small data" { + t.Errorf("Body() = %q, want %q", body, "small data") + } + }) + + t.Run("大响应切换到临时文件", func(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "100b", "10mb") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + writer := NewDynamicTempFileWriter(manager) + defer writer.Cleanup() + + // 写入大数据(超过阈值) + data := make([]byte, 200) + for i := range data { + data[i] = byte(i % 256) + } + err = writer.Write(data) + if err != nil { + t.Errorf("Write() 错误 = %v", err) + return + } + + // 验证最终化 + ctx := &fasthttp.RequestCtx{} + err = writer.Finalize(ctx, 200) + if err != nil { + t.Errorf("Finalize() 错误 = %v", err) + return + } + + // 验证内容 + body := ctx.Response.Body() + if len(body) != 200 { + t.Errorf("Body() 长度 = %d, want 200", len(body)) + } + }) + + t.Run("超过最大大小返回错误", func(t *testing.T) { + manager, err := NewTempFileManager(t.TempDir(), "1mb", "100b") + if err != nil { + t.Fatalf("创建管理器失败: %v", err) + } + + writer := NewDynamicTempFileWriter(manager) + defer writer.Cleanup() + + // 写入超过最大值的数据 + data := make([]byte, 200) + err = writer.Write(data) + if err == nil { + t.Error("Write() 应返回错误(超过最大大小)") + } + if !writer.IsExceeded() { + t.Error("IsExceeded() 应为 true") + } + }) +} + +// TestGetDefaultTempFileManager 测试默认临时文件管理器 +func TestGetDefaultTempFileManager(t *testing.T) { + manager1 := GetDefaultTempFileManager() + if manager1 == nil { + t.Fatal("GetDefaultTempFileManager() 返回 nil") + } + + // 验证单例 + manager2 := GetDefaultTempFileManager() + if manager1 != manager2 { + t.Error("GetDefaultTempFileManager() 应返回相同的实例") + } + + // 验证默认配置 + if manager1.GetThreshold() != 1<<20 { + t.Errorf("默认阈值 = %d, want %d", manager1.GetThreshold(), 1<<20) + } + if manager1.GetMaxSize() != 1<<30 { + t.Errorf("默认最大大小 = %d, want %d", manager1.GetMaxSize(), 1<<30) + } +} + +// TestTempFileCleaner 测试临时文件清理器 +func TestTempFileCleaner(t *testing.T) { + t.Run("创建和启动", func(t *testing.T) { + tempDir := t.TempDir() + cleaner := NewTempFileCleaner(tempDir, time.Second, time.Second) + + if cleaner.GetTempPath() != tempDir { + t.Errorf("GetTempPath() = %s, want %s", cleaner.GetTempPath(), tempDir) + } + + cleaner.Start() + time.Sleep(100 * time.Millisecond) + + if cleaner.IsStopped() { + t.Error("Start() 后 IsStopped() 应为 false") + } + + cleaner.Stop() + + if !cleaner.IsStopped() { + t.Error("Stop() 后 IsStopped() 应为 true") + } + }) + + t.Run("清理过期文件", func(t *testing.T) { + tempDir := t.TempDir() + + // 创建一个过期的临时文件 + oldFile := filepath.Join(tempDir, TempFilePrefix+"old") + if err := os.WriteFile(oldFile, []byte("old"), 0644); err != nil { + t.Fatalf("创建测试文件失败: %v", err) + } + + // 修改文件修改时间为过去 + oldTime := time.Now().Add(-2 * time.Hour) + if err := os.Chtimes(oldFile, oldTime, oldTime); err != nil { + t.Fatalf("修改文件时间失败: %v", err) + } + + // 创建一个非过期的临时文件 + newFile := filepath.Join(tempDir, TempFilePrefix+"new") + if err := os.WriteFile(newFile, []byte("new"), 0644); err != nil { + t.Fatalf("创建测试文件失败: %v", err) + } + + // 执行清理(1 小时过期时间) + cleaner := NewTempFileCleaner(tempDir, time.Hour, time.Hour) + cleaner.CleanupNow() + + // 验证过期文件被删除 + if _, err := os.Stat(oldFile); !os.IsNotExist(err) { + t.Error("过期文件应被删除") + } + + // 验证新文件保留 + if _, err := os.Stat(newFile); err != nil { + t.Error("新文件应被保留") + } + }) + + t.Run("不清理非 lolly 前缀文件", func(t *testing.T) { + tempDir := t.TempDir() + + // 创建一个非 lolly 前缀的文件 + otherFile := filepath.Join(tempDir, "other-file") + if err := os.WriteFile(otherFile, []byte("other"), 0644); err != nil { + t.Fatalf("创建测试文件失败: %v", err) + } + + // 修改文件修改时间为过去 + oldTime := time.Now().Add(-2 * time.Hour) + if err := os.Chtimes(otherFile, oldTime, oldTime); err != nil { + t.Fatalf("修改文件时间失败: %v", err) + } + + // 执行清理 + cleaner := NewTempFileCleaner(tempDir, time.Hour, time.Hour) + cleaner.CleanupNow() + + // 验证非 lolly 文件保留 + if _, err := os.Stat(otherFile); err != nil { + t.Error("非 lolly 前缀文件应被保留") + } + }) + + t.Run("统计孤儿文件", func(t *testing.T) { + tempDir := t.TempDir() + + // 创建孤儿文件 + orphanFile := filepath.Join(tempDir, TempFilePrefix+"orphan") + if err := os.WriteFile(orphanFile, []byte("orphan"), 0644); err != nil { + t.Fatalf("创建测试文件失败: %v", err) + } + + // 修改文件修改时间为过去 + oldTime := time.Now().Add(-2 * time.Hour) + if err := os.Chtimes(orphanFile, oldTime, oldTime); err != nil { + t.Fatalf("修改文件时间失败: %v", err) + } + + cleaner := NewTempFileCleaner(tempDir, time.Hour, time.Hour) + count := cleaner.CountOrphanFiles() + if count != 1 { + t.Errorf("CountOrphanFiles() = %d, want 1", count) + } + }) +} + +// TestGlobalTempFileCleaner 测试全局临时文件清理器 +func TestGlobalTempFileCleaner(t *testing.T) { + // 启动全局清理器 + StartGlobalTempFileCleaner(os.TempDir()) + + // 验证已启动 + cleaner := GetGlobalTempFileCleaner() + if cleaner == nil { + t.Error("GetGlobalTempFileCleaner() 不应返回 nil") + } + + // 停止全局清理器 + StopGlobalTempFileCleaner() + + // 验证已停止 + cleaner = GetGlobalTempFileCleaner() + if cleaner != nil { + t.Error("StopGlobalTempFileCleaner() 后 GetGlobalTempFileCleaner() 应返回 nil") + } +} + +// strContains 检查字符串是否包含子串 +func strContains(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(s) > 0 && strContainsHelper(s, substr)) +} + +func strContainsHelper(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +}