diff --git a/drivers/aliyundrive_share/driver.go b/drivers/aliyundrive_share/driver.go index db25065c..6a963713 100644 --- a/drivers/aliyundrive_share/driver.go +++ b/drivers/aliyundrive_share/driver.go @@ -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 diff --git a/drivers/aliyundrive_share/util.go b/drivers/aliyundrive_share/util.go index 58e0c683..a29f86f5 100644 --- a/drivers/aliyundrive_share/util.go +++ b/drivers/aliyundrive_share/util.go @@ -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 {