From 5ee5cef1fef7f60d60851e9f99235ee67019b488 Mon Sep 17 00:00:00 2001
From: Daniel Brondani <daniel.brondani@arm.com>
Date: Fri, 13 Dec 2024 16:46:43 +0100
Subject: [PATCH] [cpackget] Handle `Cookie` headers in HTTP GET requests when
 required

---
 cmd/utils/utils.go      | 14 ++++++++++++++
 cmd/utils/utils_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/cmd/utils/utils.go b/cmd/utils/utils.go
index dc4ff8e..bc2471d 100644
--- a/cmd/utils/utils.go
+++ b/cmd/utils/utils.go
@@ -152,6 +152,20 @@ func DownloadFile(URL string, timeout int) (string, error) {
 	}
 	defer resp.Body.Close()
 
+	if resp.StatusCode == http.StatusForbidden {
+		cookie := resp.Header.Get("Set-Cookie")
+		if len(cookie) > 0 {
+			// add cookie and resend GET request
+			log.Debugf("Cookie: %s", cookie)
+			req.Header.Add("Cookie", cookie)
+			resp, err = client.Do(req)
+			if err != nil {
+				log.Error(err)
+				return "", fmt.Errorf("\"%s\": %w", URL, errs.ErrFailedDownloadingFile)
+			}
+		}
+	}
+
 	if resp.StatusCode != http.StatusOK {
 		log.Debugf("bad status: %s", resp.Status)
 		return "", fmt.Errorf("\"%s\": %w", URL, errs.ErrBadRequest)
diff --git a/cmd/utils/utils_test.go b/cmd/utils/utils_test.go
index 9994ba1..226ac90 100644
--- a/cmd/utils/utils_test.go
+++ b/cmd/utils/utils_test.go
@@ -114,6 +114,35 @@ func TestDownloadFile(t *testing.T) {
 		assert.Equal(bytes, goodResponse)
 	})
 
+	t.Run("test download set cookie", func(t *testing.T) {
+		fileName := "file.txt"
+		defer os.Remove(fileName)
+		goodResponse := []byte("all good")
+		cookieContent := "cookie=test"
+		goodServer := httptest.NewServer(
+			http.HandlerFunc(
+				func(w http.ResponseWriter, r *http.Request) {
+					cookie := r.Header.Get("Cookie")
+					if len(cookie) > 0 {
+						if cookie == cookieContent {
+							fmt.Fprint(w, string(goodResponse))
+						}
+					} else {
+						w.Header().Add("Set-Cookie", cookieContent)
+						w.WriteHeader(http.StatusForbidden)
+					}
+				},
+			),
+		)
+		url := goodServer.URL + "/" + fileName
+		_, err1 := utils.DownloadFile(url, 0)
+		assert.Nil(err1)
+		assert.True(utils.FileExists(fileName))
+		bytes, err2 := os.ReadFile(fileName)
+		assert.Nil(err2)
+		assert.Equal(bytes, goodResponse)
+	})
+
 	t.Run("test download uses cache", func(t *testing.T) {
 		fileName := "file.txt"
 		defer os.Remove(fileName)