lib: arm_ffa: Add support for MSG_SEND_DIRECT_REQ2 UUID demuxing

This change allows UUID specific handlers to be installed for
MSG_SEND_DIRECT_REQ2.

Bug: 380065728
Change-Id: I54a95e8893da49c831c8c400f802550673af28d2
Signed-off-by: Jay Monkman <jmonkman@google.com>
diff --git a/lib/arm_ffa/arm_ffa.c b/lib/arm_ffa/arm_ffa.c
index 4bc440d..5c50adf 100644
--- a/lib/arm_ffa/arm_ffa.c
+++ b/lib/arm_ffa/arm_ffa.c
@@ -49,13 +49,31 @@
 static bool supports_ns_bit = false;
 static bool supports_rx_release = false;
 static bool console_log_is_unsupported;
+static mutex_t ffa_rxtx_buffer_lock = MUTEX_INITIAL_VALUE(ffa_rxtx_buffer_lock);
+#if ARCH_ARM64
 static bool send_direct_req2_is_unsupported;
 
-static mutex_t ffa_rxtx_buffer_lock = MUTEX_INITIAL_VALUE(ffa_rxtx_buffer_lock);
+static struct bst_root arm_ffa_direct_req2_handler_tree =
+        BST_ROOT_INITIAL_VALUE;
+static spin_lock_t arm_ffa_direct_req2_tree_lock = SPIN_LOCK_INITIAL_VALUE;
 
-enum arm_ffa_init_state arm_ffa_init_state(void) {
-    return ffa_init_state;
-}
+/**
+ * struct arm_ffa_direct_req2_bst_obj - Binary search tree object for
+ * ffa_direct_req2 handler
+ * @bst_node: BST node
+ * @uuid_lo_hi: Array that holds UUID as two 64 bit words
+ *              uuid_lo_hi[0] is what the FFA spec labels "Lo" - bytes [0-7]
+ *              uuid_lo_hi[1] is what the FFA spec labels "Hi" - bytes [8-15]
+ * @handler: Pointer to FFA_DIRECT_REQ2 handler function
+ */
+struct arm_ffa_direct_req2_bst_obj {
+    struct bst_node bst_node;
+    uint64_t uuid_lo_hi[2];
+    arm_ffa_direct_req2_handler_t handler;
+};
+static int arm_ffa_direct_req2_handler_compare(struct bst_node* a,
+                                               struct bst_node* b);
+#endif
 
 /**
  * uuid_to_le64_pair() - convert uuid_t to (lo, hi)-pair per FFA spec.
@@ -76,6 +94,72 @@
     }
 }
 
+#if ARCH_ARM64
+status_t arm_ffa_register_direct_req2_handler(
+        uuid_t uuid,
+        arm_ffa_direct_req2_handler_t handler) {
+    struct arm_ffa_direct_req2_bst_obj* obj;
+
+    obj = calloc(1, sizeof(*obj));
+    if (!obj) {
+        LTRACEF("ERROR: not enough memory for direct_req2 handler\n");
+        return ERR_NO_MEMORY;
+    }
+
+    uuid_to_le64_pair(uuid, obj->uuid_lo_hi);
+    obj->handler = handler;
+
+    spin_lock(&arm_ffa_direct_req2_tree_lock);
+    if (!bst_insert(&arm_ffa_direct_req2_handler_tree, &obj->bst_node,
+                    arm_ffa_direct_req2_handler_compare)) {
+        spin_unlock(&arm_ffa_direct_req2_tree_lock);
+        free(obj);
+        LTRACEF("ERROR: couldn't insert direct_req2 hander into BST\n");
+        return ERR_ALREADY_EXISTS;
+    } else {
+        spin_unlock(&arm_ffa_direct_req2_tree_lock);
+        return 0;
+    }
+}
+
+static int arm_ffa_direct_req2_handler_compare(struct bst_node* a,
+                                               struct bst_node* b) {
+    struct arm_ffa_direct_req2_bst_obj* obj_a =
+            containerof(a, struct arm_ffa_direct_req2_bst_obj, bst_node);
+    struct arm_ffa_direct_req2_bst_obj* obj_b =
+            containerof(b, struct arm_ffa_direct_req2_bst_obj, bst_node);
+
+    return memcmp(obj_a->uuid_lo_hi, obj_b->uuid_lo_hi,
+                  sizeof(obj_a->uuid_lo_hi));
+}
+
+status_t arm_ffa_handle_direct_req2(struct smc_ret18* regs) {
+    struct arm_ffa_direct_req2_bst_obj search_obj;
+    struct arm_ffa_direct_req2_bst_obj* found_obj;
+    uint16_t sender_id = (regs->r1 >> 16) & 0xffff;
+    search_obj.uuid_lo_hi[0] = regs->r2;
+    search_obj.uuid_lo_hi[1] = regs->r3;
+
+    spin_lock(&arm_ffa_direct_req2_tree_lock);
+    found_obj = bst_search_type(&arm_ffa_direct_req2_handler_tree, &search_obj,
+                                arm_ffa_direct_req2_handler_compare,
+                                struct arm_ffa_direct_req2_bst_obj, bst_node);
+    spin_unlock(&arm_ffa_direct_req2_tree_lock);
+
+    if (found_obj) {
+        return found_obj->handler(sender_id, &regs->r4);
+    } else {
+        LTRACEF("Error: No handler for UUID 0x%016lx 0x%016lx for sender %d\n",
+                regs->r2, regs->r3, sender_id);
+        return ERR_NOT_FOUND;
+    }
+}
+#endif
+
+enum arm_ffa_init_state arm_ffa_init_state(void) {
+    return ffa_init_state;
+}
+
 static status_t arm_ffa_call_id_get(uint16_t* id) {
     struct smc_ret8 smc_ret;
 
@@ -468,10 +552,12 @@
     return smc8(fid, sender_receiver_id, flags, a0, a1, a2, a3, a4);
 }
 
-status_t arm_ffa_msg_send_direct_req2(uuid_t uuid,
-                                      uint16_t receiver_id,
-                                      uint64_t args[static 14],
-                                      struct smc_ret18* resp) {
+#if ARCH_ARM64
+status_t arm_ffa_msg_send_direct_req2(
+        uuid_t uuid,
+        uint16_t receiver_id,
+        uint64_t args[static REQ2_OTHER_PARAM_REGS],
+        struct smc_ret18* resp) {
     struct smc_ret18 smc_ret;
     uint64_t uuid_lo_hi[2];
     uint32_t fid = SMC_FC64_FFA_MSG_SEND_DIRECT_REQ2;
@@ -530,6 +616,7 @@
         return ERR_NOT_VALID;
     }
 }
+#endif
 
 ssize_t arm_ffa_console_log(const char* buf, size_t len) {
     struct smc_ret8 smc_ret;
diff --git a/lib/arm_ffa/include/lib/arm_ffa/arm_ffa.h b/lib/arm_ffa/include/lib/arm_ffa/arm_ffa.h
index f446b35..57d4963 100644
--- a/lib/arm_ffa/include/lib/arm_ffa/arm_ffa.h
+++ b/lib/arm_ffa/include/lib/arm_ffa/arm_ffa.h
@@ -202,3 +202,38 @@
  * Return: the number of characters successfully printed, or an error code.
  */
 ssize_t arm_ffa_console_log(const char* buf, size_t len);
+
+#if ARCH_ARM64
+/**
+ * arm_ffa_direct_req2_handler_t - Handler function for DIRECT_REQ2 calls
+ *
+ * @sender_id: Sender's endpoint ID
+ * @regs: Contents of message - x4-x17
+ *
+ * Return: 0 on success, LK error code on failure.
+ */
+typedef status_t (*arm_ffa_direct_req2_handler_t)(
+        uint16_t sender_id,
+        uint64_t regs[static REQ2_OTHER_PARAM_REGS]);
+
+/**
+ * arm_ffa_handle_direct_req2() - Handle DIRECT_REQ2 call
+ *
+ * @regs: CPU registers for FFA call
+ *
+ * Return: 0 on success, LK error code on failure.
+ */
+status_t arm_ffa_handle_direct_req2(struct smc_ret18* regs);
+
+/**
+ * arm_ffa_register_direct_req2_handler() - Register DIRECT_REQ2 call handler
+ *
+ * @uuid: UUID of handler to register
+ * @handler: pointer to handler function
+ *
+ * Return: 0 on success, LK error code on failure.
+ */
+status_t arm_ffa_register_direct_req2_handler(
+        uuid_t uuid,
+        arm_ffa_direct_req2_handler_t handler);
+#endif