3 files changed, 334 insertions, 0 deletions
diff --git a/workhorse/internal/sendfile/sendfile.go b/workhorse/internal/sendfile/sendfile.go
new file mode 100644
index 00000000000..d009f216eb9
--- /dev/null
+++ b/workhorse/internal/sendfile/sendfile.go
@@ -0,0 +1,162 @@
+/*
+The xSendFile middleware transparently sends static files in HTTP responses
+via the X-Sendfile mechanism. All that is needed in the Rails code is the
+'send_file' method.
+*/
+
+package sendfile
+
+import (
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"regexp"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+
+	"gitlab.com/gitlab-org/labkit/log"
+	"gitlab.com/gitlab-org/labkit/mask"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/headers"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
+)
+
+var (
+	sendFileRequests = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "gitlab_workhorse_sendfile_requests",
+			Help: "How many X-Sendfile requests have been processed by gitlab-workhorse, partitioned by sendfile type.",
+		},
+		[]string{"type"},
+	)
+
+	sendFileBytes = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "gitlab_workhorse_sendfile_bytes",
+			Help: "How many X-Sendfile bytes have been sent by gitlab-workhorse, partitioned by sendfile type.",
+		},
+		[]string{"type"},
+	)
+
+	artifactsSendFile = regexp.MustCompile("builds/[0-9]+/artifacts")
+)
+
+type sendFileResponseWriter struct {
+	rw       http.ResponseWriter
+	status   int
+	hijacked bool
+	req      *http.Request
+}
+
+func SendFile(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		s := &sendFileResponseWriter{
+			rw:  rw,
+			req: req,
+		}
+		// Advertise to upstream (Rails) that we support X-Sendfile
+		req.Header.Set(headers.XSendFileTypeHeader, headers.XSendFileHeader)
+		defer s.flush()
+		h.ServeHTTP(s, req)
+	})
+}
+
+func (s *sendFileResponseWriter) Header() http.Header {
+	return s.rw.Header()
+}
+
+func (s *sendFileResponseWriter) Write(data []byte) (int, error) {
+	if s.status == 0 {
+		s.WriteHeader(http.StatusOK)
+	}
+	if s.hijacked {
+		return len(data), nil
+	}
+	return s.rw.Write(data)
+}
+
+func (s *sendFileResponseWriter) WriteHeader(status int) {
+	if s.status != 0 {
+		return
+	}
+
+	s.status = status
+	if s.status != http.StatusOK {
+		s.rw.WriteHeader(s.status)
+		return
+	}
+
+	file := s.Header().Get(headers.XSendFileHeader)
+	if file != "" && !s.hijacked {
+		// Mark this connection as hijacked
+		s.hijacked = true
+
+		// Serve the file
+		helper.DisableResponseBuffering(s.rw)
+		sendFileFromDisk(s.rw, s.req, file)
+		return
+	}
+
+	s.rw.WriteHeader(s.status)
+}
+
+func sendFileFromDisk(w http.ResponseWriter, r *http.Request, file string) {
+	log.WithContextFields(r.Context(), log.Fields{
+		"file":   file,
+		"method": r.Method,
+		"uri":    mask.URL(r.RequestURI),
+	}).Print("Send file")
+
+	contentTypeHeaderPresent := false
+
+	if headers.IsDetectContentTypeHeaderPresent(w) {
+		// Removing the GitlabWorkhorseDetectContentTypeHeader header to
+		// avoid handling the response by the senddata handler
+		w.Header().Del(headers.GitlabWorkhorseDetectContentTypeHeader)
+		contentTypeHeaderPresent = true
+	}
+
+	content, fi, err := helper.OpenFile(file)
+	if err != nil {
+		http.NotFound(w, r)
+		return
+	}
+	defer content.Close()
+
+	countSendFileMetrics(fi.Size(), r)
+
+	if contentTypeHeaderPresent {
+		data, err := ioutil.ReadAll(io.LimitReader(content, headers.MaxDetectSize))
+		if err != nil {
+			helper.Fail500(w, r, fmt.Errorf("content type detection: %v", err))
+			return
+		}
+
+		content.Seek(0, io.SeekStart)
+
+		contentType, contentDisposition := headers.SafeContentHeaders(data, w.Header().Get(headers.ContentDispositionHeader))
+		w.Header().Set(headers.ContentTypeHeader, contentType)
+		w.Header().Set(headers.ContentDispositionHeader, contentDisposition)
+	}
+
+	http.ServeContent(w, r, "", fi.ModTime(), content)
+}
+
+func countSendFileMetrics(size int64, r *http.Request) {
+	var requestType string
+	switch {
+	case artifactsSendFile.MatchString(r.RequestURI):
+		requestType = "artifacts"
+	default:
+		requestType = "other"
+	}
+
+	sendFileRequests.WithLabelValues(requestType).Inc()
+	sendFileBytes.WithLabelValues(requestType).Add(float64(size))
+}
+
+func (s *sendFileResponseWriter) flush() {
+	s.WriteHeader(http.StatusOK)
+}
diff --git a/workhorse/internal/sendfile/sendfile_test.go b/workhorse/internal/sendfile/sendfile_test.go
new file mode 100644
index 00000000000..d424814b5e5
--- /dev/null
+++ b/workhorse/internal/sendfile/sendfile_test.go
@@ -0,0 +1,171 @@
+package sendfile
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/headers"
+)
+
+func TestResponseWriter(t *testing.T) {
+	upstreamResponse := "hello world"
+
+	fixturePath := "testdata/sent-file.txt"
+	fixtureContent, err := ioutil.ReadFile(fixturePath)
+	require.NoError(t, err)
+
+	testCases := []struct {
+		desc           string
+		sendfileHeader string
+		out            string
+	}{
+		{
+			desc:           "send a file",
+			sendfileHeader: fixturePath,
+			out:            string(fixtureContent),
+		},
+		{
+			desc:           "pass through unaltered",
+			sendfileHeader: "",
+			out:            upstreamResponse,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			r, err := http.NewRequest("GET", "/foo", nil)
+			require.NoError(t, err)
+
+			rw := httptest.NewRecorder()
+			sf := &sendFileResponseWriter{rw: rw, req: r}
+			sf.Header().Set(headers.XSendFileHeader, tc.sendfileHeader)
+
+			upstreamBody := []byte(upstreamResponse)
+			n, err := sf.Write(upstreamBody)
+			require.NoError(t, err)
+			require.Equal(t, len(upstreamBody), n, "bytes written")
+
+			rw.Flush()
+
+			body := rw.Result().Body
+			data, err := ioutil.ReadAll(body)
+			require.NoError(t, err)
+			require.NoError(t, body.Close())
+
+			require.Equal(t, tc.out, string(data))
+		})
+	}
+}
+
+func TestAllowExistentContentHeaders(t *testing.T) {
+	fixturePath := "../../testdata/forgedfile.png"
+
+	httpHeaders := map[string]string{
+		headers.ContentTypeHeader:        "image/png",
+		headers.ContentDispositionHeader: "inline",
+	}
+
+	resp := makeRequest(t, fixturePath, httpHeaders)
+	require.Equal(t, "image/png", resp.Header.Get(headers.ContentTypeHeader))
+	require.Equal(t, "inline", resp.Header.Get(headers.ContentDispositionHeader))
+}
+
+func TestSuccessOverrideContentHeadersFeatureEnabled(t *testing.T) {
+	fixturePath := "../../testdata/forgedfile.png"
+
+	httpHeaders := make(map[string]string)
+	httpHeaders[headers.ContentTypeHeader] = "image/png"
+	httpHeaders[headers.ContentDispositionHeader] = "inline"
+	httpHeaders["Range"] = "bytes=1-2"
+
+	resp := makeRequest(t, fixturePath, httpHeaders)
+	require.Equal(t, "image/png", resp.Header.Get(headers.ContentTypeHeader))
+	require.Equal(t, "inline", resp.Header.Get(headers.ContentDispositionHeader))
+}
+
+func TestSuccessOverrideContentHeadersRangeRequestFeatureEnabled(t *testing.T) {
+	fixturePath := "../../testdata/forgedfile.png"
+
+	fixtureContent, err := ioutil.ReadFile(fixturePath)
+	require.NoError(t, err)
+
+	r, err := http.NewRequest("GET", "/foo", nil)
+	r.Header.Set("Range", "bytes=1-2")
+	require.NoError(t, err)
+
+	rw := httptest.NewRecorder()
+	sf := &sendFileResponseWriter{rw: rw, req: r}
+
+	sf.Header().Set(headers.XSendFileHeader, fixturePath)
+	sf.Header().Set(headers.ContentTypeHeader, "image/png")
+	sf.Header().Set(headers.ContentDispositionHeader, "inline")
+	sf.Header().Set(headers.GitlabWorkhorseDetectContentTypeHeader, "true")
+
+	upstreamBody := []byte(fixtureContent)
+	_, err = sf.Write(upstreamBody)
+	require.NoError(t, err)
+
+	rw.Flush()
+
+	resp := rw.Result()
+	body := resp.Body
+	data, err := ioutil.ReadAll(body)
+	require.NoError(t, err)
+	require.NoError(t, body.Close())
+
+	require.Len(t, data, 2)
+
+	require.Equal(t, "application/octet-stream", resp.Header.Get(headers.ContentTypeHeader))
+	require.Equal(t, "attachment", resp.Header.Get(headers.ContentDispositionHeader))
+}
+
+func TestSuccessInlineWhitelistedTypesFeatureEnabled(t *testing.T) {
+	fixturePath := "../../testdata/image.png"
+
+	httpHeaders := map[string]string{
+		headers.ContentDispositionHeader:               "inline",
+		headers.GitlabWorkhorseDetectContentTypeHeader: "true",
+	}
+
+	resp := makeRequest(t, fixturePath, httpHeaders)
+
+	require.Equal(t, "image/png", resp.Header.Get(headers.ContentTypeHeader))
+	require.Equal(t, "inline", resp.Header.Get(headers.ContentDispositionHeader))
+}
+
+func makeRequest(t *testing.T, fixturePath string, httpHeaders map[string]string) *http.Response {
+	fixtureContent, err := ioutil.ReadFile(fixturePath)
+	require.NoError(t, err)
+
+	r, err := http.NewRequest("GET", "/foo", nil)
+	require.NoError(t, err)
+
+	rw := httptest.NewRecorder()
+	sf := &sendFileResponseWriter{rw: rw, req: r}
+
+	sf.Header().Set(headers.XSendFileHeader, fixturePath)
+	for name, value := range httpHeaders {
+		sf.Header().Set(name, value)
+	}
+
+	upstreamBody := []byte("hello")
+	n, err := sf.Write(upstreamBody)
+	require.NoError(t, err)
+	require.Equal(t, len(upstreamBody), n, "bytes written")
+
+	rw.Flush()
+
+	resp := rw.Result()
+	body := resp.Body
+	data, err := ioutil.ReadAll(body)
+	require.NoError(t, err)
+	require.NoError(t, body.Close())
+
+	require.Equal(t, fixtureContent, data)
+
+	return resp
+}
diff --git a/workhorse/internal/sendfile/testdata/sent-file.txt b/workhorse/internal/sendfile/testdata/sent-file.txt
new file mode 100644
index 00000000000..40e33f8a628
--- /dev/null
+++ b/workhorse/internal/sendfile/testdata/sent-file.txt
@@ -0,0 +1 @@
+This file is sent with X-SendFile