diff --git a/drivers/local/driver.go b/drivers/local/driver.go index abe5c6b5..bf993e5d 100644 --- a/drivers/local/driver.go +++ b/drivers/local/driver.go @@ -30,6 +30,10 @@ type Local struct { model.Storage Addition mkdirPerm int32 + + // zero means no limit + thumbConcurrency int + thumbTokenBucket TokenBucket } func (d *Local) Config() driver.Config { @@ -62,6 +66,18 @@ func (d *Local) Init(ctx context.Context) error { return err } } + if d.ThumbConcurrency != "" { + v, err := strconv.ParseUint(d.ThumbConcurrency, 10, 32) + if err != nil { + return err + } + d.thumbConcurrency = int(v) + } + if d.thumbConcurrency == 0 { + d.thumbTokenBucket = NewNopTokenBucket() + } else { + d.thumbTokenBucket = NewStaticTokenBucket(d.thumbConcurrency) + } return nil } @@ -126,7 +142,6 @@ func (d *Local) FileInfoToObj(f fs.FileInfo, reqPath string, fullPath string) mo }, } return &file - } func (d *Local) GetMeta(ctx context.Context, path string) (model.Obj, error) { f, err := os.Stat(path) @@ -178,7 +193,13 @@ func (d *Local) Link(ctx context.Context, file model.Obj, args model.LinkArgs) ( fullPath := file.GetPath() var link model.Link if args.Type == "thumb" && utils.Ext(file.GetName()) != "svg" { - buf, thumbPath, err := d.getThumb(file) + var buf *bytes.Buffer + var thumbPath *string + err := d.thumbTokenBucket.Do(ctx, func() error { + var err error + buf, thumbPath, err = d.getThumb(file) + return err + }) if err != nil { return nil, err } diff --git a/drivers/local/meta.go b/drivers/local/meta.go index 51b49e64..5ffac920 100644 --- a/drivers/local/meta.go +++ b/drivers/local/meta.go @@ -9,6 +9,7 @@ type Addition struct { driver.RootPath Thumbnail bool `json:"thumbnail" required:"true" help:"enable thumbnail"` ThumbCacheFolder string `json:"thumb_cache_folder"` + ThumbConcurrency string `json:"thumb_concurrency" default:"16" required:"false" help:"Number of concurrent thumbnail generation goroutines. This controls how many thumbnails can be generated in parallel."` ShowHidden bool `json:"show_hidden" default:"true" required:"false" help:"show hidden directories and files"` MkdirPerm string `json:"mkdir_perm" default:"777"` RecycleBinPath string `json:"recycle_bin_path" default:"delete permanently" help:"path to recycle bin, delete permanently if empty or keep 'delete permanently'"` diff --git a/drivers/local/token_bucket.go b/drivers/local/token_bucket.go new file mode 100644 index 00000000..38fbe73f --- /dev/null +++ b/drivers/local/token_bucket.go @@ -0,0 +1,61 @@ +package local + +import "context" + +type TokenBucket interface { + Take() <-chan struct{} + Put() + Do(context.Context, func() error) error +} + +// StaticTokenBucket is a bucket with a fixed number of tokens, +// where the retrieval and return of tokens are manually controlled. +// In the initial state, the bucket is full. +type StaticTokenBucket struct { + bucket chan struct{} +} + +func NewStaticTokenBucket(size int) StaticTokenBucket { + bucket := make(chan struct{}, size) + for range size { + bucket <- struct{}{} + } + return StaticTokenBucket{bucket: bucket} +} + +func (b StaticTokenBucket) Take() <-chan struct{} { + return b.bucket +} + +func (b StaticTokenBucket) Put() { + b.bucket <- struct{}{} +} + +func (b StaticTokenBucket) Do(ctx context.Context, f func() error) error { + select { + case <-ctx.Done(): + return ctx.Err() + case <-b.bucket: + defer b.Put() + } + return f() +} + +// NopTokenBucket all function calls to this bucket will success immediately +type NopTokenBucket struct { + nop chan struct{} +} + +func NewNopTokenBucket() NopTokenBucket { + nop := make(chan struct{}) + close(nop) + return NopTokenBucket{nop} +} + +func (b NopTokenBucket) Take() <-chan struct{} { + return b.nop +} + +func (b NopTokenBucket) Put() {} + +func (b NopTokenBucket) Do(_ context.Context, f func() error) error { return f() } diff --git a/drivers/pikpak/driver.go b/drivers/pikpak/driver.go index 87f3bc7b..8dab4f72 100644 --- a/drivers/pikpak/driver.go +++ b/drivers/pikpak/driver.go @@ -4,6 +4,10 @@ import ( "context" "encoding/json" "fmt" + "net/http" + "strconv" + "strings" + "github.com/alist-org/alist/v3/drivers/base" "github.com/alist-org/alist/v3/internal/driver" "github.com/alist-org/alist/v3/internal/model" @@ -13,9 +17,6 @@ import ( "github.com/go-resty/resty/v2" log "github.com/sirupsen/logrus" "golang.org/x/oauth2" - "net/http" - "strconv" - "strings" ) type PikPak struct { @@ -90,54 +91,36 @@ func (d *PikPak) Init(ctx context.Context) (err error) { // 如果已经有RefreshToken，直接获取AccessToken if d.Addition.RefreshToken != "" { - // 使用 oauth2 刷新令牌 - // 初始化 oauth2Token - d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) { - return oauth2Config.TokenSource(ctx, &oauth2.Token{ - RefreshToken: d.Addition.RefreshToken, - }).Token() - })) - _, err := d.oauth2Token.Token() - if err != nil { - if err := d.login(); err != nil { + if d.RefreshTokenMethod == "oauth2" { + // 使用 oauth2 刷新令牌 + // 初始化 oauth2Token + d.initializeOAuth2Token(ctx, oauth2Config, d.Addition.RefreshToken) + if err := d.refreshTokenByOAuth2(); err != nil { + return err + } + } else { + if err := d.refreshToken(d.Addition.RefreshToken); err != nil { return err } - d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) { - return oauth2Config.TokenSource(ctx, &oauth2.Token{ - RefreshToken: d.RefreshToken, - }).Token() - })) } + } else { // 如果没有填写RefreshToken，尝试登录 获取 refreshToken if err := d.login(); err != nil { return err } - d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) { - return oauth2Config.TokenSource(ctx, &oauth2.Token{ - RefreshToken: d.RefreshToken, - }).Token() - })) - } + if d.RefreshTokenMethod == "oauth2" { + d.initializeOAuth2Token(ctx, oauth2Config, d.RefreshToken) + } - token, err := d.oauth2Token.Token() - if err != nil { - return err } - d.RefreshToken = token.RefreshToken - d.AccessToken = token.AccessToken // 获取CaptchaToken - err = d.RefreshCaptchaTokenAtLogin(GetAction(http.MethodGet, "https://api-drive.mypikpak.com/drive/v1/files"), d.Username) + err = d.RefreshCaptchaTokenAtLogin(GetAction(http.MethodGet, "https://api-drive.mypikpak.com/drive/v1/files"), d.Common.GetUserID()) if err != nil { return err } - // 获取用户ID - userID := token.Extra("sub").(string) - if userID != "" { - d.Common.SetUserID(userID) - } // 更新UserAgent if d.Platform == "android" { d.Common.UserAgent = BuildCustomUserAgent(utils.GetMD5EncodeStr(d.Username+d.Password), AndroidClientID, AndroidPackageName, AndroidSdkVersion, AndroidClientVersion, AndroidPackageName, d.Common.UserID) diff --git a/drivers/pikpak/meta.go b/drivers/pikpak/meta.go index df480fb1..9e157b6e 100644 --- a/drivers/pikpak/meta.go +++ b/drivers/pikpak/meta.go @@ -7,14 +7,15 @@ import ( type Addition struct { driver.RootID - Username string `json:"username" required:"true"` - Password string `json:"password" required:"true"` - Platform string `json:"platform" required:"true" type:"select" options:"android,web"` - RefreshToken string `json:"refresh_token" required:"true" default:""` - CaptchaToken string `json:"captcha_token" default:""` - DeviceID string `json:"device_id" required:"false" default:""` - DisableMediaLink bool `json:"disable_media_link" default:"true"` - CaptchaApi string `json:"captcha_api" default:""` + Username string `json:"username" required:"true"` + Password string `json:"password" required:"true"` + Platform string `json:"platform" required:"true" type:"select" options:"android,web"` + RefreshToken string `json:"refresh_token" required:"true" default:""` + RefreshTokenMethod string `json:"refresh_token_method" required:"true" type:"select" options:"oauth2,http"` + CaptchaToken string `json:"captcha_token" default:""` + DeviceID string `json:"device_id" required:"false" default:""` + DisableMediaLink bool `json:"disable_media_link" default:"true"` + CaptchaApi string `json:"captcha_api" default:""` } var config = driver.Config{ diff --git a/drivers/pikpak/util.go b/drivers/pikpak/util.go index bbbcd8e7..94089133 100644 --- a/drivers/pikpak/util.go +++ b/drivers/pikpak/util.go @@ -2,17 +2,11 @@ package pikpak import ( "bytes" + "context" "crypto/md5" "crypto/sha1" "encoding/hex" "fmt" - "github.com/alist-org/alist/v3/internal/driver" - "github.com/alist-org/alist/v3/internal/model" - "github.com/alist-org/alist/v3/internal/op" - "github.com/alist-org/alist/v3/pkg/utils" - "github.com/aliyun/aliyun-oss-go-sdk/oss" - jsoniter "github.com/json-iterator/go" - "github.com/pkg/errors" "io" "net/http" "path/filepath" @@ -21,6 +15,15 @@ import ( "sync" "time" + "github.com/alist-org/alist/v3/internal/driver" + "github.com/alist-org/alist/v3/internal/model" + "github.com/alist-org/alist/v3/internal/op" + "github.com/alist-org/alist/v3/pkg/utils" + "github.com/aliyun/aliyun-oss-go-sdk/oss" + jsoniter "github.com/json-iterator/go" + "github.com/pkg/errors" + "golang.org/x/oauth2" + "github.com/alist-org/alist/v3/drivers/base" "github.com/go-resty/resty/v2" ) @@ -112,39 +115,63 @@ func (d *PikPak) login() error { return nil } -//func (d *PikPak) refreshToken() error { -// url := "https://user.mypikpak.com/v1/auth/token" -// var e ErrResp -// res, err := base.RestyClient.SetRetryCount(1).R().SetError(&e). -// SetHeader("user-agent", "").SetBody(base.Json{ -// "client_id": ClientID, -// "client_secret": ClientSecret, -// "grant_type": "refresh_token", -// "refresh_token": d.RefreshToken, -// }).SetQueryParam("client_id", ClientID).Post(url) -// if err != nil { -// d.Status = err.Error() -// op.MustSaveDriverStorage(d) -// return err -// } -// if e.ErrorCode != 0 { -// if e.ErrorCode == 4126 { -// // refresh_token invalid, re-login -// return d.login() -// } -// d.Status = e.Error() -// op.MustSaveDriverStorage(d) -// return errors.New(e.Error()) -// } -// data := res.Body() -// d.Status = "work" -// d.RefreshToken = jsoniter.Get(data, "refresh_token").ToString() -// d.AccessToken = jsoniter.Get(data, "access_token").ToString() -// d.Common.SetUserID(jsoniter.Get(data, "sub").ToString()) -// d.Addition.RefreshToken = d.RefreshToken -// op.MustSaveDriverStorage(d) -// return nil -//} +func (d *PikPak) refreshToken(refreshToken string) error { + url := "https://user.mypikpak.com/v1/auth/token" + var e ErrResp + res, err := base.RestyClient.SetRetryCount(1).R().SetError(&e). + SetHeader("user-agent", "").SetBody(base.Json{ + "client_id": d.ClientID, + "client_secret": d.ClientSecret, + "grant_type": "refresh_token", + "refresh_token": refreshToken, + }).SetQueryParam("client_id", d.ClientID).Post(url) + if err != nil { + d.Status = err.Error() + op.MustSaveDriverStorage(d) + return err + } + if e.ErrorCode != 0 { + if e.ErrorCode == 4126 { + // refresh_token invalid, re-login + return d.login() + } + d.Status = e.Error() + op.MustSaveDriverStorage(d) + return errors.New(e.Error()) + } + data := res.Body() + d.Status = "work" + d.RefreshToken = jsoniter.Get(data, "refresh_token").ToString() + d.AccessToken = jsoniter.Get(data, "access_token").ToString() + d.Common.SetUserID(jsoniter.Get(data, "sub").ToString()) + d.Addition.RefreshToken = d.RefreshToken + op.MustSaveDriverStorage(d) + return nil +} + +func (d *PikPak) initializeOAuth2Token(ctx context.Context, oauth2Config *oauth2.Config, refreshToken string) { + d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) { + return oauth2Config.TokenSource(ctx, &oauth2.Token{ + RefreshToken: refreshToken, + }).Token() + })) +} + +func (d *PikPak) refreshTokenByOAuth2() error { + token, err := d.oauth2Token.Token() + if err != nil { + return err + } + d.Status = "work" + d.RefreshToken = token.RefreshToken + d.AccessToken = token.AccessToken + // 获取用户ID + userID := token.Extra("sub").(string) + d.Common.SetUserID(userID) + d.Addition.RefreshToken = d.RefreshToken + op.MustSaveDriverStorage(d) + return nil +} func (d *PikPak) request(url string, method string, callback base.ReqCallback, resp interface{}) ([]byte, error) { req := base.RestyClient.R() @@ -181,18 +208,15 @@ func (d *PikPak) request(url string, method string, callback base.ReqCallback, r return res.Body(), nil case 4122, 4121, 16: // access_token 过期 - - //if err1 := d.refreshToken(); err1 != nil { - // return nil, err1 - //} - t, err := d.oauth2Token.Token() - if err != nil { - return nil, err + if d.RefreshTokenMethod == "oauth2" { + if err1 := d.refreshTokenByOAuth2(); err1 != nil { + return nil, err1 + } + } else { + if err1 := d.refreshToken(d.RefreshToken); err1 != nil { + return nil, err1 + } } - d.AccessToken = t.AccessToken - d.RefreshToken = t.RefreshToken - d.Addition.RefreshToken = t.RefreshToken - op.MustSaveDriverStorage(d) return d.request(url, method, callback, resp) case 9: // 验证码token过期 @@ -339,6 +363,10 @@ func (c *Common) GetDeviceID() string { return c.DeviceID } +func (c *Common) GetUserID() string { + return c.UserID +} + // RefreshCaptchaTokenAtLogin 刷新验证码token(登录后) func (d *PikPak) RefreshCaptchaTokenAtLogin(action, userID string) error { metas := map[string]string{ diff --git a/internal/net/serve.go b/internal/net/serve.go index e58d7eb9..0eb8cbb8 100644 --- a/internal/net/serve.go +++ b/internal/net/serve.go @@ -87,9 +87,9 @@ func ServeHTTP(w http.ResponseWriter, r *http.Request, name string, modTime time sendSize := size var sendContent io.ReadCloser ranges, err := http_range.ParseRange(rangeReq, size) - switch err { - case nil: - case http_range.ErrNoOverlap: + switch { + case err == nil: + case errors.Is(err, http_range.ErrNoOverlap): if size == 0 { // Some clients add a Range header to all requests to // limit the size of the response. If the file is empty, @@ -105,7 +105,7 @@ func ServeHTTP(w http.ResponseWriter, r *http.Request, name string, modTime time return } - if sumRangesSize(ranges) > size || size < 0 { + if sumRangesSize(ranges) > size { // The total number of bytes in all the ranges is larger than the size of the file // or unknown file size, ignore the range request. ranges = nil @@ -174,6 +174,7 @@ func ServeHTTP(w http.ResponseWriter, r *http.Request, name string, modTime time pw.Close() }() } + defer sendContent.Close() w.Header().Set("Accept-Ranges", "bytes") if w.Header().Get("Content-Encoding") == "" { @@ -192,7 +193,6 @@ func ServeHTTP(w http.ResponseWriter, r *http.Request, name string, modTime time http.Error(w, err.Error(), http.StatusInternalServerError) } } - //defer sendContent.Close() } func ProcessHeader(origin, override http.Header) http.Header { result := http.Header{}