From e13b64869e7920b17396fd634624ccc88b698d8b Mon Sep 17 00:00:00 2001
From: Marks Polakovs <marks@markspolakovs.me>
Date: Thu, 14 Dec 2023 11:59:27 +0000
Subject: [PATCH] handler: support range requests for GET /upload

---
 pkg/handler/get_test.go         | 57 +++++++++++++++++++++++++++++++++
 pkg/handler/unrouted_handler.go |  8 +++++
 2 files changed, 65 insertions(+)

diff --git a/pkg/handler/get_test.go b/pkg/handler/get_test.go
index b3a26b548..7a51bf5f4 100644
--- a/pkg/handler/get_test.go
+++ b/pkg/handler/get_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/golang/mock/gomock"
+
 	. "github.com/tus/tusd/v2/pkg/handler"
 )
 
@@ -164,4 +165,60 @@ func TestGet(t *testing.T) {
 			ResBody: "",
 		}).Run(handler, t)
 	})
+
+	SubTest(t, "RangeRequest", func(t *testing.T, store *MockFullDataStore, _ *StoreComposer) {
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+		upload := NewMockFullUpload(ctrl)
+		locker := NewMockFullLocker(ctrl)
+		lock := NewMockFullLock(ctrl)
+
+		reader := &closingStringReader{
+			Reader: strings.NewReader("hello"),
+		}
+
+		gomock.InOrder(
+			locker.EXPECT().NewLock("yes").Return(lock, nil),
+			lock.EXPECT().Lock(gomock.Any(), gomock.Any()).Return(nil),
+			store.EXPECT().GetUpload(gomock.Any(), "yes").Return(upload, nil),
+			upload.EXPECT().GetInfo(gomock.Any()).Return(FileInfo{
+				Offset: 5,
+				Size:   20,
+				MetaData: map[string]string{
+					"filename": "file.jpg\"evil",
+					"filetype": "image/jpeg",
+				},
+				IsFinal: true,
+			}, nil),
+			upload.EXPECT().GetReader(gomock.Any()).Return(reader, nil),
+			lock.EXPECT().Unlock().Return(nil),
+		)
+
+		composer := NewStoreComposer()
+		composer.UseCore(store)
+		composer.UseLocker(locker)
+
+		handler, _ := NewHandler(Config{
+			StoreComposer: composer,
+		})
+
+		(&httpTest{
+			Method: "GET",
+			URL:    "yes",
+			ReqHeader: map[string]string{
+				"Range": "bytes=0-3",
+			},
+			ResHeader: map[string]string{
+				"Content-Length":      "4",
+				"Content-Type":        "image/jpeg",
+				"Content-Disposition": `inline;filename="file.jpg\"evil"`,
+			},
+			Code:    http.StatusOK,
+			ResBody: "hell",
+		}).Run(handler, t)
+
+		if !reader.closed {
+			t.Error("expected reader to be closed")
+		}
+	})
 }
diff --git a/pkg/handler/unrouted_handler.go b/pkg/handler/unrouted_handler.go
index 32a3237f2..08eb323ed 100644
--- a/pkg/handler/unrouted_handler.go
+++ b/pkg/handler/unrouted_handler.go
@@ -646,6 +646,7 @@ func (handler *UnroutedHandler) HeadFile(w http.ResponseWriter, r *http.Request)
 		} else {
 			resp.Header["Upload-Length"] = strconv.FormatInt(info.Size, 10)
 			resp.Header["Content-Length"] = strconv.FormatInt(info.Size, 10)
+			resp.Header["Accept-Ranges"] = "bytes"
 		}
 
 		resp.StatusCode = http.StatusOK
@@ -1002,6 +1003,13 @@ func (handler *UnroutedHandler) GetFile(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
+	if seeker, ok := src.(io.ReadSeeker); ok {
+		handler.sendResp(c, resp)
+		http.ServeContent(w, r, info.ID, time.Time{}, seeker)
+		_ = src.Close()
+		return
+	}
+
 	handler.sendResp(c, resp)
 	io.Copy(w, src)