bj_power/bj_power_mes/internal/handler/upload.go

package handler

import (
	"crypto/md5"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// ---------------- 附件统一存储方案（2026-09-19） ----------------
//
// 目录结构：<Upload.Dir>/年/月/日/<文件类型>/<uuid>.<ext>
//   - 年月日按当前日期自动生成（月/日补零）
//   - 业务类型目录全部大写（见下方枚举）
//   - 文件名统一 UUID，扩展名小写，不含中文/空格/特殊字符
//   - 数据库只存相对路径（如 2026/09/19/PROCESS_PDF/xxx.pdf），不存绝对路径
//   - 根目录从 etc/*.yaml 的 Upload.Dir 读取，不硬编码

// 文件类型枚举（全部大写）
const (
	FileTypeProcessPDF      = "PROCESS_PDF"      // 工艺流程作业指导书
	FileTypeInspectionPhoto = "INSPECTION_PHOTO" // 检验/巡检照片
	FileTypeInspectionPDF   = "INSPECTION_PDF"   // 检验报告
	FileTypeDrawing         = "DRAWING"          // 图纸
	FileTypeProcessCard     = "PROCESS_CARD"     // 流程卡
	FileTypePackagePhoto    = "PACKAGE_PHOTO"    // 包装照片
	FileTypeExcelImport     = "EXCEL_IMPORT"     // Excel 导入
	FileTypeOther           = "OTHER"            // 其他
)

var validFileTypes = map[string]bool{
	FileTypeProcessPDF: true, FileTypeInspectionPhoto: true, FileTypeInspectionPDF: true,
	FileTypeDrawing: true, FileTypeProcessCard: true, FileTypePackagePhoto: true,
	FileTypeExcelImport: true, FileTypeOther: true,
}

// allowedExt 允许的扩展名白名单（图纸/报告/图片/表格/压缩包）。白名单而非黑名单，避免上传可执行文件。
var allowedExt = map[string]bool{
	".pdf": true, ".png": true, ".jpg": true, ".jpeg": true, ".gif": true,
	".xlsx": true, ".xls": true, ".doc": true, ".docx": true,
	".dwg": true, ".dxf": true, ".zip": true, ".txt": true, ".csv": true,
}

// deriveFileType 未显式指定文件类型时，按业务类型 + 扩展名推断，兜底 OTHER。
func deriveFileType(bizType, ext string) string {
	ext = strings.ToLower(ext)
	switch {
	case ext == ".xlsx" || ext == ".xls" || ext == ".csv":
		return FileTypeExcelImport
	case bizType == "inspection":
		if ext == ".pdf" {
			return FileTypeInspectionPDF
		}
		if ext == ".png" || ext == ".jpg" || ext == ".jpeg" || ext == ".gif" {
			return FileTypeInspectionPhoto
		}
	case bizType == "material" && (ext == ".pdf" || ext == ".dwg" || ext == ".dxf"):
		return FileTypeDrawing
	case bizType == "work_order":
		if ext == ".pdf" {
			return FileTypeProcessCard
		}
	case bizType == "process_flow" && ext == ".pdf":
		return FileTypeProcessPDF
	}
	return FileTypeOther
}

// buildRelDir 相对目录 年/月/日/文件类型（月/日补零）。
func buildRelDir(fileType string, t time.Time) string {
	return fmt.Sprintf("%d/%02d/%02d/%s", t.Year(), int(t.Month()), t.Day(), fileType)
}

// safeAbsPath 把库内相对路径安全解析为根目录内的绝对路径（阻断 ../ 穿越）。
func safeAbsPath(root, rel string) (string, bool) {
	rel = strings.ReplaceAll(rel, "\\", "/")
	clean := filepath.Clean(filepath.Join(root, rel))
	rootAbs, err := filepath.Abs(root)
	if err != nil {
		return "", false
	}
	abs, err := filepath.Abs(clean)
	if err != nil {
		return "", false
	}
	if abs != rootAbs && !strings.HasPrefix(abs, rootAbs+string(os.PathSeparator)) {
		return "", false
	}
	return abs, true
}

// newUUID 生成 RFC4122 v4 UUID（crypto/rand，不引第三方依赖）。
func newUUID() string {
	b := make([]byte, 16)
	_, _ = rand.Read(b)
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16])
}

// UploadResult 上传落盘结果。
type UploadResult struct {
	RelPath  string // 相对路径：年/月/日/文件类型/uuid.ext
	FileName string // 原始文件名（含扩展名）
	Size     int64
	MD5      string
	Ext      string // 小写，不含点
	MimeType string
}

// RelPath2FileType 从相对路径反推文件类型目录名（年/月/日/类型/uuid.ext → 类型）。
// 供处理器在落盘后回填库内 fileType（避免再解析一次文件名）。
func (u *UploadResult) RelPath2FileType() string {
	parts := strings.Split(u.RelPath, "/")
	if len(parts) >= 2 {
		if ft := parts[len(parts)-2]; validFileTypes[ft] {
			return ft
		}
	}
	return FileTypeOther
}

// saveTypedUpload 流式保存上传文件（边读边算 MD5，内存占用恒定；绝不 ReadAll）。
// fileType 为空或非法时按 bizType+扩展名推断；maxBytes<=0 时用 Upload.MaxMB。
// 返回相对路径（年/月/日/文件类型/uuid.ext）与原始文件名/大小/MD5。
func saveTypedUpload(r *http.Request, field, rootDir, bizType, fileType string, maxBytes int64) (*UploadResult, error) {
	file, header, err := r.FormFile(field)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	ext := strings.ToLower(filepath.Ext(header.Filename))
	if ext == "" || !allowedExt[ext] {
		return nil, fmt.Errorf("不支持的文件类型（仅支持 PDF/图片/Office/图纸/压缩包）")
	}
	if fileType == "" || !validFileTypes[strings.ToUpper(fileType)] {
		fileType = deriveFileType(bizType, ext)
	} else {
		fileType = strings.ToUpper(fileType)
	}
	if maxBytes > 0 && header.Size > 0 && header.Size > maxBytes {
		return nil, fmt.Errorf("文件超过 %dMB 上限，请压缩后重传", maxBytes>>20)
	}

	relDir := buildRelDir(fileType, time.Now())
	absDir, ok := safeAbsPath(rootDir, relDir)
	if !ok {
		return nil, fmt.Errorf("非法存储路径")
	}
	if err := os.MkdirAll(absDir, 0o755); err != nil {
		return nil, err
	}
	storedName := newUUID() + ext
	relPath := relDir + "/" + storedName
	dst := filepath.Join(absDir, storedName)

	out, err := os.Create(dst)
	if err != nil {
		return nil, err
	}
	hash := md5.New()
	var written int64
	if maxBytes > 0 {
		written, err = io.Copy(io.MultiWriter(out, hash), io.LimitReader(file, maxBytes+1))
	} else {
		written, err = io.Copy(io.MultiWriter(out, hash), file)
	}
	_ = out.Close()
	if err != nil {
		_ = os.Remove(dst)
		return nil, err
	}
	if maxBytes > 0 && written > maxBytes {
		_ = os.Remove(dst)
		return nil, fmt.Errorf("文件超过 %dMB 上限，请压缩后重传", maxBytes>>20)
	}

	return &UploadResult{
		RelPath:  relPath,
		FileName: header.Filename,
		Size:     written,
		MD5:      hex.EncodeToString(hash.Sum(nil)),
		Ext:      strings.TrimPrefix(ext, "."),
		MimeType: header.Header.Get("Content-Type"),
	}, nil
}

// uploadMaxBytes 单文件上限：按文件类型取 Upload.MaxSize 配置，未配置用 Upload.MaxMB。
func uploadMaxBytes(maxSizes map[string]int64, maxMB int64, fileType string) int64 {
	if maxSizes != nil {
		if v, ok := maxSizes[strings.ToUpper(fileType)]; ok && v > 0 {
			return v
		}
	}
	if maxMB > 0 {
		return maxMB << 20
	}
	return 20 << 20
}

// maxUploadBytes 计算路由组请求体上限：取 MaxMB 与各文件类型 MaxSize 的最大值（兜底 20MB）。
func maxUploadBytes(maxMB int64, maxSizes map[string]int64) int64 {
	limit := maxMB << 20
	if limit <= 0 {
		limit = 20 << 20
	}
	for _, v := range maxSizes {
		if v > limit {
			limit = v
		}
	}
	return limit
}

// fileURL 由相对路径生成免鉴权直链（query 传名，天然支持多级目录）。
// 注：go-zero 路由不支持多段通配，故统一走 /api/v1/files?name=<urlencoded 相对路径>。
func fileURL(relPath string) string {
	return "/api/v1/files?name=" + url.QueryEscape(relPath)
}

// serveUploadFile 下发上传目录中的文件（相对路径可为多级「年/月/日/类型/文件名」，防路径穿越）。
func serveUploadFile(w http.ResponseWriter, r *http.Request, dir, name string) {
	name = strings.TrimSpace(name)
	if name == "" {
		http.Error(w, "非法文件名", http.StatusBadRequest)
		return
	}
	abs, ok := safeAbsPath(dir, name)
	if !ok {
		http.Error(w, "非法文件名", http.StatusBadRequest)
		return
	}
	f, err := os.Open(abs)
	if err != nil {
		http.Error(w, "文件不存在", http.StatusNotFound)
		return
	}
	defer f.Close()
	ct := "application/octet-stream"
	base := strings.ToLower(abs)
	switch {
	case strings.HasSuffix(base, ".pdf"):
		ct = "application/pdf"
	case strings.HasSuffix(base, ".jpg"), strings.HasSuffix(base, ".jpeg"):
		ct = "image/jpeg"
	case strings.HasSuffix(base, ".png"):
		ct = "image/png"
	case strings.HasSuffix(base, ".gif"):
		ct = "image/gif"
	case strings.HasSuffix(base, ".xlsx"):
		ct = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
	case strings.HasSuffix(base, ".xls"):
		ct = "application/vnd.ms-excel"
	}
	w.Header().Set("Content-Type", ct)
	http.ServeContent(w, r, filepath.Base(abs), time.Time{}, f)
}

// randHex 随机十六进制串（保留给其他场景的短随机名使用）。
func randHex(n int) string {
	b := make([]byte, n)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)[:n*2]
}