Accelerate Maven artifact repository uploads

internal/filestore/body_uploader.go

+package filestore
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"net/url"
+	"strings"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/api"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
+)
+
+type PreAuthorizer interface {
+	PreAuthorizeHandler(next api.HandleFunc, suffix string) http.Handler
+}
+
+// UploadVerifier allows to check an upload before sending it to rails
+type UploadVerifier interface {
+	// Verify can abort the upload returning an error
+	Verify(handler *FileHandler) error
+}
+
+// UploadPreparer allows to customize BodyUploader configuration
+type UploadPreparer interface {
+	// Prepare converts api.Response into a *SaveFileOpts, it can optionally return an UploadVerifier that will be
+	// invoked after the real upload, before the finalization with rails
+	Prepare(a *api.Response) (*SaveFileOpts, UploadVerifier, error)
+}
+
+type defaultPreparer struct{}
+
+func (s *defaultPreparer) Prepare(a *api.Response) (*SaveFileOpts, UploadVerifier, error) {
+	return GetOpts(a), nil, nil
+}
+
+// BodyUploader is an http.Handler that perform a pre authorization call to rails before hijacking the request body and
+// uploading it.
+// Providing an UploadPreparer allows to customize the upload process
+func BodyUploader(rails PreAuthorizer, h http.Handler, p UploadPreparer) http.Handler {
+	if p == nil {
+		p = &defaultPreparer{}
+	}
+
+	return rails.PreAuthorizeHandler(func(w http.ResponseWriter, r *http.Request, a *api.Response) {
+		opts, verifier, err := p.Prepare(a)
+		if err != nil {
+			helper.Fail500(w, r, fmt.Errorf("BodyUploader: preparation failed: %v", err))
+			return
+		}
+
+		fh, err := SaveFileFromReader(r.Context(), r.Body, r.ContentLength, opts)
+		if err != nil {
+			helper.Fail500(w, r, fmt.Errorf("BodyUploader: upload failed: %v", err))
+			return
+		}
+
+		if verifier != nil {
+			if err := verifier.Verify(fh); err != nil {
+				helper.Fail500(w, r, fmt.Errorf("BodyUploader: verification failed: %v", err))
+				return
+			}
+		}
+
+		data := url.Values{}
+		for k, v := range fh.GitLabFinalizeFields("file") {
+			data.Set(k, v)
+		}
+
+		// Hijack body
+		body := data.Encode()
+		r.Body = ioutil.NopCloser(strings.NewReader(body))
+		r.ContentLength = int64(len(body))
+		r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+
+		// And proxy the request
+		h.ServeHTTP(w, r)
+	}, "/authorize")
+}

internal/filestore/body_uploader_test.go

+package filestore_test
+
+import (
+	"fmt"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strconv"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/api"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/filestore"
+)
+
+const (
+	fileContent = "A test file content"
+	fileLen     = len(fileContent)
+)
+
+func TestBodyUploader(t *testing.T) {
+	body := strings.NewReader(fileContent)
+
+	resp := testUpload(&rails{}, nil, echoProxy(t, fileLen), body)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	uploadEcho, err := ioutil.ReadAll(resp.Body)
+
+	require.NoError(t, err, "Can't read response body")
+	require.Equal(t, fileContent, string(uploadEcho))
+}
+
+func TestBodyUploaderCustomPreparer(t *testing.T) {
+	body := strings.NewReader(fileContent)
+
+	resp := testUpload(&rails{}, &alwaysLocalPreparer{}, echoProxy(t, fileLen), body)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	uploadEcho, err := ioutil.ReadAll(resp.Body)
+	require.NoError(t, err, "Can't read response body")
+	require.Equal(t, fileContent, string(uploadEcho))
+}
+
+func TestBodyUploaderCustomVerifier(t *testing.T) {
+	body := strings.NewReader(fileContent)
+	verifier := &mockVerifier{}
+
+	resp := testUpload(&rails{}, &alwaysLocalPreparer{verifier: verifier}, echoProxy(t, fileLen), body)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	uploadEcho, err := ioutil.ReadAll(resp.Body)
+	require.NoError(t, err, "Can't read response body")
+	require.Equal(t, fileContent, string(uploadEcho))
+	require.True(t, verifier.invoked, "Verifier.Verify not invoked")
+}
+
+func TestBodyUploaderAuthorizationFailure(t *testing.T) {
+	testNoProxyInvocation(t, http.StatusUnauthorized, &rails{unauthorized: true}, nil)
+}
+
+func TestBodyUploaderErrors(t *testing.T) {
+	tests := []struct {
+		name     string
+		preparer *alwaysLocalPreparer
+	}{
+		{name: "Prepare failure", preparer: &alwaysLocalPreparer{prepareError: fmt.Errorf("")}},
+		{name: "Verify failure", preparer: &alwaysLocalPreparer{verifier: &alwaysFailsVerifier{}}},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			testNoProxyInvocation(t, http.StatusInternalServerError, &rails{}, test.preparer)
+		})
+	}
+}
+
+func testNoProxyInvocation(t *testing.T, expectedStatus int, auth filestore.PreAuthorizer, preparer filestore.UploadPreparer) {
+	proxy := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Fail(t, "request proxied upstream")
+	})
+
+	resp := testUpload(auth, preparer, proxy, nil)
+	require.Equal(t, expectedStatus, resp.StatusCode)
+}
+
+func testUpload(auth filestore.PreAuthorizer, preparer filestore.UploadPreparer, proxy http.Handler, body io.Reader) *http.Response {
+	req := httptest.NewRequest("POST", "http://example.com/upload", body)
+	w := httptest.NewRecorder()
+
+	filestore.BodyUploader(auth, proxy, preparer).ServeHTTP(w, req)
+
+	return w.Result()
+}
+
+func echoProxy(t *testing.T, expectedBodyLength int) http.Handler {
+	require := require.New(t)
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		err := r.ParseForm()
+		require.NoError(err)
+
+		require.Equal("application/x-www-form-urlencoded", r.Header.Get("Content-Type"), "Wrong Content-Type header")
+
+		require.Contains(r.PostForm, "file.md5")
+		require.Contains(r.PostForm, "file.sha1")
+		require.Contains(r.PostForm, "file.sha256")
+		require.Contains(r.PostForm, "file.sha512")
+
+		require.Contains(r.PostForm, "file.path")
+		require.Contains(r.PostForm, "file.size")
+		require.Equal(strconv.Itoa(expectedBodyLength), r.PostFormValue("file.size"))
+
+		path := r.PostFormValue("file.path")
+		uploaded, err := os.Open(path)
+		require.NoError(err, "File not uploaded")
+
+		//sending back the file for testing purpose
+		io.Copy(w, uploaded)
+	})
+}
+
+type rails struct {
+	unauthorized bool
+}
+
+func (r *rails) PreAuthorizeHandler(next api.HandleFunc, _ string) http.Handler {
+	if r.unauthorized {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusUnauthorized)
+		})
+	}
+
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		next(w, r, &api.Response{TempPath: os.TempDir()})
+	})
+}
+
+type alwaysLocalPreparer struct {
+	verifier     filestore.UploadVerifier
+	prepareError error
+}
+
+func (a *alwaysLocalPreparer) Prepare(_ *api.Response) (*filestore.SaveFileOpts, filestore.UploadVerifier, error) {
+	return filestore.GetOpts(&api.Response{TempPath: os.TempDir()}), a.verifier, a.prepareError
+}
+
+type alwaysFailsVerifier struct{}
+
+func (_ alwaysFailsVerifier) Verify(handler *filestore.FileHandler) error {
+	return fmt.Errorf("Verification failed")
+}
+
+type mockVerifier struct {
+	invoked bool
+}
+
+func (m *mockVerifier) Verify(handler *filestore.FileHandler) error {
+	m.invoked = true
+
+	return nil
+}

internal/lfs/lfs.go

@@ -6,65 +6,38 @@ package lfs
 
 import (
 	"fmt"
-	"io/ioutil"
 	"net/http"
-	"net/url"
-	"path/filepath"
-	"strings"
 
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/api"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/filestore"
-	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
 )
 
-func PutStore(a *api.API, h http.Handler) http.Handler {
-	return handleStoreLFSObject(a, h)
+type object struct {
+	size int64
+	oid  string
 }
 
-func handleStoreLFSObject(myAPI *api.API, h http.Handler) http.Handler {
-	return myAPI.PreAuthorizeHandler(func(w http.ResponseWriter, r *http.Request, a *api.Response) {
-		opts := filestore.GetOpts(a)
-		opts.TempFilePrefix = a.LfsOid
-
-		// backward compatible api check - to be removed on next release
-		if a.StoreLFSPath != "" {
-			opts.LocalTempPath = a.StoreLFSPath
-		}
-		// end of backward compatible api check
+func (l *object) Verify(fh *filestore.FileHandler) error {
+	if fh.Size != l.size {
+		return fmt.Errorf("LFSObject: expected size %d, wrote %d", l.size, fh.Size)
+	}
 
-		fh, err := filestore.SaveFileFromReader(r.Context(), r.Body, r.ContentLength, opts)
-		if err != nil {
-			helper.Fail500(w, r, fmt.Errorf("handleStoreLFSObject: copy body to tempfile: %v", err))
-			return
-		}
+	if fh.SHA256() != l.oid {
+		return fmt.Errorf("LFSObject: expected sha256 %s, got %s", l.oid, fh.SHA256())
+	}
 
-		if fh.Size != a.LfsSize {
-			helper.Fail500(w, r, fmt.Errorf("handleStoreLFSObject: expected size %d, wrote %d", a.LfsSize, fh.Size))
-			return
-		}
+	return nil
+}
 
-		if fh.SHA256() != a.LfsOid {
-			helper.Fail500(w, r, fmt.Errorf("handleStoreLFSObject: expected sha256 %s, got %s", a.LfsOid, fh.SHA256()))
-			return
-		}
+type uploadPreparer struct{}
 
-		data := url.Values{}
-		for k, v := range fh.GitLabFinalizeFields("file") {
-			data.Set(k, v)
-		}
+func (l *uploadPreparer) Prepare(a *api.Response) (*filestore.SaveFileOpts, filestore.UploadVerifier, error) {
+	opts := filestore.GetOpts(a)
+	opts.TempFilePrefix = a.LfsOid
 
-		// Hijack body
-		body := data.Encode()
-		r.Body = ioutil.NopCloser(strings.NewReader(body))
-		r.ContentLength = int64(len(body))
-		r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-		// backward compatible API header - to be removed on next release
-		if opts.IsLocal() {
-			r.Header.Set("X-GitLab-Lfs-Tmp", filepath.Base(fh.LocalPath))
-		}
-		// end of backward compatible API header
+	return opts, &object{oid: a.LfsOid, size: a.LfsSize}, nil
+}
 
-		// And proxy the request
-		h.ServeHTTP(w, r)
-	}, "/authorize")
+func PutStore(a *api.API, h http.Handler) http.Handler {
+	return filestore.BodyUploader(a, h, &uploadPreparer{})
 }

internal/upload/accelerate.go

@@ -25,11 +25,7 @@ type MultipartClaims struct {
 	jwt.StandardClaims
 }
 
-type PreAuthorizer interface {
-	PreAuthorizeHandler(next api.HandleFunc, suffix string) http.Handler
-}
-
-func Accelerate(rails PreAuthorizer, h http.Handler) http.Handler {
+func Accelerate(rails filestore.PreAuthorizer, h http.Handler) http.Handler {
 	return rails.PreAuthorizeHandler(func(w http.ResponseWriter, r *http.Request, a *api.Response) {
 		s := &savedFileTracker{request: r}
 		HandleFileUploads(w, r, h, a, s)

internal/upstream/routes.go

@@ -10,6 +10,7 @@ import (
 	apipkg "gitlab.com/gitlab-org/gitlab-workhorse/internal/api"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/artifacts"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/builds"
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/filestore"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/git"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/lfs"
@@ -145,6 +146,9 @@ func (u *Upstream) configureRoutes() {
 		route("", apiPattern+`v4/jobs/request\z`, ciAPILongPolling),
 		route("", ciAPIPattern+`v1/builds/register.json\z`, ciAPILongPolling),
 
+		// Maven Artifact Repository
+		route("PUT", apiPattern+`v4/projects/[0-9]+/packages/maven/`, filestore.BodyUploader(api, proxy, nil)),
+
 		// Explicitly proxy API requests
 		route("", apiPattern, proxy),
 		route("", ciAPIPattern, proxy),