优化结构

internal/utils/uploadpolicy/policy.go

@@ -0,0 +1,74 @@
+// Package uploadpolicy contains the stateless filename and content policy
+// shared by upload handling and local file serving.
+package uploadpolicy
+
+import (
+	"errors"
+	"path/filepath"
+	"strings"
+)
+
+var (
+	allowedExtensions = map[string]struct{}{
+		".jpg": {}, ".jpeg": {}, ".png": {}, ".gif": {}, ".webp": {}, ".bmp": {}, ".ico": {}, ".avif": {},
+		".mp3": {}, ".wav": {}, ".ogg": {}, ".m4a": {}, ".flac": {}, ".aac": {},
+		".mp4": {}, ".webm": {}, ".mov": {}, ".avi": {}, ".mkv": {},
+		".txt": {}, ".md": {}, ".csv": {}, ".json": {}, ".log": {}, ".pdf": {},
+		".doc": {}, ".docx": {}, ".xls": {}, ".xlsx": {}, ".ppt": {}, ".pptx": {},
+		".zip": {}, ".rar": {}, ".7z": {}, ".tar": {}, ".gz": {}, ".tgz": {}, ".bin": {},
+	}
+
+	inlineExtensions = map[string]struct{}{
+		".jpg": {}, ".jpeg": {}, ".png": {}, ".gif": {}, ".webp": {}, ".bmp": {}, ".ico": {}, ".avif": {},
+		".mp3": {}, ".wav": {}, ".ogg": {}, ".m4a": {}, ".flac": {}, ".aac": {},
+		".mp4": {}, ".webm": {}, ".mov": {}, ".avi": {}, ".mkv": {},
+	}
+)
+
+var errFileExtensionNotAllowed = errors.New("file extension is not allowed")
+
+// ValidateFileName rejects paths, surrounding whitespace, and extensions
+// outside the upload allowlist.
+func ValidateFileName(filename string) error {
+	if filename == "" || strings.TrimSpace(filename) != filename || strings.ContainsAny(filename, `/\`) {
+		return errFileExtensionNotAllowed
+	}
+	if _, ok := allowedExtensions[strings.ToLower(filepath.Ext(filename))]; !ok {
+		return errFileExtensionNotAllowed
+	}
+	return nil
+}
+
+// CanServeInline reports whether a stored object may be rendered by a
+// browser instead of being forced to download.
+func CanServeInline(filename string) bool {
+	if ValidateFileName(filename) != nil {
+		return false
+	}
+	_, ok := inlineExtensions[strings.ToLower(filepath.Ext(filename))]
+	return ok
+}
+
+// MIMEByExtension returns the small extension fallback table used when a
+// client did not provide a content type. Unknown extensions are left for
+// content sniffing by the caller.
+func MIMEByExtension(ext string) string {
+	switch strings.ToLower(ext) {
+	case ".json":
+		return "application/json"
+	case ".xml":
+		return "application/xml"
+	case ".pdf":
+		return "application/pdf"
+	case ".zip":
+		return "application/zip"
+	case ".mp4":
+		return "video/mp4"
+	case ".mp3":
+		return "audio/mpeg"
+	case ".csv":
+		return "text/csv"
+	default:
+		return ""
+	}
+}

internal/utils/uploadpolicy/policy_test.go

@@ -0,0 +1,53 @@
+package uploadpolicy
+
+import "testing"
+
+func TestValidateFileName(t *testing.T) {
+	tests := []struct {
+		name    string
+		value   string
+		wantErr bool
+	}{
+		{name: "known extension", value: "photo.PNG"},
+		{name: "document", value: "report.pdf"},
+		{name: "empty", value: "", wantErr: true},
+		{name: "unknown extension", value: "payload.exe", wantErr: true},
+		{name: "path", value: "../photo.png", wantErr: true},
+		{name: "backslash path", value: `dir\photo.png`, wantErr: true},
+		{name: "surrounding whitespace", value: "photo.png ", wantErr: true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := ValidateFileName(tt.value); (err != nil) != tt.wantErr {
+				t.Fatalf("ValidateFileName(%q) error = %v, wantErr %v", tt.value, err, tt.wantErr)
+			}
+		})
+	}
+}
+
+func TestCanServeInline(t *testing.T) {
+	for _, filename := range []string{"photo.png", "track.MP3", "movie.webm"} {
+		if !CanServeInline(filename) {
+			t.Errorf("CanServeInline(%q) = false", filename)
+		}
+	}
+	for _, filename := range []string{"report.pdf", "archive.zip", "payload.exe", "../photo.png"} {
+		if CanServeInline(filename) {
+			t.Errorf("CanServeInline(%q) = true", filename)
+		}
+	}
+}
+
+func TestMIMEByExtension(t *testing.T) {
+	tests := map[string]string{
+		".JSON": "application/json",
+		".mp3":  "audio/mpeg",
+		".csv":  "text/csv",
+		".exe":  "",
+	}
+	for ext, want := range tests {
+		if got := MIMEByExtension(ext); got != want {
+			t.Errorf("MIMEByExtension(%q) = %q, want %q", ext, got, want)
+		}
+	}
+}