From 8711f2a1c597c643631578b2e6ec276bfb3a8e74 Mon Sep 17 00:00:00 2001
From: Andy Hsu <i@nn.ci>
Date: Mon, 17 Apr 2023 15:31:39 +0800
Subject: [PATCH] feat(quark): shard request file (close #4175)

---
 drivers/quark/driver.go | 59 +++++++++++++++++++++++++++++++++++++----
 drivers/quark/meta.go   |  6 ++---
 internal/model/args.go  | 13 ++++-----
 server/common/proxy.go  |  2 ++
 server/webdav/webdav.go |  2 ++
 5 files changed, 68 insertions(+), 14 deletions(-)

diff --git a/drivers/quark/driver.go b/drivers/quark/driver.go
index b9928811..5882ab40 100644
--- a/drivers/quark/driver.go
+++ b/drivers/quark/driver.go
@@ -5,15 +5,18 @@ import (
 	"crypto/md5"
 	"crypto/sha1"
 	"encoding/hex"
+	"fmt"
 	"io"
 	"net/http"
 	"os"
+	"strconv"
 	"time"
 
 	"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"
+	"github.com/alist-org/alist/v3/pkg/http_range"
 	"github.com/alist-org/alist/v3/pkg/utils"
 	"github.com/go-resty/resty/v2"
 	log "github.com/sirupsen/logrus"
@@ -64,12 +67,58 @@ func (d *Quark) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (
 	if err != nil {
 		return nil, err
 	}
+	u := resp.Data[0].DownloadUrl
+	start, end := int64(0), file.GetSize()
 	return &model.Link{
-		URL: resp.Data[0].DownloadUrl,
-		Header: http.Header{
-			"Cookie":     []string{d.Cookie},
-			"Referer":    []string{"https://pan.quark.cn"},
-			"User-Agent": []string{ua},
+		Handle: func(w http.ResponseWriter, r *http.Request) error {
+			if rg := r.Header.Get("Range"); rg != "" {
+				parseRange, err := http_range.ParseRange(rg, file.GetSize())
+				if err != nil {
+					return err
+				}
+				start, end = parseRange[0].Start, parseRange[0].Start+parseRange[0].Length
+				w.Header().Set("Content-Range", parseRange[0].ContentRange(file.GetSize()))
+				w.Header().Set("Content-Length", strconv.FormatInt(parseRange[0].Length, 10))
+				w.WriteHeader(http.StatusPartialContent)
+			} else {
+				w.Header().Set("Content-Length", strconv.FormatInt(file.GetSize(), 10))
+				w.WriteHeader(http.StatusOK)
+			}
+			// request 10 MB at a time
+			chunkSize := int64(10 * 1024 * 1024)
+			for start < end {
+				_end := start + chunkSize
+				if _end > end {
+					_end = end
+				}
+				_range := "bytes=" + strconv.FormatInt(start, 10) + "-" + strconv.FormatInt(_end-1, 10)
+				start = _end
+				err = func() error {
+					req, err := http.NewRequest(r.Method, u, nil)
+					if err != nil {
+						return err
+					}
+					req.Header.Set("Range", _range)
+					req.Header.Set("User-Agent", ua)
+					req.Header.Set("Cookie", d.Cookie)
+					req.Header.Set("Referer", "https://pan.quark.cn")
+					resp, err := http.DefaultClient.Do(req)
+					if err != nil {
+						return err
+					}
+					defer resp.Body.Close()
+					if resp.StatusCode != http.StatusPartialContent {
+						return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
+					}
+					_, err = io.Copy(w, resp.Body)
+					return err
+				}()
+				if err != nil {
+					return err
+				}
+
+			}
+			return nil
 		},
 	}, nil
 }
diff --git a/drivers/quark/meta.go b/drivers/quark/meta.go
index 70021383..ba38738d 100644
--- a/drivers/quark/meta.go
+++ b/drivers/quark/meta.go
@@ -13,9 +13,9 @@ type Addition struct {
 }
 
 var config = driver.Config{
-	Name:        "Quark",
-	OnlyProxy:   true,
-	DefaultRoot: "0",
+	Name:              "Quark",
+	OnlyLocal:         true,
+	DefaultRoot:       "0",
 	NoOverwriteUpload: true,
 }
 
diff --git a/internal/model/args.go b/internal/model/args.go
index b3655476..6e224ee4 100644
--- a/internal/model/args.go
+++ b/internal/model/args.go
@@ -17,12 +17,13 @@ type LinkArgs struct {
 }
 
 type Link struct {
-	URL        string         `json:"url"`
-	Header     http.Header    `json:"header"` // needed header
-	Data       io.ReadCloser  // return file reader directly
-	Status     int            // status maybe 200 or 206, etc
-	FilePath   *string        // local file, return the filepath
-	Expiration *time.Duration // url expiration time
+	URL        string                                             `json:"url"`
+	Header     http.Header                                        `json:"header"` // needed header
+	Data       io.ReadCloser                                      // return file reader directly
+	Status     int                                                // status maybe 200 or 206, etc
+	FilePath   *string                                            // local file, return the filepath
+	Expiration *time.Duration                                     // url expiration time
+	Handle     func(w http.ResponseWriter, r *http.Request) error // custom handler
 }
 
 type OtherArgs struct {
diff --git a/server/common/proxy.go b/server/common/proxy.go
index 37c51089..397201fe 100644
--- a/server/common/proxy.go
+++ b/server/common/proxy.go
@@ -72,6 +72,8 @@ func Proxy(w http.ResponseWriter, r *http.Request, link *model.Link, file model.
 		w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"; filename*=UTF-8''%s`, filename, url.PathEscape(filename)))
 		http.ServeContent(w, r, file.GetName(), fileStat.ModTime(), f)
 		return nil
+	} else if link.Handle != nil {
+		return link.Handle(w, r)
 	} else {
 		req, err := http.NewRequest(r.Method, link.URL, nil)
 		if err != nil {
diff --git a/server/webdav/webdav.go b/server/webdav/webdav.go
index 943045e5..83918379 100644
--- a/server/webdav/webdav.go
+++ b/server/webdav/webdav.go
@@ -21,6 +21,7 @@ import (
 	"github.com/alist-org/alist/v3/internal/sign"
 	"github.com/alist-org/alist/v3/pkg/utils"
 	"github.com/alist-org/alist/v3/server/common"
+	log "github.com/sirupsen/logrus"
 )
 
 type Handler struct {
@@ -237,6 +238,7 @@ func (h *Handler) handleGetHeadPost(w http.ResponseWriter, r *http.Request) (sta
 		}
 		err = common.Proxy(w, r, link, fi)
 		if err != nil {
+			log.Errorf("webdav proxy error: %+v", err)
 			return http.StatusInternalServerError, err
 		}
 	} else if storage.GetStorage().WebdavProxy() && downProxyUrl != "" {