Prevent uploading two files as artifacts in single request

internal/artifacts/artifacts_store_test.go

@@ -136,7 +136,8 @@ func TestUploadHandlerSendingToExternalStorage(t *testing.T) {
 
 			contentBuffer, contentType := createTestMultipartForm(t, archiveData)
 			response := testUploadArtifacts(contentType, &contentBuffer, t, ts)
-			testhelper.AssertResponseCode(t, response, 200)
+			testhelper.AssertResponseCode(t, response, http.StatusOK)
+			testhelper.AssertResponseHeader(t, response, MetadataHeaderKey, MetadataHeaderPresent)
 			assert.Equal(t, 1, storeServerCalled, "store should be called only once")
 			assert.Equal(t, 1, responseProcessorCalled, "response processor should be called only once")
 		})
@@ -166,7 +167,7 @@ func TestUploadHandlerSendingToExternalStorageAndStorageServerUnreachable(t *tes
 	defer ts.Close()
 
 	response := testUploadArtifactsFromTestZip(t, ts)
-	testhelper.AssertResponseCode(t, response, 500)
+	testhelper.AssertResponseCode(t, response, http.StatusInternalServerError)
 }
 
 func TestUploadHandlerSendingToExternalStorageAndInvalidURLIsUsed(t *testing.T) {
@@ -192,7 +193,7 @@ func TestUploadHandlerSendingToExternalStorageAndInvalidURLIsUsed(t *testing.T)
 	defer ts.Close()
 
 	response := testUploadArtifactsFromTestZip(t, ts)
-	testhelper.AssertResponseCode(t, response, 500)
+	testhelper.AssertResponseCode(t, response, http.StatusInternalServerError)
 }
 
 func TestUploadHandlerSendingToExternalStorageAndItReturnsAnError(t *testing.T) {
@@ -230,7 +231,7 @@ func TestUploadHandlerSendingToExternalStorageAndItReturnsAnError(t *testing.T)
 	defer ts.Close()
 
 	response := testUploadArtifactsFromTestZip(t, ts)
-	testhelper.AssertResponseCode(t, response, 500)
+	testhelper.AssertResponseCode(t, response, http.StatusInternalServerError)
 	assert.Equal(t, 1, putCalledTimes, "upload should be called only once")
 }
 
@@ -271,7 +272,7 @@ func TestUploadHandlerSendingToExternalStorageAndSupportRequestTimeout(t *testin
 	defer ts.Close()
 
 	response := testUploadArtifactsFromTestZip(t, ts)
-	testhelper.AssertResponseCode(t, response, 500)
+	testhelper.AssertResponseCode(t, response, http.StatusInternalServerError)
 	assert.Equal(t, 1, putCalledTimes, "upload should be called only once")
 }
 

internal/artifacts/artifacts_upload.go

@@ -19,7 +19,6 @@ import (
 
 type artifactsUploadProcessor struct {
 	opts   *filestore.SaveFileOpts
-	metadataFile string
 	stored bool
 }
 
@@ -80,9 +79,10 @@ func (a *artifactsUploadProcessor) ProcessFile(ctx context.Context, formName str
 	if formName != "file" {
 		return fmt.Errorf("Invalid form field: %q", formName)
 	}
-	if a.metadataFile != "" {
-		return fmt.Errorf("Artifacts request contains more than one file!")
+	if a.stored {
+		return fmt.Errorf("Artifacts request contains more than one file")
 	}
+	a.stored = true
 
 	select {
 	case <-ctx.Done():

internal/artifacts/artifacts_upload_test.go

@@ -23,6 +23,12 @@ import (
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/zipartifacts"
 )
 
+const (
+	MetadataHeaderKey     = "Metadata-Status"
+	MetadataHeaderPresent = "present"
+	MetadataHeaderMissing = "missing"
+)
+
 func testArtifactsUploadServer(t *testing.T, authResponse api.Response, bodyProcessor func(w http.ResponseWriter, r *http.Request)) *httptest.Server {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/url/path/authorize", func(w http.ResponseWriter, r *http.Request) {
@@ -46,52 +52,54 @@ func testArtifactsUploadServer(t *testing.T, authResponse api.Response, bodyProc
 		}
 		if opts.IsLocal() {
 			if r.FormValue("file.path") == "" {
-				w.WriteHeader(501)
+				t.Fatal("Expected file to be present")
 				return
 			}
 
 			_, err := ioutil.ReadFile(r.FormValue("file.path"))
 			if err != nil {
-				w.WriteHeader(404)
+				t.Fatal("Expected file to be readable")
 				return
 			}
-		}
-
-		if opts.IsRemote() && r.FormValue("file.remote_url") == "" {
-			w.WriteHeader(501)
+		} else if opts.IsRemote() {
+			if r.FormValue("file.remote_url") == "" {
+				t.Fatal("Expected file to be remote accessible")
 				return
 			}
-
-		if r.FormValue("metadata.path") == "" {
-			w.WriteHeader(502)
-			return
 		}
 
+		if r.FormValue("metadata.path") != "" {
 			metadata, err := ioutil.ReadFile(r.FormValue("metadata.path"))
 			if err != nil {
-			w.WriteHeader(404)
+				t.Fatal("Expected metadata to be readable")
 				return
 			}
 			gz, err := gzip.NewReader(bytes.NewReader(metadata))
 			if err != nil {
-			w.WriteHeader(405)
+				t.Fatal("Expected metadata to be valid gzip")
 				return
 			}
 			defer gz.Close()
 			metadata, err = ioutil.ReadAll(gz)
 			if err != nil {
-			w.WriteHeader(404)
+				t.Fatal("Expected metadata to be valid")
 				return
 			}
 			if !bytes.HasPrefix(metadata, []byte(zipartifacts.MetadataHeaderPrefix+zipartifacts.MetadataHeader)) {
-			w.WriteHeader(400)
+				t.Fatal("Expected metadata to be of valid format")
 				return
 			}
 
+			w.Header().Set(MetadataHeaderKey, MetadataHeaderPresent)
+
+		} else {
+			w.Header().Set(MetadataHeaderKey, MetadataHeaderMissing)
+		}
+
+		w.WriteHeader(http.StatusOK)
+
 		if bodyProcessor != nil {
 			bodyProcessor(w, r)
-		} else {
-			w.WriteHeader(200)
 		}
 	})
 	return testhelper.TestServerWithHandler(nil, mux.ServeHTTP)
@@ -141,7 +149,8 @@ func TestUploadHandlerAddingMetadata(t *testing.T) {
 	writer.Close()
 
 	response := testUploadArtifacts(writer.FormDataContentType(), &buffer, t, ts)
-	testhelper.AssertResponseCode(t, response, 200)
+	testhelper.AssertResponseCode(t, response, http.StatusOK)
+	testhelper.AssertResponseHeader(t, response, MetadataHeaderKey, MetadataHeaderPresent)
 }
 
 func TestUploadHandlerForUnsupportedArchive(t *testing.T) {
@@ -164,8 +173,37 @@ func TestUploadHandlerForUnsupportedArchive(t *testing.T) {
 	writer.Close()
 
 	response := testUploadArtifacts(writer.FormDataContentType(), &buffer, t, ts)
-	// 502 is a custom response code from the mock server in testUploadArtifacts
-	testhelper.AssertResponseCode(t, response, 502)
+	testhelper.AssertResponseCode(t, response, http.StatusOK)
+	testhelper.AssertResponseHeader(t, response, MetadataHeaderKey, MetadataHeaderMissing)
+}
+
+func TestUploadHandlerForMultipleFiles(t *testing.T) {
+	tempPath, err := ioutil.TempDir("", "uploads")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tempPath)
+
+	ts := testArtifactsUploadServer(t, api.Response{TempPath: tempPath}, nil)
+	defer ts.Close()
+
+	var buffer bytes.Buffer
+	writer := multipart.NewWriter(&buffer)
+	file, err := writer.CreateFormFile("file", "my.file")
+	if err != nil {
+		t.Fatal(err)
+	}
+	fmt.Fprint(file, "test")
+
+	file, err = writer.CreateFormFile("file", "my.file")
+	if err != nil {
+		t.Fatal(err)
+	}
+	fmt.Fprint(file, "test")
+	writer.Close()
+
+	response := testUploadArtifacts(writer.FormDataContentType(), &buffer, t, ts)
+	testhelper.AssertResponseCode(t, response, http.StatusInternalServerError)
 }
 
 func TestUploadFormProcessing(t *testing.T) {
@@ -188,5 +226,5 @@ func TestUploadFormProcessing(t *testing.T) {
 	writer.Close()
 
 	response := testUploadArtifacts(writer.FormDataContentType(), &buffer, t, ts)
-	testhelper.AssertResponseCode(t, response, 500)
+	testhelper.AssertResponseCode(t, response, http.StatusInternalServerError)
 }

internal/testhelper/testhelper.go

@@ -78,8 +78,20 @@ func AssertResponseWriterHeader(t *testing.T, w http.ResponseWriter, header stri
 	assertHeaderExists(t, header, actual, expected)
 }
 
-func AssertResponseHeader(t *testing.T, w *http.Response, header string, expected ...string) {
-	actual := w.Header[http.CanonicalHeaderKey(header)]
+func AssertResponseHeader(t *testing.T, w interface{}, header string, expected ...string) {
+	var actual []string
+
+	header = http.CanonicalHeaderKey(header)
+
+	if resp, ok := w.(*http.Response); ok {
+		actual = resp.Header[header]
+	} else if resp, ok := w.(http.ResponseWriter); ok {
+		actual = resp.Header()[header]
+	} else if resp, ok := w.(*httptest.ResponseRecorder); ok {
+		actual = resp.Header()[header]
+	} else {
+		t.Fatalf("invalid type of w passed AssertResponseHeader")
+	}
 
 	assertHeaderExists(t, header, actual, expected)
 }