3 files changed, 195 insertions, 0 deletions
diff --git a/workhorse/internal/headers/content_headers.go b/workhorse/internal/headers/content_headers.go
new file mode 100644
index 00000000000..e43f10745d4
--- /dev/null
+++ b/workhorse/internal/headers/content_headers.go
@@ -0,0 +1,109 @@
+package headers
+
+import (
+	"net/http"
+	"regexp"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/utils/svg"
+)
+
+var (
+	ImageTypeRegex   = regexp.MustCompile(`^image/*`)
+	SvgMimeTypeRegex = regexp.MustCompile(`^image/svg\+xml$`)
+
+	TextTypeRegex = regexp.MustCompile(`^text/*`)
+
+	VideoTypeRegex = regexp.MustCompile(`^video/*`)
+
+	PdfTypeRegex = regexp.MustCompile(`application\/pdf`)
+
+	AttachmentRegex = regexp.MustCompile(`^attachment`)
+	InlineRegex     = regexp.MustCompile(`^inline`)
+)
+
+// Mime types that can't be inlined. Usually subtypes of main types
+var forbiddenInlineTypes = []*regexp.Regexp{SvgMimeTypeRegex}
+
+// Mime types that can be inlined. We can add global types like "image/" or
+// specific types like "text/plain". If there is a specific type inside a global
+// allowed type that can't be inlined we must add it to the forbiddenInlineTypes var.
+// One example of this is the mime type "image". We allow all images to be
+// inlined except for SVGs.
+var allowedInlineTypes = []*regexp.Regexp{ImageTypeRegex, TextTypeRegex, VideoTypeRegex, PdfTypeRegex}
+
+func SafeContentHeaders(data []byte, contentDisposition string) (string, string) {
+	contentType := safeContentType(data)
+	contentDisposition = safeContentDisposition(contentType, contentDisposition)
+	return contentType, contentDisposition
+}
+
+func safeContentType(data []byte) string {
+	// Special case for svg because DetectContentType detects it as text
+	if svg.Is(data) {
+		return "image/svg+xml"
+	}
+
+	// Override any existing Content-Type header from other ResponseWriters
+	contentType := http.DetectContentType(data)
+
+	// If the content is text type, we set to plain, because we don't
+	// want to render it inline if they're html or javascript
+	if isType(contentType, TextTypeRegex) {
+		return "text/plain; charset=utf-8"
+	}
+
+	return contentType
+}
+
+func safeContentDisposition(contentType string, contentDisposition string) string {
+	// If the existing disposition is attachment we return that. This allow us
+	// to force a download from GitLab (ie: RawController)
+	if AttachmentRegex.MatchString(contentDisposition) {
+		return contentDisposition
+	}
+
+	// Checks for mime types that are forbidden to be inline
+	for _, element := range forbiddenInlineTypes {
+		if isType(contentType, element) {
+			return attachmentDisposition(contentDisposition)
+		}
+	}
+
+	// Checks for mime types allowed to be inline
+	for _, element := range allowedInlineTypes {
+		if isType(contentType, element) {
+			return inlineDisposition(contentDisposition)
+		}
+	}
+
+	// Anything else is set to attachment
+	return attachmentDisposition(contentDisposition)
+}
+
+func attachmentDisposition(contentDisposition string) string {
+	if contentDisposition == "" {
+		return "attachment"
+	}
+
+	if InlineRegex.MatchString(contentDisposition) {
+		return InlineRegex.ReplaceAllString(contentDisposition, "attachment")
+	}
+
+	return contentDisposition
+}
+
+func inlineDisposition(contentDisposition string) string {
+	if contentDisposition == "" {
+		return "inline"
+	}
+
+	if AttachmentRegex.MatchString(contentDisposition) {
+		return AttachmentRegex.ReplaceAllString(contentDisposition, "inline")
+	}
+
+	return contentDisposition
+}
+
+func isType(contentType string, mimeType *regexp.Regexp) bool {
+	return mimeType.MatchString(contentType)
+}
diff --git a/workhorse/internal/headers/headers.go b/workhorse/internal/headers/headers.go
new file mode 100644
index 00000000000..63b39a6aa41
--- /dev/null
+++ b/workhorse/internal/headers/headers.go
@@ -0,0 +1,62 @@
+package headers
+
+import (
+	"net/http"
+	"strconv"
+)
+
+// Max number of bytes that http.DetectContentType needs to get the content type
+// Fixme: Go back to 512 bytes once https://gitlab.com/gitlab-org/gitlab-workhorse/issues/208
+// has been merged
+const MaxDetectSize = 4096
+
+// HTTP Headers
+const (
+	ContentDispositionHeader = "Content-Disposition"
+	ContentTypeHeader        = "Content-Type"
+
+	// Workhorse related headers
+	GitlabWorkhorseSendDataHeader = "Gitlab-Workhorse-Send-Data"
+	XSendFileHeader               = "X-Sendfile"
+	XSendFileTypeHeader           = "X-Sendfile-Type"
+
+	// Signal header that indicates Workhorse should detect and set the content headers
+	GitlabWorkhorseDetectContentTypeHeader = "Gitlab-Workhorse-Detect-Content-Type"
+)
+
+var ResponseHeaders = []string{
+	XSendFileHeader,
+	GitlabWorkhorseSendDataHeader,
+	GitlabWorkhorseDetectContentTypeHeader,
+}
+
+func IsDetectContentTypeHeaderPresent(rw http.ResponseWriter) bool {
+	header, err := strconv.ParseBool(rw.Header().Get(GitlabWorkhorseDetectContentTypeHeader))
+	if err != nil || !header {
+		return false
+	}
+
+	return true
+}
+
+// AnyResponseHeaderPresent checks in the ResponseWriter if there is any Response Header
+func AnyResponseHeaderPresent(rw http.ResponseWriter) bool {
+	// If this header is not present means that we want the old behavior
+	if !IsDetectContentTypeHeaderPresent(rw) {
+		return false
+	}
+
+	for _, header := range ResponseHeaders {
+		if rw.Header().Get(header) != "" {
+			return true
+		}
+	}
+	return false
+}
+
+// RemoveResponseHeaders removes any ResponseHeader from the ResponseWriter
+func RemoveResponseHeaders(rw http.ResponseWriter) {
+	for _, header := range ResponseHeaders {
+		rw.Header().Del(header)
+	}
+}
diff --git a/workhorse/internal/headers/headers_test.go b/workhorse/internal/headers/headers_test.go
new file mode 100644
index 00000000000..555406ff165
--- /dev/null
+++ b/workhorse/internal/headers/headers_test.go
@@ -0,0 +1,24 @@
+package headers
+
+import (
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsDetectContentTypeHeaderPresent(t *testing.T) {
+	rw := httptest.NewRecorder()
+
+	rw.Header().Del(GitlabWorkhorseDetectContentTypeHeader)
+	require.Equal(t, false, IsDetectContentTypeHeaderPresent(rw))
+
+	rw.Header().Set(GitlabWorkhorseDetectContentTypeHeader, "true")
+	require.Equal(t, true, IsDetectContentTypeHeaderPresent(rw))
+
+	rw.Header().Set(GitlabWorkhorseDetectContentTypeHeader, "false")
+	require.Equal(t, false, IsDetectContentTypeHeaderPresent(rw))
+
+	rw.Header().Set(GitlabWorkhorseDetectContentTypeHeader, "foobar")
+	require.Equal(t, false, IsDetectContentTypeHeaderPresent(rw))
+}