topdown: Add net.cidr_contains_matches built-in function

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

ast/builtins.go

@@ -200,6 +200,7 @@ var DefaultBuiltins = [...]*Builtin{
 	NetCIDROverlap,
 	NetCIDRIntersects,
 	NetCIDRContains,
+	NetCIDRContainsMatches,
 	NetCIDRExpand,
 
 	// Glob
@@ -1745,6 +1746,34 @@ var NetCIDRContains = &Builtin{
 	),
 }
 
+// NetCIDRContainsMatches checks if collections of cidrs or ips are contained within another collection of cidrs and returns matches.
+var NetCIDRContainsMatches = &Builtin{
+	Name: "net.cidr_contains_matches",
+	Decl: types.NewFunction(
+		types.Args(netCidrContainsMatchesOperandType, netCidrContainsMatchesOperandType),
+		types.NewSet(types.NewArray([]types.Type{types.A, types.A}, nil)),
+	),
+}
+
+var netCidrContainsMatchesOperandType = types.NewAny(
+	types.S,
+	types.NewArray(nil, types.NewAny(
+		types.S,
+		types.NewArray(nil, types.A),
+	)),
+	types.NewSet(types.NewAny(
+		types.S,
+		types.NewArray(nil, types.A),
+	)),
+	types.NewObject(nil, types.NewDynamicProperty(
+		types.S,
+		types.NewAny(
+			types.S,
+			types.NewArray(nil, types.A),
+		),
+	)),
+)
+
 /**
  * Deprecated built-ins.
  */

docs/content/policy-reference.md

@@ -604,9 +604,62 @@ The table below shows examples of calling `http.send`:
 | Built-in | Description |
 | ------- |-------------|
 | <span class="opa-keep-it-together">``net.cidr_contains(cidr, cidr_or_ip)``</span> | `output` is `true` if `cidr_or_ip` (e.g. `127.0.0.64/26` or `127.0.0.1`) is contained within `cidr` (e.g. `127.0.0.1/24`) and false otherwise. Supports both IPv4 and IPv6 notations.|
+| <span class="opa-keep-it-together">``output := net.cidr_contains_matches(cidrs, cidrs_or_ips)``</span> | `output` is a `set` of tuples identifying matches where `cidrs_or_ips` are contained within `cidrs`. This function is similar to `net.cidr_contains` except it allows callers to pass collections of CIDRs or IPs as arguments and returns the matches (as opposed to a boolean result indicating a match between two CIDRs/IPs.) See below for examples. |
 | <span class="opa-keep-it-together">``net.cidr_intersects(cidr1, cidr2)``</span> | `output` is `true` if `cidr1` (e.g. `192.168.0.0/16`) overlaps with `cidr2` (e.g. `192.168.1.0/24`) and false otherwise. Supports both IPv4 and IPv6 notations.|
 | <span class="opa-keep-it-together">``net.cidr_expand(cidr)``</span> | `output` is the set of hosts in `cidr`  (e.g., `net.cidr_expand("192.168.0.0/30")` generates 4 hosts: `{"192.168.0.0", "192.168.0.1", "192.168.0.2", "192.168.0.3"}` |
 
+**`net.cidr_contains_matches` examples**
+
+The `output := net.cidr_contains_matches(a, b)` function allows callers to supply
+strings, arrays, sets, or objects for either `a` or `b`. The `output` value in
+all cases is a set of tuples (2-element arrays) that identify matches, i.e.,
+elements of `b` contained by elements of `a`. The first tuple element refers to
+the match in `a` and the second tuple element refers to the match in `b`.
+
+| Input Type | Output Type |
+| --- | --- |
+| `string` | `string` |
+| `array` | `array` index |
+| `set` | `set` element |
+| `object` | `object` key |
+
+
+```live:netcidrcontainsmatches:module:hidden
+package netcidrcontainsmatches
+```
+
+If both operands are string values the function is similar to `net.cidr_contains`.
+
+```live:netcidrcontainsmatches/strings:query:merge_down
+net.cidr_contains_matches("1.1.1.0/24", "1.1.1.128")
+```
+```live:netcidrcontainsmatches/strings:output
+```
+
+Either (or both) operand(s) may be an array, set, or object.
+
+```live:netcidrcontainsmatches/array:query:merge_down
+net.cidr_contains_matches(["1.1.1.0/24", "1.1.2.0/24"], "1.1.1.128")
+```
+```live:netcidrcontainsmatches/array:output
+```
+
+The array/set/object elements may be arrays. In that case, the first element must be a valid CIDR/IP.
+
+```live:netcidrcontainsmatches/tuples:query:merge_down
+net.cidr_contains_matches([["1.1.0.0/16", "foo"], "1.1.2.0/24"], ["1.1.1.128", ["1.1.254.254", "bar"]])
+```
+```live:netcidrcontainsmatches/tuples:output
+```
+
+If the operand is a set, the outputs are matching elements. If the operand is an object, the outputs are matching keys.
+
+```live:netcidrcontainsmatches/sets_and_objects:query:merge_down
+net.cidr_contains_matches({["1.1.0.0/16", "foo"], "1.1.2.0/24"}, {"x": "1.1.1.128", "y": ["1.1.254.254", "bar"]})
+```
+```live:netcidrcontainsmatches/sets_and_objects:output:merge_down
+```
+
 ### Rego
 | Built-in | Description |
 | ------- |-------------|

topdown/cidr.go

@@ -1,6 +1,7 @@
 package topdown
 
 import (
+	"errors"
 	"fmt"
 	"math/big"
 	"net"
@@ -111,6 +112,75 @@ func builtinNetCIDRContains(a, b ast.Value) (ast.Value, error) {
 	return ast.Boolean(cidrContained), nil
 }
 
+var errNetCIDRContainsMatchElementType = errors.New("element must be string or non-empty array")
+
+func getCIDRMatchTerm(a *ast.Term) (*ast.Term, error) {
+	switch v := a.Value.(type) {
+	case ast.String:
+		return a, nil
+	case ast.Array:
+		if len(v) == 0 {
+			return nil, errNetCIDRContainsMatchElementType
+		}
+		return v[0], nil
+	default:
+		return nil, errNetCIDRContainsMatchElementType
+	}
+}
+
+func evalNetCIDRContainsMatchesOperand(operand int, a *ast.Term, iter func(cidr, index *ast.Term) error) error {
+	switch v := a.Value.(type) {
+	case ast.String:
+		return iter(a, a)
+	case ast.Array:
+		for i := range v {
+			cidr, err := getCIDRMatchTerm(v[i])
+			if err != nil {
+				return fmt.Errorf("operand %v: %v", operand, err)
+			}
+			if err := iter(cidr, ast.IntNumberTerm(i)); err != nil {
+				return err
+			}
+		}
+		return nil
+	case ast.Set:
+		return v.Iter(func(x *ast.Term) error {
+			cidr, err := getCIDRMatchTerm(x)
+			if err != nil {
+				return fmt.Errorf("operand %v: %v", operand, err)
+			}
+			return iter(cidr, x)
+		})
+	case ast.Object:
+		return v.Iter(func(k, v *ast.Term) error {
+			cidr, err := getCIDRMatchTerm(v)
+			if err != nil {
+				return fmt.Errorf("operand %v: %v", operand, err)
+			}
+			return iter(cidr, k)
+		})
+	}
+	return nil
+}
+
+func builtinNetCIDRContainsMatches(bctx BuiltinContext, args []*ast.Term, iter func(*ast.Term) error) error {
+	result := ast.NewSet()
+	err := evalNetCIDRContainsMatchesOperand(1, args[0], func(cidr1 *ast.Term, index1 *ast.Term) error {
+		return evalNetCIDRContainsMatchesOperand(2, args[1], func(cidr2 *ast.Term, index2 *ast.Term) error {
+			if v, err := builtinNetCIDRContains(cidr1.Value, cidr2.Value); err != nil {
+				return err
+			} else if vb, ok := v.(ast.Boolean); ok && bool(vb) {
+				result.Add(ast.ArrayTerm(index1, index2))
+			}
+			return nil
+		})
+	})
+	if err == nil {
+		return iter(ast.NewTerm(result))
+	}
+	return err
+}
+
 func builtinNetCIDRExpand(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
 
 	s, err := builtins.StringOperand(operands[0].Value, 1)
@@ -153,5 +223,6 @@ func init() {
 	RegisterFunctionalBuiltin2(ast.NetCIDROverlap.Name, builtinNetCIDRContains)
 	RegisterFunctionalBuiltin2(ast.NetCIDRIntersects.Name, builtinNetCIDRIntersects)
 	RegisterFunctionalBuiltin2(ast.NetCIDRContains.Name, builtinNetCIDRContains)
+	RegisterBuiltinFunc(ast.NetCIDRContainsMatches.Name, builtinNetCIDRContainsMatches)
 	RegisterBuiltinFunc(ast.NetCIDRExpand.Name, builtinNetCIDRExpand)
 }

topdown/cidr_test.go

@@ -75,6 +75,66 @@ func TestNetCIDRContains(t *testing.T) {
 	}
 }
 
+func TestNetCIDRContainsMatches(t *testing.T) {
+	tests := []struct {
+		note     string
+		rules    []string
+		expected interface{}
+	}{
+		{
+			note:     "strings",
+			rules:    []string{`p = x { x := net.cidr_contains_matches("1.1.1.0/24", "1.1.1.1") }`},
+			expected: `[["1.1.1.0/24", "1.1.1.1"]]`,
+		},
+		{
+			note:     "arrays",
+			rules:    []string{`p = x { x := net.cidr_contains_matches(["1.1.2.0/24", "1.1.1.0/24"], ["1.1.1.1", "1.1.2.1"]) }`},
+			expected: `[[0,1], [1,0]]`,
+		},
+		{
+			note:     "arrays of tuples",
+			rules:    []string{`p = x { x := net.cidr_contains_matches([["1.1.2.0/24", 1], "1.1.1.0/24"], ["1.1.1.1", "1.1.2.1"]) }`},
+			expected: `[[0,1], [1,0]]`,
+		},
+		{
+			note:     "bad array",
+			rules:    []string{`p = x { x := net.cidr_contains_matches(["1.1.2.0/24", "1.1.1.0/24"], ["1.1.1.1", data.a[0]]) }`},
+			expected: &Error{Code: BuiltinErr, Message: "net.cidr_contains_matches: operand 2: element must be string or non-empty array"},
+		},
+		{
+			note:     "sets of strings",
+			rules:    []string{`p = x { x := net.cidr_contains_matches({"1.1.2.0/24", "1.1.1.0/24"}, {"1.1.1.1", "1.1.2.1"}) }`},
+			expected: `[["1.1.1.0/24", "1.1.1.1"], ["1.1.2.0/24", "1.1.2.1"]]`,
+		},
+		{
+			note:     "sets of tuples",
+			rules:    []string{`p = x { x := net.cidr_contains_matches({["1.1.2.0/24", "foo"], ["1.1.1.0/24", "bar"]}, {["1.1.1.1", "baz"], ["1.1.2.1", "qux"]}) }`},
+			expected: `[[["1.1.1.0/24", "bar"], ["1.1.1.1", "baz"]], [["1.1.2.0/24", "foo"], ["1.1.2.1", "qux"]]]`,
+		},
+		{
+			note:     "bad set",
+			rules:    []string{`p = x { x := net.cidr_contains_matches({["1.1.2.0/24", "foo"], ["1.1.1.0/24", "bar"]}, {data.a[0], ["1.1.2.1", "qux"]}) }`},
+			expected: &Error{Code: BuiltinErr, Message: `net.cidr_contains_matches: operand 2: element must be string or non-empty array`},
+		},
+		{
+			note:     "bad set tuple element",
+			rules:    []string{`p = x { x := net.cidr_contains_matches({["1.1.2.0/24", "foo"], ["1.1.1.0/24", "bar"]}, {[], ["1.1.2.1", "qux"]}) }`},
+			expected: &Error{Code: BuiltinErr, Message: `net.cidr_contains_matches: operand 2: element must be string or non-empty array`},
+		},
+		{
+			note:     "objects",
+			rules:    []string{`p = x { x := net.cidr_contains_matches({"k1": "1.1.1.1/24", "k2": ["1.1.1.2/24", 1]}, "1.1.1.128") }`},
+			expected: `[["k1", "1.1.1.128"], ["k2", "1.1.1.128"]]`,
+		},
+	}
+
+	data := loadSmallTestData()
+
+	for _, tc := range tests {
+		runTopDownTestCase(t, data, tc.note, tc.rules, tc.expected)
+	}
+}
+
 func TestNetCIDRExpand(t *testing.T) {
 	tests := []struct {
 		note     string