address comments

Azure · Oct 4, 2024 · e96bf25 · e96bf25
1 parent 19638c2
commit e96bf25
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/pkg/config/draftconfig_template_test.go b/pkg/config/draftconfig_template_test.go
@@ -137,7 +137,7 @@ func loadTemplatesWithValidation() error {
 				return fmt.Errorf("template %s has a variable with cyclical reference to itself: %s", path, currVar.Name)
 			}
 
-			if isCyclicalVariableReference(currVar, refVar, allVariables, true) {
+			if isCyclicalVariableReference(currVar, refVar, allVariables, true, map[string]bool{}) {
 				return fmt.Errorf("template %s has a variable with cyclical reference to itself: %s", path, currVar.Name)
 			}
 		}
@@ -147,11 +147,15 @@ func loadTemplatesWithValidation() error {
 	})
 }
 
-func isCyclicalVariableReference(initialVar, currRefVar *BuilderVar, allVariables map[string]*BuilderVar, isFirst bool) bool {
+func isCyclicalVariableReference(initialVar, currRefVar *BuilderVar, allVariables map[string]*BuilderVar, isFirst bool, visited map[string]bool) bool {
 	if !isFirst && initialVar.Name == currRefVar.Name {
 		return true
 	}
 
+	if _, ok := visited[currRefVar.Name]; ok {
+		return true
+	}
+
 	if currRefVar.Default.ReferenceVar == "" {
 		return false
 	}
@@ -161,5 +165,6 @@ func isCyclicalVariableReference(initialVar, currRefVar *BuilderVar, allVariable
 		return false
 	}
 
-	return isCyclicalVariableReference(initialVar, refVar, allVariables, false)
+	visited[currRefVar.Name] = true
+	return isCyclicalVariableReference(initialVar, refVar, allVariables, false, visited)
 }
diff --git a/pkg/handlers/template.go b/pkg/handlers/template.go
@@ -4,13 +4,13 @@ import (
 	"bytes"
 	"fmt"
 	"io/fs"
-	"log"
 	"path/filepath"
 	"strings"
 	tmpl "text/template"
 
 	"github.com/Azure/draft/pkg/config"
 	"github.com/Azure/draft/pkg/templatewriter"
+	log "github.com/sirupsen/logrus"
 )
 
 type Template struct {

diff --git a/pkg/handlers/template_utils.go b/pkg/handlers/template_utils.go
@@ -3,13 +3,13 @@ package handlers
 import (
 	"fmt"
 	"io/fs"
-	"log"
 	"path/filepath"
 	"strings"
 
 	"github.com/Azure/draft/pkg/config"
 	"github.com/Azure/draft/template"
 	"github.com/blang/semver/v4"
+	log "github.com/sirupsen/logrus"
 )
 
 var templateConfigs map[string]*Template