From 38e5d9c6997a8ea56c69eb698c0cae4b42905999 Mon Sep 17 00:00:00 2001
From: Pam Selle <204372+pselle@users.noreply.github.com>
Date: Tue, 5 May 2020 16:44:22 -0400
Subject: [PATCH] Add more validation to expanding modules

---
 terraform/context_validate_test.go | 66 +++++++++++++++++++++++++++++-
 terraform/node_module_expand.go    | 21 +++++++++-
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/terraform/context_validate_test.go b/terraform/context_validate_test.go
index 495cad6ca..f811dfcad 100644
--- a/terraform/context_validate_test.go
+++ b/terraform/context_validate_test.go
@@ -1425,7 +1425,7 @@ module "mod2" {
 }
 
 module "mod3" {
-  count = len(module.mod2)
+  count = length(module.mod2)
   source = "./mod"
 }
 `,
@@ -1472,3 +1472,67 @@ resource "aws_instance" "foo" {
 		t.Fatal(diags.ErrWithWarnings())
 	}
 }
+
+func TestContext2Validate_expandModulesInvalidCount(t *testing.T) {
+	m := testModuleInline(t, map[string]string{
+		"main.tf": `
+module "mod1" {
+  count = -1
+  source = "./mod"
+}
+`,
+		"mod/main.tf": `
+resource "aws_instance" "foo" {
+}
+`,
+	})
+
+	p := testProvider("aws")
+	p.DiffFn = testDiffFn
+	ctx := testContext2(t, &ContextOpts{
+		Config: m,
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("aws"): testProviderFuncFixed(p),
+		},
+	})
+
+	diags := ctx.Validate()
+	if !diags.HasErrors() {
+		t.Fatal("succeeded; want errors")
+	}
+	if got, want := diags.Err().Error(), `Invalid count argument`; strings.Index(got, want) == -1 {
+		t.Fatalf("wrong error:\ngot:  %s\nwant: message containing %q", got, want)
+	}
+}
+
+func TestContext2Validate_expandModulesInvalidForEach(t *testing.T) {
+	m := testModuleInline(t, map[string]string{
+		"main.tf": `
+module "mod1" {
+  for_each = ["a", "b"]
+  source = "./mod"
+}
+`,
+		"mod/main.tf": `
+resource "aws_instance" "foo" {
+}
+`,
+	})
+
+	p := testProvider("aws")
+	p.DiffFn = testDiffFn
+	ctx := testContext2(t, &ContextOpts{
+		Config: m,
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("aws"): testProviderFuncFixed(p),
+		},
+	})
+
+	diags := ctx.Validate()
+	if !diags.HasErrors() {
+		t.Fatal("succeeded; want errors")
+	}
+	if got, want := diags.Err().Error(), `Invalid for_each argument`; strings.Index(got, want) == -1 {
+		t.Fatalf("wrong error:\ngot:  %s\nwant: message containing %q", got, want)
+	}
+}
diff --git a/terraform/node_module_expand.go b/terraform/node_module_expand.go
index 057ac2614..91fd6d23a 100644
--- a/terraform/node_module_expand.go
+++ b/terraform/node_module_expand.go
@@ -281,8 +281,27 @@ func (n *evalValidateModule) Eval(ctx EvalContext) (interface{}, error) {
 	// will be a single instance, but still get our address in the expected
 	// manner anyway to ensure they've been registered correctly.
 	for _, module := range expander.ExpandModule(n.Addr.Parent()) {
+		ctx = ctx.WithPath(module)
+
+		// Validate our for_each and count expressions at a basic level
+		// We skip validation on known, because there will be unknown values before
+		// a full expansion, presuming these errors will be caught in later steps
+		switch {
+		case n.ModuleCall.Count != nil:
+			_, diags := evaluateCountExpressionValue(n.ModuleCall.Count, ctx)
+			if diags.HasErrors() {
+				return nil, diags.Err()
+			}
+
+		case n.ModuleCall.ForEach != nil:
+			_, diags := evaluateForEachExpressionValue(n.ModuleCall.ForEach, ctx)
+			if diags.HasErrors() {
+				return nil, diags.Err()
+			}
+		}
+
 		// now set our own mode to single
-		ctx.InstanceExpander().SetModuleSingle(module, call)
+		expander.SetModuleSingle(module, call)
 	}
 	return nil, nil
 }