From 40024cebd6fbe03871ebef452cf784644a93fb0d Mon Sep 17 00:00:00 2001 From: Johan Fylling Date: Thu, 7 May 2026 14:43:08 +0200 Subject: [PATCH] ast: not-body marshaling (#8614) JSON- and pretty format marshaling of `ast.Not` Signed-off-by: Johan Fylling --- cmd/parse_test.go | 374 +++++++++++++++++++++++++++++++++++++++++ v1/ast/json/json.go | 1 + v1/ast/marshal_test.go | 250 +++++++++++++++++++++++++++ v1/ast/policy.go | 2 + v1/ast/pretty.go | 2 + v1/ast/strings.go | 2 + v1/ast/term.go | 66 +++++++- v1/ast/visit.go | 2 + 8 files changed, 695 insertions(+), 4 deletions(-) diff --git a/cmd/parse_test.go b/cmd/parse_test.go index ce036565d5..a3da2e1484 100644 --- a/cmd/parse_test.go +++ b/cmd/parse_test.go @@ -6,6 +6,7 @@ import ( "strings" "testing" + "github.com/google/go-cmp/cmp" "github.com/open-policy-agent/opa/cmd/formats" "github.com/open-policy-agent/opa/v1/util/test" ) @@ -258,6 +259,379 @@ p = 1 } } +func TestParseOutputWithNotImport(t *testing.T) { + cases := []struct { + format string + exp string + }{ + { + format: formats.JSON, + exp: `{ + "package": { + "path": [ + { + "type": "var", + "value": "data" + }, + { + "type": "string", + "value": "test" + } + ] + }, + "imports": [ + { + "path": { + "type": "ref", + "value": [ + { + "type": "var", + "value": "future" + }, + { + "type": "string", + "value": "keywords" + }, + { + "type": "string", + "value": "not" + } + ] + } + } + ], + "rules": [ + { + "body": [ + { + "index": 0, + "terms": { + "body": [ + { + "index": 0, + "terms": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "equal" + } + ] + }, + { + "type": "call", + "value": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "plus" + } + ] + }, + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "input" + }, + { + "type": "string", + "value": "x" + } + ] + }, + { + "type": "number", + "value": 2 + } + ] + }, + { + "type": "number", + "value": 42 + } + ] + } + ], + "explicit_body": false, + "type": "not" + } + } + ], + "head": { + "name": "implicit_body", + "value": { + "type": "boolean", + "value": true + }, + "ref": [ + { + "type": "var", + "value": "implicit_body" + } + ] + } + }, + { + "body": [ + { + "index": 0, + "terms": { + "body": [ + { + "index": 0, + "terms": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "assign" + } + ] + }, + { + "type": "var", + "value": "x" + }, + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "input" + }, + { + "type": "string", + "value": "x" + } + ] + } + ] + }, + { + "index": 1, + "terms": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "assign" + } + ] + }, + { + "type": "var", + "value": "y" + }, + { + "type": "number", + "value": 2 + } + ] + }, + { + "index": 2, + "terms": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "assign" + } + ] + }, + { + "type": "var", + "value": "z" + }, + { + "type": "call", + "value": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "plus" + } + ] + }, + { + "type": "var", + "value": "x" + }, + { + "type": "var", + "value": "y" + } + ] + } + ] + }, + { + "index": 3, + "terms": [ + { + "type": "ref", + "value": [ + { + "type": "var", + "value": "equal" + } + ] + }, + { + "type": "var", + "value": "z" + }, + { + "type": "number", + "value": 42 + } + ] + } + ], + "explicit_body": true, + "type": "not" + } + } + ], + "head": { + "name": "explicit_body", + "value": { + "type": "boolean", + "value": true + }, + "ref": [ + { + "type": "var", + "value": "explicit_body" + } + ] + } + } + ] +} +`, + }, + { + format: formats.Pretty, + exp: `module + package + ref + data + "test" + import + ref + future + "keywords" + "not" + + rule + head + ref + implicit_body + true + body + expr index=0 + not + body + expr index=0 + ref + equal + call + ref + plus + ref + input + "x" + 2 + 42 + rule + head + ref + explicit_body + true + body + expr index=0 + not + body + expr index=0 + ref + assign + x + ref + input + "x" + expr index=1 + ref + assign + y + 2 + expr index=2 + ref + assign + z + call + ref + plus + x + y + expr index=3 + ref + equal + z + 42 +`, + }, + } + + files := map[string]string{ + "x.rego": `package test + + import future.keywords.not + + implicit_body if { + not input.x + 2 == 42 + } + + explicit_body if { + not { + x := input.x + y := 2 + z := x + y + z == 42 + } + } + `, + } + + for _, tc := range cases { + t.Run(tc.format, func(t *testing.T) { + errc, stdout, stderr, _ := testParse(t, files, &parseParams{ + format: formats.Flag(tc.format), + }) + 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)) + } + + if diff := cmp.Diff(tc.exp, string(stdout)); diff != "" { + t.Errorf("unexpected result (-want, +got):\n%s", diff) + } + }) + } +} + func TestParseRefsJSONOutput(t *testing.T) { files := map[string]string{ diff --git a/v1/ast/json/json.go b/v1/ast/json/json.go index 9081fe7039..42949e9965 100644 --- a/v1/ast/json/json.go +++ b/v1/ast/json/json.go @@ -54,6 +54,7 @@ type NodeToggle struct { With bool Annotations bool AnnotationsRef bool + Not bool } // configuredJSONOptions synchronizes access to the global JSON options diff --git a/v1/ast/marshal_test.go b/v1/ast/marshal_test.go index a748d1a76a..a2b23cbda7 100644 --- a/v1/ast/marshal_test.go +++ b/v1/ast/marshal_test.go @@ -2,6 +2,7 @@ package ast import ( "encoding/json" + "strings" "testing" astJSON "github.com/open-policy-agent/opa/v1/ast/json" @@ -1136,3 +1137,252 @@ p = 1`, }) } } + +func TestNot_MarshalJSON(t *testing.T) { + rawModule := ` + package test + + import future.keywords.not + + implicit_body if { + not input.x + 2 == 42 + } + + explicit_body if { + not { + x := input.x + y := 2 + z := x + y + z == 42 + } + } + ` + + module, err := ParseModule("example.rego", rawModule) + if err != nil { + t.Fatal(err) + } + + testCases := map[string]struct { + Not *Not + Options astJSON.Options + ExpectedJSON string + }{ + "implicit body: base case": { + Not: module.Rules[0].Body[0].Terms.(*Not), + ExpectedJSON: `{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]},{"type":"number","value":2}]},{"type":"number","value":42}]}],"explicit_body":false,"type":"not"}`, + }, + "explicit body: base case": { + Not: module.Rules[1].Body[0].Terms.(*Not), + ExpectedJSON: `{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"x"},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]}]},{"index":1,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"y"},{"type":"number","value":2}]},{"index":2,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"z"},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"var","value":"x"},{"type":"var","value":"y"}]}]},{"index":3,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"var","value":"z"},{"type":"number","value":42}]}],"explicit_body":true,"type":"not"}`, + }, + "implicit body: location excluded": { + Not: module.Rules[0].Body[0].Terms.(*Not), + Options: astJSON.Options{ + MarshalOptions: astJSON.MarshalOptions{IncludeLocation: astJSON.NodeToggle{Not: false}}, + }, + ExpectedJSON: `{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]},{"type":"number","value":2}]},{"type":"number","value":42}]}],"explicit_body":false,"type":"not"}`, + }, + "explicit body: location excluded": { + Not: module.Rules[1].Body[0].Terms.(*Not), + Options: astJSON.Options{ + MarshalOptions: astJSON.MarshalOptions{IncludeLocation: astJSON.NodeToggle{Not: false}}, + }, + ExpectedJSON: `{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"x"},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]}]},{"index":1,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"y"},{"type":"number","value":2}]},{"index":2,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"z"},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"var","value":"x"},{"type":"var","value":"y"}]}]},{"index":3,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"var","value":"z"},{"type":"number","value":42}]}],"explicit_body":true,"type":"not"}`, + }, + "implicit body: location included": { + Not: module.Rules[0].Body[0].Terms.(*Not), + Options: astJSON.Options{MarshalOptions: astJSON.MarshalOptions{IncludeLocation: astJSON.NodeToggle{Not: true}}}, + ExpectedJSON: `{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]},{"type":"number","value":2}]},{"type":"number","value":42}]}],"explicit_body":false,"location":{"file":"example.rego","row":8,"col":3},"type":"not"}`, + }, + "explicit body: location included": { + Not: module.Rules[1].Body[0].Terms.(*Not), + Options: astJSON.Options{MarshalOptions: astJSON.MarshalOptions{IncludeLocation: astJSON.NodeToggle{Not: true}}}, + ExpectedJSON: `{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"x"},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]}]},{"index":1,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"y"},{"type":"number","value":2}]},{"index":2,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"z"},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"var","value":"x"},{"type":"var","value":"y"}]}]},{"index":3,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"var","value":"z"},{"type":"number","value":42}]}],"explicit_body":true,"location":{"file":"example.rego","row":11,"col":8},"type":"not"}`, + }, + "explicit body: location included, also for nested expressions": { + Not: module.Rules[1].Body[0].Terms.(*Not), + Options: astJSON.Options{MarshalOptions: astJSON.MarshalOptions{IncludeLocation: astJSON.NodeToggle{Not: true, Expr: true}}}, + ExpectedJSON: `{"body":[{"index":0,"location":{"file":"example.rego","row":12,"col":5},"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"x"},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]}]},{"index":1,"location":{"file":"example.rego","row":13,"col":5},"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"y"},{"type":"number","value":2}]},{"index":2,"location":{"file":"example.rego","row":14,"col":5},"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"z"},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"var","value":"x"},{"type":"var","value":"y"}]}]},{"index":3,"location":{"file":"example.rego","row":15,"col":5},"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"var","value":"z"},{"type":"number","value":42}]}],"explicit_body":true,"location":{"file":"example.rego","row":11,"col":8},"type":"not"}`, + }, + } + + for name, data := range testCases { + t.Run(name, func(t *testing.T) { + astJSON.SetOptions(data.Options) + t.Cleanup(resetJSONOptions) + + bs := util.MustMarshalJSON(data.Not) + got := string(bs) + exp := data.ExpectedJSON + + if got != exp { + t.Fatalf("expected:\n%s got\n%s", exp, got) + } + }) + } +} + +func TestNot_UnmarshalJSON(t *testing.T) { + rawModule := ` + package test + + import future.keywords.not + + implicit_body if { + not input.x + 2 == 42 + } + + explicit_body if { + not { + x := input.x + y := 2 + z := x + y + z == 42 + } + } + ` + + module, err := ParseModule("example.rego", rawModule) + if err != nil { + t.Fatal(err) + } + + implicitBodyExpr := module.Rules[0].Body[0] + // text is not marshalled to JSON so we just drop it in our examples + implicitBodyExpr.Location.Text = nil + + explicitBodyExpr := module.Rules[1].Body[0] + explicitBodyExpr.Location.Text = nil + + testCases := map[string]struct { + JSON string + ExpectedExpr *Expr + }{ + "implicit body": { + JSON: `{"index":0,"terms":{"type":"not","body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]},{"type":"number","value":2}]},{"type":"number","value":42}]}],"explicit_body":false}}`, + ExpectedExpr: func() *Expr { + e := implicitBodyExpr.Copy() + e.Location = nil + return e + }(), + }, + "explicit body": { + JSON: `{"index":0,"terms":{"body":[{"index":0,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"x"},{"type":"ref","value":[{"type":"var","value":"input"},{"type":"string","value":"x"}]}]},{"index":1,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"y"},{"type":"number","value":2}]},{"index":2,"terms":[{"type":"ref","value":[{"type":"var","value":"assign"}]},{"type":"var","value":"z"},{"type":"call","value":[{"type":"ref","value":[{"type":"var","value":"plus"}]},{"type":"var","value":"x"},{"type":"var","value":"y"}]}]},{"index":3,"terms":[{"type":"ref","value":[{"type":"var","value":"equal"}]},{"type":"var","value":"z"},{"type":"number","value":42}]}],"explicit_body":true,"type":"not"}}`, + ExpectedExpr: func() *Expr { + e := explicitBodyExpr.Copy() + e.Location = nil + return e + }(), + }, + } + + for name, data := range testCases { + t.Run(name, func(t *testing.T) { + var expr Expr + err := json.Unmarshal([]byte(data.JSON), &expr) + if err != nil { + t.Fatal(err) + } + + if !expr.Equal(data.ExpectedExpr) { + t.Fatalf("expected:\n%#v got\n%#v", data.ExpectedExpr, expr) + } + if data.ExpectedExpr.Location != nil { + if !expr.Location.Equal(data.ExpectedExpr.Location) { + t.Fatalf("expected location:\n%#v got\n%#v", data.ExpectedExpr.Location, expr.Location) + } + } + }) + } +} + +func TestNot_MarshalUnmarshalRoundTrip(t *testing.T) { + rawModule := ` + package test + + import future.keywords.not + + implicit_body if { + not input.x + 2 == 42 + } + + explicit_body if { + not { + x := input.x + y := 2 + z := x + y + z == 42 + } + } + ` + + module, err := ParseModule("example.rego", rawModule) + if err != nil { + t.Fatal(err) + } + + testCases := map[string]struct { + Expr *Expr + }{ + "implicit body": { + Expr: module.Rules[0].Body[0], + }, + "explicit body": { + Expr: module.Rules[1].Body[0], + }, + } + + for name, data := range testCases { + t.Run(name, func(t *testing.T) { + bs := util.MustMarshalJSON(data.Expr) + + var expr Expr + err := json.Unmarshal(bs, &expr) + if err != nil { + t.Fatalf("unmarshal failed: %v\njson: %s", err, string(bs)) + } + + if !expr.Equal(data.Expr) { + t.Fatalf("round-trip mismatch\noriginal: %#v\ngot: %#v", data.Expr, &expr) + } + }) + } +} + +func TestNot_UnmarshalJSON_Errors(t *testing.T) { + testCases := map[string]struct { + JSON string + expErr string + }{ + "body is not an array": { + JSON: `{"index":0,"terms":{"type":"not","body":"invalid","explicit_body":false}}`, + expErr: "invalid body field type", + }, + "body is missing": { + JSON: `{"index":0,"terms":{"type":"not","explicit_body":false}}`, + expErr: "invalid body field type", + }, + "explicit_body is not a bool": { + JSON: `{"index":0,"terms":{"type":"not","body":[],"explicit_body":"yes"}}`, + expErr: "unable to unmarshal explicit_body field", + }, + "body contains invalid expression": { + JSON: `{"index":0,"terms":{"type":"not","body":[{"index":0,"terms":"bad"}],"explicit_body":false}}`, + expErr: "unable to unmarshal not body", + }, + } + + for name, data := range testCases { + t.Run(name, func(t *testing.T) { + var expr Expr + err := json.Unmarshal([]byte(data.JSON), &expr) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), data.expErr) { + t.Fatalf("expected error containing %q, got: %v", data.expErr, err) + } + }) + } +} diff --git a/v1/ast/policy.go b/v1/ast/policy.go index 22aecd8e0a..36af41761a 100644 --- a/v1/ast/policy.go +++ b/v1/ast/policy.go @@ -1438,6 +1438,8 @@ func (expr *Expr) IsGround() bool { } case *Term: return ts.IsGround() + case *Not: + return ts.IsGround() } return true } diff --git a/v1/ast/pretty.go b/v1/ast/pretty.go index aa34f37471..1ee32ee555 100644 --- a/v1/ast/pretty.go +++ b/v1/ast/pretty.go @@ -50,6 +50,8 @@ func (pp *prettyPrinter) Before(x any) bool { pp.writeIndent("%v %v", TypeName(x), strings.Join(extras, " ")) case Null, Boolean, Number, String, Var: pp.writeValue(x) + case *Not: + pp.writeType(x) default: pp.writeType(x) } diff --git a/v1/ast/strings.go b/v1/ast/strings.go index 72ec03f8cc..8bb7e7ddfd 100644 --- a/v1/ast/strings.go +++ b/v1/ast/strings.go @@ -50,6 +50,8 @@ func ValueName(x Value) string { return "setcomprehension" case *TemplateString: return "templatestring" + case *Not: + return "not" } return TypeName(x) diff --git a/v1/ast/term.go b/v1/ast/term.go index ce51413476..ff5b2f221d 100644 --- a/v1/ast/term.go +++ b/v1/ast/term.go @@ -660,6 +660,56 @@ func (n *Not) String() string { return "not {" + n.Body.String() + "}" } +func (n *Not) MarshalJSON() ([]byte, error) { + data := map[string]any{ + "type": "not", + "body": n.Body, + "explicit_body": n.ExplicitBody, + } + + if astJSON.GetOptions().MarshalOptions.IncludeLocation.Not { + if n.Location != nil { + data["location"] = n.Location + } + } + + return json.Marshal(data) +} + +func (n *Not) UnmarshalJSON(bs []byte) error { + v := map[string]any{} + if err := util.UnmarshalJSON(bs, &v); err != nil { + return err + } + + return unmarshalNot(n, v) +} + +func unmarshalNot(n *Not, v map[string]any) error { + var eb bool + if x, ok := v["explicit_body"]; ok { + eb, ok = x.(bool) + if !ok { + return fmt.Errorf("ast: unable to unmarshal explicit_body field with type: %T (expected true or false)", v["explicit_body"]) + } + } + + b, ok := v["body"].([]any) + if !ok { + return fmt.Errorf("ast: unable to unmarshal not, invalid body field type: %T (expected list)", v["body"]) + } + + body, err := unmarshalBody(b) + if err != nil { + return fmt.Errorf("ast: unable to unmarshal not body: %w", err) + } + + n.ExplicitBody = eb + n.Body = body + + return nil +} + // Null represents the null value defined by JSON. type Null struct{} @@ -3192,11 +3242,19 @@ func unmarshalExpr(expr *Expr, v map[string]any) error { } switch ts := v["terms"].(type) { case map[string]any: - t, err := unmarshalTerm(ts) - if err != nil { - return err + if tt, ok := ts["type"]; ok && tt == "not" { + n := &Not{} + if err := unmarshalNot(n, ts); err != nil { + return err + } + expr.Terms = n + } else { + t, err := unmarshalTerm(ts) + if err != nil { + return err + } + expr.Terms = t } - expr.Terms = t case []any: terms, err := unmarshalTermSlice(ts) if err != nil { diff --git a/v1/ast/visit.go b/v1/ast/visit.go index e49ca4f31c..16d7ebd78c 100644 --- a/v1/ast/visit.go +++ b/v1/ast/visit.go @@ -736,6 +736,8 @@ func (vis *BeforeAfterVisitor) Walk(x any) { for i := range x.Symbols { vis.Walk(x.Symbols[i]) } + case *Not: + vis.Walk(x.Body) } }