Fix calling ObReferenceObjectByHandle at DISPATCH (#840)

* bugfix, add test cases

* fix build break

* cr comments

* cr comments

* cr comments

libs/execution_context/ebpf_maps.c

@@ -496,20 +496,13 @@ _check_value_type(_In_ const ebpf_core_map_t* outer_map, _In_ const ebpf_core_ob
     return allowed;
 }
 
-// Validate that a value handle is appropriate for this map,
-// and if so, return a pointer to the object with that handle.
-static ebpf_result_t
-_get_map_value_object(
-    _In_ const ebpf_core_map_t* map,
-    ebpf_handle_t value_handle,
-    ebpf_object_type_t value_type,
-    _Outptr_ ebpf_core_object_t** value_object_result)
+// Validate that a value object is appropriate for this map.
+// Also set the program type if not yet set.
+static _Requires_lock_held_(object_map->lock) ebpf_result_t _validate_map_value_object(
+    _In_ ebpf_core_object_map_t* object_map, ebpf_object_type_t value_type, _In_ const ebpf_core_object_t* value_object)
 {
-    // Convert value handle to an object pointer.
-    ebpf_core_object_t* value_object = NULL;
-    ebpf_result_t result = ebpf_reference_object_by_handle(value_handle, value_type, &value_object);
-    if (result != EBPF_SUCCESS)
-        return result;
+    ebpf_result_t result = EBPF_SUCCESS;
+    const ebpf_core_map_t* map = &object_map->core_map;
 
     const ebpf_program_type_t* value_program_type =
         (value_object->get_program_type) ? value_object->get_program_type(value_object) : NULL;
@@ -525,21 +518,18 @@ _get_map_value_object(
     // Validate that the value's program type (if any) is
     // not in conflict with the map's program type.
     if (value_program_type != NULL) {
-        ebpf_core_object_map_t* map_of_objects = (ebpf_core_object_map_t*)map;
-        if (!map_of_objects->is_program_type_set) {
-            map_of_objects->is_program_type_set = TRUE;
-            map_of_objects->program_type = *value_program_type;
-        } else if (memcmp(&map_of_objects->program_type, value_program_type, sizeof(*value_program_type)) != 0) {
+        if (!object_map->is_program_type_set) {
+            object_map->is_program_type_set = TRUE;
+            object_map->program_type = *value_program_type;
+        } else if (memcmp(&object_map->program_type, value_program_type, sizeof(*value_program_type)) != 0) {
             result = EBPF_INVALID_FD;
             goto Error;
         }
     }
 
-    *value_object_result = value_object;
     return EBPF_SUCCESS;
 
 Error:
-    ebpf_object_release_reference((ebpf_core_object_t*)value_object);
     return result;
 }
 
@@ -562,14 +552,19 @@ _update_array_map_entry_with_handle(
     ebpf_result_t result = EBPF_SUCCESS;
 
     ebpf_core_object_map_t* object_map = EBPF_FROM_FIELD(ebpf_core_object_map_t, core_map, map);
+    ebpf_core_object_t* value_object = NULL;
+
+    if (value_handle != (uintptr_t)ebpf_handle_invalid) {
+        result = ebpf_reference_object_by_handle(value_handle, value_type, &value_object);
+        if (result != EBPF_SUCCESS) {
+            return result;
+        }
+    }
 
     ebpf_lock_state_t lock_state = ebpf_lock_lock(&object_map->lock);
 
-    ebpf_core_object_t* value_object = NULL;
-    // If value_handle is valid, resolve it to an object. Else we just need to clear
-    // the existing entry from the map.
     if (value_handle != (uintptr_t)ebpf_handle_invalid) {
-        result = _get_map_value_object(map, value_handle, value_type, &value_object);
+        result = _validate_map_value_object(object_map, value_type, value_object);
         if (result != EBPF_SUCCESS) {
             goto Done;
         }
@@ -588,6 +583,9 @@ _update_array_map_entry_with_handle(
     memcpy(entry, &id, map->ebpf_map_definition.value_size);
 
 Done:
+    if (result != EBPF_SUCCESS && value_object != NULL) {
+        ebpf_object_release_reference((ebpf_core_object_t*)value_object);
+    }
     ebpf_lock_unlock(&object_map->lock, lock_state);
 
     return result;
@@ -1063,6 +1061,11 @@ _update_hash_map_entry_with_handle(
     }
 
     ebpf_core_object_map_t* object_map = EBPF_FROM_FIELD(ebpf_core_object_map_t, core_map, map);
+    ebpf_core_object_t* value_object = NULL;
+    result = ebpf_reference_object_by_handle(value_handle, value_type, &value_object);
+    if (result != EBPF_SUCCESS) {
+        return result;
+    }
 
     ebpf_lock_state_t lock_state = ebpf_lock_lock(&object_map->lock);
 
@@ -1071,7 +1074,6 @@ _update_hash_map_entry_with_handle(
     uint8_t* old_value = NULL;
     ebpf_result_t found_result = ebpf_hash_table_find((ebpf_hash_table_t*)map->data, key, &old_value);
     ebpf_id_t old_id = (old_value) ? *(ebpf_id_t*)old_value : 0;
-    ebpf_core_object_t* value_object = NULL;
 
     if ((entry_count == map->ebpf_map_definition.max_entries) && (found_result != EBPF_SUCCESS)) {
         // The hash table is already full.
@@ -1079,7 +1081,7 @@ _update_hash_map_entry_with_handle(
         goto Done;
     }
 
-    result = _get_map_value_object(map, value_handle, value_type, &value_object);
+    result = _validate_map_value_object(object_map, value_type, value_object);
     if (result != EBPF_SUCCESS) {
         goto Done;
     }

libs/platform/ebpf_handle.h

@@ -59,9 +59,8 @@ extern "C"
      * @retval EBPF_SUCCESS The operation was successful.
      * @retval EBPF_INVALID_OBJECT The provided handle is not valid.
      */
-    ebpf_result_t
-    ebpf_reference_object_by_handle(
-        ebpf_handle_t handle, ebpf_object_type_t object_type, struct _ebpf_core_object** object);
+    _IRQL_requires_max_(PASSIVE_LEVEL) ebpf_result_t ebpf_reference_object_by_handle(
+        ebpf_handle_t handle, ebpf_object_type_t object_type, _Outptr_ struct _ebpf_core_object** object);
 
 #ifdef __cplusplus
 }

libs/platform/kernel/ebpf_handle_kernel.c

@@ -87,8 +87,8 @@ ebpf_handle_close(ebpf_handle_t handle)
         EBPF_RETURN_RESULT(EBPF_SUCCESS);
 }
 
-ebpf_result_t
-ebpf_reference_object_by_handle(ebpf_handle_t handle, ebpf_object_type_t object_type, ebpf_core_object_t** object)
+_IRQL_requires_max_(PASSIVE_LEVEL) ebpf_result_t ebpf_reference_object_by_handle(
+    ebpf_handle_t handle, ebpf_object_type_t object_type, _Outptr_ ebpf_core_object_t** object)
 {
     ebpf_result_t return_value;
     NTSTATUS status;

libs/platform/user/ebpf_handle_user.c

@@ -84,8 +84,8 @@ ebpf_handle_close(ebpf_handle_t handle)
     return return_value;
 }
 
-ebpf_result_t
-ebpf_reference_object_by_handle(ebpf_handle_t handle, ebpf_object_type_t object_type, ebpf_core_object_t** object)
+_IRQL_requires_max_(PASSIVE_LEVEL) ebpf_result_t ebpf_reference_object_by_handle(
+    ebpf_handle_t handle, ebpf_object_type_t object_type, _Outptr_ ebpf_core_object_t** object)
 {
     ebpf_result_t return_value;
     ebpf_lock_state_t state;

tests/api_test/api_test.cpp

@@ -314,3 +314,84 @@ TEST_CASE("divide_by_zero_jit", "[divide_by_zero]") { divide_by_zero_test_km(EBP
 TEST_CASE("ringbuf_api_interpret", "[test_ringbuf_api]") { ring_buffer_api_test(EBPF_EXECUTION_INTERPRET); }
 TEST_CASE("divide_by_zero_interpret", "[divide_by_zero]") { divide_by_zero_test_km(EBPF_EXECUTION_INTERPRET); }
 #endif
+
+void
+_test_nested_maps(bpf_map_type type)
+{
+    // Create first inner map.
+    fd_t inner1 = bpf_create_map(BPF_MAP_TYPE_ARRAY, sizeof(uint32_t), sizeof(uint32_t), 1, 0);
+    REQUIRE(inner1 > 0);
+
+    // Create outer map.
+    fd_t outer_map_fd = bpf_create_map_in_map(type, "outer_map", sizeof(uint32_t), inner1, 10, 0);
+    REQUIRE(outer_map_fd > 0);
+
+    // Create second inner map.
+    fd_t inner2 = bpf_create_map(BPF_MAP_TYPE_ARRAY, sizeof(uint32_t), sizeof(uint32_t), 1, 0);
+    REQUIRE(inner2 > 0);
+
+    // Insert both inner maps in outer map.
+    uint32_t key = 1;
+    uint32_t result = bpf_map_update_elem(outer_map_fd, &key, &inner1, 0);
+    REQUIRE(result == ERROR_SUCCESS);
+
+    key = 2;
+    result = bpf_map_update_elem(outer_map_fd, &key, &inner1, 0);
+    REQUIRE(result == ERROR_SUCCESS);
+
+    // Remove the inner maps from outer map.
+    key = 1;
+    result = bpf_map_delete_elem(outer_map_fd, &key);
+    REQUIRE(result == ERROR_SUCCESS);
+    key = 2;
+    result = bpf_map_delete_elem(outer_map_fd, &key);
+    REQUIRE(result == ERROR_SUCCESS);
+
+    _close(inner1);
+    _close(inner2);
+    _close(outer_map_fd);
+}
+
+TEST_CASE("array_of_maps", "[map_in_map]") { _test_nested_maps(BPF_MAP_TYPE_ARRAY_OF_MAPS); }
+
+TEST_CASE("hash_of_maps", "[map_in_map]") { _test_nested_maps(BPF_MAP_TYPE_HASH_OF_MAPS); }
+
+TEST_CASE("tailcall_load_test", "[tailcall_load_test]")
+{
+    ebpf_result_t result;
+    struct bpf_object* object = nullptr;
+    fd_t program_fd;
+
+    result =
+        _program_load_helper("tail_call_multiple.o", &EBPF_PROGRAM_TYPE_XDP, EBPF_EXECUTION_ANY, &object, &program_fd);
+    REQUIRE(result == EBPF_SUCCESS);
+
+    REQUIRE(program_fd > 0);
+
+    // Set up tail calls.
+    struct bpf_program* callee0 = bpf_object__find_program_by_name(object, "callee0");
+    REQUIRE(callee0 != nullptr);
+    fd_t callee0_fd = bpf_program__fd(callee0);
+    REQUIRE(callee0_fd > 0);
+
+    struct bpf_program* callee1 = bpf_object__find_program_by_name(object, "callee1");
+    REQUIRE(callee1 != nullptr);
+    fd_t callee1_fd = bpf_program__fd(callee1);
+    REQUIRE(callee1_fd > 0);
+
+    fd_t prog_map_fd = bpf_object__find_map_fd_by_name(object, "map");
+    REQUIRE(prog_map_fd > 0);
+
+    uint32_t index = 0;
+    REQUIRE(bpf_map_update_elem(prog_map_fd, &index, &callee0_fd, 0) == 0);
+    index = 1;
+    REQUIRE(bpf_map_update_elem(prog_map_fd, &index, &callee1_fd, 0) == 0);
+
+    // Cleanup tail calls.
+    index = 0;
+    REQUIRE(bpf_map_update_elem(prog_map_fd, &index, &ebpf_fd_invalid, 0) == 0);
+    index = 1;
+    REQUIRE(bpf_map_update_elem(prog_map_fd, &index, &ebpf_fd_invalid, 0) == 0);
+
+    bpf_object__close(object);
+}