Harden link state machine (#2396)

* Harden link state machine Signed-off-by: Alan Jowett <alanjo@microsoft.com> * Fix cmake release build Signed-off-by: Alan Jowett <alanjo@microsoft.com> --------- Signed-off-by: Alan Jowett <alanjo@microsoft.com>
microsoft · Apr 27, 2023 · 0314806 · 0314806
1 parent ebeeb96
commit 0314806
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 39 deletions.
diff --git a/libs/execution_context/ebpf_link.c b/libs/execution_context/ebpf_link.c
@@ -11,18 +11,25 @@
 
 /**
  * @brief State of the link between program and provider.
- * The link starts in IDLE state. When a program is attached to a provider, the
- * state transitions to ATTACHING. If the provider is not found, the state
- * transitions to IDLE. If the provider is found, the state transitions to
- * ATTACHED. When a program is detached from a provider, the state transitions
- * to DETACHING. Once the provider is notified, the state transitions to IDLE.
+ * Legal state transitions are:
+ * 1. EBPF_LINK_STATE_INITIAL -> EBPF_LINK_STATE_ATTACHING - In ebpf_link_attach_program when the program is being
+ * linked.
+ * 2. EBPF_LINK_STATE_ATTACHING -> EBPF_LINK_STATE_ATTACHED - In _ebpf_link_extension_changed_callback when the provider
+ * attaches.
+ * 3. EBPF_LINK_STATE_ATTACHED -> EBPF_LINK_STATE_ATTACHING - In _ebpf_link_extension_changed_callback when the provider
+ * is reattaching.
+ * 4. EBPF_LINK_STATE_ATTACHED -> EBPF_LINK_STATE_DETACHING - In ebpf_link_detach_program when the program is being
+ * unlinked.
+ * 5. EBPF_LINK_STATE_DETACHING -> EBPF_LINK_STATE_DETACHED - In _ebpf_link_extension_changed_callback when the provider
+ * detaches.
  */
 typedef enum _ebpf_link_state
 {
-    EBPF_LINK_STATE_IDLE,      ///< Program is not attached to any provider.
-    EBPF_LINK_STATE_ATTACHING, ///< Program is in the process of getting attached.
+    EBPF_LINK_STATE_INITIAL,   ///< Program is not attached to any provider.
+    EBPF_LINK_STATE_ATTACHING, ///< Program is being attached to a provider.
     EBPF_LINK_STATE_ATTACHED,  ///< Program is attached to a provider.
-    EBPF_LINK_STATE_DETACHING, ///< Program is in the process of getting detached.
+    EBPF_LINK_STATE_DETACHING, ///< Program is being detached from a provider.
+    EBPF_LINK_STATE_DETACHED,  ///< Program is detached from a provider.
 } ebpf_link_state_t;
 
 typedef struct _ebpf_link
@@ -41,6 +48,9 @@ typedef struct _ebpf_link
     void* provider_binding_context;
 } ebpf_link_t;
 
+_Requires_lock_held_(link->attach_lock) static void _ebpf_link_set_state(
+    _Inout_ ebpf_link_t* link, ebpf_link_state_t new_state);
+
 static ebpf_result_t
 _ebpf_link_instance_invoke(
     _In_ const void* extension_client_binding_context, _Inout_ void* program_context, _Out_ uint32_t* result);
@@ -106,7 +116,7 @@ ebpf_link_create(
 
     memset(local_link, 0, sizeof(ebpf_link_t));
 
-    local_link->state = EBPF_LINK_STATE_IDLE;
+    local_link->state = EBPF_LINK_STATE_INITIAL;
 
     local_link->client_data.version = 0;
     local_link->client_data.size = context_data_length;
@@ -154,10 +164,23 @@ _ebpf_link_extension_changed_callback(
 
     // Complete detach.
     if (provider_data == NULL) {
+        // Check if the link is in the process of attaching.
+        if (link->state == EBPF_LINK_STATE_DETACHING) {
+            // If so, complete the detach.
+            _ebpf_link_set_state(link, EBPF_LINK_STATE_DETACHED);
+        } else if (link->state == EBPF_LINK_STATE_ATTACHED) {
+            // If not, mark the link as ready to reattach.
+            _ebpf_link_set_state(link, EBPF_LINK_STATE_ATTACHING);
+        }
         result = EBPF_SUCCESS;
         goto Done;
     }
 
+    if (link->state != EBPF_LINK_STATE_ATTACHING) {
+        result = EBPF_INVALID_ARGUMENT;
+        goto Done;
+    }
+
     if ((provider_data->version != EBPF_ATTACH_PROVIDER_DATA_VERSION) || (!provider_data->data) ||
         (provider_data->size != sizeof(ebpf_attach_provider_data_t))) {
         EBPF_LOG_MESSAGE_GUID(
@@ -185,6 +208,7 @@ _ebpf_link_extension_changed_callback(
     link->program_type = attach_provider_data->supported_program_type;
     link->bpf_attach_type = attach_provider_data->bpf_attach_type;
     link->link_type = attach_provider_data->link_type;
+    _ebpf_link_set_state(link, EBPF_LINK_STATE_ATTACHED);
     ebpf_assert(link->program != NULL);
 
     result = EBPF_SUCCESS;
@@ -215,11 +239,8 @@ ebpf_link_attach_program(_Inout_ ebpf_link_t* link, _Inout_ ebpf_program_t* prog
     state = ebpf_lock_lock(&link->attach_lock);
     attach_lock_held = true;
 
-    // If the link is not in idle state, then it is either:
-    // 1. Attaching to a program.
-    // 2. Detaching from a program.
-    // 3. Attached to a program.
-    if (link->state != EBPF_LINK_STATE_IDLE) {
+    // If the link is already attached, fail.
+    if (link->state != EBPF_LINK_STATE_INITIAL) {
         return_value = EBPF_INVALID_ARGUMENT;
         goto Done;
     }
@@ -233,7 +254,7 @@ ebpf_link_attach_program(_Inout_ ebpf_link_t* link, _Inout_ ebpf_program_t* prog
 
     link->program = program;
     ebpf_program_attach_link(program, link);
-    link->state = EBPF_LINK_STATE_ATTACHING;
+    _ebpf_link_set_state(link, EBPF_LINK_STATE_ATTACHING);
 
     ebpf_lock_unlock(&link->attach_lock, state);
     attach_lock_held = false;
@@ -271,15 +292,8 @@ ebpf_link_attach_program(_Inout_ ebpf_link_t* link, _Inout_ ebpf_program_t* prog
         link->program = NULL;
     }
 
-    if (return_value == EBPF_SUCCESS) {
-        ebpf_assert(link->state == EBPF_LINK_STATE_ATTACHING);
-        ebpf_assert(link_is_attaching);
-        ebpf_assert(link->program != NULL);
-        link->state = EBPF_LINK_STATE_ATTACHED;
-    } else {
-        if (link_is_attaching) {
-            link->state = EBPF_LINK_STATE_IDLE;
-        }
+    if (return_value != EBPF_SUCCESS) {
+        _ebpf_link_set_state(link, EBPF_LINK_STATE_DETACHED);
     }
 
     if (attach_lock_held) {
@@ -300,14 +314,15 @@ ebpf_link_detach_program(_Inout_ ebpf_link_t* link)
 
     state = ebpf_lock_lock(&link->attach_lock);
 
-    if (link->state == EBPF_LINK_STATE_ATTACHED) {
-        link->state = EBPF_LINK_STATE_DETACHING;
+    if (link->state == EBPF_LINK_STATE_ATTACHED || link->state == EBPF_LINK_STATE_ATTACHING) {
         // This thread is responsible for detaching the link from the program.
+        _ebpf_link_set_state(link, EBPF_LINK_STATE_DETACHING);
         link_is_detaching = true;
         program = link->program;
         ebpf_assert(program != NULL);
     }
-    if (link->state == EBPF_LINK_STATE_IDLE) {
+
+    if (link->state == EBPF_LINK_STATE_INITIAL || link->state == EBPF_LINK_STATE_DETACHED) {
         ebpf_assert(link->program == NULL);
     }
 
@@ -326,8 +341,6 @@ ebpf_link_detach_program(_Inout_ ebpf_link_t* link)
 
     state = ebpf_lock_lock(&link->attach_lock);
 
-    link->state = EBPF_LINK_STATE_IDLE;
-
     link->extension_client_context = NULL;
 
     ebpf_free(link->client_data.data);
@@ -393,6 +406,8 @@ _ebpf_link_instance_invoke_batch_begin(
     ((ebpf_execution_context_state_t*)state)->epoch_state = ebpf_epoch_enter();
     epoch_entered = true;
 
+    ebpf_assert(link->state == EBPF_LINK_STATE_ATTACHED);
+
     return_value = ebpf_program_reference_providers(link->program);
     if (return_value != EBPF_SUCCESS) {
         goto Done;
@@ -469,3 +484,37 @@ ebpf_link_get_info(
     *info_size = sizeof(*info);
     EBPF_RETURN_RESULT(EBPF_SUCCESS);
 }
+
+/**
+ * @brief Set and validate the link state.
+ *
+ * @param[in] link Link to set state on.
+ * @param[in] new_state New state to set.
+ */
+_Requires_lock_held_(link->attach_lock) static void _ebpf_link_set_state(
+    _Inout_ ebpf_link_t* link, ebpf_link_state_t new_state)
+{
+    ebpf_link_state_t old_state = link->state;
+    switch (new_state) {
+    case EBPF_LINK_STATE_ATTACHING:
+        // Runtime has requested that the program be linked to a provider.
+        ebpf_assert(old_state == EBPF_LINK_STATE_INITIAL || old_state == EBPF_LINK_STATE_ATTACHED);
+        break;
+    case EBPF_LINK_STATE_ATTACHED:
+        // Program is linked to a provider.
+        ebpf_assert(old_state == EBPF_LINK_STATE_ATTACHING);
+        break;
+    case EBPF_LINK_STATE_DETACHING:
+        // Runtime has requested that the program be unlinked from a provider.
+        ebpf_assert(old_state == EBPF_LINK_STATE_ATTACHED || old_state == EBPF_LINK_STATE_ATTACHING);
+        break;
+    case EBPF_LINK_STATE_DETACHED:
+        // Program is unlinked from a provider.
+        ebpf_assert(
+            old_state == EBPF_LINK_STATE_INITIAL || old_state == EBPF_LINK_STATE_DETACHING ||
+            old_state == EBPF_LINK_STATE_ATTACHING);
+        break;
+    }
+    UNREFERENCED_PARAMETER(old_state);
+    link->state = new_state;
+}
diff --git a/libs/execution_context/unit/execution_context_unit_test.cpp b/libs/execution_context/unit/execution_context_unit_test.cpp
@@ -775,35 +775,36 @@ TEST_CASE("program", "[execution_context]")
     REQUIRE(addresses[2] != 0);
 
     link_ptr link;
-    {
-        ebpf_link_t* local_link = nullptr;
-        REQUIRE(ebpf_link_create(EBPF_ATTACH_TYPE_XDP, nullptr, 0, &local_link) == EBPF_SUCCESS);
-        link.reset(local_link);
-    }
 
     // Correct attach type, but wrong program type.
     {
         single_instance_hook_t hook(EBPF_PROGRAM_TYPE_BIND, EBPF_ATTACH_TYPE_XDP);
-
+        ebpf_link_t* local_link = nullptr;
+        REQUIRE(ebpf_link_create(EBPF_ATTACH_TYPE_XDP, nullptr, 0, &local_link) == EBPF_SUCCESS);
+        link.reset(local_link);
         REQUIRE(ebpf_link_attach_program(link.get(), program.get()) == EBPF_EXTENSION_FAILED_TO_LOAD);
     }
 
     // Wrong attach type, but correct program type.
     {
         single_instance_hook_t hook(EBPF_PROGRAM_TYPE_XDP, EBPF_ATTACH_TYPE_BIND);
-
+        ebpf_link_t* local_link = nullptr;
+        REQUIRE(ebpf_link_create(EBPF_ATTACH_TYPE_XDP, nullptr, 0, &local_link) == EBPF_SUCCESS);
+        link.reset(local_link);
         REQUIRE(ebpf_link_attach_program(link.get(), program.get()) == EBPF_EXTENSION_FAILED_TO_LOAD);
     }
 
     // Correct attach type and correct program type.
     {
         single_instance_hook_t hook(EBPF_PROGRAM_TYPE_XDP, EBPF_ATTACH_TYPE_XDP);
+        ebpf_link_t* local_link = nullptr;
+        REQUIRE(ebpf_link_create(EBPF_ATTACH_TYPE_XDP, nullptr, 0, &local_link) == EBPF_SUCCESS);
+        link.reset(local_link);
 
-        // First attach should succeed.
+        // Attach should succeed.
         REQUIRE(ebpf_link_attach_program(link.get(), program.get()) == EBPF_SUCCESS);
 
-        // Second attach should fail.
-        REQUIRE(ebpf_link_attach_program(link.get(), program.get()) == EBPF_INVALID_ARGUMENT);
+        // Not possible to attach again.
 
         // First detach should succeed.
         ebpf_link_detach_program(link.get());

diff --git a/libs/platform/ebpf_extension.c b/libs/platform/ebpf_extension.c
@@ -146,6 +146,13 @@ _ebpf_extension_client_attach_provider(
 
     if (!NT_SUCCESS(status)) {
         EBPF_LOG_NTSTATUS_API_FAILURE(EBPF_TRACELOG_KEYWORD_BASE, NmrClientAttachProvider, status);
+
+        // If the attach failed, then the provider binding context is not valid.
+        // Notify client that provider is no longer available.
+        local_client_binding_context->provider_binding_context = NULL;
+        local_client_binding_context->provider_dispatch_table = NULL;
+        local_client_binding_context->provider_data = NULL;
+        (void)_ebpf_extension_client_notify_change(local_client_context, local_client_binding_context);
         goto Done;
     }