From 163af0515fee7c9ee27dda6e599fbf18fe8dc662 Mon Sep 17 00:00:00 2001
From: foxxorcat <95907542+foxxorcat@users.noreply.github.com>
Date: Mon, 27 May 2024 21:31:59 +0800
Subject: [PATCH] fix(pikpak): refresh_token contention (#6501 close #6511)

---
 drivers/pikpak/driver.go       | 16 +++++++---------
 drivers/pikpak_share/driver.go | 21 ++++++++-------------
 pkg/utils/oauth2.go            | 15 +++++++++++++++
 3 files changed, 30 insertions(+), 22 deletions(-)
 create mode 100644 pkg/utils/oauth2.go

diff --git a/drivers/pikpak/driver.go b/drivers/pikpak/driver.go
index e27263ddbb7..2dab2a9b066 100644
--- a/drivers/pikpak/driver.go
+++ b/drivers/pikpak/driver.go
@@ -41,10 +41,6 @@ func (d *PikPak) Init(ctx context.Context) (err error) {
 		d.ClientSecret = "dbw2OtmVEeuUvIptb1Coyg"
 	}
 
-	withClient := func(ctx context.Context) context.Context {
-		return context.WithValue(ctx, oauth2.HTTPClient, base.HttpClient)
-	}
-
 	oauth2Config := &oauth2.Config{
 		ClientID:     d.ClientID,
 		ClientSecret: d.ClientSecret,
@@ -55,11 +51,13 @@ func (d *PikPak) Init(ctx context.Context) (err error) {
 		},
 	}
 
-	oauth2Token, err := oauth2Config.PasswordCredentialsToken(withClient(ctx), d.Username, d.Password)
-	if err != nil {
-		return err
-	}
-	d.oauth2Token = oauth2Config.TokenSource(withClient(context.Background()), oauth2Token)
+	d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) {
+		return oauth2Config.PasswordCredentialsToken(
+			context.WithValue(context.Background(), oauth2.HTTPClient, base.HttpClient),
+			d.Username,
+			d.Password,
+		)
+	}))
 	return nil
 }
 
diff --git a/drivers/pikpak_share/driver.go b/drivers/pikpak_share/driver.go
index 58c2c8c4b55..1862db06bf4 100644
--- a/drivers/pikpak_share/driver.go
+++ b/drivers/pikpak_share/driver.go
@@ -33,10 +33,6 @@ func (d *PikPakShare) Init(ctx context.Context) error {
 		d.ClientSecret = "dbw2OtmVEeuUvIptb1Coyg"
 	}
 
-	withClient := func(ctx context.Context) context.Context {
-		return context.WithValue(ctx, oauth2.HTTPClient, base.HttpClient)
-	}
-
 	oauth2Config := &oauth2.Config{
 		ClientID:     d.ClientID,
 		ClientSecret: d.ClientSecret,
@@ -47,17 +43,16 @@ func (d *PikPakShare) Init(ctx context.Context) error {
 		},
 	}
 
-	oauth2Token, err := oauth2Config.PasswordCredentialsToken(withClient(ctx), d.Username, d.Password)
-	if err != nil {
-		return err
-	}
-	d.oauth2Token = oauth2Config.TokenSource(withClient(context.Background()), oauth2Token)
+	d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) {
+		return oauth2Config.PasswordCredentialsToken(
+			context.WithValue(context.Background(), oauth2.HTTPClient, base.HttpClient),
+			d.Username,
+			d.Password,
+		)
+	}))
 
 	if d.SharePwd != "" {
-		err = d.getSharePassToken()
-		if err != nil {
-			return err
-		}
+		return d.getSharePassToken()
 	}
 	return nil
 }
diff --git a/pkg/utils/oauth2.go b/pkg/utils/oauth2.go
new file mode 100644
index 00000000000..c1ad161245f
--- /dev/null
+++ b/pkg/utils/oauth2.go
@@ -0,0 +1,15 @@
+package utils
+
+import "golang.org/x/oauth2"
+
+type tokenSource struct {
+	fn func() (*oauth2.Token, error)
+}
+
+func (t *tokenSource) Token() (*oauth2.Token, error) {
+	return t.fn()
+}
+
+func TokenSource(fn func() (*oauth2.Token, error)) oauth2.TokenSource {
+	return &tokenSource{fn}
+}