DPE: Add custom argument cert_id to DeriveContext

Cert_id is custom argument passed by the client to DeriveContext
to assign the context to specific certificate. For now only
support the parameter as mandatory input.

Signed-off-by: Maulik Patel <maulik.patel@arm.com>
Change-Id: Ib8c00cbab8fd0ad5b49012d854ae3fc6077409b9
diff --git a/partitions/dice_protection_environment/dpe_boot_data.c b/partitions/dice_protection_environment/dpe_boot_data.c
index 3bfdda7..8540765 100644
--- a/partitions/dice_protection_environment/dpe_boot_data.c
+++ b/partitions/dice_protection_environment/dpe_boot_data.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2023-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -14,11 +14,14 @@
 
 #include "boot_hal.h"
 #include "boot_measurement.h"
+#include "dice_protection_environment.h"
 #include "dpe_context_mngr.h"
 #include "service_api.h"
 #include "tfm_boot_status.h"
 #include "tfm_plat_otp.h"
 
+#define DPE_PLATFORM_CERT_ID 0x200
+
 /* Maximum measurement size is size of SHA-512 hash */
 #define MEASUREMENT_VALUE_MAX_SIZE 64
 
@@ -296,6 +299,7 @@
 
     /* Derive RoT layer */
     err = derive_context_request(rot_ctx_handle,
+                                 DPE_ROT_CERT_ID, /* cert_id */
                                  false, /* retain_parent_context */
                                  true, /* allow_new_context_to_derive */
                                  true, /* create certificate */
@@ -327,6 +331,7 @@
 
     /* Derive BL2 context */
     err = derive_context_request(plat_ctx_handle,
+                                 DPE_PLATFORM_CERT_ID, /* cert_id */
                                  false, /* close parent context */
                                  true, /* allow BL2 to derive further */
                                  false, /* create_certificate */
@@ -354,6 +359,7 @@
                                                 &dice_inputs)) == 1) {
         /* Derive rest of platform contexts from retained BL2 context */
         err = derive_context_request(plat_ctx_handle,
+                                     DPE_CERT_ID_SAME_AS_PARENT, /* cert_id */
                                      true, /* retain parent context */
                                      false, /* do not allow derived context to derive */
                                      false, /* create_certificate */
@@ -394,9 +400,10 @@
      * caller in the new_ctx_handle output parameter.
      */
     return derive_context_request(plat_ctx_handle,
+                                  DPE_CERT_ID_SAME_AS_PARENT, /* cert_id */
                                   false, /* close parent context */
                                   true, /* allow AP to derive */
-                                  false, /* create_certificate */
+                                  true, /* create_certificate */
                                   &dice_inputs,
                                   0, /* client_id */
                                   0, /* target_locality */
diff --git a/partitions/dice_protection_environment/dpe_cmd_decode.c b/partitions/dice_protection_environment/dpe_cmd_decode.c
index 63e23ea..c4f2623 100644
--- a/partitions/dice_protection_environment/dpe_cmd_decode.c
+++ b/partitions/dice_protection_environment/dpe_cmd_decode.c
@@ -150,6 +150,7 @@
     int new_parent_context_handle;
     uint8_t *new_certificate_buf = REUSE_CMD_BUF(DICE_CERT_SIZE);
     uint8_t exported_cdi_buf[DICE_MAX_ENCODED_CDI_SIZE];
+    uint32_t cert_id;
     size_t new_certificate_actual_size = 0;
     size_t exported_cdi_actual_size = 0;
 
@@ -158,11 +159,19 @@
 
     QCBORDecode_GetByteStringInMapN(decode_ctx, DPE_DERIVE_CONTEXT_CONTEXT_HANDLE,
                                     &out);
-    if (out.len != sizeof(context_handle)) {
+    qcbor_err = QCBORDecode_GetAndResetError(decode_ctx);
+    if ((qcbor_err != QCBOR_SUCCESS) || (out.len != sizeof(context_handle))) {
         return DPE_INVALID_COMMAND;
     }
     memcpy(&context_handle, out.ptr, out.len);
 
+    QCBORDecode_GetUInt64InMapN(decode_ctx, DPE_DERIVE_CONTEXT_CERT_ID, &cert_id);
+    /* Check if cert_id was encoded in the received command buffer */
+    qcbor_err = QCBORDecode_GetAndResetError(decode_ctx);
+    if (qcbor_err != QCBOR_SUCCESS) {
+        cert_id = DPE_CERT_ID_INVALID;
+    }
+
     QCBORDecode_GetBoolInMapN(decode_ctx, DPE_DERIVE_CONTEXT_RETAIN_PARENT_CONTEXT,
                               &retain_parent_context);
 
@@ -204,7 +213,7 @@
         return DPE_INVALID_COMMAND;
     }
 
-    dpe_err = derive_context_request(context_handle, retain_parent_context,
+    dpe_err = derive_context_request(context_handle, cert_id, retain_parent_context,
                                      allow_new_context_to_derive, create_certificate,
                                      &dice_inputs, client_id,
                                      target_locality,
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.c b/partitions/dice_protection_environment/dpe_context_mngr.c
index d750cad..86e67f1 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.c
+++ b/partitions/dice_protection_environment/dpe_context_mngr.c
@@ -10,6 +10,7 @@
 #include <string.h>
 #include "dice_protection_environment.h"
 #include "dpe_certificate.h"
+#include "dpe_client.h"
 #include "dpe_crypto_interface.h"
 #include "dpe_log.h"
 #include "dpe_plat.h"
@@ -82,6 +83,7 @@
     layer_ctx_array[i].state = LAYER_STATE_CLOSED;
     layer_ctx_array[i].parent_layer_idx = INVALID_LAYER_IDX;
     layer_ctx_array[i].is_cdi_to_be_exported = false;
+    layer_ctx_array[i].cert_id = DPE_CERT_ID_INVALID;
     (void)memset(&layer_ctx_array[i].attest_cdi_hash_input, 0,
                  sizeof(layer_ctx_array[i].attest_cdi_hash_input));
     (void)psa_destroy_key(layer_ctx_array[i].data.cdi_key_id);
@@ -378,11 +380,6 @@
     return MAX_NUM_OF_LAYERS - 1;
 }
 
-static inline void link_layer(uint16_t derived_ctx_layer, uint16_t parent_ctx_layer)
-{
-    layer_ctx_array[derived_ctx_layer].parent_layer_idx = parent_ctx_layer;
-}
-
 static inline bool is_input_client_id_valid(int32_t client_id, int32_t target_locality)
 {
     // TODO: FIXME
@@ -390,31 +387,68 @@
     return true;
 }
 
-static dpe_error_t assign_layer_to_context(struct component_context_t *new_ctx)
+static bool is_cert_id_used(uint32_t cert_id, uint16_t *layer_idx)
 {
-    uint16_t new_layer_idx, parent_layer_idx;
+    int i;
+
+    for (i = 0; i < MAX_NUM_OF_LAYERS; i++) {
+        if (layer_ctx_array[i].cert_id == cert_id) {
+            *layer_idx = i;
+            return true;
+        }
+    }
+
+    /* No certificate ID match found */
+    return false;
+}
+
+static dpe_error_t assign_layer_to_context(struct component_context_t *new_ctx,
+                                           uint32_t cert_id)
+{
+    uint16_t parent_layer_idx, layer_idx_to_link;
 
     assert(new_ctx->parent_idx < MAX_NUM_OF_COMPONENTS);
 
     parent_layer_idx = component_ctx_array[new_ctx->parent_idx].linked_layer_idx;
     assert(parent_layer_idx < MAX_NUM_OF_LAYERS);
 
-    if (layer_ctx_array[parent_layer_idx].state == LAYER_STATE_FINALISED) {
-        /* Parent comp's layer of new derived context is finalised; open a new layer */
-        new_layer_idx = open_new_layer();
-        if (new_layer_idx == INVALID_LAYER_IDX) {
-            return DPE_INTERNAL_ERROR;
+    if (cert_id != DPE_CERT_ID_INVALID) {
+        /* cert id was sent by the client */
+        if (cert_id == DPE_CERT_ID_SAME_AS_PARENT) {
+            if (layer_ctx_array[parent_layer_idx].state == LAYER_STATE_FINALISED) {
+                /* Cannot add to the layer which is already finalised */
+                return DPE_INTERNAL_ERROR;
+            }
+            /* Derived context belongs to the same certificate as its parent component */
+            new_ctx->linked_layer_idx = parent_layer_idx;
+
+        } else if (is_cert_id_used(cert_id, &layer_idx_to_link)) {
+            /* Cert ID is already in use */
+            if (layer_ctx_array[layer_idx_to_link].state == LAYER_STATE_FINALISED) {
+                /* Cannot add to the layer which is already finalised */
+                return DPE_INTERNAL_ERROR;
+            }
+            /* Use the same layer that is associated with cert_id */
+            new_ctx->linked_layer_idx = layer_idx_to_link;
+            /* Linked layer's parent is already assigned when it was opened */
+
+        } else {
+            /* Open new layer and link derived context to new layer */
+            layer_idx_to_link = open_new_layer();
+            if (layer_idx_to_link == INVALID_LAYER_IDX) {
+                return DPE_INTERNAL_ERROR;
+            }
+            /* Link this context to the new layer */
+            new_ctx->linked_layer_idx = layer_idx_to_link;
+            /* New layer's parent is parent component's layer */
+            layer_ctx_array[layer_idx_to_link].parent_layer_idx = parent_layer_idx;
+            layer_ctx_array[layer_idx_to_link].cert_id = cert_id;
         }
-        /* Link this context to the new layer */
-        new_ctx->linked_layer_idx = new_layer_idx;
-        /* New layer's parent is current layer */
-        link_layer(new_layer_idx, parent_layer_idx);
 
     } else {
-        /* Parent comp's layer is not yet finalised, link
-         * new component to the same layer as parent
-         */
-        new_ctx->linked_layer_idx = parent_layer_idx;
+        /* cert id was not sent by the client */
+        //TODO: To be implemented; return error for now.
+        return DPE_INVALID_ARGUMENT;
     }
 
     return DPE_NO_ERROR;
@@ -439,8 +473,6 @@
     struct component_context_t *rot_comp_ctx = &component_ctx_array[0];
     struct layer_context_t *rot_layer_ctx = &layer_ctx_array[DPE_ROT_LAYER_IDX];
 
-    /* Open RoT layer */
-    rot_layer_ctx->state = LAYER_STATE_OPEN;
     /* Parent layer for RoT context's layer is same */
     rot_layer_ctx->parent_layer_idx = DPE_ROT_LAYER_IDX;
 
@@ -489,6 +521,7 @@
 }
 
 dpe_error_t derive_context_request(int input_ctx_handle,
+                                   uint32_t cert_id,
                                    bool retain_parent_context,
                                    bool allow_new_context_to_derive,
                                    bool create_certificate,
@@ -514,7 +547,7 @@
     struct layer_context_t *layer_ctx;
     psa_status_t status;
 
-    log_derive_context(input_ctx_handle, retain_parent_context,
+    log_derive_context(input_ctx_handle, cert_id, retain_parent_context,
                        allow_new_context_to_derive, create_certificate, dice_inputs,
                        client_id);
 
@@ -593,7 +626,7 @@
     derived_ctx->parent_idx = parent_ctx_idx;
     /* Mark new derived component index as in use */
     derived_ctx->in_use = true;
-    err = assign_layer_to_context(derived_ctx);
+    err = assign_layer_to_context(derived_ctx, cert_id);
     if (err != DPE_NO_ERROR) {
         return err;
     }
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.h b/partitions/dice_protection_environment/dpe_context_mngr.h
index b1bd815..b34fb9a 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.h
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -98,6 +98,7 @@
     enum layer_state_t state;
     bool is_external_pub_key_provided;
     bool is_cdi_to_be_exported;
+    uint32_t cert_id;
 };
 
 /**
@@ -114,6 +115,8 @@
  *        chain.
  *
  * \param[in]  input_context_handle        Input handle to parent component context.
+ * \param[in]  cert_id                     Logical certificate id to which derived
+ *                                         context belongs to.
  * \param[in]  retain_parent_context       Flag to indicate if parent context need
  *                                         to be retained. TRUE only if a client
  *                                         is calling DPE commands multiple times.
@@ -149,6 +152,7 @@
  * \return Returns error code of type dpe_error_t
  */
 dpe_error_t derive_context_request(int input_ctx_handle,
+                                   uint32_t cert_id,
                                    bool retain_parent_context,
                                    bool allow_new_context_to_derive,
                                    bool create_certificate,
diff --git a/partitions/dice_protection_environment/dpe_log.c b/partitions/dice_protection_environment/dpe_log.c
index 4986432..82e86cd 100644
--- a/partitions/dice_protection_environment/dpe_log.c
+++ b/partitions/dice_protection_environment/dpe_log.c
@@ -58,6 +58,7 @@
 }
 
 void log_derive_context(int context_handle,
+                        uint32_t cert_id,
                         bool retain_parent_context,
                         bool allow_new_context_to_derive,
                         bool create_certificate,
@@ -67,6 +68,7 @@
     LOG_DBGFMT("DPE DeriveContext:\r\n");
     LOG_DBGFMT(" - context_handle index = %d\r\n", GET_IDX(context_handle));
     LOG_DBGFMT(" - context_handle nonce = %d\r\n", GET_NONCE(context_handle));
+    LOG_DBGFMT(" - cert_id = 0x%x\r\n", cert_id);
     LOG_DBGFMT(" - retain_parent_context = %d\r\n", retain_parent_context);
     LOG_DBGFMT(" - allow_new_context_to_derive = %d\r\n", allow_new_context_to_derive);
     LOG_DBGFMT(" - create_certificate = %d\r\n", create_certificate);
diff --git a/partitions/dice_protection_environment/dpe_log.h b/partitions/dice_protection_environment/dpe_log.h
index f7b207c..0bb95f1 100644
--- a/partitions/dice_protection_environment/dpe_log.h
+++ b/partitions/dice_protection_environment/dpe_log.h
@@ -26,6 +26,7 @@
  * \brief Log the derive context command parameters.
  */
 void log_derive_context(int context_handle,
+                        uint32_t cert_id,
                         bool retain_parent_context,
                         bool allow_new_context_to_derive,
                         bool create_certificate,
diff --git a/partitions/dice_protection_environment/interface/include/dice_protection_environment.h b/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
index 5fed3a9..cddb9ed 100644
--- a/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
+++ b/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2023-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -34,6 +34,10 @@
 #define DPE_INSUFFICIENT_MEMORY       ((dpe_error_t)128)
 #define DPE_ERR_CBOR_FORMATTING       ((dpe_error_t)129)
 
+#define DPE_CERT_ID_INVALID 0
+#define DPE_ROT_CERT_ID 0x100
+#define DPE_CERT_ID_SAME_AS_PARENT 0xFFFFFFFF
+
 /**
  * \brief Performs the DICE computation to derive a new context and optionally
  *        creates an intermediate certificate. Software component measurement
@@ -41,6 +45,8 @@
  *
  * \param[in]  context_handle              Input context handle for the DPE
  *                                         context.
+ * \param[in]  cert_id                     Logical certificate id to which derived
+ *                                         context belongs to.
  * \param[in]  retain_parent_context       Flag to indicate whether to retain the
  *                                         parent context. True only if a client
  *                                         will call further DPE commands on the
@@ -78,6 +84,7 @@
  */
 dpe_error_t
 dpe_derive_context(int                    context_handle,
+                   uint32_t               cert_id,
                    bool                   retain_parent_context,
                    bool                   allow_new_context_to_derive,
                    bool                   create_certificate,
diff --git a/partitions/dice_protection_environment/interface/include/dpe_client.h b/partitions/dice_protection_environment/interface/include/dpe_client.h
index eed9598..cef190a 100644
--- a/partitions/dice_protection_environment/interface/include/dpe_client.h
+++ b/partitions/dice_protection_environment/interface/include/dpe_client.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2023-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -59,6 +59,8 @@
     DPE_DERIVE_CONTEXT_RETURN_CERTIFICATE = 9,
     DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_EXPORT = 10,
     DPE_DERIVE_CONTEXT_EXPORT_CDI = 11,
+    /* enum values 256 and onwards are reserved for custom arguments */
+    DPE_DERIVE_CONTEXT_CERT_ID = 256,
 };
 
 enum dpe_destroy_context_input_labels_t {
diff --git a/partitions/dice_protection_environment/interface/include/dpe_cmd_encode.h b/partitions/dice_protection_environment/interface/include/dpe_cmd_encode.h
index 64c2644..b06a56d 100644
--- a/partitions/dice_protection_environment/interface/include/dpe_cmd_encode.h
+++ b/partitions/dice_protection_environment/interface/include/dpe_cmd_encode.h
@@ -18,6 +18,7 @@
 
 struct derive_context_input_t {
     int context_handle;
+    uint32_t cert_id;
     bool retain_parent_context;
     bool allow_new_context_to_derive;
     bool create_certificate;
diff --git a/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c b/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
index c0ef7cb..c929f61 100644
--- a/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
+++ b/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
@@ -79,6 +79,8 @@
     QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CONTEXT_HANDLE,
                                (UsefulBufC){ &args->context_handle,
                                              sizeof(args->context_handle) });
+    QCBOREncode_AddUInt64ToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CERT_ID,
+                                args->cert_id);
     QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_RETAIN_PARENT_CONTEXT,
                               args->retain_parent_context);
     QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_DERIVE,
@@ -351,6 +353,7 @@
 
 dpe_error_t
 dpe_derive_context(int                    context_handle,
+                   uint32_t               cert_id,
                    bool                   retain_parent_context,
                    bool                   allow_new_context_to_derive,
                    bool                   create_certificate,
@@ -376,6 +379,7 @@
 
     const struct derive_context_input_t in_args = {
         context_handle,
+        cert_id,
         retain_parent_context,
         allow_new_context_to_derive,
         create_certificate,