From 5fa70e40107180816b0c9b1c1a2325e0b5318bdb Mon Sep 17 00:00:00 2001
From: seiuneko <25706824+seiuneko@users.noreply.github.com>
Date: Thu, 25 Jul 2024 20:08:59 +0800
Subject: [PATCH] perf(123pan): optimize rate limiting (#6859)

- eliminating fixed 200 ms delay in getFiles to prevent thread starvation
- allowing cancellation via context to mitigate potential DoS attacks by immediately cancelling excessive requests
---
 drivers/123/driver.go       | 16 +++++++---------
 drivers/123/util.go         | 10 +++++-----
 drivers/123_share/driver.go | 13 +++++++------
 drivers/123_share/util.go   |  8 ++++----
 4 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/drivers/123/driver.go b/drivers/123/driver.go
index bd7089e85d2..aeda7fcf742 100644
--- a/drivers/123/driver.go
+++ b/drivers/123/driver.go
@@ -53,7 +53,7 @@ func (d *Pan123) Drop(ctx context.Context) error {
 }
 
 func (d *Pan123) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
-	files, err := d.getFiles(dir.GetID(), dir.GetName())
+	files, err := d.getFiles(ctx, dir.GetID(), dir.GetName())
 	if err != nil {
 		return nil, err
 	}
@@ -247,9 +247,6 @@ func (d *Pan123) Put(ctx context.Context, dstDir model.Obj, stream model.FileStr
 		}
 		_, err = uploader.UploadWithContext(ctx, input)
 	}
-	if err != nil {
-		return err
-	}
 	_, err = d.request(UploadComplete, http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{
 			"fileId": resp.Data.FileId,
@@ -258,11 +255,12 @@ func (d *Pan123) Put(ctx context.Context, dstDir model.Obj, stream model.FileStr
 	return err
 }
 
-func (d *Pan123) APIRateLimit(api string) bool {
-	limiter, _ := d.apiRateLimit.LoadOrStore(api,
-		rate.NewLimiter(rate.Every(time.Millisecond*700), 1))
-	ins := limiter.(*rate.Limiter)
-	return ins.Allow()
+func (d *Pan123) APIRateLimit(ctx context.Context, api string) error {
+	value, _ := d.apiRateLimit.LoadOrStore(api,
+		rate.NewLimiter(rate.Every(700*time.Millisecond), 1))
+	limiter := value.(*rate.Limiter)
+
+	return limiter.Wait(ctx)
 }
 
 var _ driver.Driver = (*Pan123)(nil)
diff --git a/drivers/123/util.go b/drivers/123/util.go
index 540201f7035..73c73b3b3b3 100644
--- a/drivers/123/util.go
+++ b/drivers/123/util.go
@@ -1,6 +1,7 @@
 package _123
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"hash/crc32"
@@ -14,7 +15,7 @@ import (
 
 	"github.com/alist-org/alist/v3/drivers/base"
 	"github.com/alist-org/alist/v3/pkg/utils"
-	resty "github.com/go-resty/resty/v2"
+	"github.com/go-resty/resty/v2"
 	jsoniter "github.com/json-iterator/go"
 	log "github.com/sirupsen/logrus"
 )
@@ -233,15 +234,14 @@ func (d *Pan123) request(url string, method string, callback base.ReqCallback, r
 	return body, nil
 }
 
-func (d *Pan123) getFiles(parentId string, name string) ([]File, error) {
+func (d *Pan123) getFiles(ctx context.Context, parentId string, name string) ([]File, error) {
 	page := 1
 	total := 0
 	res := make([]File, 0)
 	// 2024-02-06 fix concurrency by 123pan
 	for {
-		if !d.APIRateLimit(FileList) {
-			time.Sleep(time.Millisecond * 200)
-			continue
+		if err := d.APIRateLimit(ctx, FileList); err != nil {
+			return nil, err
 		}
 		var resp Files
 		query := map[string]string{
diff --git a/drivers/123_share/driver.go b/drivers/123_share/driver.go
index 7fca7cc145e..9c1f3803710 100644
--- a/drivers/123_share/driver.go
+++ b/drivers/123_share/driver.go
@@ -45,7 +45,7 @@ func (d *Pan123Share) Drop(ctx context.Context) error {
 
 func (d *Pan123Share) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
 	// TODO return the files list, required
-	files, err := d.getFiles(dir.GetID())
+	files, err := d.getFiles(ctx, dir.GetID())
 	if err != nil {
 		return nil, err
 	}
@@ -150,11 +150,12 @@ func (d *Pan123Share) Put(ctx context.Context, dstDir model.Obj, stream model.Fi
 //	return nil, errs.NotSupport
 //}
 
-func (d *Pan123Share) APIRateLimit(api string) bool {
-	limiter, _ := d.apiRateLimit.LoadOrStore(api,
-		rate.NewLimiter(rate.Every(time.Millisecond*700), 1))
-	ins := limiter.(*rate.Limiter)
-	return ins.Allow()
+func (d *Pan123Share) APIRateLimit(ctx context.Context, api string) error {
+	value, _ := d.apiRateLimit.LoadOrStore(api,
+		rate.NewLimiter(rate.Every(700*time.Millisecond), 1))
+	limiter := value.(*rate.Limiter)
+
+	return limiter.Wait(ctx)
 }
 
 var _ driver.Driver = (*Pan123Share)(nil)
diff --git a/drivers/123_share/util.go b/drivers/123_share/util.go
index a192993bf19..80ea8f0ca46 100644
--- a/drivers/123_share/util.go
+++ b/drivers/123_share/util.go
@@ -1,6 +1,7 @@
 package _123Share
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"hash/crc32"
@@ -80,13 +81,12 @@ func (d *Pan123Share) request(url string, method string, callback base.ReqCallba
 	return body, nil
 }
 
-func (d *Pan123Share) getFiles(parentId string) ([]File, error) {
+func (d *Pan123Share) getFiles(ctx context.Context, parentId string) ([]File, error) {
 	page := 1
 	res := make([]File, 0)
 	for {
-		if !d.APIRateLimit(FileList) {
-			time.Sleep(time.Millisecond * 200)
-			continue
+		if err := d.APIRateLimit(ctx, FileList); err != nil {
+			return nil, err
 		}
 		var resp Files
 		query := map[string]string{