fix(aliyundrive_share): add limit rate and lift rate limit restrictions (#4587)

pull/4665/head
wenmig 2023-06-26 14:49:21 +08:00 committed by GitHub
parent 2ae9cd8634
commit 21b8e7f6e5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 1 deletions

View File

@ -22,6 +22,9 @@ type AliyundriveShare struct {
ShareToken string
DriveId string
cron *cron.Cron
limitList func(ctx context.Context, dir model.Obj) ([]model.Obj, error)
limitLink func(ctx context.Context, file model.Obj) (*model.Link, error)
}
func (d *AliyundriveShare) Config() driver.Config {
@ -48,6 +51,8 @@ func (d *AliyundriveShare) Init(ctx context.Context) error {
log.Errorf("%+v", err)
}
})
d.limitList = utils.LimitRateCtx(d.list, time.Second/4)
d.limitLink = utils.LimitRateCtx(d.link, time.Second)
return nil
}
@ -60,6 +65,10 @@ func (d *AliyundriveShare) Drop(ctx context.Context) error {
}
func (d *AliyundriveShare) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
return d.limitList(ctx, dir)
}
func (d *AliyundriveShare) list(ctx context.Context, dir model.Obj) ([]model.Obj, error) {
files, err := d.getFiles(dir.GetID())
if err != nil {
return nil, err
@ -70,6 +79,10 @@ func (d *AliyundriveShare) List(ctx context.Context, dir model.Obj, args model.L
}
func (d *AliyundriveShare) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*model.Link, error) {
return d.limitLink(ctx, file)
}
func (d *AliyundriveShare) link(ctx context.Context, file model.Obj) (*model.Link, error) {
data := base.Json{
"drive_id": d.DriveId,
"file_id": file.GetID(),
@ -79,7 +92,7 @@ func (d *AliyundriveShare) Link(ctx context.Context, file model.Obj, args model.
}
var resp ShareLinkResp
_, err := d.request("https://api.aliyundrive.com/v2/file/get_share_link_download_url", http.MethodPost, func(req *resty.Request) {
req.SetBody(data).SetResult(&resp)
req.SetHeader(CanaryHeaderKey, CanaryHeaderValue).SetBody(data).SetResult(&resp)
})
if err != nil {
return nil, err

View File

@ -9,6 +9,12 @@ import (
log "github.com/sirupsen/logrus"
)
const (
// CanaryHeaderKey CanaryHeaderValue for lifting rate limit restrictions
CanaryHeaderKey = "X-Canary"
CanaryHeaderValue = "client=web,app=share,version=v2.3.1"
)
func (d *AliyundriveShare) refreshToken() error {
url := "https://auth.aliyundrive.com/v2/account/token"
var resp base.TokenResp
@ -58,6 +64,7 @@ func (d *AliyundriveShare) request(url, method string, callback base.ReqCallback
SetError(&e).
SetHeader("content-type", "application/json").
SetHeader("Authorization", "Bearer\t"+d.AccessToken).
SetHeader(CanaryHeaderKey, CanaryHeaderValue).
SetHeader("x-share-token", d.ShareToken)
if callback != nil {
callback(req)
@ -107,6 +114,7 @@ func (d *AliyundriveShare) getFiles(fileId string) ([]File, error) {
var resp ListResp
res, err := base.RestyClient.R().
SetHeader("x-share-token", d.ShareToken).
SetHeader(CanaryHeaderKey, CanaryHeaderValue).
SetResult(&resp).SetError(&e).SetBody(data).
Post("https://api.aliyundrive.com/adrive/v3/file/list")
if err != nil {