From 4ae81b5a79f853e47c8ff741ba297b4fae44ebcb Mon Sep 17 00:00:00 2001 From: foxxorcat <95907542+foxxorcat@users.noreply.github.com> Date: Sat, 5 Mar 2022 13:14:57 +0800 Subject: [PATCH] feat: aliyundrive fast upload (#652) --- drivers/alidrive/driver.go | 132 +++++++++++++++++++++++++------------ utils/md5.go | 7 ++ 2 files changed, 98 insertions(+), 41 deletions(-) diff --git a/drivers/alidrive/driver.go b/drivers/alidrive/driver.go index a870422e..0ddcf68b 100644 --- a/drivers/alidrive/driver.go +++ b/drivers/alidrive/driver.go @@ -2,18 +2,25 @@ package alidrive import ( "bytes" + "crypto/sha1" + "encoding/base64" + "encoding/hex" "errors" "fmt" + "io" + "io/ioutil" + "math" + "math/big" + "net/http" + "os" + "path/filepath" + "github.com/Xhofe/alist/conf" "github.com/Xhofe/alist/drivers/base" "github.com/Xhofe/alist/model" "github.com/Xhofe/alist/utils" "github.com/robfig/cron/v3" log "github.com/sirupsen/logrus" - "io" - "math" - "net/http" - "path/filepath" ) type AliDrive struct{} @@ -376,15 +383,16 @@ type UploadResp struct { PartInfoList []struct { UploadUrl string `json:"upload_url"` } `json:"part_info_list"` + + RapidUpload bool `json:"rapid_upload"` } func (driver AliDrive) Upload(file *model.FileStream, account *model.Account) error { if file == nil { return base.ErrEmptyFile } - const DEFAULT uint64 = 10485760 + const DEFAULT int64 = 10485760 var count = int64(math.Ceil(float64(file.GetSize()) / float64(DEFAULT))) - var finish uint64 = 0 parentFile, err := driver.File(file.ParentPath, account) if err != nil { return err @@ -392,32 +400,38 @@ func (driver AliDrive) Upload(file *model.FileStream, account *model.Account) er if !parentFile.IsDir() { return base.ErrNotFolder } - var resp UploadResp - var e AliRespError - partInfoList := make([]base.Json, 0) + + partInfoList := make([]base.Json, 0, count) var i int64 for i = 0; i < count; i++ { partInfoList = append(partInfoList, base.Json{ "part_number": i + 1, }) } - _, err = aliClient.R().SetResult(&resp).SetError(&e). - SetHeader("authorization", "Bearer\t"+account.AccessToken). - SetBody(base.Json{ - "check_name_mode": "auto_rename", - // content_hash - "content_hash_name": "none", - "drive_id": account.DriveId, - "name": file.GetFileName(), - "parent_file_id": parentFile.Id, - "part_info_list": partInfoList, - //proof_code - "proof_version": "v1", - "size": file.GetSize(), - "type": "file", - }).Post("https://api.aliyundrive.com/adrive/v2/file/createWithFolders") // /v2/file/create_with_proof - //log.Debugf("%+v\n%+v", resp, e) - if e.Code != "" { + + buf := make([]byte, 1024) + n, _ := file.Read(buf[:]) + reqBody := base.Json{ + "check_name_mode": "auto_rename", + "drive_id": account.DriveId, + "name": file.GetFileName(), + "parent_file_id": parentFile.Id, + "part_info_list": partInfoList, + "size": file.GetSize(), + "type": "file", + "pre_hash": utils.GetSHA1Encode(string(buf[:n])), + } + fileReader := io.MultiReader(bytes.NewReader(buf[:n]), file.File) + + var resp UploadResp + var e AliRespError + client := aliClient.R().SetResult(&resp).SetError(&e).SetHeader("authorization", "Bearer\t"+account.AccessToken).SetBody(reqBody) + + _, err = client.Post("https://api.aliyundrive.com/adrive/v2/file/createWithFolders") + if err != nil { + return err + } + if e.Code != "" && e.Code != "PreHashMatched" { if e.Code == "AccessTokenInvalid" { err = driver.RefreshToken(account) if err != nil { @@ -429,26 +443,59 @@ func (driver AliDrive) Upload(file *model.FileStream, account *model.Account) er } return fmt.Errorf("%s", e.Message) } - var byteSize uint64 - for i = 0; i < count; i++ { - byteSize = file.GetSize() - finish - if DEFAULT < byteSize { - byteSize = DEFAULT - } - log.Debugf("%d,%d", byteSize, finish) - byteData := make([]byte, byteSize) - n, err := io.ReadFull(file, byteData) - //n, err := file.Read(byteData) - //byteData, err := io.ReadAll(file) - //n := len(byteData) - log.Debug(err, n) + + if e.Code == "PreHashMatched" { + tempFile, err := ioutil.TempFile(conf.Conf.TempDir, "file-*") if err != nil { return err } - finish += uint64(n) + defer tempFile.Close() + defer os.Remove(tempFile.Name()) - req, err := http.NewRequest("PUT", resp.PartInfoList[i].UploadUrl, bytes.NewBuffer(byteData)) + delete(reqBody, "pre_hash") + h := sha1.New() + if _, err = io.Copy(tempFile, io.TeeReader(fileReader, h)); err != nil { + return err + } + reqBody["content_hash"] = hex.EncodeToString(h.Sum(nil)) + reqBody["content_hash_name"] = "sha1" + reqBody["proof_version"] = "v1" + + /* + js 隐性转换太坑不知道有没有bug + var n = e.access_token， + r = new BigNumber('0x'.concat(md5(n).slice(0, 16)))， + i = new BigNumber(t.file.size)， + o = i ? r.mod(i) : new gt.BigNumber(0); + (t.file.slice(o.toNumber(), Math.min(o.plus(8).toNumber(), t.file.size))) + */ + r, _ := new(big.Int).SetString(utils.GetMD5Encode(account.AccessToken)[:16], 16) + i := new(big.Int).SetUint64(file.Size) + o := r.Mod(r, i) + n, _ = io.NewSectionReader(tempFile, o.Int64(), 8).Read(buf[:8]) + reqBody["proof_code"] = base64.StdEncoding.EncodeToString(buf[:n]) + + _, err = client.Post("https://api.aliyundrive.com/adrive/v2/file/createWithFolders") + if err != nil { + return err + } + if e.Code != "" && e.Code != "PreHashMatched" { + return fmt.Errorf("%s", e.Message) + } + + if resp.RapidUpload { + return nil + } + + if _, err = tempFile.Seek(0, io.SeekStart); err != nil { + return err + } + fileReader = tempFile + } + + for i = 0; i < count; i++ { + req, err := http.NewRequest("PUT", resp.PartInfoList[i].UploadUrl, io.LimitReader(fileReader, DEFAULT)) if err != nil { return err } @@ -473,6 +520,9 @@ func (driver AliDrive) Upload(file *model.FileStream, account *model.Account) er "file_id": resp.FileId, "upload_id": resp.UploadId, }).Post("https://api.aliyundrive.com/v2/file/complete") + if err != nil { + return err + } if e.Code != "" { //if e.Code == "AccessTokenInvalid" { // err = driver.RefreshToken(account) diff --git a/utils/md5.go b/utils/md5.go index 67f962da..2cd85e63 100644 --- a/utils/md5.go +++ b/utils/md5.go @@ -2,10 +2,17 @@ package utils import ( "crypto/md5" + "crypto/sha1" "encoding/hex" "fmt" ) +func GetSHA1Encode(data string) string { + h := sha1.New() + h.Write([]byte(data)) + return hex.EncodeToString(h.Sum(nil)) +} + // GetMD5Encode func GetMD5Encode(data string) string { h := md5.New()