From ffeded20a4d87c949de9f9624cdc192015ca4355 Mon Sep 17 00:00:00 2001
From: Christian Mesh <christianmesh1@gmail.com>
Date: Thu, 29 Aug 2024 10:32:01 -0400
Subject: [PATCH] Better handling of key_provider references (#1921)

Signed-off-by: Christian Mesh <christianmesh1@gmail.com>
---
 internal/encryption/keyprovider.go  | 20 +++++++++++++----
 internal/encryption/targets_test.go | 34 +++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/internal/encryption/keyprovider.go b/internal/encryption/keyprovider.go
index 9c2b490700..c0977bea94 100644
--- a/internal/encryption/keyprovider.go
+++ b/internal/encryption/keyprovider.go
@@ -127,15 +127,27 @@ func (e *targetBuilder) setupKeyProvider(cfg config.KeyProviderConfig, stack []c
 			continue
 		}
 
-		// TODO this should be more defensive
+		// This will always be a TraverseRoot, panic is OK if that's not the case
 		depRoot := (dep[0].(hcl.TraverseRoot)).Name
-		depType := (dep[1].(hcl.TraverseAttr)).Name
-		depName := (dep[2].(hcl.TraverseAttr)).Name
-
 		if depRoot != "key_provider" {
 			nonKeyProviderDeps = append(nonKeyProviderDeps, dep)
 			continue
 		}
+		depTypeAttr, typeOk := dep[1].(hcl.TraverseAttr)
+		depNameAttr, nameOk := dep[2].(hcl.TraverseAttr)
+
+		if !typeOk || !nameOk {
+			diags = append(diags, &hcl.Diagnostic{
+				Severity: hcl.DiagError,
+				Summary:  "Invalid Key Provider expression format",
+				Detail:   "Expected key_provider.<type>.<name>",
+				Subject:  dep.SourceRange().Ptr(),
+			})
+			continue
+		}
+
+		depType := depTypeAttr.Name
+		depName := depNameAttr.Name
 
 		kpc, ok := e.cfg.GetKeyProvider(depType, depName)
 		if !ok {
diff --git a/internal/encryption/targets_test.go b/internal/encryption/targets_test.go
index f48050cc16..ec87e18f83 100644
--- a/internal/encryption/targets_test.go
+++ b/internal/encryption/targets_test.go
@@ -131,6 +131,22 @@ func TestBaseEncryption_buildTargetMethods(t *testing.T) {
 				aesgcm.Is,
 			},
 		},
+		"key-from-complex-vars": {
+			rawConfig: `
+				key_provider "static" "basic" {
+					key = var.obj[0].key
+				}
+				method "aes_gcm" "example" {
+					keys = key_provider.static.basic
+				}
+				state {
+					method = method.aes_gcm.example
+				}
+			`,
+			wantMethods: []func(method.Method) bool{
+				aesgcm.Is,
+			},
+		},
 		"undefined-key-from-vars": {
 			rawConfig: `
 				key_provider "static" "basic" {
@@ -145,6 +161,20 @@ func TestBaseEncryption_buildTargetMethods(t *testing.T) {
 			`,
 			wantErr: "Test Config Source:3,12-28: Undefined variable; Undefined variable var.undefinedkey",
 		},
+		"bad-keyprovider-format": {
+			rawConfig: `
+				key_provider "static" "basic" {
+					key = key_provider.static[0]
+				}
+				method "aes_gcm" "example" {
+					keys = key_provider.static.basic
+				}
+				state {
+					method = method.aes_gcm.example
+				}
+			`,
+			wantErr: "Test Config Source:3,12-34: Invalid Key Provider expression format; Expected key_provider.<type>.<name>",
+		},
 	}
 
 	reg := lockingencryptionregistry.New()
@@ -165,6 +195,10 @@ func TestBaseEncryption_buildTargetMethods(t *testing.T) {
 				Default: cty.StringVal("6f6f706830656f67686f6834616872756f3751756165686565796f6f72653169"),
 				Type:    cty.String,
 			},
+			"obj": {
+				Name:    "obj",
+				Default: cty.ListVal([]cty.Value{cty.ObjectVal(map[string]cty.Value{"key": cty.StringVal("6f6f706830656f67686f6834616872756f3751756165686565796f6f72653169")})}),
+			},
 		},
 	}