From 6639cab1ae697d463474d535a70966965332ede9 Mon Sep 17 00:00:00 2001
From: Noah Hsu <i@nn.ci>
Date: Mon, 7 Nov 2022 20:58:52 +0800
Subject: [PATCH] feat(google_drive): chunk upload (close #2241)

---
 drivers/google_drive/driver.go | 13 ++++++++++---
 drivers/google_drive/meta.go   |  1 +
 drivers/google_drive/util.go   | 26 ++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/drivers/google_drive/driver.go b/drivers/google_drive/driver.go
index 0bbca06d..57d47ae9 100644
--- a/drivers/google_drive/driver.go
+++ b/drivers/google_drive/driver.go
@@ -36,6 +36,9 @@ func (d *GoogleDrive) Init(ctx context.Context, storage model.Storage) error {
 	if err != nil {
 		return err
 	}
+	if d.ChunkSize == 0 {
+		d.ChunkSize = 5
+	}
 	return d.refreshToken()
 }
 
@@ -160,9 +163,13 @@ func (d *GoogleDrive) Put(ctx context.Context, dstDir model.Obj, stream model.Fi
 		return fmt.Errorf("%s: %v", e.Error.Message, e.Error.Errors)
 	}
 	putUrl := res.Header().Get("location")
-	_, err = d.request(putUrl, http.MethodPut, func(req *resty.Request) {
-		req.SetHeader("Content-Length", strconv.FormatInt(stream.GetSize(), 10)).SetBody(stream.GetReadCloser())
-	}, nil)
+	if stream.GetSize() < d.ChunkSize*1024*1024 {
+		_, err = d.request(putUrl, http.MethodPut, func(req *resty.Request) {
+			req.SetHeader("Content-Length", strconv.FormatInt(stream.GetSize(), 10)).SetBody(stream.GetReadCloser())
+		}, nil)
+	} else {
+		err = d.chunkUpload(ctx, stream, putUrl)
+	}
 	return err
 }
 
diff --git a/drivers/google_drive/meta.go b/drivers/google_drive/meta.go
index cb8f274e..b98a93e5 100644
--- a/drivers/google_drive/meta.go
+++ b/drivers/google_drive/meta.go
@@ -12,6 +12,7 @@ type Addition struct {
 	OrderDirection string `json:"order_direction" type:"select" options:"asc,desc"`
 	ClientID       string `json:"client_id" required:"true" default:"202264815644.apps.googleusercontent.com"`
 	ClientSecret   string `json:"client_secret" required:"true" default:"X4Z3ca8xfWDb1Voo-F9a7ZxJ"`
+	ChunkSize      int64  `json:"chunk_size" default:"5" help:"chunk size while uploading (unit: MB)"`
 }
 
 var config = driver.Config{
diff --git a/drivers/google_drive/util.go b/drivers/google_drive/util.go
index eaae584f..5eaca9e0 100644
--- a/drivers/google_drive/util.go
+++ b/drivers/google_drive/util.go
@@ -1,10 +1,14 @@
 package google_drive
 
 import (
+	"context"
 	"fmt"
+	"io"
 	"net/http"
+	"strconv"
 
 	"github.com/alist-org/alist/v3/drivers/base"
+	"github.com/alist-org/alist/v3/internal/model"
 	"github.com/go-resty/resty/v2"
 	log "github.com/sirupsen/logrus"
 )
@@ -95,3 +99,25 @@ func (d *GoogleDrive) getFiles(id string) ([]File, error) {
 	}
 	return res, nil
 }
+
+func (d *GoogleDrive) chunkUpload(ctx context.Context, stream model.FileStreamer, url string) error {
+	var defaultChunkSize = d.ChunkSize * 1024 * 1024
+	var finish int64 = 0
+	for finish < stream.GetSize() {
+		chunkSize := stream.GetSize() - finish
+		if chunkSize > defaultChunkSize {
+			chunkSize = defaultChunkSize
+		}
+		_, err := d.request(url, http.MethodPut, func(req *resty.Request) {
+			req.SetHeaders(map[string]string{
+				"Content-Length": strconv.FormatInt(chunkSize, 10),
+				"Content-Range":  fmt.Sprintf("bytes %d-%d/%d", finish, finish+chunkSize-1, stream.GetSize()),
+			}).SetBody(io.LimitReader(stream.GetReadCloser(), chunkSize))
+		}, nil)
+		if err != nil {
+			return err
+		}
+		finish += chunkSize
+	}
+	return nil
+}