Fix headers forwarding (#928)

We were doing the headers forwarding after writing the body, but
`ResponseWriter.Write()` can set its own headers when others are not found.
Set the headers before doing any write in the response.

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Bugfixes
 
 * Update Go runtime to 1.19.4. [#924](https://github.com/elastic/package-registry/pull/924)
+* Fix headers forwarding when forwarding artifacts requests to the package storage. [#926](https://github.com/elastic/package-registry/pull/926)
 
 ### Added
 

package_storage_test.go

@@ -214,6 +214,7 @@ func TestPackageStorage_ResolverHeadersResponse(t *testing.T) {
 	webServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Foo", "bar")
 		w.Header().Set("Last-Modified", "time")
+		w.Header().Set("Content-Type", "image/svg+xml")
 		fmt.Fprintf(w, "%s\n%s\n%+v\n", r.Method, r.RequestURI, r.Header)
 	}))
 	defer webServer.Close()
@@ -236,11 +237,14 @@ func TestPackageStorage_ResolverHeadersResponse(t *testing.T) {
 		handler         func(w http.ResponseWriter, r *http.Request)
 	}{
 		{
-			endpoint:        "/package/1password/0.1.1/img/1password-logo-light-bg.svg",
-			path:            staticRouterPath,
-			file:            "1password-logo-light-bg.svg.response",
-			responseHeaders: map[string]string{"Last-Modified": "time"},
-			handler:         staticHandler,
+			endpoint: "/package/1password/0.1.1/img/1password-logo-light-bg.svg",
+			path:     staticRouterPath,
+			file:     "1password-logo-light-bg.svg.response",
+			responseHeaders: map[string]string{
+				"Last-Modified": "time",
+				"Content-Type":  "image/svg+xml",
+			},
+			handler: staticHandler,
 		},
 	}
 

storage/resolver.go

@@ -27,7 +27,7 @@ var acceptedHeaders = map[string]string{
 	"Date":           "",
 }
 
-func (resolver storageResolver) pipeRequestProxy(w http.ResponseWriter, r *http.Request, remoteURL string) error {
+func (resolver storageResolver) pipeRequestProxy(w http.ResponseWriter, r *http.Request, remoteURL string) {
 	client := &http.Client{}
 
 	forwardRequest, err := http.NewRequestWithContext(r.Context(), r.Method, remoteURL, nil)
@@ -36,17 +36,23 @@ func (resolver storageResolver) pipeRequestProxy(w http.ResponseWriter, r *http.
 	resp, err := client.Do(forwardRequest)
 	if err != nil {
 		http.Error(w, "error from package-storage server", http.StatusInternalServerError)
+		return
 	}
 	defer resp.Body.Close()
 
+	// Set headers before setting the body. If not, first call to w.Write will
+	// add some default values.
+	addRequestHeadersToResponse(w, resp)
+
 	_, err = io.Copy(w, resp.Body)
 	if err != nil {
 		http.Error(w, "error writing response", http.StatusInternalServerError)
+		return
 	}
 
-	addRequestHeadersToResponse(w, resp)
 	w.WriteHeader(resp.StatusCode)
-	return nil
+
+	return
 }
 
 func addRequestHeadersToRequest(orig, forward *http.Request) {