prov/efa: Fix use-after-free in DC packet completion handling

DC packets had a race condition where RECEIPT acknowledgments could
arrive before all TX completions, causing premature txe release and
use-after-free errors when remaining TX completions accessed the
freed memory.

Fix by using efa_outstanding_tx_ops for completion tracking:
- DC packets use efa_outstanding_tx_ops to track TX completion status
- RECEIPT handler writes completion immediately (preserves DC semantics)
- txe release delayed until both TX completions and receipt received
- Added EFA_RDM_TXE_RECEIPT_RECEIVED flag to track receipt arrival
- Added efa_rdm_txe_dc_ready_for_release() with proper documentation

This prevents the race condition while maintaining existing DC
completion behavior for applications.

Signed-off-by: Shi Jin <sjina@amazon.com>

Author: shijin-aws

prov/efa/src/rdm/efa_rdm_ope.c

@@ -1971,7 +1971,7 @@ ssize_t efa_rdm_ope_repost_ope_queued_before_handshake(struct efa_rdm_ope *ope)
 	}
 }
 
-int efa_rdm_ope_process_queued_ope(struct efa_rdm_ope *ope, uint16_t flag)
+int efa_rdm_ope_process_queued_ope(struct efa_rdm_ope *ope, uint32_t flag)
 {
 	int ret = 0;
 

prov/efa/src/rdm/efa_rdm_ope.h

@@ -103,7 +103,7 @@ struct efa_rdm_ope {
 	 * This flag is different from #cq_entry.flags, which is
 	 * applied to CQ entry's returned to user.
 	 */
-	uint16_t internal_flags;
+	uint32_t internal_flags;
 
 	size_t iov_count;
 	struct iovec iov[EFA_RDM_IOV_LIMIT];
@@ -281,13 +281,18 @@ void efa_rdm_rxe_release_internal(struct efa_rdm_ope *rxe);
  * @brief flag to indicate that the ope was created
  * for internal operations, so it should not generate
  * any cq entry or err entry.
- * NOTICE: the ope->internal_flags is uint16_t, so
- * to introduce more bits for internal flags, the
- * internal_flags needs to be changed to uint32_t
- * or larger.
  */
 #define EFA_RDM_OPE_INTERNAL			BIT_ULL(15)
 
+/**
+ * @brief flag to indicate that a DC txe has received its receipt packet
+ *
+ * This flag is used to track when a delivery complete operation has
+ * received acknowledgment from the receiver, preventing premature
+ * completion before all TX operations finish.
+ */
+#define EFA_RDM_TXE_RECEIPT_RECEIVED		BIT_ULL(16)
+
 #define EFA_RDM_OPE_QUEUED_FLAGS (EFA_RDM_OPE_QUEUED_RNR | EFA_RDM_OPE_QUEUED_CTRL | EFA_RDM_OPE_QUEUED_READ | EFA_RDM_OPE_QUEUED_BEFORE_HANDSHAKE)
 
 void efa_rdm_ope_try_fill_desc(struct efa_rdm_ope *ope, int mr_iov_start, uint64_t access);
@@ -311,6 +316,27 @@ void efa_rdm_ope_handle_recv_completed(struct efa_rdm_ope *ope);
 
 void efa_rdm_ope_handle_send_completed(struct efa_rdm_ope *ope);
 
+/**
+ * @brief Check if a delivery complete (DC) TXE is ready for release
+ *
+ * @details
+ * For DC packets, this function prevents use-after-free race conditions by
+ * ensuring the TXE is only released when both conditions are met:
+ * 1. All TX operations have completed (efa_outstanding_tx_ops == 0)
+ * 2. Receipt packet has been received (EFA_RDM_TXE_RECEIPT_RECEIVED flag set)
+ *
+ * This dual-condition check ensures proper synchronization between send
+ * completions and receipt acknowledgments in the delivery complete protocol.
+ *
+ * @param[in] txe TX operation entry to check
+ * @return true if TXE is ready for release, false otherwise
+ */
+static inline bool efa_rdm_txe_dc_ready_for_release(struct efa_rdm_ope *txe)
+{
+	return (txe->efa_outstanding_tx_ops == 0) &&
+	       (txe->internal_flags & EFA_RDM_TXE_RECEIPT_RECEIVED);
+}
+
 int efa_rdm_ope_prepare_to_post_read(struct efa_rdm_ope *ope);
 
 void efa_rdm_ope_prepare_to_post_write(struct efa_rdm_ope *ope);
@@ -342,6 +368,6 @@ ssize_t efa_rdm_ope_repost_ope_queued_before_handshake(struct efa_rdm_ope *ope);
 
 ssize_t efa_rdm_txe_prepare_local_read_pkt_entry(struct efa_rdm_ope *txe);
 
-int efa_rdm_ope_process_queued_ope(struct efa_rdm_ope *ope, uint16_t flag);
+int efa_rdm_ope_process_queued_ope(struct efa_rdm_ope *ope, uint32_t flag);
 
 #endif

prov/efa/src/rdm/efa_rdm_pke_cmd.c

@@ -646,16 +646,16 @@ void efa_rdm_pke_handle_send_completion(struct efa_rdm_pke *pkt_entry)
 	case EFA_RDM_DC_LONGCTS_MSGRTM_PKT:
 	case EFA_RDM_DC_LONGCTS_TAGRTM_PKT:
 	case EFA_RDM_DC_LONGCTS_RTW_PKT:
-		/* no action to be taken here
-		 * For non-dc version of these packet types,
-		 * this is the place to increase bytes_acked or
-		 * write tx completion.
-		 * For dc, tx completion will always be
-		 * written upon receving the receipt packet
-		 * Moreoever, because receipt can arrive
-		 * before send completion, we cannot take
-		 * any action on txe here.
+		/* For DC packets, use efa_outstanding_tx_ops to track TX completions
+		 * instead of bytes_acked to avoid issues with unset payload_size.
+		 * Note: efa_rdm_ep_record_tx_op_completed() above decrements efa_outstanding_tx_ops,
+		 * so this check must come after that call.
+		 * Only release TXE when both TX ops complete and receipt is received.
 		 */
+		assert(pkt_entry->ope);
+		if (efa_rdm_txe_dc_ready_for_release(pkt_entry->ope))
+			efa_rdm_txe_release(pkt_entry->ope);
+		break;
 	case EFA_RDM_READ_NACK_PKT:
 		/* no action needed for NACK packet */
 		break;

prov/efa/src/rdm/efa_rdm_pke_nonreq.c

@@ -783,7 +783,16 @@ void efa_rdm_pke_handle_receipt_recv(struct efa_rdm_pke *pkt_entry)
 		return;
 	}
 
-	efa_rdm_ope_handle_send_completed(txe);
+	/* Write send completion immediately to preserve DC semantics */
+	efa_rdm_txe_report_completion(txe);
+
+	/* Set receipt received flag for DC operations */
+	txe->internal_flags |= EFA_RDM_TXE_RECEIPT_RECEIVED;
+
+	/* Only release txe if both conditions are met */
+	if (efa_rdm_txe_dc_ready_for_release(txe))
+		efa_rdm_txe_release(txe);
+
 	efa_rdm_pke_release_rx(pkt_entry);
 }
 

prov/efa/test/efa_unit_test_ope.c

@@ -2,6 +2,8 @@
 /* SPDX-FileCopyrightText: Copyright Amazon.com, Inc. or its affiliates. All rights reserved. */
 
 #include "efa_unit_tests.h"
+#include "rdm/efa_rdm_pke_cmd.h"
+#include "rdm/efa_rdm_pke_nonreq.h"
 
 typedef void (*efa_rdm_ope_handle_error_func_t)(struct efa_rdm_ope *ope, int err, int prov_errno);
 
@@ -1176,3 +1178,97 @@ void test_efa_rdm_atomic_compare_desc_persistence(struct efa_resource **state)
 	efa_unit_test_buff_destruct(&result_buff);
 	efa_unit_test_buff_destruct(&compare_buff);
 }
+/**
+ * @brief Common helper for DC packet TXE release testing
+ *
+ * Sets up test environment and packet entries for DC packet testing.
+ * Tests the specified completion order and verifies TXE release behavior.
+ *
+ * @param[in] resource Test resource structure
+ * @param[in] send_first If true, send completion happens first; if false, receipt first
+ */
+static void test_efa_rdm_txe_dc_release_common(struct efa_resource *resource, bool send_first)
+{
+	struct efa_rdm_ep *efa_rdm_ep;
+	struct efa_rdm_ope *txe;
+	struct efa_rdm_pke *dc_pkt_entry, *receipt_pkt_entry;
+
+	efa_unit_test_resource_construct(resource, FI_EP_RDM, EFA_FABRIC_NAME);
+
+	efa_rdm_ep = container_of(resource->ep, struct efa_rdm_ep, base_ep.util_ep.ep_fid);
+
+	/* Allocate TXE and set up for DC operation */
+	txe = efa_unit_test_alloc_txe(resource, ofi_op_msg);
+	assert_non_null(txe);
+	txe->internal_flags |= EFA_RDM_TXE_DELIVERY_COMPLETE_REQUESTED;
+	txe->efa_outstanding_tx_ops = 1;
+
+	/* Create fake DC packet entry */
+	dc_pkt_entry = efa_rdm_pke_alloc(efa_rdm_ep, efa_rdm_ep->efa_tx_pkt_pool, EFA_RDM_PKE_FROM_EFA_TX_POOL);
+	assert_non_null(dc_pkt_entry);
+	dc_pkt_entry->ope = txe;
+	dc_pkt_entry->ep = efa_rdm_ep;
+	dc_pkt_entry->peer = txe->peer;
+	/* Set DC packet type in wiredata */
+	struct efa_rdm_base_hdr *base_hdr = (struct efa_rdm_base_hdr *)dc_pkt_entry->wiredata;
+	base_hdr->type = EFA_RDM_DC_EAGER_MSGRTM_PKT;
+
+	/* Create fake receipt packet entry */
+	receipt_pkt_entry = efa_rdm_pke_alloc(efa_rdm_ep, efa_rdm_ep->efa_rx_pkt_pool, EFA_RDM_PKE_FROM_EFA_RX_POOL);
+	assert_non_null(receipt_pkt_entry);
+	receipt_pkt_entry->ope = txe;
+	receipt_pkt_entry->ep = efa_rdm_ep;
+
+	/* Verify TXE is not ready for release initially */
+	assert_false(efa_rdm_txe_dc_ready_for_release(txe));
+	assert_int_equal(efa_unit_test_get_dlist_length(&efa_rdm_ep->txe_list), 1);
+
+	if (send_first) {
+		/* Send completion first - should not release TXE yet */
+		efa_rdm_pke_handle_send_completion(dc_pkt_entry);
+		assert_int_equal(efa_unit_test_get_dlist_length(&efa_rdm_ep->txe_list), 1);
+		assert_false(efa_rdm_txe_dc_ready_for_release(txe));
+
+		/* Receipt handling - should now release TXE */
+		efa_rdm_pke_handle_receipt_recv(receipt_pkt_entry);
+	} else {
+		/* Receipt handling first - should not release TXE yet */
+		efa_rdm_pke_handle_receipt_recv(receipt_pkt_entry);
+		assert_int_equal(efa_unit_test_get_dlist_length(&efa_rdm_ep->txe_list), 1);
+		assert_true(txe->internal_flags & EFA_RDM_TXE_RECEIPT_RECEIVED);
+		assert_false(efa_rdm_txe_dc_ready_for_release(txe));
+
+		/* Send completion - should now release TXE */
+		efa_rdm_pke_handle_send_completion(dc_pkt_entry);
+	}
+
+	/* Verify TXE is released */
+	assert_int_equal(efa_unit_test_get_dlist_length(&efa_rdm_ep->txe_list), 0);
+}
+
+/**
+ * @brief Test DC packet TXE release with send completion first
+ *
+ * This test verifies the DC (Delivery Complete) TXE release logic when
+ * send completion arrives before receipt acknowledgment.
+ *
+ * @param[in] state cmocka state variable
+ */
+void test_efa_rdm_txe_dc_send_first(struct efa_resource **state)
+{
+	test_efa_rdm_txe_dc_release_common(*state, true);
+}
+
+/**
+ * @brief Test DC packet TXE release with receipt completion first
+ *
+ * This test verifies the race condition fix where receipt acknowledgment
+ * arrives before send completion. The TXE should only be released when
+ * both conditions are met, regardless of order.
+ *
+ * @param[in] state cmocka state variable
+ */
+void test_efa_rdm_txe_dc_receipt_first(struct efa_resource **state)
+{
+	test_efa_rdm_txe_dc_release_common(*state, false);
+}

prov/efa/test/efa_unit_tests.c

@@ -333,6 +333,8 @@ int main(void)
 		cmocka_unit_test_setup_teardown(test_efa_rdm_ope_eor_packet_failed_posting, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		cmocka_unit_test_setup_teardown(test_efa_rdm_ope_eor_packet_tracking_unresponsive_wait_send, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		cmocka_unit_test_setup_teardown(test_efa_rdm_atomic_compare_desc_persistence, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+		cmocka_unit_test_setup_teardown(test_efa_rdm_txe_dc_send_first, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
+		cmocka_unit_test_setup_teardown(test_efa_rdm_txe_dc_receipt_first, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),
 		/* end of efa_unit_test_ope.c */
 
 		cmocka_unit_test_setup_teardown(test_efa_rdm_msg_send_to_local_peer_with_null_desc, efa_unit_test_mocks_setup, efa_unit_test_mocks_teardown),

prov/efa/test/efa_unit_tests.h

@@ -291,6 +291,8 @@ void test_efa_rdm_ope_eor_packet_tracking_wait_send();
 void test_efa_rdm_ope_eor_packet_failed_posting();
 void test_efa_rdm_ope_eor_packet_tracking_unresponsive_wait_send();
 void test_efa_rdm_atomic_compare_desc_persistence();
+void test_efa_rdm_txe_dc_send_first();
+void test_efa_rdm_txe_dc_receipt_first();
 
 
 /* end of efa_unit_test_ope.c */