From 03bbe4f9127227d7800e90d0d135184378a56697 Mon Sep 17 00:00:00 2001
From: fanpei-tes <pei.fan@target-energysolutions.com>
Date: Fri, 12 Jun 2020 14:58:42 +0800
Subject: [PATCH] extend tus protocol

---
 .gitignore                      |  1 +
 pkg/handler/unrouted_handler.go | 10 +++++++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 12ba2985f..cf80e60e6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,4 @@ cover.out
 data/
 node_modules/
 .DS_Store
+.idea
diff --git a/pkg/handler/unrouted_handler.go b/pkg/handler/unrouted_handler.go
index ef403dbe4..bb1f1afaa 100644
--- a/pkg/handler/unrouted_handler.go
+++ b/pkg/handler/unrouted_handler.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/base64"
 	"errors"
+	"fmt"
 	"io"
 	"log"
 	"math"
@@ -414,7 +415,7 @@ func (handler *UnroutedHandler) PostFile(w http.ResponseWriter, r *http.Request)
 		// Directly finish the upload if the upload is empty (i.e. has a size of 0).
 		// This statement is in an else-if block to avoid causing duplicate calls
 		// to finishUploadIfComplete if an upload is empty and contains a chunk.
-		if err := handler.finishUploadIfComplete(ctx, upload, info, r); err != nil {
+		if err := handler.finishUploadIfComplete(ctx, upload, info, r, w); err != nil {
 			handler.sendError(w, r, err)
 			return
 		}
@@ -676,13 +677,13 @@ func (handler *UnroutedHandler) writeChunk(ctx context.Context, upload Upload, i
 	handler.Metrics.incBytesReceived(uint64(bytesWritten))
 	info.Offset = newOffset
 
-	return handler.finishUploadIfComplete(ctx, upload, info, r)
+	return handler.finishUploadIfComplete(ctx, upload, info, r, w)
 }
 
 // finishUploadIfComplete checks whether an upload is completed (i.e. upload offset
 // matches upload size) and if so, it will call the data store's FinishUpload
 // function and send the necessary message on the CompleteUpload channel.
-func (handler *UnroutedHandler) finishUploadIfComplete(ctx context.Context, upload Upload, info FileInfo, r *http.Request) error {
+func (handler *UnroutedHandler) finishUploadIfComplete(ctx context.Context, upload Upload, info FileInfo, r *http.Request, rw http.ResponseWriter) error {
 	// If the upload is completed, ...
 	if !info.SizeIsDeferred && info.Offset == info.Size {
 		// ... allow custom mechanism to finish and cleanup the upload
@@ -690,6 +691,9 @@ func (handler *UnroutedHandler) finishUploadIfComplete(ctx context.Context, uplo
 			return err
 		}
 
+		rw.Header().Set("Tus-Finished", "true")
+		rw.Header().Set("Tus-Filesize", fmt.Sprintf("%d", info.Size))
+
 		// ... send the info out to the channel
 		if handler.config.NotifyCompleteUploads {
 			handler.CompleteUploads <- newHookEvent(info, r)