Merge pull request #132 from whywaita/fix/registration-token-cache-bug

Org token and Repo token has a same installation ID
whywaita · Jan 31, 2022 · f082923 · f082923
2 parents e1eeca0 + 51b1cb7
commit f082923
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/pkg/gh/token_registration.go b/pkg/gh/token_registration.go
@@ -3,7 +3,6 @@ package gh
 import (
 	"context"
 	"fmt"
-	"strconv"
 	"time"
 
 	"github.com/patrickmn/go-cache"
@@ -16,7 +15,7 @@ var (
 // GetRunnerRegistrationToken get token for register runner
 // clientInstallation needs to response of `NewClientInstallation()`
 func GetRunnerRegistrationToken(ctx context.Context, gheDomain string, installationID int64, scope string) (string, error) {
-	cachedToken := getRunnerRegisterTokenFromCache(installationID)
+	cachedToken := getRunnerRegisterTokenFromCache(installationID, scope)
 	if cachedToken != "" {
 		return cachedToken, nil
 	}
@@ -25,7 +24,7 @@ func GetRunnerRegistrationToken(ctx context.Context, gheDomain string, installat
 	if err != nil {
 		return "", fmt.Errorf("failed to generate runner register token: %w", err)
 	}
-	setRunnerRegisterTokenCache(installationID, rrToken, *expiresAt)
+	setRunnerRegisterTokenCache(installationID, scope, rrToken, *expiresAt)
 	return rrToken, nil
 }
 
@@ -56,14 +55,14 @@ func generateRunnerRegisterToken(ctx context.Context, gheDomain string, installa
 	}
 }
 
-func setRunnerRegisterTokenCache(installationID int64, token string, expiresAt time.Time) {
-	expiresDuration := time.Until(expiresAt.Add(-time.Minute))
+func setRunnerRegisterTokenCache(installationID int64, scope, token string, expiresAt time.Time) {
+	expiresDuration := time.Until(expiresAt.Add(-6 * time.Minute))
 
-	cacheRegistrationToken.Set(strconv.FormatInt(installationID, 10), token, expiresDuration)
+	cacheRegistrationToken.Set(getCacheKeyRegistrationToken(installationID, scope), token, expiresDuration)
 }
 
-func getRunnerRegisterTokenFromCache(installationID int64) string {
-	got, found := cacheRegistrationToken.Get(strconv.FormatInt(installationID, 10))
+func getRunnerRegisterTokenFromCache(installationID int64, scope string) string {
+	got, found := cacheRegistrationToken.Get(getCacheKeyRegistrationToken(installationID, scope))
 	if !found {
 		return ""
 	}
@@ -73,3 +72,7 @@ func getRunnerRegisterTokenFromCache(installationID int64) string {
 	}
 	return token
 }
+
+func getCacheKeyRegistrationToken(installationID int64, scope string) string {
+	return fmt.Sprintf("%s-%d", scope, installationID)
+}