1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  //
  3  // Licensed under the Apache License, Version 2.0 (the "License");
  4  // you may not use this file except in compliance with the License.
  5  // You may obtain a copy of the License at
  6  //
  7  //     http://www.apache.org/licenses/LICENSE-2.0
  8  //
  9  // Unless required by applicable law or agreed to in writing, software
 10  // distributed under the License is distributed on an "AS IS" BASIS,
 11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  // See the License for the specific language governing permissions and
 13  // limitations under the License.
 14  //
 15  // Requirement: Any integration or derivative work must explicitly attribute
 16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  // documentation or user interface, as detailed in the NOTICE file.
 18  
 19  package websocket
 20  
 21  import (
 22  	"fmt"
 23  	"os"
 24  	"path/filepath"
 25  )
 26  
 27  // FileUploadConfig 文件上传配置
 28  type FileUploadConfig struct {
 29  	UploadDir string `json:"upload_dir"` // 文件上传目录
 30  }
 31  
 32  // DefaultFileUploadConfig 默认文件上传配置
 33  func DefaultFileUploadConfig() *FileUploadConfig {
 34  	return &FileUploadConfig{
 35  		UploadDir: "./uploads",
 36  	}
 37  }
 38  
 39  // LoadFileUploadConfigFromEnv 从环境变量加载文件上传配置
 40  func LoadFileUploadConfigFromEnv() *FileUploadConfig {
 41  	config := DefaultFileUploadConfig()
 42  
 43  	// 从环境变量读取配置
 44  	if uploadDir := os.Getenv("FILE_UPLOAD_DIR"); uploadDir != "" {
 45  		config.UploadDir = uploadDir
 46  	}
 47  
 48  	return config
 49  }
 50  
 51  // ValidateConfig 验证配置的有效性
 52  func (c *FileUploadConfig) ValidateConfig() error {
 53  	// 检查存储目录是否存在或可以创建
 54  	if err := c.ensureUploadDir(); err != nil {
 55  		return fmt.Errorf("存储目录配置错误: %v", err)
 56  	}
 57  	return nil
 58  }
 59  
 60  // ensureUploadDir 确保上传目录存在并且可写
 61  func (c *FileUploadConfig) ensureUploadDir() error {
 62  	// 获取绝对路径
 63  	absPath, err := filepath.Abs(c.UploadDir)
 64  	if err != nil {
 65  		return fmt.Errorf("无法解析存储路径: %v", err)
 66  	}
 67  
 68  	// 检查目录是否存在
 69  	if _, err := os.Stat(absPath); os.IsNotExist(err) {
 70  		// 目录不存在，尝试创建
 71  		if err := os.MkdirAll(absPath, 0755); err != nil {
 72  			return fmt.Errorf("无法创建存储目录 %s: %v", absPath, err)
 73  		}
 74  	}
 75  
 76  	// 检查目录是否可写
 77  	if err := c.checkDirWritable(absPath); err != nil {
 78  		return fmt.Errorf("存储目录不可写: %v", err)
 79  	}
 80  
 81  	// 更新为绝对路径
 82  	c.UploadDir = absPath
 83  	return nil
 84  }
 85  
 86  // checkDirWritable 检查目录是否可写
 87  func (c *FileUploadConfig) checkDirWritable(dir string) error {
 88  	// 尝试创建一个临时文件来测试写权限
 89  	testFile := filepath.Join(dir, ".test_write")
 90  	file, err := os.Create(testFile)
 91  	if err != nil {
 92  		return err
 93  	}
 94  	defer func() {
 95  		file.Close()
 96  		os.Remove(testFile) // 清理测试文件
 97  	}()
 98  	return nil
 99  }
100  
101  // GetFileURL 根据文件名生成完整的文件访问URL
102  func (c *FileUploadConfig) GetFileURL(fileName string) string {
103  	return fileName
104  }