Skip to content

Commit

Permalink
ast: Add location to single entry rule head ref (open-policy-agent#6212)
Browse files Browse the repository at this point in the history 
Fixes: open-policy-agent#6199
Signed-off-by: Ronnie Personal <[email protected]>
  • Loading branch information
@Ronnie-personal
Ronnie-personal authored Sep 18, 2023
1 parent 00877d6 commit 38733ed
Show file tree
Hide file tree
Showing 5 changed files with 238 additions and 5 deletions.
4 changes: 2 additions & 2 deletions ast/marshal_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -937,8 +937,8 @@ p = 1`,
},
expected: []string{
`{"annotations":{"authors":[{"name":"pkg"}],"custom":{"pkg":"pkg"},"description":"pkg","location":{"file":"","row":1,"col":1},"organizations":["pkg"],"related_resources":[{"ref":"https://pkg"}],"schemas":[{"path":[{"type":"var","value":"input"},{"type":"string","value":"foo"}],"definition":{"type":"boolean"}}],"scope":"package","title":"pkg"},"location":{"file":"","row":14,"col":1},"path":[{"location":{"file":"","row":14,"col":9},"type":"var","value":"data"},{"location":{"file":"","row":14,"col":9},"type":"string","value":"test"}]}`,
`{"annotations":{"authors":[{"name":"doc"}],"custom":{"doc":"doc"},"description":"doc","location":{"file":"","row":16,"col":1},"organizations":["doc"],"related_resources":[{"ref":"https://doc"}],"schemas":[{"path":[{"type":"var","value":"input"},{"type":"string","value":"bar"}],"definition":{"type":"integer"}}],"scope":"document","title":"doc"},"location":{"file":"","row":44,"col":1},"path":[{"location":{"file":"","row":14,"col":9},"type":"var","value":"data"},{"location":{"file":"","row":14,"col":9},"type":"string","value":"test"},{"type":"string","value":"p"}]}`,
`{"annotations":{"authors":[{"name":"rule"}],"custom":{"rule":"rule"},"description":"rule","location":{"file":"","row":31,"col":1},"organizations":["rule"],"related_resources":[{"ref":"https://rule"}],"schemas":[{"path":[{"type":"var","value":"input"},{"type":"string","value":"baz"}],"definition":{"type":"string"}}],"scope":"rule","title":"rule"},"location":{"file":"","row":44,"col":1},"path":[{"location":{"file":"","row":14,"col":9},"type":"var","value":"data"},{"location":{"file":"","row":14,"col":9},"type":"string","value":"test"},{"type":"string","value":"p"}]}`,
`{"annotations":{"authors":[{"name":"doc"}],"custom":{"doc":"doc"},"description":"doc","location":{"file":"","row":16,"col":1},"organizations":["doc"],"related_resources":[{"ref":"https://doc"}],"schemas":[{"path":[{"type":"var","value":"input"},{"type":"string","value":"bar"}],"definition":{"type":"integer"}}],"scope":"document","title":"doc"},"location":{"file":"","row":44,"col":1},"path":[{"location":{"file":"","row":14,"col":9},"type":"var","value":"data"},{"location":{"file":"","row":14,"col":9},"type":"string","value":"test"},{"location":{"file":"","row":44,"col":1},"type":"string","value":"p"}]}`,
`{"annotations":{"authors":[{"name":"rule"}],"custom":{"rule":"rule"},"description":"rule","location":{"file":"","row":31,"col":1},"organizations":["rule"],"related_resources":[{"ref":"https://rule"}],"schemas":[{"path":[{"type":"var","value":"input"},{"type":"string","value":"baz"}],"definition":{"type":"string"}}],"scope":"rule","title":"rule"},"location":{"file":"","row":44,"col":1},"path":[{"location":{"file":"","row":14,"col":9},"type":"var","value":"data"},{"location":{"file":"","row":14,"col":9},"type":"string","value":"test"},{"location":{"file":"","row":44,"col":1},"type":"string","value":"p"}]}`,
},
},
}
Expand Down
4 changes: 3 additions & 1 deletion ast/parser.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -856,7 +856,9 @@ func (p *Parser) parseHead(defaultRule bool) (*Head, bool) {

switch x := ref.Value.(type) {
case Var:
head = NewHead(x)
// Modify the code to add the location to the head ref
// and set the head ref's jsonOptions.
head = VarHead(x, ref.Location, p.po.JSONOptions)
case Ref:
head = RefHead(x)
case Call:
Expand Down
8 changes: 6 additions & 2 deletions ast/parser_ext.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,9 @@ func ParseCompleteDocRuleFromEqExpr(module *Module, lhs, rhs *Term) (*Rule, erro
var head *Head

if v, ok := lhs.Value.(Var); ok {
head = NewHead(v)
// Modify the code to add the location to the head ref
// and set the head ref's jsonOptions.
head = VarHead(v, lhs.Location, &lhs.jsonOptions)
} else if r, ok := lhs.Value.(Ref); ok { // groundness ?
if _, ok := r[0].Value.(Var); !ok {
return nil, fmt.Errorf("invalid rule head: %v", r)
Expand Down Expand Up @@ -350,7 +352,9 @@ func ParsePartialSetDocRuleFromTerm(module *Module, term *Term) (*Rule, error) {
if !ok {
return nil, fmt.Errorf("%vs cannot be used for rule head", TypeName(term.Value))
}
head = NewHead(v)
// Modify the code to add the location to the head ref
// and set the head ref's jsonOptions.
head = VarHead(v, ref[0].Location, &ref[0].jsonOptions)
head.Key = ref[1]
}
head.Location = term.Location
Expand Down
11 changes: 11 additions & 0 deletions ast/policy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -769,6 +769,17 @@ func NewHead(name Var, args ...*Term) *Head {
return head
}

// VarHead creates a head object, initializes its Name, Location, and JSONOptions,
// and returns the new head.
func VarHead(name Var, location *Location, jsonOpts *JSONOptions) *Head {
h := NewHead(name)
h.Reference[0].Location = location
if jsonOpts != nil {
h.Reference[0].setJSONOptions(*jsonOpts)
}
return h
}

// RefHead returns a new Head object with the passed Ref. If args are provided,
// the first will be used for the value.
func RefHead(ref Ref, args ...*Term) *Head {
Expand Down
216 changes: 216 additions & 0 deletions cmd/parse_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,11 @@ func TestParseJSONOutputWithLocations(t *testing.T) {
},
"ref": [
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 3
},
"type": "var",
"value": "p"
}
Expand All @@ -231,6 +236,207 @@ func TestParseJSONOutputWithLocations(t *testing.T) {
}
}

func TestParseRefsJSONOutput(t *testing.T) {

files := map[string]string{
"x.rego": `package x
a.b.c := true
`,
}
errc, stdout, stderr, _ := testParse(t, files, &parseParams{
format: util.NewEnumFlag(parseFormatJSON, []string{parseFormatPretty, parseFormatJSON}),
})
if errc != 0 {
t.Fatalf("Expected exit code 0, got %v", errc)
}
if len(stderr) > 0 {
t.Fatalf("Expected no stderr output, got:\n%s\n", string(stderr))
}

expectedOutput := `{
"package": {
"path": [
{
"type": "var",
"value": "data"
},
{
"type": "string",
"value": "x"
}
]
},
"rules": [
{
"body": [
{
"index": 0,
"terms": {
"type": "boolean",
"value": true
}
}
],
"head": {
"value": {
"type": "boolean",
"value": true
},
"assign": true,
"ref": [
{
"type": "var",
"value": "a"
},
{
"type": "string",
"value": "b"
},
{
"type": "string",
"value": "c"
}
]
}
}
]
}
`

if got, want := string(stdout), expectedOutput; got != want {
t.Fatalf("Expected output\n%v\n, got\n%v", want, got)
}
}

func TestParseRefsJSONOutputWithLocations(t *testing.T) {

files := map[string]string{
"x.rego": `package x
a.b.c := true
`,
}
errc, stdout, stderr, tempDirPath := testParse(t, files, &parseParams{
format: util.NewEnumFlag(parseFormatJSON, []string{parseFormatPretty, parseFormatJSON}),
jsonInclude: "locations",
})
if errc != 0 {
t.Fatalf("Expected exit code 0, got %v", errc)
}
if len(stderr) > 0 {
t.Fatalf("Expected no stderr output, got:\n%s\n", string(stderr))
}

expectedOutput := strings.Replace(`{
"package": {
"location": {
"file": "TEMPDIR/x.rego",
"row": 1,
"col": 1
},
"path": [
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 1,
"col": 9
},
"type": "var",
"value": "data"
},
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 1,
"col": 9
},
"type": "string",
"value": "x"
}
]
},
"rules": [
{
"body": [
{
"index": 0,
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 14
},
"terms": {
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 14
},
"type": "boolean",
"value": true
}
}
],
"head": {
"value": {
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 14
},
"type": "boolean",
"value": true
},
"assign": true,
"ref": [
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 5
},
"type": "var",
"value": "a"
},
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 7
},
"type": "string",
"value": "b"
},
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 9
},
"type": "string",
"value": "c"
}
],
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 5
}
},
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 5
}
}
]
}
`, "TEMPDIR", tempDirPath, -1)

if got, want := string(stdout), expectedOutput; got != want {
t.Fatalf("Expected output\n%v\n, got\n%v", want, got)
}
}
func TestParseRulesBlockJSONOutputWithLocations(t *testing.T) {

files := map[string]string{
Expand Down Expand Up @@ -318,6 +524,11 @@ func TestParseRulesBlockJSONOutputWithLocations(t *testing.T) {
},
"ref": [
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 3,
"col": 11
},
"type": "var",
"value": "allow"
}
Expand Down Expand Up @@ -567,6 +778,11 @@ func TestParseRulesBlockJSONOutputWithLocations(t *testing.T) {
},
"ref": [
{
"location": {
"file": "TEMPDIR/x.rego",
"row": 4,
"col": 5
},
"type": "var",
"value": "allow"
}
Expand Down

0 comments on commit 38733ed

Please sign in to comment.