Use interfaces

internal/git/archive.go

@@ -6,6 +6,7 @@ package git
 
 import (
 	"../helper"
+	"../senddata"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -19,16 +20,18 @@ import (
 	"time"
 )
 
-const SendArchivePrefix = "git-archive:"
+type archive struct{ senddata.Prefix }
 
-func SendArchive(w http.ResponseWriter, r *http.Request, sendData string) {
+var SendArchive = &archive{"git-archive:"}
+
+func (a *archive) Handle(w http.ResponseWriter, r *http.Request, sendData string) {
 	var params struct {
 		RepoPath      string
 		ArchivePath   string
 		ArchivePrefix string
 		CommitId      string
 	}
-	if err := unpackSendData(&params, sendData, SendArchivePrefix); err != nil {
+	if err := a.Unpack(&params, sendData); err != nil {
 		helper.Fail500(w, fmt.Errorf("SendArchive: unpack sendData: %v", err))
 		return
 	}

internal/git/blob.go

@@ -2,20 +2,20 @@ package git
 
 import (
 	"../helper"
-	"encoding/base64"
-	"encoding/json"
+	"../senddata"
 	"fmt"
 	"io"
 	"log"
 	"net/http"
-	"strings"
 )
 
-const SendBlobPrefix = "git-blob:"
+type blob struct{ senddata.Prefix }
 
-func SendBlob(w http.ResponseWriter, r *http.Request, sendData string) {
+var SendBlob = &blob{"git-blob:"}
+
+func (b *blob) Handle(w http.ResponseWriter, r *http.Request, sendData string) {
 	var params struct{ RepoPath, BlobId string }
-	if err := unpackSendData(&params, sendData, SendBlobPrefix); err != nil {
+	if err := b.Unpack(&params, sendData); err != nil {
 		helper.Fail500(w, fmt.Errorf("SendBlob: unpack sendData: %v", err))
 		return
 	}
@@ -43,14 +43,3 @@ func SendBlob(w http.ResponseWriter, r *http.Request, sendData string) {
 		return
 	}
 }
-
-func unpackSendData(result interface{}, sendData string, prefix string) error {
-	jsonBytes, err := base64.URLEncoding.DecodeString(strings.TrimPrefix(sendData, prefix))
-	if err != nil {
-		return err
-	}
-	if err := json.Unmarshal([]byte(jsonBytes), result); err != nil {
-		return err
-	}
-	return nil
-}

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

@@ -4,20 +4,17 @@ via the X-Sendfile mechanism. All that is needed in the Rails code is the
 'send_file' method.
 */
 
-package senddata
+package inject
 
 import (
 	"../git"
 	"../helper"
+	"../senddata"
 	"log"
 	"net/http"
-	"strings"
 )
 
-const (
-	sendDataResponseHeader = "Gitlab-Workhorse-Send-Data"
-	sendFileResponseHeader = "X-Sendfile"
-)
+const sendFileResponseHeader = "X-Sendfile"
 
 type sendFileResponseWriter struct {
 	rw       http.ResponseWriter
@@ -71,18 +68,15 @@ func (s *sendFileResponseWriter) WriteHeader(status int) {
 		return
 	}
 
-	sendData := s.Header().Get(sendDataResponseHeader)
-	s.Header().Del(sendDataResponseHeader)
-	for _, handler := range []struct {
-		prefix string
-		f      func(http.ResponseWriter, *http.Request, string)
-	}{
-		{git.SendBlobPrefix, git.SendBlob},
-		{git.SendArchivePrefix, git.SendArchive},
+	header := s.Header().Get(senddata.Header)
+	s.Header().Del(senddata.Header)
+	for _, handler := range []senddata.Handler{
+		git.SendBlob,
+		git.SendArchive,
 	} {
-		if strings.HasPrefix(sendData, handler.prefix) {
+		if handler.Match(header) {
 			s.hijacked = true
-			handler.f(s.rw, s.req, sendData)
+			handler.Handle(s.rw, s.req, header)
 			return
 		}
 	}

internal/proxy/proxy.go

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

internal/senddata/handler.go

+package senddata
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"net/http"
+	"strings"
+)
+
+type Handler interface {
+	Match(string) bool
+	Handle(http.ResponseWriter, *http.Request, string)
+}
+
+type Prefix string
+
+const Header = "Gitlab-Workhorse-Send-Data"
+
+func (p Prefix) Match(s string) bool {
+	return strings.HasPrefix(s, string(p))
+}
+
+func (p Prefix) Unpack(result interface{}, sendData string) error {
+	jsonBytes, err := base64.URLEncoding.DecodeString(strings.TrimPrefix(sendData, string(p)))
+	if err != nil {
+		return err
+	}
+	if err := json.Unmarshal([]byte(jsonBytes), result); err != nil {
+		return err
+	}
+	return nil
+}