2 files changed, 32 insertions, 26 deletions

diff --git a/workhorse/internal/dependencyproxy/dependencyproxy.go b/workhorse/internal/dependencyproxy/dependencyproxy.go
index cfb3045544f..90f3042a342 100644
--- a/workhorse/internal/dependencyproxy/dependencyproxy.go
+++ b/workhorse/internal/dependencyproxy/dependencyproxy.go
@@ -4,37 +4,17 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"net"
 	"net/http"
-	"time"
 
-	"gitlab.com/gitlab-org/labkit/correlation"
 	"gitlab.com/gitlab-org/labkit/log"
-	"gitlab.com/gitlab-org/labkit/tracing"
 
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper"
+	"gitlab.com/gitlab-org/gitlab/workhorse/internal/helper/httptransport"
 	"gitlab.com/gitlab-org/gitlab/workhorse/internal/senddata"
 )
 
-// httpTransport defines a http.Transport with values
-// that are more restrictive than for http.DefaultTransport,
-// they define shorter TLS Handshake, and more aggressive connection closing
-// to prevent the connection hanging and reduce FD usage
-var httpTransport = tracing.NewRoundTripper(correlation.NewInstrumentedRoundTripper(&http.Transport{
-	Proxy: http.ProxyFromEnvironment,
-	DialContext: (&net.Dialer{
-		Timeout:   30 * time.Second,
-		KeepAlive: 10 * time.Second,
-	}).DialContext,
-	MaxIdleConns:          2,
-	IdleConnTimeout:       30 * time.Second,
-	TLSHandshakeTimeout:   10 * time.Second,
-	ExpectContinueTimeout: 10 * time.Second,
-	ResponseHeaderTimeout: 30 * time.Second,
-}))
-
 var httpClient = &http.Client{
-	Transport: httpTransport,
+	Transport: httptransport.New(),
 }
 
 type Injector struct {
@@ -87,15 +67,28 @@ func (p *Injector) Inject(w http.ResponseWriter, r *http.Request, sendData strin
 		return
 	}
 
+	w.Header().Set("Content-Length", dependencyResponse.Header.Get("Content-Length"))
+
 	teeReader := io.TeeReader(dependencyResponse.Body, w)
 	saveFileRequest, err := http.NewRequestWithContext(r.Context(), "POST", r.URL.String()+"/upload", teeReader)
 	if err != nil {
 		helper.Fail500(w, r, fmt.Errorf("dependency proxy: failed to create request: %w", err))
 	}
 	saveFileRequest.Header = helper.HeaderClone(r.Header)
-	saveFileRequest.ContentLength = dependencyResponse.ContentLength
 
-	w.Header().Del("Content-Length")
+	// forward headers from dependencyResponse to rails and client
+	for key, values := range dependencyResponse.Header {
+		saveFileRequest.Header.Del(key)
+		w.Header().Del(key)
+		for _, value := range values {
+			saveFileRequest.Header.Add(key, value)
+			w.Header().Add(key, value)
+		}
+	}
+
+	// workhorse hijack overwrites the Content-Type header, but we need this header value
+	saveFileRequest.Header.Set("Workhorse-Proxy-Content-Type", dependencyResponse.Header.Get("Content-Type"))
+	saveFileRequest.ContentLength = dependencyResponse.ContentLength
 
 	nrw := &nullResponseWriter{header: make(http.Header)}
 	p.uploadHandler.ServeHTTP(nrw, saveFileRequest)
diff --git a/workhorse/internal/dependencyproxy/dependencyproxy_test.go b/workhorse/internal/dependencyproxy/dependencyproxy_test.go
index 37e54c0b756..d9169b2b4ce 100644
--- a/workhorse/internal/dependencyproxy/dependencyproxy_test.go
+++ b/workhorse/internal/dependencyproxy/dependencyproxy_test.go
@@ -33,7 +33,7 @@ func (f *fakeUploadHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 type errWriter struct{ writes int }
 
-func (w *errWriter) Header() http.Header { return nil }
+func (w *errWriter) Header() http.Header { return make(http.Header) }
 func (w *errWriter) WriteHeader(h int)   {}
 
 // First call of Write function succeeds while all the subsequent ones fail
@@ -112,8 +112,15 @@ func TestInject(t *testing.T) {
 
 func TestSuccessfullRequest(t *testing.T) {
 	content := []byte("result")
+	contentLength := strconv.Itoa(len(content))
+	contentType := "foo"
+	dockerContentDigest := "sha256:asdf1234"
+	overriddenHeader := "originResourceServer"
 	originResourceServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Length", strconv.Itoa(len(content)))
+		w.Header().Set("Content-Length", contentLength)
+		w.Header().Set("Content-Type", contentType)
+		w.Header().Set("Docker-Content-Digest", dockerContentDigest)
+		w.Header().Set("Overridden-Header", overriddenHeader)
 		w.Write(content)
 	}))
 
@@ -130,11 +137,16 @@ func TestSuccessfullRequest(t *testing.T) {
 
 	require.Equal(t, "/target/upload", uploadHandler.request.URL.Path)
 	require.Equal(t, int64(6), uploadHandler.request.ContentLength)
+	require.Equal(t, contentType, uploadHandler.request.Header.Get("Workhorse-Proxy-Content-Type"))
+	require.Equal(t, dockerContentDigest, uploadHandler.request.Header.Get("Docker-Content-Digest"))
+	require.Equal(t, overriddenHeader, uploadHandler.request.Header.Get("Overridden-Header"))
 
 	require.Equal(t, content, uploadHandler.body)
 
 	require.Equal(t, 200, response.Code)
 	require.Equal(t, string(content), response.Body.String())
+	require.Equal(t, contentLength, response.Header().Get("Content-Length"))
+	require.Equal(t, dockerContentDigest, response.Header().Get("Docker-Content-Digest"))
 }
 
 func TestIncorrectSendData(t *testing.T) {
@@ -175,6 +187,7 @@ func TestFailedOriginServer(t *testing.T) {
 func makeRequest(injector *Injector, data string) *httptest.ResponseRecorder {
 	w := httptest.NewRecorder()
 	r := httptest.NewRequest("GET", "/target", nil)
+	r.Header.Set("Overridden-Header", "request")
 
 	sendData := base64.StdEncoding.EncodeToString([]byte(data))
 	injector.Inject(w, r, sendData)