fix(baidu_netdisk): improve upload handling

Author: j2rong4cn

drivers/baidu_netdisk/driver.go

@@ -1,30 +1,31 @@
 package baidu_netdisk
 
 import (
+	"bytes"
 	"context"
 	"crypto/md5"
 	"encoding/hex"
 	"errors"
-	"fmt"
 	"io"
+	"mime/multipart"
+	"net/http"
 	"net/url"
 	"os"
 	stdpath "path"
 	"strconv"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/OpenListTeam/OpenList/v4/drivers/base"
+	"github.com/OpenListTeam/OpenList/v4/internal/cache"
 	"github.com/OpenListTeam/OpenList/v4/internal/conf"
 	"github.com/OpenListTeam/OpenList/v4/internal/driver"
 	"github.com/OpenListTeam/OpenList/v4/internal/errs"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
+	"github.com/OpenListTeam/OpenList/v4/internal/net"
 	"github.com/OpenListTeam/OpenList/v4/pkg/errgroup"
-	"github.com/OpenListTeam/OpenList/v4/pkg/singleflight"
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
 	"github.com/avast/retry-go"
-	"github.com/go-resty/resty/v2"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -35,11 +36,8 @@ type BaiduNetdisk struct {
 	uploadThread int
 	vipType      int // 会员类型，0普通用户(4G/4M)、1普通会员(10G/16M)、2超级会员(20G/32M)
 
-	upClient            *resty.Client // 上传文件使用的http客户端
-	uploadUrlG          singleflight.Group[string]
-	uploadUrlMu         sync.RWMutex
-	uploadUrl           string    // 上传域名
-	uploadUrlUpdateTime time.Time // 上传域名上次更新时间
+	uploadUrl               string // 上传域名
+	uploadUrlExpirationTime cache.Expirable
 }
 
 var ErrUploadIDExpired = errors.New("uploadid expired")
@@ -53,11 +51,6 @@ func (d *BaiduNetdisk) GetAddition() driver.Additional {
 }
 
 func (d *BaiduNetdisk) Init(ctx context.Context) error {
-	d.upClient = base.NewRestyClient().
-		SetTimeout(UPLOAD_TIMEOUT).
-		SetRetryCount(UPLOAD_RETRY_COUNT).
-		SetRetryWaitTime(UPLOAD_RETRY_WAIT_TIME).
-		SetRetryMaxWaitTime(UPLOAD_RETRY_MAX_WAIT_TIME)
 	d.uploadThread, _ = strconv.Atoi(d.UploadThread)
 	if d.uploadThread < 1 {
 		d.uploadThread, d.UploadThread = 1, "1"
@@ -95,9 +88,10 @@ func (d *BaiduNetdisk) List(ctx context.Context, dir model.Obj, args model.ListA
 }
 
 func (d *BaiduNetdisk) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*model.Link, error) {
-	if d.DownloadAPI == "crack" {
+	switch d.DownloadAPI {
+	case "crack":
 		return d.linkCrack(file, args)
-	} else if d.DownloadAPI == "crack_video" {
+	case "crack_video":
 		return d.linkCrackVideo(file, args)
 	}
 	return d.linkOfficial(file, args)
@@ -214,7 +208,7 @@ func (d *BaiduNetdisk) Put(ctx context.Context, dstDir model.Obj, stream model.F
 		tmpF  *os.File
 		err   error
 	)
-	if _, ok := cache.(io.ReaderAt); !ok {
+	if cache == nil {
 		tmpF, err = os.CreateTemp(conf.Conf.TempDir, "file-*")
 		if err != nil {
 			return nil, err
@@ -301,27 +295,25 @@ func (d *BaiduNetdisk) Put(ctx context.Context, dstDir model.Obj, stream model.F
 
 	// step.2 上传分片
 uploadLoop:
-	for attempt := 0; attempt < 2; attempt++ {
+	for range 2 {
 		// 获取上传域名
 		uploadUrl := d.getUploadUrl(path, precreateResp.Uploadid)
 		// 并发上传
 		threadG, upCtx := errgroup.NewGroupWithContext(ctx, d.uploadThread,
-			retry.Attempts(1),
-			retry.Delay(time.Second),
+			retry.Attempts(UPLOAD_RETRY_COUNT),
+			retry.Delay(UPLOAD_RETRY_WAIT_TIME),
+			retry.MaxDelay(UPLOAD_RETRY_MAX_WAIT_TIME),
 			retry.DelayType(retry.BackOffDelay))
 
-		cacheReaderAt, okReaderAt := cache.(io.ReaderAt)
-		if !okReaderAt {
-			return nil, fmt.Errorf("cache object must implement io.ReaderAt interface for upload operations")
-		}
-
 		totalParts := len(precreateResp.BlockList)
 
 		for i, partseq := range precreateResp.BlockList {
-			if utils.IsCanceled(upCtx) || partseq < 0 {
+			if utils.IsCanceled(upCtx) {
+				break
+			}
+			if partseq < 0 {
 				continue
 			}
-			i, partseq := i, partseq
 			offset, size := int64(partseq)*sliceSize, sliceSize
 			if partseq+1 == count {
 				size = lastBlockSize
@@ -335,15 +327,13 @@ uploadLoop:
 					"uploadid":     precreateResp.Uploadid,
 					"partseq":      strconv.Itoa(partseq),
 				}
-				section := io.NewSectionReader(cacheReaderAt, offset, size)
-				err := d.uploadSlice(ctx, uploadUrl, params, stream.GetName(), driver.NewLimitedUploadStream(ctx, section))
+				section := io.NewSectionReader(cache, offset, size)
+				err := d.uploadSlice(ctx, uploadUrl, params, stream.GetName(), section)
 				if err != nil {
 					return err
 				}
 				precreateResp.BlockList[i] = -1
-				// 当前goroutine还没退出，+1才是真正成功的数量
-				success := threadG.Success() + 1
-				progress := float64(success) * 100 / float64(totalParts)
+				progress := float64(threadG.Success()+1) * 100 / float64(totalParts+1)
 				up(progress)
 				return nil
 			})
@@ -430,27 +420,58 @@ func (d *BaiduNetdisk) precreate(ctx context.Context, path string, streamSize in
 }
 
 func (d *BaiduNetdisk) uploadSlice(ctx context.Context, uploadUrl string, params map[string]string, fileName string, file io.Reader) error {
-	res, err := d.upClient.R().
-		SetContext(ctx).
-		SetQueryParams(params).
-		SetFileReader("file", fileName, file).
-		Post(uploadUrl + "/rest/2.0/pcs/superfile2")
+	b := bytes.NewBuffer(nil)
+	mw := multipart.NewWriter(b)
+	_, err := mw.CreateFormFile("file", fileName)
+	if err != nil {
+		return err
+	}
+	headSize := b.Len()
+	err = mw.Close()
+	if err != nil {
+		return err
+	}
+	head := bytes.NewReader(b.Bytes()[:headSize])
+	tail := bytes.NewReader(b.Bytes()[headSize:])
+	rateLimitedRd := driver.NewLimitedUploadStream(ctx, io.MultiReader(head, file, tail))
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadUrl+"/rest/2.0/pcs/superfile2", rateLimitedRd)
+	if err != nil {
+		return err
+	}
+	query := req.URL.Query()
+	for k, v := range params {
+		query.Set(k, v)
+	}
+	req.URL.RawQuery = query.Encode()
+	req.Header.Set("Content-Type", mw.FormDataContentType())
+
+	client := net.NewHttpClient()
+	client.Timeout = UPLOAD_TIMEOUT
+	resp, err := client.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	b.Reset()
+	_, err = b.ReadFrom(resp.Body)
 	if err != nil {
 		return err
 	}
-	log.Debugln(res.RawResponse.Status + res.String())
-	errCode := utils.Json.Get(res.Body(), "error_code").ToInt()
-	errNo := utils.Json.Get(res.Body(), "errno").ToInt()
-	respStr := res.String()
+	body := b.Bytes()
+	respStr := string(body)
+	log.Debugln(respStr)
 	lower := strings.ToLower(respStr)
 	// 合并 uploadid 过期检测逻辑
 	if strings.Contains(lower, "uploadid") &&
 		(strings.Contains(lower, "invalid") || strings.Contains(lower, "expired") || strings.Contains(lower, "not found")) {
 		return ErrUploadIDExpired
 	}
 
+	errCode := utils.Json.Get(body, "error_code").ToInt()
+	errNo := utils.Json.Get(body, "errno").ToInt()
 	if errCode != 0 || errNo != 0 {
-		return errs.NewErr(errs.StreamIncomplete, "error uploading to baidu, response=%s", res.String())
+		return errs.NewErr(errs.StreamIncomplete, "error uploading to baidu, response=%s", respStr)
 	}
 	return nil
 }

drivers/baidu_netdisk/util.go

@@ -12,10 +12,12 @@ import (
 	"unicode"
 
 	"github.com/OpenListTeam/OpenList/v4/drivers/base"
+	"github.com/OpenListTeam/OpenList/v4/internal/cache"
 	"github.com/OpenListTeam/OpenList/v4/internal/driver"
 	"github.com/OpenListTeam/OpenList/v4/internal/errs"
 	"github.com/OpenListTeam/OpenList/v4/internal/model"
 	"github.com/OpenListTeam/OpenList/v4/internal/op"
+	"github.com/OpenListTeam/OpenList/v4/pkg/singleflight"
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
 	"github.com/avast/retry-go"
 	"github.com/go-resty/resty/v2"
@@ -122,7 +124,7 @@ func (d *BaiduNetdisk) request(furl string, method string, callback base.ReqCall
 				}
 			}
 
-			if 31023 == errno && d.DownloadAPI == "crack_video" {
+			if errno == 31023 && d.DownloadAPI == "crack_video" {
 				result = res.Body()
 				return nil
 			}
@@ -247,7 +249,7 @@ func (d *BaiduNetdisk) linkCrack(file model.Obj, _ model.LinkArgs) (*model.Link,
 func (d *BaiduNetdisk) linkCrackVideo(file model.Obj, _ model.LinkArgs) (*model.Link, error) {
 	param := map[string]string{
 		"type":       "VideoURL",
-		"path":       fmt.Sprintf("%s", file.GetPath()),
+		"path":       file.GetPath(),
 		"fs_id":      file.GetID(),
 		"devuid":     "0%1",
 		"clienttype": "1",
@@ -400,45 +402,26 @@ func (d *BaiduNetdisk) getUploadUrl(path, uploadId string) string {
 	if !d.UseDynamicUploadAPI {
 		return d.UploadAPI
 	}
-	getCachedUrlFunc := func() string {
-		d.uploadUrlMu.RLock()
-		defer d.uploadUrlMu.RUnlock()
-		if d.uploadUrl != "" && time.Since(d.uploadUrlUpdateTime) < UPLOAD_URL_EXPIRE_TIME {
-			uploadUrl := d.uploadUrl
-			return uploadUrl
-		}
-		return ""
-	}
 	// 检查地址缓存
-	if uploadUrl := getCachedUrlFunc(); uploadUrl != "" {
-		return uploadUrl
+	if d.uploadUrlExpirationTime != nil && !d.uploadUrlExpirationTime.Expired() {
+		return d.uploadUrl
 	}
 
-	uploadUrlGetFunc := func() (string, error) {
-		// 双重检查缓存
-		if uploadUrl := getCachedUrlFunc(); uploadUrl != "" {
-			return uploadUrl, nil
-		}
-
+	_, err, _ := singleflight.AnyGroup.Do(fmt.Sprintf("BaiduNetdisk.getUploadUrl_%p", d), func() (any, error) {
 		uploadUrl, err := d.requestForUploadUrl(path, uploadId)
 		if err != nil {
-			return "", err
+			log.Warnf("[baidu_netdisk] get upload URL failed (%v), will use fallback URL: %s", err, d.UploadAPI)
+			return nil, err
 		}
 
-		d.uploadUrlMu.Lock()
-		defer d.uploadUrlMu.Unlock()
 		d.uploadUrl = uploadUrl
-		d.uploadUrlUpdateTime = time.Now()
-		return uploadUrl, nil
-	}
-
-	uploadUrl, err, _ := d.uploadUrlG.Do("", uploadUrlGetFunc)
+		d.uploadUrlExpirationTime = cache.ExpirationTime(time.Now().Add(UPLOAD_URL_EXPIRE_TIME))
+		return nil, nil
+	})
 	if err != nil {
-		fallback := d.UploadAPI
-		log.Warnf("[baidu_netdisk] get upload URL failed (%v), will use fallback URL: %s", err, fallback)
-		return fallback
+		return d.UploadAPI
 	}
-	return uploadUrl
+	return d.uploadUrl
 }
 
 // requestForUploadUrl 请求获取上传地址。