Handle correctly parse existing tags with comments (#406)

* Handle correctly parse existing tags with comments * Update src/terraform/structure/terraform_parser.go Co-authored-by: Anton Grübel <anton.gruebel@gmail.com> --------- Co-authored-by: Anton Grübel <anton.gruebel@gmail.com>
bridgecrewio · Jul 12, 2023 · 283d738 · 283d738
1 parent ddf8dd7
commit 283d738
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 2 deletions.
diff --git a/src/terraform/structure/terraform_parser.go b/src/terraform/structure/terraform_parser.go
@@ -725,7 +725,7 @@ func (p *TerraformParser) getHclMapsContents(tokens hclwrite.Tokens) []hclwrite.
 func (p *TerraformParser) extractTagPairs(tokens hclwrite.Tokens) []hclwrite.Tokens {
 	// The function gets tokens and returns an array of tokens that represent key and value
 	// example: tokens: "a=1\n b=2, c=3", returns: ["a=1", "b=2", "c=3"]
-	separatorTokens := []hclsyntax.TokenType{hclsyntax.TokenComma, hclsyntax.TokenNewline}
+	separatorTokens := []hclsyntax.TokenType{hclsyntax.TokenComma, hclsyntax.TokenNewline, hclsyntax.TokenComment}
 
 	bracketsCounters := map[hclsyntax.TokenType]int{
 		hclsyntax.TokenOParen: 0,
@@ -746,7 +746,11 @@ func (p *TerraformParser) extractTagPairs(tokens hclwrite.Tokens) []hclwrite.Tok
 	for i, token := range tokens {
 		if utils.InSlice(separatorTokens, token.Type) && getUncloseBracketsCount(bracketsCounters) == 0 {
 			if hasEq {
-				tagPairs = append(tagPairs, tokens[startIndex:i])
+				endIndex := i
+				if token.Type == hclsyntax.TokenComment {
+					endIndex = i + 1
+				}
+				tagPairs = append(tagPairs, tokens[startIndex:endIndex])
 			}
 			startIndex = i + 1
 			hasEq = false

diff --git a/src/terraform/structure/terraform_parser_test.go b/src/terraform/structure/terraform_parser_test.go
@@ -377,6 +377,32 @@ func TestTerraformParser_Module(t *testing.T) {
 		assert.Equal(t, string(resultStr), string(expectedStr))
 	})
 
+	t.Run("Test parsing of existing tags with comments", func(t *testing.T) {
+		p := &TerraformParser{}
+		p.Init("../../../tests/terraform/data", nil)
+		defer p.Close()
+		sourceFilePath := "../../../tests/terraform/tags_with_comments/main.tf"
+		expectedFileName := "../../../tests/terraform/tags_with_comments/expected.txt"
+		blocks, err := p.ParseFile(sourceFilePath)
+		if err != nil {
+			t.Fail()
+		}
+
+		mb := blocks[0]
+		mb.AddNewTags([]tags.ITag{
+			&tags.Tag{Key: "mock_tag_2", Value: "mock_value"},
+		})
+
+		resultFileName := "result.txt"
+		defer func() {
+			_ = os.Remove(resultFileName)
+		}()
+		_ = p.WriteFile(sourceFilePath, blocks, resultFileName)
+		resultStr, _ := os.ReadFile(resultFileName)
+		expectedStr, _ := os.ReadFile(expectedFileName)
+		assert.Equal(t, string(resultStr), string(expectedStr))
+	})
+
 	t.Run("Test parsing of unsupported resources", func(t *testing.T) {
 		p := &TerraformParser{}
 		p.Init("../../../tests/terraform/supported", nil)

diff --git a/tests/terraform/tags_with_comments/expected.txt b/tests/terraform/tags_with_comments/expected.txt
@@ -0,0 +1,14 @@
+data "aws_s3_bucket" "primary" {
+  count  = var.create_bucket == true ? 0 : 1
+  bucket = "externally-created-bucket"
+}
+
+resource "aws_s3_bucket" "primary" {
+  count  = var.create_bucket == true ? 1 : 0
+  bucket = "yor-bug-test-bucket"
+  tags = {
+    mock_tag   = "mock_value" # This is for test
+    yor_trace  = "some-uuid"
+    mock_tag_2 = "mock_value"
+  }
+}
diff --git a/tests/terraform/tags_with_comments/main.tf b/tests/terraform/tags_with_comments/main.tf
@@ -0,0 +1,13 @@
+data "aws_s3_bucket" "primary" {
+  count  = var.create_bucket == true ? 0 : 1
+  bucket = "externally-created-bucket"
+}
+
+resource "aws_s3_bucket" "primary" {
+  count  = var.create_bucket == true ? 1 : 0
+  bucket = "yor-bug-test-bucket"
+  tags = {
+      mock_tag  = "mock_value" # This is for test
+      yor_trace = "some-uuid"
+    }
+}