From 7201ab18cc2d3392d699500c7e75c30d5562d12e Mon Sep 17 00:00:00 2001
From: wongearl <vamliang@163.com>
Date: Fri, 30 Jun 2023 09:31:52 +0800
Subject: [PATCH] fix: parse filename from resp header

---
 pkg/net/http.go      | 20 +++++++++++------
 pkg/net/http_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/pkg/net/http.go b/pkg/net/http.go
index ac2b6d4..8a4ea53 100644
--- 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,15 @@ 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=`):]
+			if filename == filepath {
+				filename = ""
+			}
+		}
+	}
+	return
+}
diff --git a/pkg/net/http_test.go b/pkg/net/http_test.go
index ebbd252..f788979 100644
--- 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)
+			}
+		})
+	}
+}