From 25f4cb6db7e64df5f9c9526406adb6bb249ac437 Mon Sep 17 00:00:00 2001
From: Anders Eknert <anders@styra.com>
Date: Tue, 20 Aug 2024 11:38:18 +0200
Subject: [PATCH] Include term locations in rule heads when requested (#6943)

When the JSON option to include term locations have been set, they should be
included in all parts of the head where terms appear.

Fixes #6860

Signed-off-by: Anders Eknert <anders@styra.com>
---
 ast/marshal_test.go | 36 ++++++++++++++++++++++++++++++++++++
 ast/policy.go       | 10 ++++++++--
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/ast/marshal_test.go b/ast/marshal_test.go
index 2b5952745c..63b7c9ae91 100644
--- a/ast/marshal_test.go
+++ b/ast/marshal_test.go
@@ -543,6 +543,42 @@ func TestHead_MarshalJSON(t *testing.T) {
 	}
 }
 
+func TestRuleHeadRefWithTermLocations_MarshalJSON(t *testing.T) {
+	policy := `package test
+
+import rego.v1
+
+ref.head[rule].test contains "value" if {
+	rule := "rule"
+}`
+
+	jsonOptions := &astJSON.Options{
+		MarshalOptions: astJSON.MarshalOptions{
+			IncludeLocation: astJSON.NodeToggle{
+				Head: true,
+				Term: true,
+			},
+		},
+	}
+
+	module, err := ParseModuleWithOpts("test.rego", policy, ParserOptions{JSONOptions: jsonOptions})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	bs, err := json.Marshal(module.Rules[0].Head)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Ensure marshalled JSON includes location for any term
+	expectedJSON := `{"key":{"location":{"file":"test.rego","row":5,"col":30},"type":"string","value":"value"},"ref":[{"location":{"file":"test.rego","row":5,"col":1},"type":"var","value":"ref"},{"location":{"file":"test.rego","row":5,"col":5},"type":"string","value":"head"},{"location":{"file":"test.rego","row":5,"col":10},"type":"var","value":"rule"},{"location":{"file":"test.rego","row":5,"col":16},"type":"string","value":"test"}],"location":{"file":"test.rego","row":5,"col":1}}`
+
+	if string(bs) != expectedJSON {
+		t.Errorf("expected %s but got %s", expectedJSON, string(bs))
+	}
+}
+
 func TestExpr_MarshalJSON(t *testing.T) {
 	rawModule := `
 	package foo
diff --git a/ast/policy.go b/ast/policy.go
index ee6e14171f..09bcea6487 100644
--- a/ast/policy.go
+++ b/ast/policy.go
@@ -1034,16 +1034,22 @@ func (head *Head) setJSONOptions(opts astJSON.Options) {
 
 func (head *Head) MarshalJSON() ([]byte, error) {
 	var loc *Location
-	if head.jsonOptions.MarshalOptions.IncludeLocation.Head {
+	includeLoc := head.jsonOptions.MarshalOptions.IncludeLocation
+	if includeLoc.Head {
 		if head.Location != nil {
 			loc = head.Location
 		}
+
+		for _, term := range head.Reference {
+			if term.Location != nil {
+				term.jsonOptions.MarshalOptions.IncludeLocation.Term = includeLoc.Term
+			}
+		}
 	}
 
 	// NOTE(sr): we do this to override the rendering of `head.Reference`.
 	// It's still what'll be used via the default means of encoding/json
 	// for unmarshaling a json object into a Head struct!
-	// NOTE(charlieegan3): we also need to optionally include the location
 	type h Head
 	return json.Marshal(struct {
 		h