Configure sendfile and injecters in routes.go

internal/lfs/lfs.go

@@ -7,7 +7,6 @@ package lfs
 import (
 	"../api"
 	"../helper"
-	"../proxy"
 	"bytes"
 	"crypto/sha256"
 	"encoding/hex"
@@ -20,8 +19,8 @@ import (
 	"path/filepath"
 )
 
-func PutStore(a *api.API, p *proxy.Proxy) http.Handler {
-	return lfsAuthorizeHandler(a, handleStoreLfsObject(p))
+func PutStore(a *api.API, h http.Handler) http.Handler {
+	return lfsAuthorizeHandler(a, handleStoreLfsObject(h))
 }
 
 func lfsAuthorizeHandler(myAPI *api.API, handleFunc api.HandleFunc) http.Handler {

internal/proxy/proxy.go

@@ -3,7 +3,6 @@ package proxy
 import (
 	"../badgateway"
 	"../helper"
-	"../inject"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -34,8 +33,6 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	// Set Workhorse version
 	req.Header.Set("Gitlab-Workhorse", p.Version)
-	rw := inject.NewSendFileResponseWriter(w, &req)
-	defer rw.Flush()
 
-	p.reverseProxy.ServeHTTP(&rw, &req)
+	p.reverseProxy.ServeHTTP(w, &req)
 }

internal/senddata/injecter.go

@@ -14,7 +14,7 @@ type Injecter interface {
 
 type Prefix string
 
-const Header = "Gitlab-Workhorse-Send-Data"
+const HeaderKey = "Gitlab-Workhorse-Send-Data"
 
 func (p Prefix) Match(s string) bool {
 	return strings.HasPrefix(s, string(p))

internal/senddata/senddata.go

+package senddata
+
+import (
+	"net/http"
+)
+
+type sendDataResponseWriter struct {
+	rw        http.ResponseWriter
+	status    int
+	hijacked  bool
+	req       *http.Request
+	injecters []Injecter
+}
+
+func SendData(h http.Handler, injecters ...Injecter) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		s := sendDataResponseWriter{
+			rw:        w,
+			req:       r,
+			injecters: injecters,
+		}
+		defer s.Flush()
+		h.ServeHTTP(&s, r)
+	})
+}
+
+func (s *sendDataResponseWriter) Header() http.Header {
+	return s.rw.Header()
+}
+
+func (s *sendDataResponseWriter) Write(data []byte) (n int, err error) {
+	if s.status == 0 {
+		s.WriteHeader(http.StatusOK)
+	}
+	if s.hijacked {
+		return
+	}
+	return s.rw.Write(data)
+}
+
+func (s *sendDataResponseWriter) WriteHeader(status int) {
+	if s.status != 0 {
+		return
+	}
+
+	s.status = status
+	if s.status != http.StatusOK {
+		s.rw.WriteHeader(s.status)
+		return
+	}
+
+	if header := s.Header().Get(HeaderKey); header != "" {
+		s.Header().Del(HeaderKey)
+		for _, injecter := range s.injecters {
+			if injecter.Match(header) {
+				s.hijacked = true
+				injecter.Inject(s.rw, s.req, header)
+				return
+			}
+		}
+	}
+
+	s.rw.WriteHeader(s.status)
+	return
+}
+
+func (s *sendDataResponseWriter) Flush() {
+	s.WriteHeader(http.StatusOK)
+}

internal/inject/sendfile.go → internal/sendfile/sendfile.go

@@ -4,12 +4,10 @@ via the X-Sendfile mechanism. All that is needed in the Rails code is the
 'send_file' method.
 */
 
-package inject
+package sendfile
 
 import (
-	"../git"
 	"../helper"
-	"../senddata"
 	"log"
 	"net/http"
 )
@@ -23,14 +21,17 @@ type sendFileResponseWriter struct {
 	req      *http.Request
 }
 
-func NewSendFileResponseWriter(rw http.ResponseWriter, req *http.Request) sendFileResponseWriter {
-	s := sendFileResponseWriter{
+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("X-Sendfile-Type", "X-Sendfile")
-	return s
+		defer s.Flush()
+		h.ServeHTTP(s, req)
+	})
 }
 
 func (s *sendFileResponseWriter) Header() http.Header {
@@ -68,20 +69,6 @@ func (s *sendFileResponseWriter) WriteHeader(status int) {
 		return
 	}
 
-	if header := s.Header().Get(senddata.Header); header != "" {
-		s.Header().Del(senddata.Header)
-		for _, handler := range []senddata.Injecter{
-			git.SendBlob,
-			git.SendArchive,
-		} {
-			if handler.Match(header) {
-				s.hijacked = true
-				handler.Inject(s.rw, s.req, header)
-				return
-			}
-		}
-	}
-
 	s.rw.WriteHeader(s.status)
 	return
 }

internal/upstream/routes.go

@@ -6,6 +6,8 @@ import (
 	"../git"
 	"../lfs"
 	proxypkg "../proxy"
+	"../senddata"
+	"../sendfile"
 	"../staticpages"
 	"net/http"
 	"regexp"
@@ -37,10 +39,14 @@ func (u *Upstream) configureRoutes() {
 		u.RoundTripper,
 	)
 	static := &staticpages.Static{u.DocumentRoot}
-	proxy := proxypkg.NewProxy(
+	proxy := senddata.SendData(
+		sendfile.SendFile(proxypkg.NewProxy(
 			u.Backend,
 			u.Version,
 			u.RoundTripper,
+		)),
+		git.SendArchive,
+		git.SendBlob,
 	)
 
 	u.Routes = []route{