fix: parse filename from response header

LinuxSuRen · Jun 30, 2023 · 920dfc7 · 920dfc7
1 parent 957c061
commit 920dfc7
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 7 deletions.
diff --git a/pkg/net/http.go b/pkg/net/http.go
@@ -156,13 +156,7 @@ func (h *HTTPDownloader) DownloadFile() error {
 		}
 	}
 
-	if disposition, ok := resp.Header["Content-Disposition"]; ok && len(disposition) >= 1 {
-		h.suggestedFilename = strings.TrimPrefix(disposition[0], `filename="`)
-		h.suggestedFilename = strings.TrimSuffix(h.suggestedFilename, `"`)
-		if h.suggestedFilename == filepath {
-			h.suggestedFilename = ""
-		}
-	}
+	h.suggestedFilename = ParseSuggestedFilename(resp.Header, filepath)
 
 	// pre-hook before get started to download file
 	if h.PreStart != nil && !h.PreStart(resp) {
@@ -336,3 +330,16 @@ func DetectSizeWithRoundTripper(targetURL, output string, showProgress, noProxy,
 	}
 	return
 }
+
+func ParseSuggestedFilename(header http.Header, filepath string) (filename string) {
+	if disposition, ok := header["Content-Disposition"]; ok && len(disposition) >= 1 {
+		if index := strings.LastIndex(disposition[0], `filename="`); index != -1 {
+			filename = disposition[0][index+len(`filename="`):]
+			filename = strings.TrimSuffix(filename, `"`)
+			if filename == filepath {
+				filename = ""
+			}
+		}
+	}
+	return
+}
diff --git a/pkg/net/http_test.go b/pkg/net/http_test.go
@@ -402,3 +402,55 @@ func TestMultiThreadDownloader(t *testing.T) {
 		})
 	}
 }
+
+func Test_getSuggestedFilename(t *testing.T) {
+	type args struct {
+		header   http.Header
+		filepath string
+	}
+	tests := []struct {
+		name         string
+		args         args
+		wantFilename string
+	}{
+		{
+			name: "test1",
+			args: args{
+				header:   http.Header{"Content-Disposition": []string{`attachment; filename="harbor-helm-1.3.18.tar.gz"`}},
+				filepath: "harbor.tar.gz",
+			},
+			wantFilename: "harbor-helm-1.3.18.tar.gz",
+		},
+		{
+			name: "test2",
+			args: args{
+				header:   http.Header{"Content-Disposition": []string{`attachment; filename="harbor-helm-1.3.18.tar.gz"`}},
+				filepath: "harbor-helm-1.3.18.tar.gz",
+			},
+			wantFilename: "",
+		},
+		{
+			name: "test3",
+			args: args{
+				header:   http.Header{"Content-Disposition": []string{`filename="harbor-helm-1.3.18.tar.gz"`}},
+				filepath: "harbor.tar.gz",
+			},
+			wantFilename: "harbor-helm-1.3.18.tar.gz",
+		},
+		{
+			name: "test4",
+			args: args{
+				header:   http.Header{"Content-Disposition": []string{`filename="harbor-helm-1.3.18.tar.gz"`}},
+				filepath: "harbor-helm-1.3.18.tar.gz",
+			},
+			wantFilename: "",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if gotFilename := net.ParseSuggestedFilename(tt.args.header, tt.args.filepath); gotFilename != tt.wantFilename {
+				t.Errorf("getSuggestedFilename() = %v, want %v", gotFilename, tt.wantFilename)
+			}
+		})
+	}
+}