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

* 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 <[email protected]>

Author: SasSwart

go.mod

@@ -6,6 +6,7 @@ require (
 .com/docker/docker v26.1.5+incompatible
 .com/google/uuid v1.6.0
 .com/hashicorp/go-cty v1.4.1-0.20200414143053-d3edf31b6320
+.com/hashicorp/terraform-plugin-log v0.9.0
 .com/hashicorp/terraform-plugin-sdk/v2 v2.36.1
 .com/masterminds/semver v1.5.0
 .com/mitchellh/mapstructure v1.5.0
@@ -50,7 +51,6 @@ require (
 .com/hashicorp/terraform-exec v0.22.0 // indirect
 .com/hashicorp/terraform-json v0.24.0 // indirect
 .com/hashicorp/terraform-plugin-go v0.26.0 // indirect
-.com/hashicorp/terraform-plugin-log v0.9.0 // indirect
 .com/hashicorp/terraform-registry-address v0.2.4 // indirect
 .com/hashicorp/terraform-svchost v0.1.1 // indirect
 .com/hashicorp/yamux v0.1.1 // indirect

provider/agent.go

@@ -2,13 +2,16 @@ package provider
 
 import (
 "context"
+"crypto/sha256"
+"encoding/hex"
 "fmt"
 "path/filepath"
 "reflect"
 "strings"
 
 ".com/google/uuid"
 ".com/hashicorp/go-cty/cty"
+".com/hashicorp/terraform-plugin-log/tflog"
 ".com/hashicorp/terraform-plugin-sdk/v2/diag"
 ".com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 ".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[:])
+}

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"
 }