feat: reuse agent tokens when a prebuilt agent reinitializes (#374)

SasSwart · Emyrk · web-flow · commit 5648efbf6db0 · 2025-05-02T08:57:05.000+02:00
* feat: allow presets to define prebuilds

* document prebuild parameters

* remove todo

* make gen

* feat: reuse agent tokens when a prebuilt agent reinitializes

* WIP: get agent.go ready to be merged with support for prebuilds

* fix: ensure the agent token is reused for prebuilds

* lint and make gen

* simplify function

* test: rbac role test assertion to handle site wide roles

Site wide roles have empty string org_ids

---------

Co-authored-by: Steven Masley &lt;stevenmasley@gmail.com&gt;
diff --git a/go.mod b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/docker/docker v26.1.5+incompatible
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/go-cty v1.4.1-0.20200414143053-d3edf31b6320
+	github.com/hashicorp/terraform-plugin-log v0.9.0
 	github.com/hashicorp/terraform-plugin-sdk/v2 v2.36.1
 	github.com/masterminds/semver v1.5.0
 	github.com/mitchellh/mapstructure v1.5.0
@@ -50,7 +51,6 @@ require (
 	github.com/hashicorp/terraform-exec v0.22.0 // indirect
 	github.com/hashicorp/terraform-json v0.24.0 // indirect
 	github.com/hashicorp/terraform-plugin-go v0.26.0 // indirect
-	github.com/hashicorp/terraform-plugin-log v0.9.0 // indirect
 	github.com/hashicorp/terraform-registry-address v0.2.4 // indirect
 	github.com/hashicorp/terraform-svchost v0.1.1 // indirect
 	github.com/hashicorp/yamux v0.1.1 // indirect
diff --git a/provider/agent.go b/provider/agent.go
@@ -2,13 +2,16 @@ package provider
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"fmt"
 	"path/filepath"
 	"reflect"
 	"strings"
 
 	"github.com/google/uuid"
 	"github.com/hashicorp/go-cty/cty"
+	"github.com/hashicorp/terraform-plugin-log/tflog"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
@@ -22,10 +25,12 @@ func agentResource() *schema.Resource {
 		SchemaVersion: 1,
 
 		Description: "Use this resource to associate an agent.",
-		CreateContext: func(_ context.Context, resourceData *schema.ResourceData, i interface{}) diag.Diagnostics {
-			// This should be a real authentication token!
-			resourceData.SetId(uuid.NewString())
-			err := resourceData.Set("token", uuid.NewString())
+		CreateContext: func(ctx context.Context, resourceData *schema.ResourceData, i interface{}) diag.Diagnostics {
+			agentID := uuid.NewString()
+			resourceData.SetId(agentID)
+
+			token := agentAuthToken(ctx, "")
+			err := resourceData.Set("token", token)
 			if err != nil {
 				return diag.FromErr(err)
 			}
@@ -48,10 +53,12 @@ func agentResource() *schema.Resource {
 			return updateInitScript(resourceData, i)
 		},
 		ReadWithoutTimeout: func(ctx context.Context, resourceData *schema.ResourceData, i interface{}) diag.Diagnostics {
-			err := resourceData.Set("token", uuid.NewString())
+			token := agentAuthToken(ctx, "")
+			err := resourceData.Set("token", token)
 			if err != nil {
 				return diag.FromErr(err)
 			}
+
 			if _, ok := resourceData.GetOk("display_apps"); !ok {
 				err = resourceData.Set("display_apps", []interface{}{
 					map[string]bool{
@@ -469,3 +476,37 @@ func updateInitScript(resourceData *schema.ResourceData, i interface{}) diag.Dia
 	}
 	return nil
 }
+
+func agentAuthToken(ctx context.Context, agentID string) string {
+	existingToken := helpers.OptionalEnv(RunningAgentTokenEnvironmentVariable(agentID))
+	if existingToken == "" {
+		// Most of the time, we will generate a new token for the agent.
+		// In the case of a prebuilt workspace being claimed, we will override with
+		// an existing token provided below.
+		token := uuid.NewString()
+		return token
+	}
+
+	// An existing token was provided for this agent. That means that this
+	// is a prebuilt workspace in the process of being claimed.
+	// We should reuse the token.
+	tflog.Info(ctx, "using provided agent token for prebuild", map[string]interface{}{
+		"agent_id": agentID,
+	})
+	return existingToken
+}
+
+// RunningAgentTokenEnvironmentVariable returns the name of an environment variable
+// that contains the token to use for the running agent. This is used for prebuilds,
+// where we want to reuse the same token for the next iteration of a workspace agent
+// before and after the workspace was claimed by a user.
+//
+// By reusing an existing token, we can avoid the need to change a value that may have been
+// used immutably. Thus, allowing us to avoid reprovisioning resources that may take a long time
+// to replace.
+//
+// agentID is unused for now, but will be used as soon as we support multiple agents.
+func RunningAgentTokenEnvironmentVariable(agentID string) string {
+	sum := sha256.Sum256([]byte(agentID))
+	return "CODER_RUNNING_WORKSPACE_AGENT_TOKEN_" + hex.EncodeToString(sum[:])
+}
diff --git a/provider/workspace.go b/provider/workspace.go
@@ -27,13 +27,13 @@ func workspaceDataSource() *schema.Resource {
 			}
 			_ = rd.Set("start_count", count)
 
-			prebuild := helpers.OptionalEnv(IsPrebuildEnvironmentVariable())
-			prebuildCount := 0
-			if prebuild == "true" {
-				prebuildCount = 1
+			if isPrebuiltWorkspace() {
+				_ = rd.Set("prebuild_count", 1)
 				_ = rd.Set("is_prebuild", true)
+			} else {
+				_ = rd.Set("prebuild_count", 0)
+				_ = rd.Set("is_prebuild", false)
 			}
-			_ = rd.Set("prebuild_count", prebuildCount)
 
 			name := helpers.OptionalEnvOrDefault("CODER_WORKSPACE_NAME", "default")
 			rd.Set("name", name)
@@ -140,6 +140,24 @@ func workspaceDataSource() *schema.Resource {
 	}
 }
 
+// isPrebuiltWorkspace returns true if the workspace is an unclaimed prebuilt workspace.
+func isPrebuiltWorkspace() bool {
+	return helpers.OptionalEnv(IsPrebuildEnvironmentVariable()) == "true"
+}
+
+// IsPrebuildEnvironmentVariable returns the name of the environment variable that
+// indicates whether the workspace is an unclaimed prebuilt workspace.
+//
+// Knowing whether the workspace is an unclaimed prebuilt workspace allows template
+// authors to conditionally execute code in the template based on whether the workspace
+// has been assigned to a user or not. This allows identity specific configuration to
+// be applied only after the workspace is claimed, while the rest of the workspace can
+// be pre-configured.
+//
+// The value of this environment variable should be set to "true" if the workspace is prebuilt
+// and it has not yet been claimed by a user. Any other values, including "false"
+// and "" will be interpreted to mean that the workspace is not prebuilt, or was
+// prebuilt but has since been claimed by a user.
 func IsPrebuildEnvironmentVariable() string {
 	return "CODER_WORKSPACE_IS_PREBUILD"
 }
