From 65d77383d0097f7d800612354c1b92c9e5eb9542 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Wed, 20 Apr 2022 12:28:48 -0500
Subject: [PATCH] fix: Allow nested Terraform resources (#1093)

This fixes the dependency tree by adding recursion. It
now finds indirect connections and associates it with
an agent.

An example is attached which surfaced this issue.
---
 .vscode/settings.json                   |   1 +
 examples/gcp-vm-container/README.md     |   5 +
 examples/gcp-vm-container/main.tf       | 103 ++++++++++++++++++
 go.mod                                  |   2 +-
 provisioner/terraform/provision.go      | 136 +++++++++++++-----------
 provisioner/terraform/provision_test.go |  87 +++++++++++++++
 provisionersdk/agent.go                 |  18 ++--
 7 files changed, 278 insertions(+), 74 deletions(-)
 create mode 100644 examples/gcp-vm-container/README.md
 create mode 100644 examples/gcp-vm-container/main.tf

diff --git a/.vscode/settings.json b/.vscode/settings.json
index 5319d54048..160a8cd98a 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -51,6 +51,7 @@
     "tcpip",
     "TCSETS",
     "tfexec",
+    "tfjson",
     "tfstate",
     "trimprefix",
     "unconvert",
diff --git a/examples/gcp-vm-container/README.md b/examples/gcp-vm-container/README.md
new file mode 100644
index 0000000000..75148e14d6
--- /dev/null
+++ b/examples/gcp-vm-container/README.md
@@ -0,0 +1,5 @@
+---
+name: Develop in a container on a Google Cloud VM
+description: Get started with Linux development on Google Cloud.
+tags: [cloud, google, container]
+---
diff --git a/examples/gcp-vm-container/main.tf b/examples/gcp-vm-container/main.tf
new file mode 100644
index 0000000000..03fad536ef
--- /dev/null
+++ b/examples/gcp-vm-container/main.tf
@@ -0,0 +1,103 @@
+terraform {
+  required_providers {
+    coder = {
+      source  = "coder/coder"
+      version = "~> 0.3.1"
+    }
+    google = {
+      source  = "hashicorp/google"
+      version = "~> 4.15"
+    }
+  }
+}
+
+variable "service_account" {
+  description = <<EOF
+Coder requires a Google Cloud Service Account to provision workspaces.
+
+1. Create a service account:
+   https://console.cloud.google.com/projectselector/iam-admin/serviceaccounts/create
+2. Add the roles:
+   - Compute Admin
+   - Service Account User
+3. Click on the created key, and navigate to the "Keys" tab.
+4. Click "Add key", then "Create new key".
+5. Generate a JSON private key, and paste the contents below.
+EOF
+  sensitive   = true
+}
+
+variable "zone" {
+  description = "What region should your workspace live in?"
+  default     = "us-central1-a"
+  validation {
+    condition     = contains(["northamerica-northeast1-a", "us-central1-a", "us-west2-c", "europe-west4-b", "southamerica-east1-a"], var.zone)
+    error_message = "Invalid zone!"
+  }
+}
+
+provider "google" {
+  zone        = var.zone
+  credentials = var.service_account
+  project     = jsondecode(var.service_account).project_id
+}
+
+data "google_compute_default_service_account" "default" {
+}
+
+data "coder_workspace" "me" {
+}
+
+resource "coder_agent" "dev" {
+  auth = "google-instance-identity"
+  arch = "amd64"
+  os   = "linux"
+}
+
+module "gce-container" {
+  source  = "terraform-google-modules/container-vm/google"
+  version = "3.0.0"
+
+  container = {
+    image   = "mcr.microsoft.com/vscode/devcontainers/go:1"
+    command = ["sh"]
+    args    = ["-c", coder_agent.dev.init_script]
+    securityContext = {
+      privileged : true
+    }
+  }
+}
+
+resource "google_compute_instance" "dev" {
+  zone         = var.zone
+  count        = data.coder_workspace.me.start_count
+  name         = "coder-${data.coder_workspace.me.owner}-${data.coder_workspace.me.name}"
+  machine_type = "e2-medium"
+  network_interface {
+    network = "default"
+    access_config {
+      // Ephemeral public IP
+    }
+  }
+  boot_disk {
+    initialize_params {
+      image = module.gce-container.source_image
+    }
+  }
+  service_account {
+    email  = data.google_compute_default_service_account.default.email
+    scopes = ["cloud-platform"]
+  }
+  metadata = {
+    "gce-container-declaration" = module.gce-container.metadata_value
+  }
+  labels = {
+    container-vm = module.gce-container.vm_container_label
+  }
+}
+
+resource "coder_agent_instance" "dev" {
+  count       = data.coder_workspace.me.start_count
+  agent_id    = coder_agent.dev.id
+  instance_id = google_compute_instance.dev[0].instance_id
+}
diff --git a/go.mod b/go.mod
index f832074c61..2e98572e65 100644
--- a/go.mod
+++ b/go.mod
@@ -67,6 +67,7 @@ require (
 	github.com/hashicorp/hcl/v2 v2.11.1
 	github.com/hashicorp/terraform-config-inspect v0.0.0-20211115214459-90acf1ca460f
 	github.com/hashicorp/terraform-exec v0.15.0
+	github.com/hashicorp/terraform-json v0.13.0
 	github.com/hashicorp/yamux v0.0.0-20211028200310-0bc27b27de87
 	github.com/jedib0t/go-pretty/v6 v6.3.1
 	github.com/justinas/nosurf v1.1.1
@@ -162,7 +163,6 @@ require (
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
-	github.com/hashicorp/terraform-json v0.13.0 // indirect
 	github.com/imdario/mergo v0.3.12 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
diff --git a/provisioner/terraform/provision.go b/provisioner/terraform/provision.go
index b9f6fdaa7e..4ca20ba830 100644
--- a/provisioner/terraform/provision.go
+++ b/provisioner/terraform/provision.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/awalterschulze/gographviz"
 	"github.com/hashicorp/terraform-exec/tfexec"
+	tfjson "github.com/hashicorp/terraform-json"
 	"github.com/mitchellh/mapstructure"
 	"golang.org/x/xerrors"
 
@@ -88,15 +89,24 @@ func (t *terraform) Provision(stream proto.DRPCProvisioner_ProvisionStream) erro
 			})
 		}
 	}()
+	terraformEnv := map[string]string{}
+	// Required for "terraform init" to find "git" to
+	// clone Terraform modules.
+	for _, env := range os.Environ() {
+		parts := strings.SplitN(env, "=", 2)
+		if len(parts) < 2 {
+			continue
+		}
+		terraformEnv[parts[0]] = parts[1]
+	}
 	// Only Linux reliably works with the Terraform plugin
 	// cache directory. It's unknown why this is.
 	if t.cachePath != "" && runtime.GOOS == "linux" {
-		err = terraform.SetEnv(map[string]string{
-			"TF_PLUGIN_CACHE_DIR": t.cachePath,
-		})
-		if err != nil {
-			return xerrors.Errorf("set terraform plugin cache dir: %w", err)
-		}
+		terraformEnv["TF_PLUGIN_CACHE_DIR"] = t.cachePath
+	}
+	err = terraform.SetEnv(terraformEnv)
+	if err != nil {
+		return xerrors.Errorf("set terraform env: %w", err)
 	}
 	terraform.SetStdout(writer)
 	t.logger.Debug(shutdown, "running initialization")
@@ -320,40 +330,22 @@ func parseTerraformPlan(ctx context.Context, terraform *tfexec.Terraform, planfi
 				agent.StartupScript = startupScript
 			}
 		}
-		if _, has := resource.Expressions["instance_id"]; has {
-			// This is a dynamic value. If it's expressed, we know
-			// it's at least an instance ID, which is better than nothing.
-			agent.Auth = &proto.Agent_InstanceId{
-				InstanceId: "",
-			}
-		}
 
 		agents[resource.Address] = agent
 	}
+
 	for _, resource := range plan.PlannedValues.RootModule.Resources {
-		if resource.Type == "coder_agent" {
+		if resource.Mode == tfjson.DataResourceMode {
+			continue
+		}
+		if resource.Type == "coder_agent" || resource.Type == "coder_agent_instance" {
 			continue
 		}
 		resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
-		resourceNode, exists := resourceDependencies[resourceKey]
-		if !exists {
-			continue
-		}
-		// Associate resources that depend on an agent.
-		resourceAgents := make([]*proto.Agent, 0)
-		for _, dep := range resourceNode {
-			var has bool
-			agent, has := agents[dep]
-			if !has {
-				continue
-			}
-			resourceAgents = append(resourceAgents, agent)
-		}
-
 		resources = append(resources, &proto.Resource{
 			Name:   resource.Name,
 			Type:   resource.Type,
-			Agents: resourceAgents,
+			Agents: findAgents(resourceDependencies, agents, resourceKey),
 		})
 	}
 
@@ -460,32 +452,25 @@ func parseTerraformApply(ctx context.Context, terraform *tfexec.Terraform, state
 		}
 
 		for _, resource := range state.Values.RootModule.Resources {
+			if resource.Mode == tfjson.DataResourceMode {
+				continue
+			}
 			if resource.Type == "coder_agent" || resource.Type == "coder_agent_instance" {
 				continue
 			}
 			resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
-			resourceNode, exists := resourceDependencies[resourceKey]
-			if !exists {
-				continue
-			}
-			// Associate resources that depend on an agent.
-			resourceAgents := make([]*proto.Agent, 0)
-			for _, dep := range resourceNode {
-				var has bool
-				agent, has := agents[dep]
-				if !has {
-					continue
-				}
-				resourceAgents = append(resourceAgents, agent)
-
+			resourceAgents := findAgents(resourceDependencies, agents, resourceKey)
+			for _, agent := range resourceAgents {
 				// Didn't use instance identity.
 				if agent.GetToken() != "" {
 					continue
 				}
 
 				key, isValid := map[string]string{
-					"google_compute_instance": "instance_id",
-					"aws_instance":            "id",
+					"google_compute_instance":         "instance_id",
+					"aws_instance":                    "id",
+					"azurerm_linux_virtual_machine":   "id",
+					"azurerm_windows_virtual_machine": "id",
 				}[resource.Type]
 				if !isValid {
 					// The resource type doesn't support
@@ -571,21 +556,50 @@ func findDirectDependencies(rawGraph string) (map[string][]string, error) {
 			continue
 		}
 		label = strings.Trim(label, `"`)
-
-		dependencies := make([]string, 0)
-		for destination := range graph.Edges.SrcToDsts[node.Name] {
-			dependencyNode, exists := graph.Nodes.Lookup[destination]
-			if !exists {
-				continue
-			}
-			label, exists := dependencyNode.Attrs["label"]
-			if !exists {
-				continue
-			}
-			label = strings.Trim(label, `"`)
-			dependencies = append(dependencies, label)
-		}
-		direct[label] = dependencies
+		direct[label] = findDependenciesWithLabels(graph, node.Name)
 	}
+
 	return direct, nil
 }
+
+// findDependenciesWithLabels recursively finds nodes with labels (resource and data nodes)
+// to build a dependency tree.
+func findDependenciesWithLabels(graph *gographviz.Graph, nodeName string) []string {
+	dependencies := make([]string, 0)
+	for destination := range graph.Edges.SrcToDsts[nodeName] {
+		dependencyNode, exists := graph.Nodes.Lookup[destination]
+		if !exists {
+			continue
+		}
+		label, exists := dependencyNode.Attrs["label"]
+		if !exists {
+			dependencies = append(dependencies, findDependenciesWithLabels(graph, dependencyNode.Name)...)
+			continue
+		}
+		label = strings.Trim(label, `"`)
+		dependencies = append(dependencies, label)
+	}
+	return dependencies
+}
+
+// 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]
+	if !exists {
+		return []*proto.Agent{}
+	}
+	// Associate resources that depend on an agent.
+	resourceAgents := make([]*proto.Agent, 0)
+	for _, dep := range resourceNode {
+		var has bool
+		agent, has := agents[dep]
+		if !has {
+			resourceAgents = append(resourceAgents, findAgents(resourceDependencies, agents, dep)...)
+			continue
+		}
+		resourceAgents = append(resourceAgents, agent)
+	}
+	return resourceAgents
+}
diff --git a/provisioner/terraform/provision_test.go b/provisioner/terraform/provision_test.go
index 645c87ba03..317db3bb49 100644
--- a/provisioner/terraform/provision_test.go
+++ b/provisioner/terraform/provision_test.go
@@ -337,6 +337,93 @@ provider "coder" {
 				},
 			},
 		},
+	}, {
+		Name: "dryrun-resource-separated-from-agent",
+		Files: map[string]string{
+			"main.tf": provider + `
+			resource "coder_agent" "A" {
+				os = "darwin"
+				arch = "amd64"
+			}
+			data "null_data_source" "values" {
+				inputs = {
+					script = coder_agent.A.init_script
+				}
+			}
+			resource "null_resource" "A" {
+				depends_on = [
+					data.null_data_source.values
+				]
+			}
+			`,
+		},
+		Request: &proto.Provision_Request{
+			Type: &proto.Provision_Request_Start{
+				Start: &proto.Provision_Start{
+					Metadata: &proto.Provision_Metadata{},
+					DryRun:   true,
+				},
+			},
+		},
+		Response: &proto.Provision_Response{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					Resources: []*proto.Resource{{
+						Name: "A",
+						Type: "null_resource",
+						Agents: []*proto.Agent{{
+							Name:            "A",
+							OperatingSystem: "darwin",
+							Architecture:    "amd64",
+							Auth:            &proto.Agent_Token{},
+						}},
+					}},
+				},
+			},
+		},
+	}, {
+		Name: "resource-separated-from-agent",
+		Files: map[string]string{
+			"main.tf": provider + `
+			resource "coder_agent" "A" {
+				os = "darwin"
+				arch = "amd64"
+			}
+			data "null_data_source" "values" {
+				inputs = {
+					script = coder_agent.A.init_script
+				}
+			}
+			resource "null_resource" "A" {
+				depends_on = [
+					data.null_data_source.values
+				]
+			}
+			`,
+		},
+		Request: &proto.Provision_Request{
+			Type: &proto.Provision_Request_Start{
+				Start: &proto.Provision_Start{
+					Metadata: &proto.Provision_Metadata{},
+				},
+			},
+		},
+		Response: &proto.Provision_Response{
+			Type: &proto.Provision_Response_Complete{
+				Complete: &proto.Provision_Complete{
+					Resources: []*proto.Resource{{
+						Name: "A",
+						Type: "null_resource",
+						Agents: []*proto.Agent{{
+							Name:            "A",
+							OperatingSystem: "darwin",
+							Architecture:    "amd64",
+							Auth:            &proto.Agent_Token{},
+						}},
+					}},
+				},
+			},
+		},
 	}} {
 		testCase := testCase
 		t.Run(testCase.Name, func(t *testing.T) {
diff --git a/provisionersdk/agent.go b/provisionersdk/agent.go
index 5ffaee48b2..9d5389e87c 100644
--- a/provisionersdk/agent.go
+++ b/provisionersdk/agent.go
@@ -11,38 +11,32 @@ var (
 		// executing shell to be named "sshd", otherwise it fails. See:
 		// https://github.com/microsoft/vscode-remote-release/issues/5699
 		"windows": {
-			"amd64": `
-$ProgressPreference = "SilentlyContinue"
+			"amd64": `$ProgressPreference = "SilentlyContinue"
 Invoke-WebRequest -Uri ${ACCESS_URL}bin/coder-windows-amd64.exe -OutFile $env:TEMP\sshd.exe
 Set-MpPreference -DisableRealtimeMonitoring $true -ExclusionPath $env:TEMP\sshd.exe
 $env:CODER_AUTH = "${AUTH_TYPE}"
 $env:CODER_URL = "${ACCESS_URL}"
-Start-Process -FilePath $env:TEMP\sshd.exe -ArgumentList "agent" -PassThru
-`,
+Start-Process -FilePath $env:TEMP\sshd.exe -ArgumentList "agent" -PassThru`,
 		},
 		"linux": {
-			"amd64": `
-#!/usr/bin/env sh
+			"amd64": `#!/usr/bin/env sh
 set -eu pipefail
 export BINARY_LOCATION=$(mktemp -d -t tmp.coderXXXXX)/coder
 curl -fsSL ${ACCESS_URL}bin/coder-linux-amd64 -o $BINARY_LOCATION
 chmod +x $BINARY_LOCATION
 export CODER_AUTH="${AUTH_TYPE}"
 export CODER_URL="${ACCESS_URL}"
-exec $BINARY_LOCATION agent
-`,
+exec $BINARY_LOCATION agent`,
 		},
 		"darwin": {
-			"amd64": `
-#!/usr/bin/env sh
+			"amd64": `#!/usr/bin/env sh
 set -eu pipefail
 export BINARY_LOCATION=$(mktemp -d -t tmp.coderXXXXX)/coder
 curl -fsSL ${ACCESS_URL}bin/coder-darwin-amd64 -o $BINARY_LOCATION
 chmod +x $BINARY_LOCATION
 export CODER_AUTH="${AUTH_TYPE}"
 export CODER_URL="${ACCESS_URL}"
-exec $BINARY_LOCATION agent
-`,
+exec $BINARY_LOCATION agent`,
 		},
 	}
 )