Fix port_quota app usage (#1184)

* port quota Signed-off-by: Dave Thaler <[email protected]> * Fix port_quota load Signed-off-by: Dave Thaler <[email protected]> Co-authored-by: Alan Jowett <[email protected]>
microsoft · Jun 9, 2022 · 25abe8f · 25abe8f
1 parent e870a29
commit 25abe8f
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 10 deletions.
diff --git a/libs/execution_context/ebpf_core.c b/libs/execution_context/ebpf_core.c
@@ -353,14 +353,14 @@ ebpf_core_resolve_maps(
 {
     EBPF_LOG_ENTRY();
     ebpf_program_t* program = NULL;
-    memset(map_addresses, 0, sizeof(uintptr_t) * count_of_maps);
+    uint32_t map_index = 0;
 
     ebpf_result_t return_value =
         ebpf_reference_object_by_handle(program_handle, EBPF_OBJECT_PROGRAM, (ebpf_core_object_t**)&program);
     if (return_value != EBPF_SUCCESS)
         goto Done;
 
-    for (uint32_t map_index = 0; map_index < count_of_maps; map_index++) {
+    for (map_index = 0; map_index < count_of_maps; map_index++) {
         ebpf_map_t* map;
         return_value =
             ebpf_reference_object_by_handle(map_handles[map_index], EBPF_OBJECT_MAP, (ebpf_core_object_t**)&map);
@@ -375,10 +375,8 @@ ebpf_core_resolve_maps(
 
 Done:
     // Release our reference only after the map has been associated with the program.
-    for (uint32_t map_index = 0; map_index < count_of_maps; map_index++) {
-        if (map_addresses[map_index]) {
-            ebpf_object_release_reference((ebpf_core_object_t*)map_addresses[map_index]);
-        }
+    for (uint32_t map_index2 = 0; map_index2 < map_index; map_index2++) {
+        ebpf_object_release_reference((ebpf_core_object_t*)map_addresses[map_index2]);
     }
 
     ebpf_object_release_reference((ebpf_core_object_t*)program);

diff --git a/scripts/deploy-ebpf.ps1 b/scripts/deploy-ebpf.ps1
@@ -68,6 +68,7 @@ $build_directory=".\x64\Debug"
     "droppacket_unsafe.o",
     "EbpfApi.pdb",
     "ebpfnetsh.pdb",
+    "ebpfsvc.pdb",
     "encap_reflect_packet.o",
     "encap_reflect_packet_um.dll",
     "encap_reflect_packet_um.pdb",
@@ -94,6 +95,8 @@ $build_directory=".\x64\Debug"
     "map.sys",
     "pidtgid.o",
     "pidtgid.sys",
+    "port_quota.exe",
+    "port_quota.pdb",
     "printk.o",
     "printk_legacy.o",
     "printk_legacy_um.dll",

diff --git a/tests/end_to_end/test_helper.cpp b/tests/end_to_end/test_helper.cpp
@@ -344,6 +344,8 @@ GlueDeviceIoControl(
     size_t minimum_request_size = 0;
     size_t minimum_reply_size = 0;
     bool async = false;
+    DWORD sharedBufferSize = (nInBufferSize > nOutBufferSize) ? nInBufferSize : nOutBufferSize;
+    std::vector<uint8_t> sharedBuffer(sharedBufferSize);
 
     result = ebpf_core_get_protocol_handler_properties(request_id, &minimum_request_size, &minimum_reply_size, &async);
     if (result != EBPF_SUCCESS)
@@ -369,15 +371,24 @@ GlueDeviceIoControl(
     // Intercept the call to perform any IOCTL specific _pre_ tasks.
     _preprocess_ioctl(user_request);
 
+    // In the kernel execution context, the request and reply share
+    // the same memory.  So to catch bugs that only show up in that
+    // case, we force the same here.
+    memcpy(sharedBuffer.data(), user_request, nInBufferSize);
+
     result = ebpf_core_invoke_protocol_handler(
         request_id,
-        user_request,
+        sharedBuffer.data(),
         static_cast<uint16_t>(nInBufferSize),
-        user_reply,
+        (minimum_reply_size > 0) ? sharedBuffer.data() : nullptr,
         static_cast<uint16_t>(nOutBufferSize),
         lpOverlapped,
         _complete_overlapped);
 
+    if (minimum_reply_size > 0) {
+        memcpy(user_reply, sharedBuffer.data(), nOutBufferSize);
+    }
+
     if (result != EBPF_SUCCESS)
         goto Fail;
 

diff --git a/tools/port_quota/port_quota.cpp b/tools/port_quota/port_quota.cpp
@@ -37,7 +37,7 @@ load(int argc, char** argv)
         return 1;
     }
 
-    ebpf_object_set_execution_type(object, EBPF_EXECUTION_INTERPRET);
+    ebpf_object_set_execution_type(object, EBPF_EXECUTION_JIT);
     program = bpf_object__next_program(object, nullptr);
     if (bpf_object__load(object) < 0) {
         fprintf(stderr, "Failed to load port quota eBPF program\n");
@@ -202,4 +202,4 @@ main(int argc, char** argv)
     }
     print_usage(argv[0]);
     return 1;
-}
+}