package handle import ( "ai_scheduler/internal/data/constants" "ai_scheduler/internal/entitys" "ai_scheduler/internal/pkg/l_request" "bytes" "errors" "fmt" "io" "net/http" "net/url" "path/filepath" "strings" "github.com/gabriel-vasile/mimetype" ) // HandleRecognizeFile 这里的目的是无论将什么类型的file都转为二进制格式 // 最终输出:1.将 files.FileData 填充为文件的二进制数据 2.将 files.FileType 填充为文件的类型(当前为 constants.Caller,兼容写入其字符串值) // 判断文件大小(统一限制为10MB);判断文件类型;判断文件是否合法(类型在白名单映射中);无法识别/非法/超限→填充unknown并兼容返回 // 若 FileData 不存在 且 FileUrl 不存在, 则直接退出 // 若 FileData 存在 FileType 存在, 则直接退出 // 若 FileData 存在 FileType 不存在, 则根据 FileData 推断文件类型并填充 FileType // 若 FileUrl 存在, 则下载文件并填充 FileData 和 FileType func HandleRecognizeFile(files *entitys.RecognizeFile) { if files == nil { return } const maxSize = 10 * 1024 * 1024 // 10MB 上限 // 工具:根据 MIME 或扩展名映射到 FileType mapToFileType := func(s string) constants.FileType { if len(s) == 0 { return constants.FileTypeUnknown } s = strings.ToLower(strings.TrimSpace(s)) for ft, items := range constants.FileTypeMappings { for _, item := range items { if !strings.HasPrefix(item, ".") { // MIME if s == item { return ft } } else { // 扩展名 if s == item { return ft } } } } return constants.FileTypeUnknown } // 分支1:无数据、无URL→直接返回 if len(files.FileData) == 0 && len(files.FileUrl) == 0 { return } // 分支2:已有数据且已有类型→直接返回 if len(files.FileData) > 0 && len(strings.TrimSpace(files.FileType.String())) > 0 { return } // 分支3:仅有数据、无类型→内容检测并填充 if len(files.FileData) > 0 && len(strings.TrimSpace(files.FileType.String())) == 0 { if len(files.FileData) > maxSize { files.FileType = constants.FileTypeUnknown return } reader := bytes.NewReader(files.FileData) detected, fileRealMime := detectFileType(reader, "") files.FileType = detected files.FileRealMime = fileRealMime return } // 分支4:存在URL→下载并填充数据与类型 if len(files.FileUrl) > 0 { fileBytes, contentType, err := downloadFile(files.FileUrl) if err != nil || len(fileBytes) == 0 { files.FileType = constants.FileTypeUnknown return } if len(fileBytes) > maxSize { // 超限:不写入数据,类型置 unknown files.FileType = constants.FileTypeUnknown return } // 优先使用响应头的 Content-Type 映射 detected := mapToFileType(contentType) fileRealMime := contentType if detected == constants.FileTypeUnknown { // 回退:内容检测 + URL 文件名扩展名辅助 var fname string if u, perr := url.Parse(files.FileUrl); perr == nil { fname = filepath.Base(u.Path) } reader := bytes.NewReader(fileBytes) detected, fileRealMime = detectFileType(reader, fname) } // 写入数据 files.FileData = fileBytes files.FileType = detected files.FileRealMime = fileRealMime return } } // 下载文件并返回二进制数据、MIME 类型 func downloadFile(fileUrl string) (fileBytes []byte, contentType string, err error) { if len(fileUrl) == 0 { return } req := l_request.Request{ Method: "GET", Url: fileUrl, Headers: map[string]string{ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", "Accept": "image/webp,image/apng,image/*,*/*;q=0.8", }, } res, err := req.Send() if err != nil { return } var ex bool if contentType, ex = res.Headers["Content-Type"]; !ex { err = errors.New("Content-Type不存在") return } if res.StatusCode != http.StatusOK { err = fmt.Errorf("server returned non-200 status: %d", res.StatusCode) } fileBytes = res.Content return fileBytes, contentType, nil } // detectFileType 判断文件类型 func detectFileType(file io.ReadSeeker, filename string) (constants.FileType, string) { // 1. 读取文件头检测 MIME buffer := make([]byte, 512) n, _ := file.Read(buffer) file.Seek(0, io.SeekStart) // 重置读取位置 detectedMIME := mimetype.Detect(buffer[:n]).String() for fileType, items := range constants.FileTypeMappings { for _, item := range items { if !strings.HasPrefix(item, ".") && item == detectedMIME { return fileType, detectedMIME } } } // 2. 备用:通过扩展名检测 ext := strings.ToLower(filepath.Ext(filename)) for fileType, items := range constants.FileTypeMappings { for _, item := range items { if strings.HasPrefix(item, ".") && item == ext { return fileType, ext } } } return constants.FileTypeUnknown, "" }