ast: Fix deep copy on Module to include comments

Previously comments were not being deep copied which could lead to data races from the transfomer and other places. Fixes open-policy-agent#3793 Signed-off-by: Torin Sandall <[email protected]>
tsandall · Sep 16, 2021 · 467fad3 · 467fad3
1 parent 4747e22
commit 467fad3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/ast/policy.go b/ast/policy.go
@@ -411,6 +411,11 @@ func (mod *Module) Copy() *Module {
 		cpy.Annotations[i] = mod.Annotations[i].Copy(nodes[mod.Annotations[i].node])
 	}
 
+	cpy.Comments = make([]*Comment, len(mod.Comments))
+	for i := range mod.Comments {
+		cpy.Comments[i] = mod.Comments[i].Copy()
+	}
+
 	return &cpy
 }
 

diff --git a/ast/policy_test.go b/ast/policy_test.go
@@ -5,6 +5,7 @@
 package ast
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"reflect"
@@ -484,6 +485,22 @@ wildcard = true { bar[_] = 1 }`
 	}
 }
 
+func TestModuleCopy(t *testing.T) {
+
+	input := `package foo
+
+	# comment
+	p := 7`
+
+	mod := MustParseModule(input)
+	cpy := mod.Copy()
+	cpy.Comments[0].Text[0] = 'X'
+
+	if !bytes.Equal(mod.Comments[0].Text, []byte(" comment")) {
+		t.Fatal("expected comment text to be unchanged")
+	}
+}
+
 func TestWithString(t *testing.T) {
 
 	with1 := &With{