From 1ab523164399dd06f174cc23bf0ed4fb2c99cc02 Mon Sep 17 00:00:00 2001
From: Adam Jensen <acjensen@gmail.com>
Date: Mon, 23 Apr 2018 17:10:23 -0400
Subject: [PATCH] Add basic Upload-Defer-Length header handling

---
 s3store/s3store_test.go |  4 ++--
 unrouted_handler.go     | 42 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/s3store/s3store_test.go b/s3store/s3store_test.go
index f9729e0..ecd8f7f 100644
--- a/s3store/s3store_test.go
+++ b/s3store/s3store_test.go
@@ -52,8 +52,8 @@ func TestNewUpload(t *testing.T) {
 		s3obj.EXPECT().PutObject(&s3.PutObjectInput{
 			Bucket:        aws.String("bucket"),
 			Key:           aws.String("uploadId.info"),
-			Body:          bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"Offset":0,"MetaData":{"bar":"menü","foo":"hello"},"IsPartial":false,"IsFinal":false,"PartialUploads":null}`)),
-			ContentLength: aws.Int64(int64(148)),
+			Body:          bytes.NewReader([]byte(`{"ID":"uploadId+multipartId","Size":500,"SizeIsDeferred":false,"Offset":0,"MetaData":{"bar":"menü","foo":"hello"},"IsPartial":false,"IsFinal":false,"PartialUploads":null}`)),
+			ContentLength: aws.Int64(int64(171)),
 		}),
 	)
 
diff --git a/unrouted_handler.go b/unrouted_handler.go
index f35adae..8376ca5 100644
--- a/unrouted_handler.go
+++ b/unrouted_handler.go
@@ -15,6 +15,8 @@ import (
 	"time"
 )
 
+const UploadLengthDeferred = "1"
+
 var (
 	reExtractFileID  = regexp.MustCompile(`([^/]+)\/?$`)
 	reForwardedHost  = regexp.MustCompile(`host=([^,]+)`)
@@ -243,6 +245,7 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request)
 	// uploads the size is sum of all sizes of these files (no need for
 	// Upload-Length header)
 	var size int64
+	var sizeIsDeferred bool
 	if isFinal {
 		// A final upload must not contain a chunk within the creation request
 		if containsChunk {
@@ -256,9 +259,11 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request)
 			return
 		}
 	} else {
-		size, err = strconv.ParseInt(r.Header.Get("Upload-Length"), 10, 64)
-		if err != nil || size < 0 {
-			handler.sendError(w, r, ErrInvalidUploadLength)
+		uploadLengthHeader := r.Header.Get("Upload-Length")
+		uploadDeferLengthHeader := r.Header.Get("Upload-Defer-Length")
+		size, sizeIsDeferred, err = validateNewUploadLengthHeaders(uploadLengthHeader, uploadDeferLengthHeader)
+		if err != nil {
+			handler.sendError(w, r, err)
 			return
 		}
 	}
@@ -274,6 +279,7 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request)
 
 	info := FileInfo{
 		Size:           size,
+		SizeIsDeferred: sizeIsDeferred,
 		MetaData:       meta,
 		IsPartial:      isPartial,
 		IsFinal:        isFinal,
@@ -382,8 +388,13 @@ func (handler *UnroutedHandler) HeadFile(w http.ResponseWriter, r *http.Request)
 		w.Header().Set("Upload-Metadata", SerializeMetadataHeader(info.MetaData))
 	}
 
+	if info.SizeIsDeferred {
+		w.Header().Set("Upload-Defer-Length", UploadLengthDeferred)
+	} else {
+		w.Header().Set("Upload-Length", strconv.FormatInt(info.Size, 10))
+	}
+
 	w.Header().Set("Cache-Control", "no-store")
-	w.Header().Set("Upload-Length", strconv.FormatInt(info.Size, 10))
 	w.Header().Set("Upload-Offset", strconv.FormatInt(info.Offset, 10))
 	handler.sendResp(w, r, http.StatusOK)
 }
@@ -849,6 +860,29 @@ func (handler *UnroutedHandler) sizeOfUploads(ids []string) (size int64, err err
 	return
 }
 
+// Verify that the Upload-Length and Upload-Defer-Length headers are acceptable for creating a
+// new upload
+func validateNewUploadLengthHeaders(uploadLengthHeader string, uploadDeferLengthHeader string) (uploadLength int64, uploadLengthDeferred bool, err error) {
+	haveBothLengthHeaders := uploadLengthHeader != "" && uploadDeferLengthHeader != ""
+	haveInvalidDeferHeader := uploadDeferLengthHeader != "" && uploadDeferLengthHeader != UploadLengthDeferred
+	lengthIsDeferred := uploadDeferLengthHeader == UploadLengthDeferred
+
+	if haveBothLengthHeaders {
+		err = ErrUploadLengthAndUploadDeferLength
+	} else if haveInvalidDeferHeader {
+		err = ErrInvalidUploadDeferLength
+	} else if lengthIsDeferred {
+		uploadLengthDeferred = true
+	} else {
+		uploadLength, err = strconv.ParseInt(uploadLengthHeader, 10, 64)
+		if err != nil || uploadLength < 0 {
+			err = ErrInvalidUploadLength
+		}
+	}
+
+	return
+}
+
 // ParseMetadataHeader parses the Upload-Metadata header as defined in the
 // File Creation extension.
 // e.g. Upload-Metadata: name bHVucmpzLnBuZw==,type aW1hZ2UvcG5n