From 8ff85cbc34ab37191a3459737ac57255e04d53b0 Mon Sep 17 00:00:00 2001
From: Kris Kennaway <95245117+KrisKennawayDD@users.noreply.github.com>
Date: Tue, 11 Mar 2025 18:07:50 +0000
Subject: [PATCH] ir: Fix nil pointer deref in Unmarshal() when handling
 IsSetStmt (#7430)

Fixes: #7415

Signed-off-by: Kris Kennaway <kris.kennaway@datadoghq.com>
---
 v1/ir/encoding/encoding_test.go | 125 +++++++++++++++++++-------------
 v1/ir/marshal.go                |  14 +++-
 2 files changed, 86 insertions(+), 53 deletions(-)

diff --git a/v1/ir/encoding/encoding_test.go b/v1/ir/encoding/encoding_test.go
index 38395add47..5712b012a0 100644
--- a/v1/ir/encoding/encoding_test.go
+++ b/v1/ir/encoding/encoding_test.go
@@ -11,62 +11,85 @@ import (
 )
 
 func TestRoundTrip(t *testing.T) {
-
-	// Note: v1 module
-	c, err := ast.CompileModules(map[string]string{
-		"test.rego": `
-			package test
-
-			p if {
-				input.foo == 7
-			}
-		`,
-	})
-
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	modules := []*ast.Module{}
-
-	for _, m := range c.Modules {
-		modules = append(modules, m)
-	}
-
-	planner := planner.New().
-		WithQueries([]planner.QuerySet{
-			{
-				Name: "main",
-				Queries: []ast.Body{
-					ast.MustParseBody("data.test.p = true"),
-				},
+	tests := []struct {
+		note    string
+		modules map[string]string
+	}{
+		{
+			note: "simple",
+			modules: map[string]string{
+				"test.rego": `
+					package test
+					p if {
+						input.foo == 7
+					}
+				`,
 			},
-		}).
-		WithModules(modules).
-		WithBuiltinDecls(ast.BuiltinMap)
-
-	plan, err := planner.Plan()
-	if err != nil {
-		t.Fatal(err)
+		},
+		{
+			note: "every",
+			modules: map[string]string{
+				"test.rego": `
+					package test
+					p if {
+						every i in input.foo { i > 0 } 
+					}
+				`,
+			},
+		},
 	}
 
-	bs, err := json.MarshalIndent(plan, "", "  ")
-	if err != nil {
-		t.Fatal(err)
-	}
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			// Note: v1 module
+			c, err := ast.CompileModules(tc.modules)
 
-	var cpy ir.Policy
-	err = json.Unmarshal(bs, &cpy)
-	if err != nil {
-		t.Fatal(err)
-	}
+			if err != nil {
+				t.Fatal(err)
+			}
 
-	bs2, err := json.MarshalIndent(plan, "", "  ")
-	if err != nil {
-		t.Fatal(err)
-	}
+			modules := []*ast.Module{}
 
-	if !bytes.Equal(bs, bs2) {
-		t.Fatal("expected bytes to be equal")
+			for _, m := range c.Modules {
+				modules = append(modules, m)
+			}
+
+			planner := planner.New().
+				WithQueries([]planner.QuerySet{
+					{
+						Name: "main",
+						Queries: []ast.Body{
+							ast.MustParseBody("data.test.p = true"),
+						},
+					},
+				}).
+				WithModules(modules).
+				WithBuiltinDecls(ast.BuiltinMap)
+
+			plan, err := planner.Plan()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			bs, err := json.MarshalIndent(plan, "", "  ")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			var cpy ir.Policy
+			err = json.Unmarshal(bs, &cpy)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			bs2, err := json.MarshalIndent(plan, "", "  ")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if !bytes.Equal(bs, bs2) {
+				t.Fatal("expected bytes to be equal")
+			}
+		})
 	}
 }
diff --git a/v1/ir/marshal.go b/v1/ir/marshal.go
index 69f4b5caf6..f792e2c1b6 100644
--- a/v1/ir/marshal.go
+++ b/v1/ir/marshal.go
@@ -6,6 +6,7 @@ package ir
 
 import (
 	"encoding/json"
+	"fmt"
 	"reflect"
 )
 
@@ -50,7 +51,11 @@ func (a *Operand) UnmarshalJSON(bs []byte) error {
 	if err := json.Unmarshal(bs, &typed); err != nil {
 		return err
 	}
-	x := valFactories[typed.Type]()
+	f, ok := valFactories[typed.Type]
+	if !ok {
+		return fmt.Errorf("unrecognized value type %q", typed.Type)
+	}
+	x := f()
 	if err := json.Unmarshal(typed.Value, &x); err != nil {
 		return err
 	}
@@ -77,7 +82,11 @@ type rawTypedStmt struct {
 }
 
 func (raw rawTypedStmt) Unmarshal() (Stmt, error) {
-	x := stmtFactories[raw.Type]()
+	f, ok := stmtFactories[raw.Type]
+	if !ok {
+		return nil, fmt.Errorf("unrecognized statement type %q", raw.Type)
+	}
+	x := f()
 	if err := json.Unmarshal(raw.Stmt, &x); err != nil {
 		return nil, err
 	}
@@ -119,6 +128,7 @@ var stmtFactories = map[string]func() Stmt{
 	"IsArrayStmt":          func() Stmt { return &IsArrayStmt{} },
 	"IsObjectStmt":         func() Stmt { return &IsObjectStmt{} },
 	"IsDefinedStmt":        func() Stmt { return &IsDefinedStmt{} },
+	"IsSetStmt":            func() Stmt { return &IsSetStmt{} },
 	"IsUndefinedStmt":      func() Stmt { return &IsUndefinedStmt{} },
 	"ArrayAppendStmt":      func() Stmt { return &ArrayAppendStmt{} },
 	"ObjectInsertStmt":     func() Stmt { return &ObjectInsertStmt{} },