From 99631490b377b2e2b6f83184f18c44b673bee17e Mon Sep 17 00:00:00 2001
From: kylecarbs <kyle@carberry.com>
Date: Tue, 24 May 2022 20:48:33 +0000
Subject: [PATCH] fix: Use Terraform address to index resource + agent
 association

Closes #1705.

There was an issue in the implementation brought by #1577 by not trimming
the array value when resources use counts. This should fix it, and adds
a test to be sure!
---
 provisioner/terraform/provision.go      | 21 ++++++++++++---------
 provisioner/terraform/provision_test.go |  6 +++++-
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/provisioner/terraform/provision.go b/provisioner/terraform/provision.go
index 8c1c96d55e57a..acffae785b759 100644
--- a/provisioner/terraform/provision.go
+++ b/provisioner/terraform/provision.go
@@ -390,7 +390,7 @@ func parseTerraformPlan(ctx context.Context, terraform *tfexec.Terraform, planfi
 			}
 		}
 
-		agents[resource.Address] = agent
+		agents[convertAddressToLabel(resource.Address)] = agent
 	}
 
 	for _, resource := range tfResources {
@@ -400,11 +400,10 @@ func parseTerraformPlan(ctx context.Context, terraform *tfexec.Terraform, planfi
 		if resource.Type == "coder_agent" || resource.Type == "coder_agent_instance" {
 			continue
 		}
-		resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
 		resources = append(resources, &proto.Resource{
 			Name:   resource.Name,
 			Type:   resource.Type,
-			Agents: findAgents(resourceDependencies, agents, resourceKey),
+			Agents: findAgents(resourceDependencies, agents, convertAddressToLabel(resource.Address)),
 		})
 	}
 
@@ -485,8 +484,7 @@ func parseTerraformApply(ctx context.Context, terraform *tfexec.Terraform, state
 			default:
 				agent.Auth = &proto.Agent_InstanceId{}
 			}
-			resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
-			agents[resourceKey] = agent
+			agents[convertAddressToLabel(resource.Address)] = agent
 		}
 
 		// Manually associate agents with instance IDs.
@@ -529,8 +527,7 @@ func parseTerraformApply(ctx context.Context, terraform *tfexec.Terraform, state
 			if resource.Type == "coder_agent" || resource.Type == "coder_agent_instance" {
 				continue
 			}
-			resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
-			resourceAgents := findAgents(resourceDependencies, agents, resourceKey)
+			resourceAgents := findAgents(resourceDependencies, agents, convertAddressToLabel(resource.Address))
 			for _, agent := range resourceAgents {
 				// Didn't use instance identity.
 				if agent.GetToken() != "" {
@@ -696,8 +693,8 @@ func findDependenciesWithLabels(graph *gographviz.Graph, nodeName string) []stri
 // findAgents recursively searches through resource dependencies
 // to find associated agents. Nested is required for indirect
 // dependency matching.
-func findAgents(resourceDependencies map[string][]string, agents map[string]*proto.Agent, resourceKey string) []*proto.Agent {
-	resourceNode, exists := resourceDependencies[resourceKey]
+func findAgents(resourceDependencies map[string][]string, agents map[string]*proto.Agent, resourceLabel string) []*proto.Agent {
+	resourceNode, exists := resourceDependencies[resourceLabel]
 	if !exists {
 		return []*proto.Agent{}
 	}
@@ -714,3 +711,9 @@ func findAgents(resourceDependencies map[string][]string, agents map[string]*pro
 	}
 	return resourceAgents
 }
+
+// convertAddressToLabel returns the Terraform address without the count
+// specifier. eg. "module.ec2_dev.ec2_instance.dev[0]" becomes "module.ec2_dev.ec2_instance.dev"
+func convertAddressToLabel(address string) string {
+	return strings.Split(address, "[")[0]
+}
diff --git a/provisioner/terraform/provision_test.go b/provisioner/terraform/provision_test.go
index d02a08fbcf2eb..05d781dc8d8ff 100644
--- a/provisioner/terraform/provision_test.go
+++ b/provisioner/terraform/provision_test.go
@@ -206,6 +206,7 @@ provider "coder" {
 		Name: "dryrun-resource-associated-with-agent",
 		Files: map[string]string{
 			"main.tf": provider + `
+			data "coder_workspace" "me" {}
 			resource "coder_agent" "A" {
 				count = 1
 				os = "linux"
@@ -216,7 +217,10 @@ provider "coder" {
 				startup_script = "code-server"
 			}
 			resource "null_resource" "A" {
-				count = length(coder_agent.A)
+				depends_on = [
+					coder_agent.A[0]
+				]
+				count = data.coder_workspace.me.start_count
 			}`,
 		},
 		Request: &proto.Provision_Request{