From 95b3b8767262f49d7f5489e747efb7d00ee10ec9 Mon Sep 17 00:00:00 2001
From: Andy Hsu <i@nn.ci>
Date: Mon, 20 Feb 2023 16:57:52 +0800
Subject: [PATCH] feat(sftp): support range header

---
 drivers/base/util.go   | 30 ++++++++++++++++++++++++++++++
 drivers/sftp/driver.go |  7 +++++--
 drivers/smb/driver.go  | 22 ++--------------------
 go.sum                 |  2 --
 4 files changed, 37 insertions(+), 24 deletions(-)
 create mode 100644 drivers/base/util.go

diff --git a/drivers/base/util.go b/drivers/base/util.go
new file mode 100644
index 00000000000..f26af147326
--- /dev/null
+++ b/drivers/base/util.go
@@ -0,0 +1,30 @@
+package base
+
+import (
+	"io"
+	"net/http"
+	"strconv"
+
+	"github.com/alist-org/alist/v3/internal/model"
+	"github.com/alist-org/alist/v3/pkg/http_range"
+	"github.com/alist-org/alist/v3/pkg/utils"
+)
+
+func HandleRange(link *model.Link, file io.ReadSeekCloser, header http.Header, size int64) {
+	if header.Get("Range") != "" {
+		r, err := http_range.ParseRange(header.Get("Range"), size)
+		if err == nil && len(r) > 0 {
+			_, err := file.Seek(r[0].Start, io.SeekStart)
+			if err == nil {
+				link.Data = utils.NewLimitReadCloser(file, func() error {
+					return file.Close()
+				}, r[0].Length)
+				link.Status = http.StatusPartialContent
+				link.Header = http.Header{
+					"Content-Range":  []string{r[0].ContentRange(size)},
+					"Content-Length": []string{strconv.FormatInt(r[0].Length, 10)},
+				}
+			}
+		}
+	}
+}
diff --git a/drivers/sftp/driver.go b/drivers/sftp/driver.go
index b9d3467b98f..24e83daa887 100644
--- a/drivers/sftp/driver.go
+++ b/drivers/sftp/driver.go
@@ -5,6 +5,7 @@ import (
 	"os"
 	"path"
 
+	"github.com/alist-org/alist/v3/drivers/base"
 	"github.com/alist-org/alist/v3/internal/driver"
 	"github.com/alist-org/alist/v3/internal/errs"
 	"github.com/alist-org/alist/v3/internal/model"
@@ -52,9 +53,11 @@ func (d *SFTP) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*
 	if err != nil {
 		return nil, err
 	}
-	return &model.Link{
+	link := &model.Link{
 		Data: remoteFile,
-	}, nil
+	}
+	base.HandleRange(link, remoteFile, args.Header, file.GetSize())
+	return link, nil
 }
 
 func (d *SFTP) MakeDir(ctx context.Context, parentDir model.Obj, dirName string) error {
diff --git a/drivers/smb/driver.go b/drivers/smb/driver.go
index 5afd0e8ebb9..17875fe16f1 100644
--- a/drivers/smb/driver.go
+++ b/drivers/smb/driver.go
@@ -3,16 +3,13 @@ package smb
 import (
 	"context"
 	"errors"
-	"io"
-	"net/http"
 	"path/filepath"
-	"strconv"
 	"strings"
 	"time"
 
+	"github.com/alist-org/alist/v3/drivers/base"
 	"github.com/alist-org/alist/v3/internal/driver"
 	"github.com/alist-org/alist/v3/internal/model"
-	"github.com/alist-org/alist/v3/pkg/http_range"
 	"github.com/alist-org/alist/v3/pkg/utils"
 
 	"github.com/hirochachacha/go-smb2"
@@ -86,22 +83,7 @@ func (d *SMB) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*m
 	link := &model.Link{
 		Data: remoteFile,
 	}
-	if args.Header.Get("Range") != "" {
-		r, err := http_range.ParseRange(args.Header.Get("Range"), file.GetSize())
-		if err == nil && len(r) > 0 {
-			_, err := remoteFile.Seek(r[0].Start, io.SeekStart)
-			if err == nil {
-				link.Data = utils.NewLimitReadCloser(remoteFile, func() error {
-					return remoteFile.Close()
-				}, r[0].Length)
-				link.Status = 206
-				link.Header = http.Header{
-					"Content-Range":  []string{r[0].ContentRange(file.GetSize())},
-					"Content-Length": []string{strconv.FormatInt(r[0].Length, 10)},
-				}
-			}
-		}
-	}
+	base.HandleRange(link, remoteFile, args.Header, file.GetSize())
 	d.updateLastConnTime()
 	return link, nil
 }
diff --git a/go.sum b/go.sum
index 6e0b0d05439..191d4d48f7b 100644
--- a/go.sum
+++ b/go.sum
@@ -93,8 +93,6 @@ github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9
 github.com/goccy/go-json v0.9.7/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
 github.com/goccy/go-json v0.9.11 h1:/pAaQDLHEoCq/5FFmSKBswWmK6H0e8g4159Kc/X/nqk=
 github.com/goccy/go-json v0.9.11/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
-github.com/golang-jwt/jwt/v4 v4.4.3 h1:Hxl6lhQFj4AnOX6MLrsCb/+7tCj7DxP7VA+2rDIq5AU=
-github.com/golang-jwt/jwt/v4 v4.4.3/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/golang-jwt/jwt/v4 v4.5.0 h1:7cYmW1XlMY7h7ii7UhUyChSgS5wUJEnm9uZVTGqOWzg=
 github.com/golang-jwt/jwt/v4 v4.5.0/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
 github.com/golang/geo v0.0.0-20210211234256-740aa86cb551 h1:gtexQ/VGyN+VVFRXSFiguSNcXmS6rkKT+X7FdIrTtfo=