From 35d047db0152fe5058c7ea1eafa77dac1892c995 Mon Sep 17 00:00:00 2001
From: Yusuke Kuoka <ykuoka@gmail.com>
Date: Tue, 16 Feb 2021 12:56:52 +0900
Subject: [PATCH] Fix enterprise runners misusing cached token (#314)

Follow-up for #290
---
 github/fake/fake.go | 36 ++++++++++++++++++++++++++++++++++++
 github/github.go    |  9 +++------
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/github/fake/fake.go b/github/fake/fake.go
index cdea9cb6..3d86e219 100644
--- a/github/fake/fake.go
+++ b/github/fake/fake.go
@@ -102,6 +102,18 @@ func NewServer(opts ...Option) *httptest.Server {
 			Status: http.StatusBadRequest,
 			Body:   "",
 		},
+		"/enterprises/test/actions/runners/registration-token": &Handler{
+			Status: http.StatusCreated,
+			Body:   fmt.Sprintf("{\"token\": \"%s\", \"expires_at\": \"%s\"}", RegistrationToken, time.Now().Add(time.Hour*1).Format(time.RFC3339)),
+		},
+		"/enterprises/invalid/actions/runners/registration-token": &Handler{
+			Status: http.StatusOK,
+			Body:   fmt.Sprintf("{\"token\": \"%s\", \"expires_at\": \"%s\"}", RegistrationToken, time.Now().Add(time.Hour*1).Format(time.RFC3339)),
+		},
+		"/enterprises/error/actions/runners/registration-token": &Handler{
+			Status: http.StatusBadRequest,
+			Body:   "",
+		},
 
 		// For ListRunners
 		"/repos/test/valid/actions/runners": config.FixedResponses.ListRunners,
@@ -125,6 +137,18 @@ func NewServer(opts ...Option) *httptest.Server {
 			Status: http.StatusBadRequest,
 			Body:   "",
 		},
+		"/enterprises/test/actions/runners": &Handler{
+			Status: http.StatusOK,
+			Body:   RunnersListBody,
+		},
+		"/enterprises/invalid/actions/runners": &Handler{
+			Status: http.StatusNoContent,
+			Body:   "",
+		},
+		"/enterprises/error/actions/runners": &Handler{
+			Status: http.StatusBadRequest,
+			Body:   "",
+		},
 
 		// For RemoveRunner
 		"/repos/test/valid/actions/runners/1": &Handler{
@@ -151,6 +175,18 @@ func NewServer(opts ...Option) *httptest.Server {
 			Status: http.StatusBadRequest,
 			Body:   "",
 		},
+		"/enterprises/test/actions/runners/1": &Handler{
+			Status: http.StatusNoContent,
+			Body:   "",
+		},
+		"/enterprises/invalid/actions/runners/1": &Handler{
+			Status: http.StatusOK,
+			Body:   "",
+		},
+		"/enterprises/error/actions/runners/1": &Handler{
+			Status: http.StatusBadRequest,
+			Body:   "",
+		},
 
 		// For auto-scaling based on the number of queued(pending) workflow runs
 		"/repos/test/valid/actions/runs": config.FixedResponses.ListRepositoryWorkflowRuns,
diff --git a/github/github.go b/github/github.go
index 3e92e023..f258c220 100644
--- a/github/github.go
+++ b/github/github.go
@@ -82,7 +82,7 @@ func (c *Client) GetRegistrationToken(ctx context.Context, enterprise, org, repo
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	key := getRegistrationKey(org, repo)
+	key := getRegistrationKey(org, repo, enterprise)
 	rt, ok := c.regTokens[key]
 
 	if ok && rt.GetExpiresAt().After(time.Now()) {
@@ -250,11 +250,8 @@ func getEnterpriseOrganisationAndRepo(enterprise, org, repo string) (string, str
 	return "", "", "", fmt.Errorf("enterprise, organization and repository are all empty")
 }
 
-func getRegistrationKey(org, repo string) string {
-	if len(org) > 0 {
-		return org
-	}
-	return repo
+func getRegistrationKey(org, repo, enterprise string) string {
+	return fmt.Sprintf("org=%s,repo=%s,enterprise=%s", org, repo, enterprise)
 }
 
 func splitOwnerAndRepo(repo string) (string, string, error) {