From 6aee31341fc95fe8b845f34e19ca68ee63e44dd5 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Tue, 21 Jan 2020 11:49:46 +0800
Subject: [PATCH] Feat: async monitor OneDrive resume upload progress

---
 middleware/auth.go                         |   4 +
 models/init.go                             |   2 +-
 models/migration.go                        |   2 +
 models/policy.go                           |   3 +-
 models/policy_test.go                      |  10 +-
 pkg/filesystem/driver/onedrive/api.go      | 264 +++++++++++++++++----
 pkg/filesystem/driver/onedrive/handller.go |  18 +-
 pkg/filesystem/driver/onedrive/types.go    |  77 ++++++
 8 files changed, 321 insertions(+), 59 deletions(-)
 create mode 100644 pkg/filesystem/driver/onedrive/types.go

diff --git a/middleware/auth.go b/middleware/auth.go
index eb7f466..2167b3e 100644
--- a/middleware/auth.go
+++ b/middleware/auth.go
@@ -8,6 +8,7 @@ import (
 	"github.com/HFO4/cloudreve/models"
 	"github.com/HFO4/cloudreve/pkg/auth"
 	"github.com/HFO4/cloudreve/pkg/cache"
+	"github.com/HFO4/cloudreve/pkg/filesystem/driver/onedrive"
 	"github.com/HFO4/cloudreve/pkg/filesystem/driver/oss"
 	"github.com/HFO4/cloudreve/pkg/filesystem/driver/upyun"
 	"github.com/HFO4/cloudreve/pkg/serializer"
@@ -286,6 +287,9 @@ func OneDriveCallbackAuth() gin.HandlerFunc {
 			return
 		}
 
+		// 发送回调结束信号
+		onedrive.FinishCallback(c.Param("key"))
+
 		c.Next()
 	}
 }
diff --git a/models/init.go b/models/init.go
index c782745..bafc977 100644
--- a/models/init.go
+++ b/models/init.go
@@ -47,7 +47,7 @@ func Init() {
 
 	// Debug模式下，输出所有 SQL 日志
 	if conf.SystemConfig.Debug {
-		db.LogMode(true)
+		db.LogMode(false)
 	}
 
 	//db.SetLogger(util.Log())
diff --git a/models/migration.go b/models/migration.go
index 3051f70..c1c3620 100644
--- a/models/migration.go
+++ b/models/migration.go
@@ -107,6 +107,8 @@ solid #e9e9e9;"bgcolor="#fff"><tbody><tr style="font-family: 'Helvetica Neue',He
 		{Name: "upload_credential_timeout", Value: `1800`, Type: "timeout"},
 		{Name: "upload_session_timeout", Value: `86400`, Type: "timeout"},
 		{Name: "slave_api_timeout", Value: `60`, Type: "timeout"},
+		{Name: "onedrive_monitor_timeout", Value: `600`, Type: "timeout"},
+		{Name: "onedrive_callback_check", Value: `20`, Type: "timeout"},
 		{Name: "allowdVisitorDownload", Value: `false`, Type: "share"},
 		{Name: "login_captcha", Value: `0`, Type: "login"},
 		{Name: "qq_login", Value: `0`, Type: "login"},
diff --git a/models/policy.go b/models/policy.go
index 3941223..eeb9733 100644
--- a/models/policy.go
+++ b/models/policy.go
@@ -180,7 +180,8 @@ func (policy *Policy) GetUploadURL() string {
 
 	var controller *url.URL
 	switch policy.Type {
-	case "local":
+	case "local", "onedrive":
+		server = GetSiteURL()
 		controller, _ = url.Parse("/api/v3/file/upload")
 	case "remote":
 		controller, _ = url.Parse("/api/v3/slave/upload")
diff --git a/models/policy_test.go b/models/policy_test.go
index 7a73243..8913cbc 100644
--- a/models/policy_test.go
+++ b/models/policy_test.go
@@ -3,6 +3,7 @@ package model
 import (
 	"encoding/json"
 	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/stretchr/testify/assert"
 	"strconv"
 	"testing"
@@ -15,13 +16,13 @@ func TestGetPolicyByID(t *testing.T) {
 	// 缓存未命中
 	{
 		rows := sqlmock.NewRows([]string{"name", "type", "options"}).
-			AddRow("默认存储策略", "local", "{\"op_name\":\"123\"}")
+			AddRow("默认存储策略", "local", "{\"od_redirect\":\"123\"}")
 		mock.ExpectQuery("^SELECT(.+)").WillReturnRows(rows)
 		policy, err := GetPolicyByID(uint(22))
 		asserts.NoError(err)
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.Equal("默认存储策略", policy.Name)
-		asserts.Equal("123", policy.OptionsSerialized.OPName)
+		asserts.Equal("123", policy.OptionsSerialized.OdRedirect)
 
 		rows = sqlmock.NewRows([]string{"name", "type", "options"})
 		mock.ExpectQuery("^SELECT(.+)").WillReturnRows(rows)
@@ -35,7 +36,7 @@ func TestGetPolicyByID(t *testing.T) {
 		policy, err := GetPolicyByID(uint(22))
 		asserts.NoError(err)
 		asserts.Equal("默认存储策略", policy.Name)
-		asserts.Equal("123", policy.OptionsSerialized.OPName)
+		asserts.Equal("123", policy.OptionsSerialized.OdRedirect)
 
 	}
 
@@ -46,7 +47,7 @@ func TestPolicy_BeforeSave(t *testing.T) {
 
 	testPolicy := Policy{
 		OptionsSerialized: PolicyOption{
-			OPName: "123",
+			OdRedirect: "123",
 		},
 	}
 	expected, _ := json.Marshal(testPolicy.OptionsSerialized)
@@ -163,6 +164,7 @@ func TestPolicy_GetUploadURL(t *testing.T) {
 
 	// 本地
 	{
+		cache.Set("setting_siteURL", "http://127.0.0.1", 0)
 		policy := Policy{Type: "local", Server: "http://127.0.0.1"}
 		asserts.Equal("http://127.0.0.1/api/v3/file/upload", policy.GetUploadURL())
 	}
diff --git a/pkg/filesystem/driver/onedrive/api.go b/pkg/filesystem/driver/onedrive/api.go
index 0d8dd31..34a27cb 100644
--- a/pkg/filesystem/driver/onedrive/api.go
+++ b/pkg/filesystem/driver/onedrive/api.go
@@ -3,52 +3,21 @@ package onedrive
 import (
 	"context"
 	"encoding/json"
+	"errors"
+	model "github.com/HFO4/cloudreve/models"
+	"github.com/HFO4/cloudreve/pkg/cache"
 	"github.com/HFO4/cloudreve/pkg/request"
+	"github.com/HFO4/cloudreve/pkg/util"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/url"
 	"path"
+	"strconv"
 	"strings"
+	"time"
 )
 
-// RespError 接口返回错误
-type RespError struct {
-	APIError APIError `json:"error"`
-}
-
-// APIError 接口返回的错误内容
-type APIError struct {
-	Code    string `json:"code"`
-	Message string `json:"message"`
-}
-
-// UploadSessionResponse 分片上传会话
-type UploadSessionResponse struct {
-	DataContext        string   `json:"@odata.context"`
-	ExpirationDateTime string   `json:"expirationDateTime"`
-	NextExpectedRanges []string `json:"nextExpectedRanges"`
-	UploadURL          string   `json:"uploadUrl"`
-}
-
-// FileInfo 文件元信息
-type FileInfo struct {
-	Name            string          `json:"name"`
-	Size            uint64          `json:"size"`
-	Image           imageInfo       `json:"image"`
-	ParentReference parentReference `json:"parentReference"`
-}
-
-type imageInfo struct {
-	Height int `json:"height"`
-	Width  int `json:"width"`
-}
-
-type parentReference struct {
-	Path string `json:"path"`
-	Name string `json:"name"`
-	ID   string `json:"id"`
-}
-
 // GetSourcePath 获取文件的绝对路径
 func (info *FileInfo) GetSourcePath() string {
 	res, err := url.PathUnescape(
@@ -84,7 +53,7 @@ func (client *Client) getRequestURL(api string) string {
 func (client *Client) Meta(ctx context.Context, id string) (*FileInfo, error) {
 
 	requestURL := client.getRequestURL("/me/drive/items/" + id)
-	res, err := client.request(ctx, "GET", requestURL+"?expand=thumbnails", "")
+	res, err := client.requestWithStr(ctx, "GET", requestURL+"?expand=thumbnails", "", 200)
 	if err != nil {
 		return nil, err
 	}
@@ -119,7 +88,7 @@ func (client *Client) CreateUploadSession(ctx context.Context, dst string, opts
 	}
 	bodyBytes, _ := json.Marshal(body)
 
-	res, err := client.request(ctx, "POST", requestURL, string(bodyBytes))
+	res, err := client.requestWithStr(ctx, "POST", requestURL, string(bodyBytes), 200)
 	if err != nil {
 		return "", err
 	}
@@ -136,6 +105,192 @@ func (client *Client) CreateUploadSession(ctx context.Context, dst string, opts
 	return uploadSession.UploadURL, nil
 }
 
+// GetUploadSessionStatus 查询上传会话状态
+func (client *Client) GetUploadSessionStatus(ctx context.Context, uploadURL string) (*UploadSessionResponse, error) {
+	res, err := client.requestWithStr(ctx, "GET", uploadURL, "", 200)
+	if err != nil {
+		return nil, err
+	}
+
+	var (
+		decodeErr     error
+		uploadSession UploadSessionResponse
+	)
+	decodeErr = json.Unmarshal([]byte(res), &uploadSession)
+	if decodeErr != nil {
+		return nil, decodeErr
+	}
+
+	return &uploadSession, nil
+}
+
+// DeleteUploadSession 删除上传会话
+func (client *Client) DeleteUploadSession(ctx context.Context, uploadURL string) error {
+	_, err := client.requestWithStr(ctx, "DELETE", uploadURL, "", 204)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// PutFile 上传小文件到dst
+func (client *Client) PutFile(ctx context.Context, dst string, body io.Reader) (*UploadResult, error) {
+	dst = strings.TrimPrefix(dst, "/")
+	requestURL := client.getRequestURL("me/drive/root:/" + dst + ":/content")
+
+	res, err := client.request(ctx, "PUT", requestURL, body, 201)
+	if err != nil {
+		return nil, err
+	}
+
+	var (
+		decodeErr error
+		uploadRes UploadResult
+	)
+	decodeErr = json.Unmarshal([]byte(res), &uploadRes)
+	if decodeErr != nil {
+		return nil, decodeErr
+	}
+
+	return &uploadRes, nil
+}
+
+// Delete 并行删除文件，返回删除失败的文件，及第一个遇到的错误
+func (client *Client) Delete(ctx context.Context, dst []string) ([]string, error) {
+	body := client.makeBatchDeleteRequestsBody(dst)
+	res, err := client.requestWithStr(ctx, "POST", client.getRequestURL("$batch"), body, 200)
+	if err != nil {
+		return dst, err
+	}
+
+	var (
+		decodeErr error
+		deleteRes BatchResponses
+	)
+	decodeErr = json.Unmarshal([]byte(res), &deleteRes)
+	if decodeErr != nil {
+		return dst, decodeErr
+	}
+
+	// 取得删除失败的文件
+	failed := getDeleteFailed(&deleteRes)
+	if len(failed) != 0 {
+		return failed, errors.New("无法删除文件")
+	}
+	return failed, nil
+}
+
+func getDeleteFailed(res *BatchResponses) []string {
+	var failed = make([]string, 0, len(res.Responses))
+	for _, v := range res.Responses {
+		if v.Status != 204 {
+			failed = append(failed, v.ID)
+		}
+	}
+	return failed
+}
+
+// makeBatchDeleteRequestsBody 生成批量删除请求正文
+func (client *Client) makeBatchDeleteRequestsBody(files []string) string {
+	req := BatchRequests{
+		Requests: make([]BatchRequest, len(files)),
+	}
+	for i, v := range files {
+		v = strings.TrimPrefix(v, "/")
+		req.Requests[i] = BatchRequest{
+			ID:     v,
+			Method: "DELETE",
+			URL:    "me/drive/root:/" + v,
+		}
+	}
+
+	res, _ := json.Marshal(req)
+	return string(res)
+}
+
+// MonitorUpload 监控客户端分片上传进度
+func (client *Client) MonitorUpload(uploadURL, callbackKey, path string, size uint64, ttl int64) {
+	// 回调完成通知chan
+	callbackChan := make(chan bool)
+	callbackSignal.Store(callbackKey, callbackChan)
+	timeout := model.GetIntSetting("onedrive_monitor_timeout", 600)
+	interval := model.GetIntSetting("onedrive_callback_check", 20)
+
+	for {
+		select {
+		case <-callbackChan:
+			util.Log().Debug("客户端完成回调")
+			return
+		case <-time.After(time.Duration(ttl) * time.Second):
+			// 上传会话到期，仍未完成上传，创建占位符
+			client.DeleteUploadSession(context.Background(), uploadURL)
+			_, err := client.PutFile(context.Background(), path, strings.NewReader(""))
+			if err != nil {
+				util.Log().Debug("无法创建占位文件，%s", err)
+			}
+			return
+		case <-time.After(time.Duration(timeout) * time.Second):
+			util.Log().Debug("检查上传情况")
+			status, err := client.GetUploadSessionStatus(context.Background(), uploadURL)
+
+			if err != nil {
+				if resErr, ok := err.(*RespError); ok {
+					if resErr.APIError.Code == "itemNotFound" {
+						util.Log().Debug("上传会话已完成，稍后检查回调")
+						time.Sleep(time.Duration(interval) * time.Second)
+						util.Log().Debug("开始检查回调")
+						_, ok := cache.Get("callback_" + callbackKey)
+						if ok {
+							util.Log().Warning("未发送回调，删除文件")
+							cache.Deletes([]string{callbackKey}, "callback_")
+							_, err = client.Delete(context.Background(), []string{path})
+							if err != nil {
+								util.Log().Warning("无法删除未回掉的文件，%s", err)
+							}
+						}
+						return
+					}
+					util.Log().Debug("无法获取上传会话状态，%s", err.Error())
+					return
+				}
+				util.Log().Debug("无法获取上传会话状态，继续下一轮，%s", err.Error())
+			}
+
+			// 成功获取分片上传状态，检查文件大小
+			sizeRange := strings.Split(
+				status.NextExpectedRanges[len(status.NextExpectedRanges)-1],
+				"-",
+			)
+			if len(sizeRange) != 2 {
+				continue
+			}
+			uploadFullSize, _ := strconv.ParseUint(sizeRange[1], 10, 64)
+			if (sizeRange[0] == "0" && sizeRange[1] == "") || uploadFullSize+1 != size {
+				util.Log().Debug("未开始上传或文件大小不一致，取消上传会话")
+				// 取消上传会话，实测OneDrive取消上传会话后，客户端还是可以上传，
+				// 所以上传一个空文件占位，阻止客户端上传
+				client.DeleteUploadSession(context.Background(), uploadURL)
+				_, err := client.PutFile(context.Background(), path, strings.NewReader(""))
+				if err != nil {
+					util.Log().Debug("无法创建占位文件，%s", err)
+				}
+				return
+			}
+
+		}
+	}
+}
+
+// FinishCallback 向Monitor发送回调结束信号
+func FinishCallback(key string) {
+	if signal, ok := callbackSignal.Load(key); ok {
+		if signalChan, ok := signal.(chan bool); ok {
+			close(signalChan)
+		}
+	}
+}
+
 func sysError(err error) *RespError {
 	return &RespError{APIError: APIError{
 		Code:    "system",
@@ -143,21 +298,14 @@ func sysError(err error) *RespError {
 	}}
 }
 
-func (client *Client) request(ctx context.Context, method string, url string, body string) (string, *RespError) {
-
+func (client *Client) request(ctx context.Context, method string, url string, body io.Reader, expectedCode int, option ...request.Option) (string, *RespError) {
 	// 获取凭证
 	err := client.UpdateCredential(ctx)
 	if err != nil {
 		return "", sysError(err)
 	}
 
-	// 发送请求
-	bodyReader := ioutil.NopCloser(strings.NewReader(body))
-	res := client.Request.Request(
-		method,
-		url,
-		bodyReader,
-		request.WithContentLength(int64(len(body))),
+	option = append(option,
 		request.WithHeader(http.Header{
 			"Authorization": {"Bearer " + client.Credential.AccessToken},
 			"Content-Type":  {"application/json"},
@@ -165,6 +313,14 @@ func (client *Client) request(ctx context.Context, method string, url string, bo
 		request.WithContext(ctx),
 	)
 
+	// 发送请求
+	res := client.Request.Request(
+		method,
+		url,
+		body,
+		option...,
+	)
+
 	if res.Err != nil {
 		return "", sysError(res.Err)
 	}
@@ -180,7 +336,7 @@ func (client *Client) request(ctx context.Context, method string, url string, bo
 		decodeErr error
 	)
 	// 如果有错误
-	if res.Response.StatusCode != 200 {
+	if res.Response.StatusCode != expectedCode {
 		decodeErr = json.Unmarshal([]byte(respBody), &errResp)
 		if decodeErr != nil {
 			return "", sysError(err)
@@ -190,3 +346,11 @@ func (client *Client) request(ctx context.Context, method string, url string, bo
 
 	return respBody, nil
 }
+
+func (client *Client) requestWithStr(ctx context.Context, method string, url string, body string, expectedCode int) (string, *RespError) {
+	// 发送请求
+	bodyReader := ioutil.NopCloser(strings.NewReader(body))
+	return client.request(ctx, method, url, bodyReader, expectedCode,
+		request.WithContentLength(int64(len(body))),
+	)
+}
diff --git a/pkg/filesystem/driver/onedrive/handller.go b/pkg/filesystem/driver/onedrive/handller.go
index e66427b..6b2176d 100644
--- a/pkg/filesystem/driver/onedrive/handller.go
+++ b/pkg/filesystem/driver/onedrive/handller.go
@@ -30,7 +30,7 @@ func (handler Driver) Put(ctx context.Context, file io.ReadCloser, dst string, s
 // Delete 删除一个或多个文件，
 // 返回未删除的文件，及遇到的最后一个错误
 func (handler Driver) Delete(ctx context.Context, files []string) ([]string, error) {
-	return []string{}, errors.New("未实现")
+	return handler.Client.Delete(ctx, files)
 }
 
 // Thumb 获取文件缩略图
@@ -50,14 +50,23 @@ func (handler Driver) Source(
 	return "", errors.New("未实现")
 }
 
-// Token 获取上传策略和认证Token
+// Token 获取上传会话URL
 func (handler Driver) Token(ctx context.Context, TTL int64, key string) (serializer.UploadCredential, error) {
 
-	// 读取上下文中生成的存储路径
+	// 读取上下文中生成的存储路径和文件大小
 	savePath, ok := ctx.Value(fsctx.SavePathCtx).(string)
 	if !ok {
 		return serializer.UploadCredential{}, errors.New("无法获取存储路径")
 	}
+	fileSize, ok := ctx.Value(fsctx.FileSizeCtx).(uint64)
+	if !ok {
+		return serializer.UploadCredential{}, errors.New("无法获取文件大小")
+	}
+
+	// 如果小于10MB，则由服务端中转
+	if fileSize <= 10*1024*1024 {
+		return serializer.UploadCredential{}, nil
+	}
 
 	// 生成回调地址
 	siteURL := model.GetSiteURL()
@@ -69,6 +78,9 @@ func (handler Driver) Token(ctx context.Context, TTL int64, key string) (seriali
 		return serializer.UploadCredential{}, err
 	}
 
+	// 监控回调及上传
+	go handler.Client.MonitorUpload(uploadURL, key, savePath, fileSize, TTL)
+
 	return serializer.UploadCredential{
 		Policy: uploadURL,
 		Token:  apiURL.String(),
diff --git a/pkg/filesystem/driver/onedrive/types.go b/pkg/filesystem/driver/onedrive/types.go
new file mode 100644
index 0000000..de41eee
--- /dev/null
+++ b/pkg/filesystem/driver/onedrive/types.go
@@ -0,0 +1,77 @@
+package onedrive
+
+import (
+	"sync"
+)
+
+// RespError 接口返回错误
+type RespError struct {
+	APIError APIError `json:"error"`
+}
+
+// APIError 接口返回的错误内容
+type APIError struct {
+	Code    string `json:"code"`
+	Message string `json:"message"`
+}
+
+// UploadSessionResponse 分片上传会话
+type UploadSessionResponse struct {
+	DataContext        string   `json:"@odata.context"`
+	ExpirationDateTime string   `json:"expirationDateTime"`
+	NextExpectedRanges []string `json:"nextExpectedRanges"`
+	UploadURL          string   `json:"uploadUrl"`
+}
+
+// FileInfo 文件元信息
+type FileInfo struct {
+	Name            string          `json:"name"`
+	Size            uint64          `json:"size"`
+	Image           imageInfo       `json:"image"`
+	ParentReference parentReference `json:"parentReference"`
+}
+
+type imageInfo struct {
+	Height int `json:"height"`
+	Width  int `json:"width"`
+}
+
+type parentReference struct {
+	Path string `json:"path"`
+	Name string `json:"name"`
+	ID   string `json:"id"`
+}
+
+// UploadResult 上传结果
+type UploadResult struct {
+	ID   string `json:"id"`
+	Name string `json:"name"`
+	Size uint64 `json:"size"`
+}
+
+// BatchRequests 批量操作请求
+type BatchRequests struct {
+	Requests []BatchRequest `json:"requests"`
+}
+
+// BatchRequest 批量操作单个请求
+type BatchRequest struct {
+	ID      string            `json:"id"`
+	Method  string            `json:"method"`
+	URL     string            `json:"url"`
+	Body    interface{}       `json:"body,omitempty"`
+	Headers map[string]string `json:"headers,omitempty"`
+}
+
+// BatchResponses 批量操作响应
+type BatchResponses struct {
+	Responses []BatchResponse `json:"responses"`
+}
+
+// BatchResponse 批量操作单个响应
+type BatchResponse struct {
+	ID     string `json:"id"`
+	Status int    `json:"status"`
+}
+
+var callbackSignal sync.Map