wasm: Add native support for json.filter builtin function

Signed-off-by: Ashutosh Narkar <[email protected]>

docs/content/policy-reference.md

@@ -314,7 +314,7 @@ complex types.
 | <span class="opa-keep-it-together">`output := object.remove(object, keys)`</span> | `output` is a new object which is the result of removing the specified `keys` from `object`. `keys` must be either an array, object, or set of keys. | ✅ |
 | <span class="opa-keep-it-together">`output := object.union(objectA, objectB)`</span> | `output` is a new object which is the result of an asymmetric recursive union of two objects where conflicts are resolved by choosing the key from the right-hand object (`objectB`). For example: `object.union({"a": 1, "b": 2, "c": {"d": 3}}, {"a": 7, "c": {"d": 4, "e": 5}})` will result in `{"a": 7, "b": 2, "c": {"d": 4, "e": 5}}`  | ✅ |
 | <span class="opa-keep-it-together">`filtered := object.filter(object, keys)`</span> | `filtered` is a new object with the remaining data from `object` with only keys specified in `keys` which is an array, object, or set of keys. For example: `object.filter({"a": {"b": "x", "c": "y"}, "d": "z"}, ["a"])` will result in `{"a": {"b": "x", "c": "y"}}`). | ✅ |
-| <span class="opa-keep-it-together">`filtered := json.filter(object, paths)`</span> | `filtered` is the remaining data from `object` with only keys specified in `paths` which is an array or set of JSON string paths. For example: `json.filter({"a": {"b": "x", "c": "y"}}, ["a/b"])` will result in `{"a": {"b": "x"}}`). Paths are not filtered in-order and are deduplicated before being evaluated. | ``SDK-dependent`` |
+| <span class="opa-keep-it-together">`filtered := json.filter(object, paths)`</span> | `filtered` is the remaining data from `object` with only keys specified in `paths` which is an array or set of JSON string paths. For example: `json.filter({"a": {"b": "x", "c": "y"}}, ["a/b"])` will result in `{"a": {"b": "x"}}`). Paths are not filtered in-order and are deduplicated before being evaluated. | ✅ |
 | <span class="opa-keep-it-together">`output := json.remove(object, paths)`</span> | `output` is a new object which is the result of removing all keys specified in `paths` which is an array or set of JSON string paths. For example: `json.remove({"a": {"b": "x", "c": "y"}}, ["a/b"])` will result in `{"a": {"c": "y"}}`. Paths are not removed in-order and are deduplicated before being evaluated. | ✅ |
 | <span class="opa-keep-it-together">`output := json.patch(object, patches)`</span> | `output` is a the object obtained after consecutively applying all [JSON Patch](https://tools.ietf.org/html/rfc6902) operations in the array `patches`. For example: `json.patch({"a": {"foo": 1}}, [{"op": "add", "path": "/a/bar", "value": 2}])` results in `{"a": {"foo": 1, "bar": 2}`.  The patches are applied atomically: if any of them fails, the result will be undefined. | ``SDK-dependent`` |
 

internal/compiler/wasm/wasm.go

@@ -144,6 +144,7 @@ var builtinsFunctions = map[string]string{
 	ast.RegexMatchDeprecated.Name:       "opa_regex_match",
 	ast.RegexFindAllStringSubmatch.Name: "opa_regex_find_all_string_submatch",
 	ast.JSONRemove.Name:                 "builtin_json_remove",
+	ast.JSONFilter.Name:                 "builtin_json_filter",
 }
 
 var builtinDispatchers = [...]string{

test/wasm/assets/018_builtins.yaml

@@ -179,6 +179,12 @@ cases:
   - note: json.remove built-in
     query: 'json.remove({"a": {"b": {"c": 7, "d": 8}}, "e": 9}, {"a/b/c", "e"}, x)'
     want_result: [{'x': {"a": {"b": {"d": 8}}}}]
+  - note: json.filter built-in
+    query: 'json.filter({"a": {"b": {"c": 7, "d": 8}}, "e": 9}, {"a/b/c"}, x)'
+    want_result: [{'x': {"a": {"b": {"c": 7}}}}]
+  - note: json.filter built-in
+    query: 'json.filter({"a": {"b": {"c": 7, "d": 8}}, "e": 9}, {["a", "b", "c"], ["e"]}, x)'
+    want_result: [{'x': {"a": {"b": {"c": 7}}, "e": 9}}]
   - note: concat built-in
     query: concat(",",["a","b"],x)
     want_result: [{'x': "a,b"}]

wasm/src/object.c

@@ -8,6 +8,7 @@ opa_array_t *__get_json_paths(opa_value *a);
 opa_object_t *__paths_to_object(opa_value *a);
 opa_array_t *__parse_path(opa_value *a);
 opa_value *__json_remove(opa_value *a, opa_value *b);
+opa_value *__json_filter(opa_value *a, opa_value *b);
 
 opa_value *__merge(opa_value *a, opa_value *b)
 {
@@ -295,6 +296,101 @@ opa_value *__json_remove(opa_value *a, opa_value *b)
     return NULL;
 }
 
+opa_value *__json_filter(opa_value *a, opa_value *b)
+{
+
+    if (opa_value_compare(b, opa_null()) == 0)
+    {
+        return a;
+    }
+
+    if (opa_value_type(b) != OPA_OBJECT)
+    {
+        return NULL;
+    }
+
+    switch (opa_value_type(a))
+    {
+        case OPA_STRING:
+        case OPA_NUMBER:
+        case OPA_BOOLEAN:
+        case OPA_NULL:
+        {
+            return a;
+        }
+        case OPA_OBJECT:
+        {
+            opa_object_t *new_obj = opa_cast_object(opa_object());
+
+            opa_object_t *iter_obj = opa_cast_object(a);
+            opa_object_t *other = opa_cast_object(b);
+
+            if (iter_obj->len < other->len)
+            {
+                iter_obj = opa_cast_object(b);
+                other = opa_cast_object(a);
+            }
+
+            for (opa_value *key = opa_value_iter(&iter_obj->hdr, NULL); key != NULL; key = opa_value_iter(&iter_obj->hdr, key))
+            {
+
+                if (opa_value_get(&other->hdr, key) != NULL)
+                {
+                    opa_value *filtered_value = __json_filter(opa_value_get(a, key), opa_value_get(b, key));
+
+                    if (filtered_value != NULL)
+                    {
+                        opa_object_insert(new_obj, key, filtered_value);
+                    }
+                }
+            }
+            return &new_obj->hdr;
+        }
+        case OPA_SET:
+        {
+            opa_set_t *new_set = opa_cast_set(opa_set());
+            opa_set_t *set = opa_cast_set(a);
+
+            for (int i = 0; i < set->n; i++)
+            {
+                opa_set_elem_t *elem = set->buckets[i];
+
+                while (elem != NULL)
+                {
+                    opa_value *filtered_value = __json_filter(elem->v, opa_value_get(b, elem->v));
+
+                    if (filtered_value != NULL)
+                    {
+                        opa_set_add(new_set, filtered_value);
+                    }
+                    elem = elem->next;
+                }
+            }
+            return &new_set->hdr;
+        }
+        case OPA_ARRAY:
+        {
+            opa_array_t *new_array = opa_cast_array(opa_array());
+            opa_array_t *array = opa_cast_array(a);
+
+            for (int i = 0; i < array->len; i++)
+            {
+                opa_value *value = array->elems[i].v;
+
+                opa_value *filtered_value = __json_filter(value, opa_value_get(b, opa_strings_format_int(opa_number_int(i), opa_number_int(10))));
+
+                if (filtered_value != NULL)
+                {
+                    opa_array_append(new_array, filtered_value);
+                }
+            }
+            return &new_array->hdr;
+        }
+    }
+
+    return NULL;
+}
+
 OPA_BUILTIN
 opa_value *builtin_object_filter(opa_value *obj, opa_value *keys)
 {
@@ -443,3 +539,31 @@ opa_value *builtin_json_remove(opa_value *obj, opa_value *paths)
 
     return r;
 }
+
+OPA_BUILTIN
+opa_value *builtin_json_filter(opa_value *obj, opa_value *paths)
+{
+    if (opa_value_type(obj) != OPA_OBJECT)
+    {
+        return NULL;
+    }
+
+    if (opa_value_type(paths) != OPA_ARRAY && opa_value_type(paths) != OPA_SET)
+    {
+        return NULL;
+    }
+
+    // Build a list of filter strings
+    opa_array_t *json_paths = __get_json_paths(paths);
+
+    if (json_paths == NULL)
+    {
+        return NULL;
+    }
+
+    opa_object_t *json_paths_obj = __paths_to_object(&json_paths->hdr);
+
+    opa_value *r = __json_filter(obj, &json_paths_obj->hdr);
+
+    return r;
+}

wasm/src/object.h

@@ -9,5 +9,6 @@ opa_value *builtin_object_get(opa_value *obj, opa_value *key, opa_value *value);
 opa_value *builtin_object_remove(opa_value *obj, opa_value *keys);
 opa_value *builtin_object_union(opa_value *a, opa_value *b);
 opa_value *builtin_json_remove(opa_value *obj, opa_value *paths);
+opa_value *builtin_json_filter(opa_value *obj, opa_value *paths);
 
 #endif