fix: enhance filename extraction logic for downloads and add unit tests

krau · krau · commit 4d2c34500355 · 2026-01-29T17:11:25.000+08:00
diff --git a/core/tasks/directlinks/execute.go b/core/tasks/directlinks/execute.go
@@ -45,9 +45,17 @@ func (t *Task) Execute(ctx context.Context) error {
 			fetchedTotalBytes.Add(resp.ContentLength)
 			file.Size = resp.ContentLength
 			if name := resp.Header.Get("Content-Disposition"); name != "" {
-				// Set file name
 				filename := parseFilename(name)
-				file.Name = filename
+				if filename != "" {
+					file.Name = filename
+				}
+			}
+			// extract filename from URL if Content-Disposition is empty or invalid
+			if file.Name == "" {
+				file.Name = parseFilenameFromURL(file.URL)
+			}
+			if file.Name == "" {
+				return fmt.Errorf("failed to determine filename for %s: Content-Disposition header is empty and URL does not contain a valid filename", file.URL)
 			}
 
 			return nil
diff --git a/core/tasks/directlinks/util.go b/core/tasks/directlinks/util.go
@@ -144,6 +144,41 @@ func tryDecodeGBK(s string) string {
 	return ""
 }
 
+// parseFilenameFromURL extracts filename from URL path
+// This is used as a fallback when Content-Disposition is not available
+func parseFilenameFromURL(rawURL string) string {
+	parsed, err := url.Parse(rawURL)
+	if err != nil {
+		return ""
+	}
+
+	// Get the path part and extract the last segment
+	path := parsed.Path
+	if path == "" {
+		return ""
+	}
+
+	// URL decode the path first
+	decodedPath, err := url.PathUnescape(path)
+	if err != nil {
+		decodedPath = path
+	}
+
+	// Get the last segment of the path
+	lastSlash := strings.LastIndex(decodedPath, "/")
+	if lastSlash == -1 {
+		return decodedPath
+	}
+	filename := decodedPath[lastSlash+1:]
+
+	// Remove query string if somehow still present
+	if idx := strings.Index(filename, "?"); idx != -1 {
+		filename = filename[:idx]
+	}
+
+	return filename
+}
+
 // parseFilenameFallback manually parses filename= when mime.ParseMediaType fails
 func parseFilenameFallback(cd string) string {
 	// Look for filename= (case-insensitive)
diff --git a/core/tasks/directlinks/util_test.go b/core/tasks/directlinks/util_test.go
@@ -0,0 +1,73 @@
+package directlinks
+
+import (
+	"testing"
+)
+
+func TestParseFilenameFromURL(t *testing.T) {
+	tests := []struct {
+		name     string
+		url      string
+		expected string
+	}{
+		{
+			name:     "simple filename",
+			url:      "https://example.com/files/document.pdf",
+			expected: "document.pdf",
+		},
+		{
+			name:     "filename with encoded characters",
+			url:      "https://example.com/files/%E6%B5%8B%E8%AF%95.zip",
+			expected: "测试.zip",
+		},
+		{
+			name:     "filename with query string in URL",
+			url:      "https://example.com/files/image.png?token=abc123",
+			expected: "image.png",
+		},
+		{
+			name:     "nested path",
+			url:      "https://example.com/a/b/c/file.txt",
+			expected: "file.txt",
+		},
+		{
+			name:     "URL with port",
+			url:      "https://example.com:8080/downloads/archive.tar.gz",
+			expected: "archive.tar.gz",
+		},
+		{
+			name:     "empty path",
+			url:      "https://example.com",
+			expected: "",
+		},
+		{
+			name:     "root path only",
+			url:      "https://example.com/",
+			expected: "",
+		},
+		{
+			name:     "filename with spaces encoded",
+			url:      "https://example.com/my%20file%20name.pdf",
+			expected: "my file name.pdf",
+		},
+		{
+			name:     "complex encoded filename",
+			url:      "https://example.com/downloads/%E4%B8%AD%E6%96%87%E6%96%87%E4%BB%B6.docx",
+			expected: "中文文件.docx",
+		},
+		{
+			name:     "invalid URL",
+			url:      "://invalid-url",
+			expected: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := parseFilenameFromURL(tt.url)
+			if result != tt.expected {
+				t.Errorf("parseFilenameFromURL(%q) = %q, want %q", tt.url, result, tt.expected)
+			}
+		})
+	}
+}
