[draft] [test] Fix for concurrent multi-attach test

libs/execution_context/ebpf_core.c

@@ -1523,12 +1523,15 @@ _ebpf_core_protocol_unlink_program(_In_ const ebpf_operation_unlink_program_requ
     EBPF_LOG_ENTRY();
     ebpf_result_t retval = EBPF_SUCCESS;
     ebpf_link_t* link = NULL;
+    bool detach_single_link = false;
+    bool detached_any_link = false;
 
     if (request->link_handle != ebpf_handle_invalid) {
         retval = EBPF_OBJECT_REFERENCE_BY_HANDLE(request->link_handle, EBPF_OBJECT_LINK, (ebpf_core_object_t**)&link);
         if (retval != EBPF_SUCCESS) {
             goto Done;
         }
+        detach_single_link = true;
     } else if (request->attach_data_present) {
         // This path will be taken for bpf_prog_detach and bpf_prog_detach2 APIs.
         // Find the link object matching the unlink request parameters.
@@ -1549,17 +1552,18 @@ _ebpf_core_protocol_unlink_program(_In_ const ebpf_operation_unlink_program_requ
             // Detach the link. Since _ebpf_core_find_matching_link takes a reference on the link object,
             // the detach function will not free the link object.
             ebpf_link_detach_program(link);
+            detached_any_link = true;
             // Pass the link object as the previous object parameter to the _ebpf_core_find_matching_link function,
             // which will release the reference from it.
             previous_link = link;
         }
         if (retval == EBPF_NO_MORE_KEYS) {
             // No more matching links to detach.
-            retval = EBPF_SUCCESS;
+            retval = detached_any_link ? EBPF_SUCCESS : EBPF_OBJECT_NOT_FOUND;
         }
     }
 
-    if (link != NULL) {
+    if (detach_single_link && link != NULL) {
         ebpf_link_detach_program(link);
     }
 

tests/connect_redirect/connect_redirect_tests.cpp

@@ -251,7 +251,7 @@ _initialize_test_globals()
 
     // Load and attach the programs.
     native_module_helper_t helper;
-    helper.initialize("cgroup_sock_addr2");
+    helper.initialize("cgroup_sock_addr2", EBPF_EXECUTION_ANY, true);
     _globals.bpf_object.reset(bpf_object__open(helper.get_file_name().c_str()));
     SAFE_REQUIRE(_globals.bpf_object.get() != nullptr);
     SAFE_REQUIRE(bpf_object__load(_globals.bpf_object.get()) == 0);

tests/libs/util/native_helper.cpp

@@ -23,20 +23,6 @@ _native_module_helper::initialize(
 #else
     ebpf_execution_type_t system_default = ebpf_execution_type_t::EBPF_EXECUTION_ANY;
 #endif
-    // Clean up any previous state.
-    // Detach all bpf links.
-    uint32_t link_id;
-    while (bpf_link_get_next_id(0, &link_id) == 0) {
-        fd_t link_fd = bpf_link_get_fd_by_id(link_id);
-        if (link_fd < 0) {
-            break;
-        }
-        bpf_link_detach(link_fd);
-        if (link_fd >= 0) {
-            (void)ebpf_close_fd(link_fd);
-        }
-    }
-
     // Set _is_main_thread before any REQUIRE is invoked.
     _is_main_thread = is_main_thread;
 

tests/socket/socket_tests.cpp

@@ -457,7 +457,8 @@ execute_connection_test(_In_ const connection_test_case& test_case)
                 ebpf_attach_type_t attach_type_guid{};
                 SAFE_REQUIRE(
                     ebpf_get_ebpf_attach_type(loaded_program.spec.attach_type, &attach_type_guid) == EBPF_SUCCESS);
-                SAFE_REQUIRE(ebpf_program_attach(program, &attach_type_guid, nullptr, 0, nullptr) == EBPF_SUCCESS);
+                SAFE_REQUIRE(
+                    ebpf_program_attach(program, &attach_type_guid, nullptr, 0, &loaded_program.link) == EBPF_SUCCESS);
             }
         }
     }
@@ -590,6 +591,18 @@ execute_connection_test(_In_ const connection_test_case& test_case)
 
         ++test_index;
     }
+
+    // Detach and clean up all attached programs before unloading objects.
+    for (auto& mod : loaded_modules) {
+        for (auto& loaded_program : mod.programs) {
+            if (loaded_program.link != nullptr) {
+                bpf_link__destroy(loaded_program.link);
+                loaded_program.link = nullptr;
+            } else if (loaded_program.spec.attach_method == attach_method_t::bpf_prog_attach) {
+                bpf_prog_detach2(bpf_program__fd(loaded_program.program), 0, loaded_program.spec.attach_type);
+            }
+        }
+    }
 }
 
 // Type tuples for TEMPLATE_TEST_CASE: (address_family, protocol).
@@ -1077,6 +1090,10 @@ sock_addr_bind_unknown_verdict_test(ADDRESS_FAMILY address_family, IPPROTO proto
     int rc = bind(sock, reinterpret_cast<sockaddr*>(&bind_addr), sizeof(bind_addr));
     int err = (rc == 0) ? 0 : WSAGetLastError();
     closesocket(sock);
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(program)), 0, attach_type);
+
     SAFE_REQUIRE(rc != 0);
     SAFE_REQUIRE(err == WSAEACCES);
 }
@@ -1226,6 +1243,10 @@ bind_helper_functions_validation_test(ADDRESS_FAMILY address_family)
     SAFE_REQUIRE(results.socket_cookie == 0);
 
     closesocket(sock);
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(bind_program)), 0, attach_type);
+
     printf(
         "Bind helper functions validation test completed successfully for %s\n",
         (address_family == AF_INET) ? "IPv4" : "IPv6");
@@ -1348,6 +1369,12 @@ helper_functions_validation_test(
     printf(
         "Helper functions validation test completed successfully for %s\n",
         (address_family == AF_INET) ? "IPv4" : "IPv6");
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(
+        bpf_program__fd(const_cast<const bpf_program*>(connect_authorization_program)),
+        0,
+        connect_authorization_attach_type);
 }
 
 TEST_CASE("connect_authorization_helper_functions_validation_tcp_v4", "[sock_addr_tests][helper_validation]")
@@ -1424,6 +1451,12 @@ TEST_CASE(
         SAFE_REQUIRE(tunnel_count == 0);
     }
 
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(
+        bpf_program__fd(const_cast<const bpf_program*>(conditional_program)),
+        0,
+        BPF_CGROUP_INET4_CONNECT_AUTHORIZATION);
+
     printf("Conditional policy validation test completed successfully\n");
 }
 
@@ -1878,6 +1911,10 @@ TEST_CASE("listen_helper_functions_validation_tcp_v4", "[sock_addr_tests][helper
     SAFE_REQUIRE(results.socket_cookie == 0);
 
     closesocket(sock);
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(listen_program)), 0, BPF_CGROUP_INET4_LISTEN);
+
     printf("Listen helper functions validation test completed successfully for IPv4\n");
 }
 
@@ -1979,6 +2016,10 @@ TEST_CASE("listen_helper_functions_validation_tcp_v6", "[sock_addr_tests][helper
     SAFE_REQUIRE(results.socket_cookie == 0);
 
     closesocket(sock);
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(listen_program)), 0, BPF_CGROUP_INET6_LISTEN);
+
     printf("Listen helper functions validation test completed successfully for IPv6\n");
 }
 
@@ -2217,6 +2258,25 @@ TEST_CASE("attach_sock_addr_programs", "[sock_addr_tests]")
         0);
     SAFE_REQUIRE(result == 0);
 
+    ZeroMemory(&program_info, program_info_size);
+    SAFE_REQUIRE(
+        bpf_obj_get_info_by_fd(
+            bpf_program__fd(const_cast<const bpf_program*>(connect6_program)), &program_info, &program_info_size) == 0);
+    SAFE_REQUIRE(program_info.link_count == 1);
+    SAFE_REQUIRE(program_info.map_ids == 0);
+
+    result = bpf_prog_detach2(
+        bpf_program__fd(const_cast<const bpf_program*>(connect6_program)),
+        DEFAULT_COMPARTMENT_ID,
+        BPF_CGROUP_INET6_CONNECT);
+    SAFE_REQUIRE(result == 0);
+
+    ZeroMemory(&program_info, program_info_size);
+    SAFE_REQUIRE(
+        bpf_obj_get_info_by_fd(
+            bpf_program__fd(const_cast<const bpf_program*>(connect6_program)), &program_info, &program_info_size) == 0);
+    SAFE_REQUIRE(program_info.link_count == 0);
+
     bpf_program* recv_accept6_program = bpf_object__find_program_by_name(object, "authorize_recv_accept6");
     SAFE_REQUIRE(recv_accept6_program != nullptr);
 
@@ -2226,6 +2286,27 @@ TEST_CASE("attach_sock_addr_programs", "[sock_addr_tests]")
         BPF_CGROUP_INET6_RECV_ACCEPT,
         0);
     SAFE_REQUIRE(result == 0);
+
+    ZeroMemory(&program_info, program_info_size);
+    SAFE_REQUIRE(
+        bpf_obj_get_info_by_fd(
+            bpf_program__fd(const_cast<const bpf_program*>(recv_accept6_program)), &program_info, &program_info_size) ==
+        0);
+    SAFE_REQUIRE(program_info.link_count == 1);
+    SAFE_REQUIRE(program_info.map_ids == 0);
+
+    result = bpf_prog_detach2(
+        bpf_program__fd(const_cast<const bpf_program*>(recv_accept6_program)),
+        DEFAULT_COMPARTMENT_ID,
+        BPF_CGROUP_INET6_RECV_ACCEPT);
+    SAFE_REQUIRE(result == 0);
+
+    ZeroMemory(&program_info, program_info_size);
+    SAFE_REQUIRE(
+        bpf_obj_get_info_by_fd(
+            bpf_program__fd(const_cast<const bpf_program*>(recv_accept6_program)), &program_info, &program_info_size) ==
+        0);
+    SAFE_REQUIRE(program_info.link_count == 0);
 }
 
 void
@@ -2372,6 +2453,9 @@ connection_monitor_test(
 
     // Unsubscribe.
     context->unsubscribe();
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(_program)), 0, BPF_CGROUP_SOCK_OPS);
 }
 
 TEST_CASE("connection_monitor_test_udp_v4", "[sock_ops_tests]")
@@ -2450,6 +2534,9 @@ TEST_CASE("attach_sockops_programs", "[sock_ops_tests]")
 
     int result = bpf_prog_attach(bpf_program__fd(const_cast<const bpf_program*>(_program)), 0, BPF_CGROUP_SOCK_OPS, 0);
     SAFE_REQUIRE(result == 0);
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(_program)), 0, BPF_CGROUP_SOCK_OPS);
 }
 
 // Custom event handler for flow ID validation
@@ -2592,6 +2679,9 @@ TEST_CASE("sock_ops_flow_id_helper_test", "[sock_ops_tests]")
     // Verify we get a non-zero flow ID.
     REQUIRE(result == 0);
     REQUIRE(stored_flow_id != 0);
+
+    // Detach the program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(_program)), 0, BPF_CGROUP_SOCK_OPS);
 }
 
 // This function populates map policies for multi-attach tests.
@@ -2907,6 +2997,15 @@ multi_attach_test(uint32_t compartment_id, socket_family_t family, ADDRESS_FAMIL
     // Validate that the connection is allowed.
     validate_connection_multi_attach(
         family, address_family, SOCKET_TEST_PORT, SOCKET_TEST_PORT, protocol, RESULT_ALLOW);
+
+    // Detach all programs before unloading to avoid leaving stale legacy links.
+    for (uint32_t i = 0; i < MULTIPLE_ATTACH_PROGRAM_COUNT; i++) {
+        bpf_program* prog = bpf_object__find_program_by_name(objects[i], connect_program_name);
+        if (prog != nullptr) {
+            bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(prog)), compartment_id, attach_type);
+        }
+    }
+    bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(connect_program)), compartment_id + 2, attach_type);
 }
 
 void
@@ -3075,6 +3174,14 @@ multi_attach_test_redirection(
     for (uint32_t i = 0; i < MULTIPLE_ATTACH_PROGRAM_COUNT; i++) {
         validate_program_redirection(i);
     }
+
+    // Detach all programs before unloading to avoid leaving stale legacy links.
+    for (uint32_t i = 0; i < MULTIPLE_ATTACH_PROGRAM_COUNT; i++) {
+        bpf_program* program = bpf_object__find_program_by_name(objects[i], connect_program_name);
+        if (program != nullptr) {
+            bpf_prog_detach2(bpf_program__fd(const_cast<const bpf_program*>(program)), compartment_id, attach_type);
+        }
+    }
 }
 
 TEST_CASE("multi_attach_test_TCP_IPv4", "[sock_addr_tests][multi_attach_tests]")
@@ -3282,6 +3389,17 @@ test_multi_attach_combined(socket_family_t family, ADDRESS_FAMILY address_family
             break;
         }
     }
+
+    // Detach all programs before unloading to avoid leaving stale legacy links.
+    for (uint32_t i = 0; i < program_count_per_hook * 2; i++) {
+        bpf_program* connect_program = bpf_object__find_program_by_name(objects[i], connect_program_name);
+        if (connect_program != nullptr) {
+            bpf_prog_detach2(
+                bpf_program__fd(const_cast<const bpf_program*>(connect_program)),
+                i < program_count_per_hook ? 1 : UNSPECIFIED_COMPARTMENT_ID,
+                attach_type);
+        }
+    }
 }
 
 TEST_CASE("multi_attach_test_combined_TCP_IPV4", "[sock_addr_tests][multi_attach_tests]")
@@ -3554,6 +3672,12 @@ TEST_CASE("multi_attach_test_invocation_order", "[sock_addr_tests][multi_attach_
     // Since the specific program is now detached, the connection should be correctly redirected by wildcard program.
     validate_connection_multi_attach(
         family, address_family, SOCKET_TEST_PORT, destination_port, IPPROTO_TCP, RESULT_ALLOW);
+
+    // Detach the wildcard program before unloading to avoid leaving stale legacy links.
+    bpf_prog_detach2(
+        bpf_program__fd(const_cast<const bpf_program*>(connect_program_wildcard)),
+        UNSPECIFIED_COMPARTMENT_ID,
+        attach_type);
 }
 
 /**
@@ -3681,6 +3805,7 @@ thread_function_allow_block_connection(
     SAFE_REQUIRE(bpf_object__load(object) == 0);
 
     fd_t prog_fd = bpf_program__fd(const_cast<const bpf_program*>(connect_program));
+    bool attached = false;
 
     // Attach the program at BPF_CGROUP_INET4_CONNECT / BPF_CGROUP_INET6_CONNECT.
     int result = bpf_prog_attach(prog_fd, compartment_id, attach_type, 0);
@@ -3692,6 +3817,7 @@ thread_function_allow_block_connection(
             << " result=" << result << " errno=" << saved_errno;
         throw test_failure(oss.str());
     }
+    attached = true;
 
     // Configure policy map to allow the connection.
     bpf_map* policy_map = bpf_object__find_map_by_name(object, "policy_map");
@@ -3714,27 +3840,45 @@ thread_function_allow_block_connection(
     _update_map_entry_multi_attach(
         map_fd, address_family, htons(destination_port), htons(destination_port), protocol, true);
 
-    while (!token.stop_requested()) {
-        // Block the connection.
-        _update_map_entry_multi_attach(
-            map_fd, address_family, htons(destination_port), htons(destination_port), protocol, false);
+    try {
+        while (!token.stop_requested()) {
+            // Block the connection.
+            _update_map_entry_multi_attach(
+                map_fd, address_family, htons(destination_port), htons(destination_port), protocol, false);
 
-        // The connection should be blocked. Due to race, it can sometimes be allowed, so we don't care about the
-        // result.
-        validate_connection_multi_attach(
-            family, address_family, destination_port, destination_port, protocol, RESULT_DONT_CARE);
+            // The connection should be blocked. Due to race, it can sometimes be allowed, so we don't care about the
+            // result.
+            validate_connection_multi_attach(
+                family, address_family, destination_port, destination_port, protocol, RESULT_DONT_CARE);
 
-        // Allow the connection.
-        _update_map_entry_multi_attach(
-            map_fd, address_family, htons(destination_port), htons(destination_port), protocol, true);
+            // Allow the connection.
+            _update_map_entry_multi_attach(
+                map_fd, address_family, htons(destination_port), htons(destination_port), protocol, true);
 
-        // The connection should be allowed. Due to race, it can sometimes be blocked, so we don't care about the
-        // result.
-        validate_connection_multi_attach(
-            family, address_family, destination_port, destination_port, protocol, RESULT_DONT_CARE);
+            // The connection should be allowed. Due to race, it can sometimes be blocked, so we don't care about the
+            // result.
+            validate_connection_multi_attach(
+                family, address_family, destination_port, destination_port, protocol, RESULT_DONT_CARE);
 
-        count++;
+            count++;
+        }
+    } catch (...) {
+        if (attached) {
+            (void)bpf_prog_detach2(prog_fd, compartment_id, attach_type);
+        }
+        throw;
+    }
+
+    result = bpf_prog_detach2(prog_fd, compartment_id, attach_type);
+    if (result != 0) {
+        int saved_errno = errno;
+        std::ostringstream oss;
+        oss << "ALLOW_BLOCK DETACH FAILED: thread=" << std::this_thread::get_id() << " compartment=" << compartment_id
+            << " prog_fd=" << prog_fd << " attach_type=" << static_cast<int>(attach_type) << " protocol=" << protocol
+            << " result=" << result << " errno=" << saved_errno;
+        throw test_failure(oss.str());
     }
+    attached = false;
 
     std::cout << "Thread (allow_block)" << std::this_thread::get_id() << " executed " << count << " times."
               << std::endl;