From b22590703d0b516ce3c916ea89d860175d41ab65 Mon Sep 17 00:00:00 2001
From: alfred <alfred@stokespace.com>
Date: Fri, 3 May 2024 14:24:29 -0700
Subject: [PATCH] Populate labels in the generated script.

---
 pkg/config/config.go   |  2 ++
 pkg/starter/scripts.go | 18 ++++++++++++------
 pkg/starter/starter.go | 11 +++++------
 pkg/web/webhook.go     |  2 +-
 4 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 8149095..13e4254 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -103,3 +103,5 @@ func marshalModeWebhookType(in string) ModeWebhookType {
 func (c Conf) IsGHES() bool {
 	return !strings.EqualFold(c.GitHubURL, "https://github.com")
 }
+
+const RequiredMyShoesLabel = "myshoes"
diff --git a/pkg/starter/scripts.go b/pkg/starter/scripts.go
index d6d44e0..888c7cf 100644
--- a/pkg/starter/scripts.go
+++ b/pkg/starter/scripts.go
@@ -22,8 +22,8 @@ func getPatchedFiles() (string, error) {
 	return runnerService, nil
 }
 
-func (s *Starter) getSetupScript(ctx context.Context, targetScope, runnerName string) (string, error) {
-	rawScript, err := s.getSetupRawScript(ctx, targetScope, runnerName)
+func (s *Starter) getSetupScript(ctx context.Context, targetScope, runnerName string, labels []string) (string, error) {
+	rawScript, err := s.getSetupRawScript(ctx, targetScope, runnerName, labels)
 	if err != nil {
 		return "", fmt.Errorf("failed to get raw setup scripts: %w", err)
 	}
@@ -44,7 +44,7 @@ func (s *Starter) getSetupScript(ctx context.Context, targetScope, runnerName st
 	return fmt.Sprintf(templateCompressedScript, encoded), nil
 }
 
-func (s *Starter) getSetupRawScript(ctx context.Context, targetScope, runnerName string) (string, error) {
+func (s *Starter) getSetupRawScript(ctx context.Context, targetScope, runnerName string, labels []string) (string, error) {
 	runnerUser := config.Config.RunnerUser
 	githubURL := config.Config.GitHubURL
 
@@ -76,9 +76,15 @@ func (s *Starter) getSetupRawScript(ctx context.Context, targetScope, runnerName
 		return "", fmt.Errorf("failed to generate runner register token: %w", err)
 	}
 
-	var labels []string
+	var sanitizedLabels []string
+
+	for _, l := range labels {
+		if !strings.EqualFold(l, config.RequiredMyShoesLabel) {
+			sanitizedLabels = append(sanitizedLabels, l)
+		}
+	}
 	if githubURL != "" && githubURL != "https://github.com" {
-		labels = append(labels, "dependabot")
+		sanitizedLabels = append(sanitizedLabels, "dependabot")
 	}
 
 	v := templateCreateLatestRunnerOnceValue{
@@ -90,7 +96,7 @@ func (s *Starter) getSetupRawScript(ctx context.Context, targetScope, runnerName
 		RunnerVersion:           runnerVersion,
 		RunnerServiceJS:         runnerServiceJs,
 		RunnerArg:               runnerTemporaryMode.StringFlag(),
-		AdditionalLabels:        labelsToOneLine(labels),
+		AdditionalLabels:        labelsToOneLine(sanitizedLabels),
 	}
 
 	t, err := template.New("templateCreateLatestRunnerOnce").Parse(templateCreateLatestRunnerOnce)
diff --git a/pkg/starter/starter.go b/pkg/starter/starter.go
index a778914..6a66e58 100644
--- a/pkg/starter/starter.go
+++ b/pkg/starter/starter.go
@@ -288,7 +288,11 @@ func (s *Starter) bung(ctx context.Context, job datastore.Job, target datastore.
 	runnerName := runner.ToName(job.UUID.String())
 
 	targetScope := getTargetScope(target, job)
-	script, err := s.getSetupScript(ctx, targetScope, runnerName)
+	labels, err := gh.ExtractRunsOnLabels([]byte(job.CheckEventJSON))
+	if err != nil {
+		return "", "", "", datastore.ResourceTypeUnknown, fmt.Errorf("failed to extract labels: %w", err)
+	}
+	script, err := s.getSetupScript(ctx, targetScope, runnerName, labels)
 	if err != nil {
 		return "", "", "", datastore.ResourceTypeUnknown, fmt.Errorf("failed to get setup scripts: %w", err)
 	}
@@ -299,11 +303,6 @@ func (s *Starter) bung(ctx context.Context, job datastore.Job, target datastore.
 	}
 	defer teardown()
 
-	labels, err := gh.ExtractRunsOnLabels([]byte(job.CheckEventJSON))
-	if err != nil {
-		return "", "", "", datastore.ResourceTypeUnknown, fmt.Errorf("failed to extract labels: %w", err)
-	}
-
 	cloudID, ipAddress, shoesType, resourceType, err := client.AddInstance(ctx, runnerName, script, target.ResourceType, labels)
 	if err != nil {
 		if stat, _ := status.FromError(err); stat.Code() == codes.InvalidArgument {
diff --git a/pkg/web/webhook.go b/pkg/web/webhook.go
index eb36e7d..8f19384 100644
--- a/pkg/web/webhook.go
+++ b/pkg/web/webhook.go
@@ -204,7 +204,7 @@ func receiveWorkflowJobWebhook(ctx context.Context, event *github.WorkflowJobEve
 
 func isRequestedMyshoesLabel(labels []string) bool {
 	for _, label := range labels {
-		if strings.EqualFold(label, "myshoes") || strings.EqualFold(label, "self-hosted") {
+		if strings.EqualFold(label, config.RequiredMyShoesLabel) || strings.EqualFold(label, "self-hosted") {
 			return true
 		}
 	}