Add remote file upload v2

storage/file_persister.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"mime/multipart"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -63,6 +64,17 @@ type RemoteFilePersister struct {
 	httpClient *http.Client
 }
 
+// PresignedURLResponse holds the response from a presigned generation request.
+type PresignedURLResponse struct {
+	Service string `json:"service"`
+	URLs    []struct {
+		Name         string            `json:"name"`
+		PreSignedURL string            `json:"pre_signed_url"` //nolint:tagliatelle
+		Method       string            `json:"method"`
+		FormFields   map[string]string `json:"form_fields"` //nolint:tagliatelle
+	} `json:"urls"`
+}
+
 // NewRemoteFilePersister creates a new instance of RemoteFilePersister.
 func NewRemoteFilePersister(
 	preSignedURLGetterURL string,
@@ -81,12 +93,12 @@ func NewRemoteFilePersister(
 
 // Persist will upload the contents of data to a remote location.
 func (r *RemoteFilePersister) Persist(ctx context.Context, path string, data io.Reader) (err error) {
-	pURL, err := r.getPreSignedURL(ctx, path)
+	psResp, err := r.getPreSignedURL(ctx, path)
 	if err != nil {
 		return fmt.Errorf("getting presigned url: %w", err)
 	}
 
-	req, err := http.NewRequestWithContext(ctx, http.MethodPut, pURL, data)
+	req, err := newUploadRequest(ctx, psResp, data)
 	if err != nil {
 		return fmt.Errorf("creating upload request: %w", err)
 	}
@@ -116,16 +128,17 @@ func checkStatusCode(resp *http.Response) error {
 	return nil
 }
 
-// getPreSignedURL will retrieve the presigned url for the current file.
-func (r *RemoteFilePersister) getPreSignedURL(ctx context.Context, path string) (string, error) {
+// getPreSignedURL will request a new presigned URL from the remote server for the given path.
+// Returns a [PresignedURLResponse] that contains the presigned URL details.
+func (r *RemoteFilePersister) getPreSignedURL(ctx context.Context, path string) (PresignedURLResponse, error) {
 	b, err := buildPresignedRequestBody(r.basePath, path)
 	if err != nil {
-		return "", fmt.Errorf("building request body: %w", err)
+		return PresignedURLResponse{}, fmt.Errorf("building request body: %w", err)
 	}
 
 	req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.preSignedURLGetterURL, bytes.NewReader(b))
 	if err != nil {
-		return "", fmt.Errorf("creating request: %w", err)
+		return PresignedURLResponse{}, fmt.Errorf("creating request: %w", err)
 	}
 
 	for k, v := range r.headers {
@@ -134,12 +147,12 @@ func (r *RemoteFilePersister) getPreSignedURL(ctx context.Context, path string)
 
 	resp, err := r.httpClient.Do(req)
 	if err != nil {
-		return "", fmt.Errorf("performing request: %w", err)
+		return PresignedURLResponse{}, fmt.Errorf("performing request: %w", err)
 	}
 	defer resp.Body.Close() //nolint:errcheck
 
 	if err := checkStatusCode(resp); err != nil {
-		return "", err
+		return PresignedURLResponse{}, err
 	}
 
 	return readResponseBody(resp)
@@ -154,7 +167,7 @@ func buildPresignedRequestBody(basePath, path string) ([]byte, error) {
 		} `json:"files"`
 	}{
 		Service:   "aws_s3",
-		Operation: "upload",
+		Operation: "upload_post",
 		Files: []struct {
 			Name string `json:"name"`
 		}{
@@ -172,24 +185,63 @@ func buildPresignedRequestBody(basePath, path string) ([]byte, error) {
 	return bb, nil
 }
 
-func readResponseBody(resp *http.Response) (string, error) {
-	rb := struct {
-		Service string `json:"service"`
-		URLs    []struct {
-			Name         string `json:"name"`
-			PreSignedURL string `json:"pre_signed_url"` //nolint:tagliatelle
-		} `json:"urls"`
-	}{}
+func readResponseBody(resp *http.Response) (PresignedURLResponse, error) {
+	var rb PresignedURLResponse
 
 	decoder := json.NewDecoder(resp.Body)
 	err := decoder.Decode(&rb)
 	if err != nil {
-		return "", fmt.Errorf("decoding response body: %w", err)
+		return PresignedURLResponse{}, fmt.Errorf("decoding response body: %w", err)
 	}
 
 	if len(rb.URLs) == 0 {
-		return "", errors.New("missing presigned url in response body")
+		return PresignedURLResponse{}, errors.New("missing presigned url in response body")
+	}
+
+	return rb, nil
+}
+
+// newUploadRequest creates a new HTTP request to upload a file as a multipart
+// form to the presigned URL received from the server.
+func newUploadRequest(
+	ctx context.Context,
+	resp PresignedURLResponse,
+	data io.Reader,
+) (*http.Request, error) {
+	// we don't support multiple presigned URLs at the moment.
+	psu := resp.URLs[0]
+
+	// copy all form fields received from a presigned URL
+	// response to the multipart form fields.
+	var form bytes.Buffer
+	fw := multipart.NewWriter(&form)
+	for k, v := range psu.FormFields {
+		if err := fw.WriteField(k, v); err != nil {
+			return nil, fmt.Errorf("writing form field key %q and value %q: %w", k, v, err)
+		}
+	}
+	// attach the file data to the form.
+	ff, err := fw.CreateFormFile("file", psu.Name)
+	if err != nil {
+		return nil, fmt.Errorf("creating multipart form file: %w", err)
+	}
+	if _, err := io.Copy(ff, data); err != nil {
+		return nil, fmt.Errorf("copying file data to multipart form: %w", err)
+	}
+	if err := fw.Close(); err != nil {
+		return nil, fmt.Errorf("closing multipart form writer: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(
+		ctx,
+		psu.Method,
+		psu.PreSignedURL,
+		&form,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("creating new request: %w", err)
 	}
+	req.Header.Set("Content-Type", fw.FormDataContentType())
 
-	return rb.URLs[0].PreSignedURL, nil
+	return req, nil
 }

storage/file_persister_test.go

@@ -83,16 +83,20 @@ func TestLocalFilePersister(t *testing.T) {
 func TestRemoteFilePersister(t *testing.T) {
 	t.Parallel()
 
-	basePath := "screenshots"
-	presignedEndpoint := "/presigned"
-	uploadEndpoint := "/upload"
+	const (
+		basePath          = "screenshots"
+		presignedEndpoint = "/presigned"
+		uploadEndpoint    = "/upload"
+	)
 
 	tests := []struct {
 		name                    string
 		path                    string
 		dataToUpload            string
+		multipartFormFields     map[string]string
 		wantPresignedURLBody    string
 		wantPresignedHeaders    map[string]string
+		wantPresignedURLMethod  string
 		uploadResponse          int
 		getPresignedURLResponse int
 		wantError               string
@@ -101,15 +105,41 @@ func TestRemoteFilePersister(t *testing.T) {
 			name:         "upload_file",
 			path:         "some/path/file.png",
 			dataToUpload: "here's some data",
+			multipartFormFields: map[string]string{
+				"fooKey": "foo",
+				"barKey": "bar",
+			},
+			wantPresignedURLBody: `{
+					"service":"aws_s3",
+					"operation": "upload_post",
+					"files":[{"name":"%s"}]
+				}`,
+			wantPresignedHeaders: map[string]string{
+				"Authorization": "token asd123",
+				"Run_id":        "123456",
+			},
+			wantPresignedURLMethod:  http.MethodPost,
+			uploadResponse:          http.StatusOK,
+			getPresignedURLResponse: http.StatusOK,
+		},
+		{
+			name:         "upload_file",
+			path:         "some/path/file.png",
+			dataToUpload: "here's some data",
+			multipartFormFields: map[string]string{ // provide different form fields then the previous test
+				"bazKey": "baz",
+				"quxKey": "qux",
+			},
 			wantPresignedURLBody: `{
 					"service":"aws_s3",
-					"operation": "upload",
-					"files":[{"name":"screenshots/some/path/file.png"}]
+					"operation": "upload_post",
+					"files":[{"name":"%s"}]
 				}`,
 			wantPresignedHeaders: map[string]string{
 				"Authorization": "token asd123",
 				"Run_id":        "123456",
 			},
+			wantPresignedURLMethod:  http.MethodPut, // accepts dynamic methods
 			uploadResponse:          http.StatusOK,
 			getPresignedURLResponse: http.StatusOK,
 		},
@@ -119,13 +149,14 @@ func TestRemoteFilePersister(t *testing.T) {
 			dataToUpload: "here's some data",
 			wantPresignedURLBody: `{
 					"service":"aws_s3",
-					"operation": "upload",
-					"files":[{"name":"screenshots/some/path/file.png"}]
+					"operation": "upload_post",
+					"files":[{"name":"%s"}]
 				}`,
 			wantPresignedHeaders: map[string]string{
 				"Authorization": "token asd123",
 				"Run_id":        "123456",
 			},
+			wantPresignedURLMethod:  http.MethodPost,
 			getPresignedURLResponse: http.StatusTooManyRequests,
 			wantError:               "getting presigned url: server returned 429 (too many requests)",
 		},
@@ -135,13 +166,14 @@ func TestRemoteFilePersister(t *testing.T) {
 			dataToUpload: "here's some data",
 			wantPresignedURLBody: `{
 					"service":"aws_s3",
-					"operation": "upload",
-					"files":[{"name":"screenshots/some/path/file.png"}]
+					"operation": "upload_post",
+					"files":[{"name":"%s"}]
 				}`,
 			wantPresignedHeaders: map[string]string{
 				"Authorization": "token asd123",
 				"Run_id":        "123456",
 			},
+			wantPresignedURLMethod:  http.MethodPost,
 			getPresignedURLResponse: http.StatusInternalServerError,
 			wantError:               "getting presigned url: server returned 500 (internal server error)",
 		},
@@ -151,13 +183,14 @@ func TestRemoteFilePersister(t *testing.T) {
 			dataToUpload: "here's some data",
 			wantPresignedURLBody: `{
 					"service":"aws_s3",
-					"operation": "upload",
-					"files":[{"name":"screenshots/some/path/file.png"}]
+					"operation": "upload_post",
+					"files":[{"name":"%s"}]
 				}`,
 			wantPresignedHeaders: map[string]string{
 				"Authorization": "token asd123",
 				"Run_id":        "123456",
 			},
+			wantPresignedURLMethod:  http.MethodPost,
 			uploadResponse:          http.StatusTooManyRequests,
 			getPresignedURLResponse: http.StatusOK,
 			wantError:               "uploading: server returned 429 (too many requests)",
@@ -168,13 +201,14 @@ func TestRemoteFilePersister(t *testing.T) {
 			dataToUpload: "here's some data",
 			wantPresignedURLBody: `{
 					"service":"aws_s3",
-					"operation": "upload",
-					"files":[{"name":"screenshots/some/path/file.png"}]
+					"operation": "upload_post",
+					"files":[{"name":"%s"}]
 				}`,
 			wantPresignedHeaders: map[string]string{
 				"Authorization": "token asd123",
 				"Run_id":        "123456",
 			},
+			wantPresignedURLMethod:  http.MethodPost,
 			uploadResponse:          http.StatusInternalServerError,
 			getPresignedURLResponse: http.StatusOK,
 			wantError:               "uploading: server returned 500 (internal server error)",
@@ -198,24 +232,39 @@ func TestRemoteFilePersister(t *testing.T) {
 					bb, err := io.ReadAll(r.Body)
 					require.NoError(t, err)
 
-					// Ensures that the body of the request matches the
-					// expected format.
-					assert.JSONEq(t, tt.wantPresignedURLBody, string(bb))
+					// Does the response match the expected format?
+					wantPresignedURLBody := fmt.Sprintf(
+						tt.wantPresignedURLBody,
+						filepath.Join(basePath, tt.path),
+					)
+					assert.JSONEq(t, wantPresignedURLBody, string(bb))
 
-					// Ensures that the headers are sent to the server from
-					// the browser module.
+					// Do the HTTP headers are sent to the server from the browser module?
 					for k, v := range tt.wantPresignedHeaders {
 						assert.Equal(t, v, r.Header[k][0])
 					}
 
+					var formFields string
+					for k, v := range tt.multipartFormFields {
+						formFields += fmt.Sprintf(`"%s":"%s",`, k, v)
+					}
+					formFields = strings.TrimRight(formFields, ",")
+
 					w.WriteHeader(tt.getPresignedURLResponse)
 					_, err = fmt.Fprintf(w, `{
 							"service": "aws_s3",
 							"urls": [{
 								"name": "%s",
-								"pre_signed_url": "%s"
+								"pre_signed_url": "%s",
+								"method": "%s",
+								"form_fields": {%s}
 							}]
-							}`, basePath, s.URL+uploadEndpoint)
+							}`,
+						basePath+"/"+tt.path,
+						s.URL+uploadEndpoint,
+						tt.wantPresignedURLMethod,
+						formFields,
+					)
 
 					require.NoError(t, err)
 				},
@@ -227,12 +276,24 @@ func TestRemoteFilePersister(t *testing.T) {
 				func(w http.ResponseWriter, r *http.Request) {
 					defer r.Body.Close() //nolint:errcheck
 
-					bb, err := io.ReadAll(r.Body)
+					assert.Equal(t, tt.wantPresignedURLMethod, r.Method)
+
+					// Does the multipart form data contain the file to upload?
+					file, header, err := r.FormFile("file")
+					require.NoError(t, err)
+					t.Cleanup(func() {
+						_ = file.Close()
+					})
+					cd := header.Header.Get("Content-Disposition")
+					assert.Equal(t, cd, `form-data; name="file"; filename="`+basePath+`/`+tt.path+`"`)
+
+					// Does the file content match the expected data?
+					bb, err := io.ReadAll(file)
 					require.NoError(t, err)
+					assert.Equal(t, string(bb), tt.dataToUpload)
 
-					// Ensures that the data is uploaded to the server and matches
-					// what was sent.
-					assert.Equal(t, tt.dataToUpload, string(bb))
+					// Is the content type set correctly to the binary data?
+					assert.Equal(t, "application/octet-stream", header.Header.Get("Content-Type"))
 
 					w.WriteHeader(tt.uploadResponse)
 				}))