Agent-System/internal/upload/validate.go

package upload

import (
	"fmt"
	"net/http"
	"path/filepath"
	"strings"
	"unicode/utf8"
)

const maxFilenameRunes = 200

// 工作流模板上传允许的扩展名 → 期望 MIME。
var allowedTemplateExt = map[string]string{
	".md":   "text/markdown",
	".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
	".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
	".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
}

// 项目材料上传允许的扩展名 → 期望 MIME。
var allowedMaterialExt = map[string]string{
	".txt":  "text/plain",
	".md":   "text/plain",
	".csv":  "text/plain",
	".json": "application/json",
	".xml":  "application/xml",
	".yaml": "text/plain",
	".yml":  "text/plain",
	".log":  "text/plain",
	".pdf":  "application/pdf",
	".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
	".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
}

var blockedDetectedMIME = map[string]struct{}{
	"text/html":              {},
	"application/javascript": {},
	"application/x-javascript": {},
}

// ValidateProjectMaterial 校验文件名、扩展名白名单与内容魔数/MIME。
func ValidateProjectMaterial(filename, declaredMIME string, data []byte) (string, error) {
	base := filepath.Base(strings.TrimSpace(filename))
	if base == "" || base == "." || base == ".." {
		return "", fmt.Errorf("invalid filename")
	}
	if utf8.RuneCountInString(base) > maxFilenameRunes {
		return "", fmt.Errorf("filename too long")
	}
	ext := strings.ToLower(filepath.Ext(base))
	expectedMIME, ok := allowedMaterialExt[ext]
	if !ok {
		return "", fmt.Errorf("file type not allowed: %s", ext)
	}
	if len(data) == 0 {
		return "", fmt.Errorf("empty file")
	}
	if err := rejectDangerousContent(data); err != nil {
		return "", err
	}
	if looksLikeHTML(data) {
		return "", fmt.Errorf("html content rejected")
	}
	detected := http.DetectContentType(data)
	if _, blocked := blockedDetectedMIME[detected]; blocked {
		return "", fmt.Errorf("dangerous content type detected: %s", detected)
	}
	if err := assertExtensionMatchesContent(ext, data, detected); err != nil {
		return "", err
	}
	declaredMIME = strings.TrimSpace(strings.Split(declaredMIME, ";")[0])
	if declaredMIME != "" && !mimeCompatible(declaredMIME, expectedMIME, detected) {
		return "", fmt.Errorf("declared content type mismatch")
	}
	return expectedMIME, nil
}

// ValidateTemplate 校验工作流模板文件名、扩展名白名单与内容魔数/MIME。
func ValidateTemplate(filename, declaredMIME string, data []byte) (string, error) {
	base := filepath.Base(strings.TrimSpace(filename))
	if base == "" || base == "." || base == ".." {
		return "", fmt.Errorf("invalid filename")
	}
	if utf8.RuneCountInString(base) > maxFilenameRunes {
		return "", fmt.Errorf("filename too long")
	}
	ext := strings.ToLower(filepath.Ext(base))
	expectedMIME, ok := allowedTemplateExt[ext]
	if !ok {
		return "", fmt.Errorf("file type not allowed: %s", ext)
	}
	if len(data) == 0 {
		return "", fmt.Errorf("empty file")
	}
	if len(data) > 32*1024*1024 {
		return "", fmt.Errorf("file too large")
	}
	if err := rejectDangerousContent(data); err != nil {
		return "", err
	}
	if looksLikeHTML(data) {
		return "", fmt.Errorf("html content rejected")
	}
	detected := http.DetectContentType(data)
	if _, blocked := blockedDetectedMIME[detected]; blocked {
		return "", fmt.Errorf("dangerous content type detected: %s", detected)
	}
	if err := assertExtensionMatchesContent(ext, data, detected); err != nil {
		return "", err
	}
	declaredMIME = strings.TrimSpace(strings.Split(declaredMIME, ";")[0])
	if declaredMIME != "" && !mimeCompatible(declaredMIME, expectedMIME, detected) {
		return "", fmt.Errorf("declared content type mismatch")
	}
	return expectedMIME, nil
}

func looksLikeHTML(data []byte) bool {
	sample := strings.ToLower(strings.TrimSpace(string(data[:min(len(data), 512)])))
	return strings.HasPrefix(sample, "<!doctype") ||
		strings.HasPrefix(sample, "<html") ||
		strings.Contains(sample, "<script") ||
		strings.Contains(sample, "<iframe")
}

func rejectDangerousContent(data []byte) error {
	if len(data) >= 2 && data[0] == 'M' && data[1] == 'Z' {
		return fmt.Errorf("executable content rejected")
	}
	if len(data) >= 4 && data[0] == 0x7f && data[1] == 'E' && data[2] == 'L' && data[3] == 'F' {
		return fmt.Errorf("executable content rejected")
	}
	head := strings.ToLower(string(data[:min(len(data), 256)]))
	if strings.HasPrefix(head, "#!/") || strings.Contains(head, "<?php") {
		return fmt.Errorf("script content rejected")
	}
	return nil
}

func assertExtensionMatchesContent(ext string, data []byte, detected string) error {
	switch ext {
	case ".pdf":
		if !bytesHasPrefix(data, []byte("%PDF")) {
			return fmt.Errorf("pdf magic mismatch")
		}
	case ".docx", ".xlsx", ".pptx":
		if !bytesHasPrefix(data, []byte("PK\x03\x04")) {
			return fmt.Errorf("office document magic mismatch")
		}
	case ".txt", ".md", ".csv", ".log", ".yaml", ".yml":
		if containsNullByte(data[:min(len(data), 8192)]) {
			return fmt.Errorf("text file contains binary content")
		}
	case ".json":
		if !strings.HasPrefix(strings.TrimSpace(string(data[:min(len(data), 32)])), "{") &&
			!strings.HasPrefix(strings.TrimSpace(string(data[:min(len(data), 32)])), "[") {
			return fmt.Errorf("json magic mismatch")
		}
	}
	_ = detected
	return nil
}

func mimeCompatible(declared, expected, detected string) bool {
	if declared == expected {
		return true
	}
	if strings.HasPrefix(detected, "text/") && strings.HasPrefix(expected, "text/") {
		return true
	}
	if detected == "application/zip" {
		switch expected {
		case allowedMaterialExt[".docx"], allowedMaterialExt[".xlsx"],
			allowedTemplateExt[".docx"], allowedTemplateExt[".xlsx"], allowedTemplateExt[".pptx"]:
			return true
		}
	}
	if detected == "application/octet-stream" {
		return declared == expected || declared == "application/octet-stream"
	}
	return declared == detected
}

func bytesHasPrefix(data, prefix []byte) bool {
	return len(data) >= len(prefix) && string(data[:len(prefix)]) == string(prefix)
}

func containsNullByte(data []byte) bool {
	for _, b := range data {
		if b == 0 {
			return true
		}
	}
	return false
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}