Skip to content

Commit

Permalink
perf(123pan): optimize rate limiting (#6859)
Browse files Browse the repository at this point in the history
- 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
  • Loading branch information
seiuneko authored Jul 25, 2024
1 parent d4e3355 commit 5fa70e4
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 24 deletions.
16 changes: 7 additions & 9 deletions drivers/123/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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,
Expand All @@ -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)
10 changes: 5 additions & 5 deletions drivers/123/util.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package _123

import (
"context"
"errors"
"fmt"
"hash/crc32"
Expand All @@ -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"
)
Expand Down Expand Up @@ -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{
Expand Down
13 changes: 7 additions & 6 deletions drivers/123_share/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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)
8 changes: 4 additions & 4 deletions drivers/123_share/util.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package _123Share

import (
"context"
"errors"
"fmt"
"hash/crc32"
Expand Down Expand Up @@ -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{
Expand Down

0 comments on commit 5fa70e4

Please sign in to comment.