DPE: Add new parameters for DeriveContext command

As per spec updates in v1.0 r9, add below new parameters to
DeriveContext command:
* return_certificate [input]
* allow_new_context_to_export [input]
* export_cdi [input]
* new_certificate [output]
* exported_cdi [output]

Signed-off-by: Maulik Patel <maulik.patel@arm.com>
Change-Id: If646caad314f38abcfcca267123aa5d3d477f156
diff --git a/partitions/dice_protection_environment/dpe_boot_data.c b/partitions/dice_protection_environment/dpe_boot_data.c
index 8aa56dd..3bfdda7 100644
--- a/partitions/dice_protection_environment/dpe_boot_data.c
+++ b/partitions/dice_protection_environment/dpe_boot_data.c
@@ -296,13 +296,23 @@
 
     /* Derive RoT layer */
     err = derive_context_request(rot_ctx_handle,
-                                 false,
-                                 true,
+                                 false, /* retain_parent_context */
+                                 true, /* allow_new_context_to_derive */
                                  true, /* create certificate */
                                  &dice_inputs,
-                                 0,
-                                 &plat_ctx_handle,
-                                 &invalid_ctx_handle);
+                                 0, /* client_id */
+                                 0, /* target_locality */
+                                 false, /* return_certificate */
+                                 true, /* allow_new_context_to_export */
+                                 false, /* export_cdi */
+                                 &plat_ctx_handle, /* new_ctx_handle */
+                                 &invalid_ctx_handle, /* new_parent_ctx_handle */
+                                 NULL, /* new_certificate_buf */
+                                 0, /* new_certificate_buf_size */
+                                 NULL, /* new_certificate_actual_size */
+                                 NULL, /* exported_cdi_buf */
+                                 0, /* exported_cdi_buf_size */
+                                 NULL); /* exported_cdi_actual_size */
     if (err != DPE_NO_ERROR) {
         return err;
     }
@@ -319,11 +329,21 @@
     err = derive_context_request(plat_ctx_handle,
                                  false, /* close parent context */
                                  true, /* allow BL2 to derive further */
-                                 false,
+                                 false, /* create_certificate */
                                  &dice_inputs,
-                                 0,
-                                 &plat_ctx_handle,
-                                 &invalid_ctx_handle);
+                                 0, /* client_id */
+                                 0, /* target_locality */
+                                 false, /* return_certificate */
+                                 true, /* allow_new_context_to_export */
+                                 false, /* export_cdi */
+                                 &plat_ctx_handle, /* new_ctx_handle */
+                                 &invalid_ctx_handle, /* new_parent_ctx_handle */
+                                 NULL, /* new_certificate_buf */
+                                 0, /* new_certificate_buf_size */
+                                 NULL, /* new_certificate_actual_size */
+                                 NULL, /* exported_cdi_buf */
+                                 0, /* exported_cdi_buf_size */
+                                 NULL); /* exported_cdi_actual_size */
     if (err != DPE_NO_ERROR) {
         return err;
     }
@@ -336,11 +356,22 @@
         err = derive_context_request(plat_ctx_handle,
                                      true, /* retain parent context */
                                      false, /* do not allow derived context to derive */
-                                     false,
+                                     false, /* create_certificate */
                                      &dice_inputs,
-                                     0,
-                                     &invalid_ctx_handle,
-                                     &plat_ctx_handle);
+                                     0, /* client_id */
+                                     0, /* target_locality */
+                                     false, /* return_certificate */
+                                     true, /* allow_new_context_to_export */
+                                     false, /* export_cdi */
+                                     &invalid_ctx_handle, /* new_ctx_handle */
+                                     &plat_ctx_handle, /* new_parent_ctx_handle */
+                                     NULL, /* new_certificate_buf */
+                                     0, /* new_certificate_buf_size */
+                                     NULL, /* new_certificate_actual_size */
+                                     NULL, /* exported_cdi_buf */
+                                     0, /* exported_cdi_buf_size */
+                                     NULL); /* exported_cdi_actual_size */
+
         if (err != DPE_NO_ERROR) {
             return err;
         }
@@ -365,9 +396,20 @@
     return derive_context_request(plat_ctx_handle,
                                   false, /* close parent context */
                                   true, /* allow AP to derive */
-                                  false,
+                                  false, /* create_certificate */
                                   &dice_inputs,
-                                  0,
-                                  new_ctx_handle,
-                                  &invalid_ctx_handle);
+                                  0, /* client_id */
+                                  0, /* target_locality */
+                                  false, /* return_certificate */
+                                  true, /* allow_new_context_to_export */
+                                  false, /* export_cdi */
+                                  new_ctx_handle, /* new_ctx_handle */
+                                  &invalid_ctx_handle, /* new_parent_ctx_handle */
+                                  NULL, /* new_certificate_buf */
+                                  0, /* new_certificate_buf_size */
+                                  NULL, /* new_certificate_actual_size */
+                                  NULL, /* exported_cdi_buf */
+                                  0, /* exported_cdi_buf_size */
+                                  NULL); /* exported_cdi_actual_size */
+
 }
diff --git a/partitions/dice_protection_environment/dpe_certificate.c b/partitions/dice_protection_environment/dpe_certificate.c
index 2564f5a..9524ccd 100644
--- a/partitions/dice_protection_environment/dpe_certificate.c
+++ b/partitions/dice_protection_environment/dpe_certificate.c
@@ -124,6 +124,14 @@
     }
 }
 
+static void add_cdi_export_claim(struct dpe_cert_encode_ctx *me,
+                                 struct layer_context_t *layer_ctx)
+{
+    QCBOREncode_AddBoolToMapN(&me->cbor_enc_ctx,
+                              DPE_CERT_LABEL_CDI_EXPORT,
+                              layer_ctx->is_cdi_to_be_exported);
+}
+
 static void add_subject_claim(struct dpe_cert_encode_ctx *me,
                               struct layer_context_t *layer_ctx)
 {
@@ -439,6 +447,11 @@
     /* Add key usage claim */
     add_key_usage_claim(&dpe_cert_ctx);
 
+    /* Add CDI exported claim */
+    if (layer_ctx->is_cdi_to_be_exported) {
+        add_cdi_export_claim(&dpe_cert_ctx, layer_ctx);
+    }
+
     /* Finish up creating the token. This is where the actual signature
      * is generated. This finishes up the CBOR encoding too.
      */
@@ -575,3 +588,69 @@
     return close_certificate_chain(&dpe_cert_chain_ctx,
                                    cert_chain_actual_size);
 }
+
+dpe_error_t encode_cdi(const uint8_t *cdi,
+                       size_t cdi_size,
+                       uint8_t *encoded_cdi_buf,
+                       size_t encoded_cdi_buf_size,
+                       size_t *encoded_cdi_actual_size)
+{
+    QCBOREncodeContext encode_ctx;
+    QCBORError encode_err;
+    UsefulBufC out;
+
+    QCBOREncode_Init(&encode_ctx, (UsefulBuf){ encoded_cdi_buf, encoded_cdi_buf_size });
+    QCBOREncode_OpenMap(&encode_ctx);
+
+    /* Encode CDI value as byte string */
+    QCBOREncode_AddBytesToMapN(&encode_ctx,
+                               DPE_LABEL_CDI_ATTEST,
+                               (UsefulBufC){ cdi, cdi_size });
+
+    QCBOREncode_CloseMap(&encode_ctx);
+    encode_err = QCBOREncode_Finish(&encode_ctx, &out);
+
+    /* Check for any encoding errors. */
+    if (encode_err == QCBOR_ERR_BUFFER_TOO_SMALL) {
+        return DPE_INSUFFICIENT_MEMORY;
+    } else if (encode_err != QCBOR_SUCCESS) {
+        return DPE_INTERNAL_ERROR;
+    }
+
+    *encoded_cdi_actual_size = out.len;
+
+    return DPE_NO_ERROR;
+}
+
+dpe_error_t add_encoded_layer_certificate(const uint8_t *cert_buf,
+                                          size_t cert_buf_size,
+                                          uint8_t *encoded_cert_buf,
+                                          size_t encoded_cert_buf_size,
+                                          size_t *encoded_cert_actual_size)
+{
+    QCBOREncodeContext encode_ctx;
+    QCBORError encode_err;
+    UsefulBufC out;
+
+    QCBOREncode_Init(&encode_ctx, (UsefulBuf){ encoded_cert_buf, encoded_cert_buf_size });
+    QCBOREncode_OpenMap(&encode_ctx);
+
+    /* Encode CDI value as byte string */
+    QCBOREncode_AddBytesToMapN(&encode_ctx,
+                               DPE_LABEL_CERT,
+                               (UsefulBufC){ cert_buf, cert_buf_size });
+
+    QCBOREncode_CloseMap(&encode_ctx);
+    encode_err = QCBOREncode_Finish(&encode_ctx, &out);
+
+    /* Check for any encoding errors. */
+    if (encode_err == QCBOR_ERR_BUFFER_TOO_SMALL) {
+        return DPE_INSUFFICIENT_MEMORY;
+    } else if (encode_err != QCBOR_SUCCESS) {
+        return DPE_INTERNAL_ERROR;
+    }
+
+    *encoded_cert_actual_size = out.len;
+
+    return DPE_NO_ERROR;
+}
diff --git a/partitions/dice_protection_environment/dpe_certificate.h b/partitions/dice_protection_environment/dpe_certificate.h
index d591098..eb925db 100644
--- a/partitions/dice_protection_environment/dpe_certificate.h
+++ b/partitions/dice_protection_environment/dpe_certificate.h
@@ -41,9 +41,14 @@
 #define DPE_CERT_LABEL_SUBJECT_PUBLIC_KEY        (DPE_CERT_LABEL_RANGE_BASE - 7)
 #define DPE_CERT_LABEL_KEY_USAGE                 (DPE_CERT_LABEL_RANGE_BASE - 8)
 
-/* Below label is custom and not specified in DICE profile */
+/* Below labels are custom and not specified in DICE profile */
 #define DPE_CERT_LABEL_SW_COMPONENTS             (DPE_CERT_LABEL_RANGE_BASE - 9)
 #define DPE_CERT_LABEL_EXTERNAL_LABEL            (DPE_CERT_LABEL_RANGE_BASE - 10)
+#define DPE_CERT_LABEL_CDI_EXPORT                (DPE_CERT_LABEL_RANGE_BASE - 11)
+#define DPE_LABEL_CDI_ATTEST                     (1)
+#define DPE_LABEL_CDI_SEAL                       (2)
+#define DPE_LABEL_CERT_CHAIN                     (3)
+#define DPE_LABEL_CERT                           (4)
 
 /* Key usage constant per RFC 5280 */
 #define DPE_CERT_KEY_USAGE_CERT_SIGN             (1 << 5);
@@ -86,6 +91,40 @@
                                   size_t cert_chain_buf_size,
                                   size_t *cert_chain_actual_size);
 
+/**
+ * \brief Returns the encoded CDI from raw value.
+ *
+ * \param[in]  cdi                       Pointer to the input CDI buffer.
+ * \param[in]  cdi_size                  Size of the input CDI buffer.
+ * \param[out] encoded_cdi_buf           Pointer to the output encoded CDI buffer.
+ * \param[in]  encoded_cdi_buf_size      Size of the encoded CDI buffer.
+ * \param[out] exported_cdi_actual_size  Actual size of the encoded CDI.
+ *
+ * \return Returns error code of type dpe_error_t
+ */
+dpe_error_t encode_cdi(const uint8_t *cdi,
+                       size_t cdi_size,
+                       uint8_t *encoded_cdi_buf,
+                       size_t encoded_cdi_buf_size,
+                       size_t *encoded_cdi_actual_size);
+
+/**
+ * \brief Adds already encoded certificate to the array.
+ *
+ * \param[in]  cert_buf                  Pointer to the input cert buffer.
+ * \param[in]  cert_buf_size             Size of the input cert buffer.
+ * \param[out] encoded_cert_buf          Pointer to the output encoded cert buffer.
+ * \param[in]  encoded_cert_buf_size     Size of the encoded cert buffer.
+ * \param[out] encoded_cert_actual_size  Actual size of the encoded cert byte array.
+ *
+ * \return Returns error code of type dpe_error_t
+ */
+dpe_error_t add_encoded_layer_certificate(const uint8_t *cert_buf,
+                                          size_t cert_buf_size,
+                                          uint8_t *encoded_cert_buf,
+                                          size_t encoded_cert_buf_size,
+                                          size_t *encoded_cert_actual_size);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/partitions/dice_protection_environment/dpe_cmd_decode.c b/partitions/dice_protection_environment/dpe_cmd_decode.c
index b6dda5a..cda48fa 100644
--- a/partitions/dice_protection_environment/dpe_cmd_decode.c
+++ b/partitions/dice_protection_environment/dpe_cmd_decode.c
@@ -103,6 +103,7 @@
     return DPE_NO_ERROR;
 }
 
+//TODO: Handle the omission of parameters from DPE commands.
 static dpe_error_t decode_derive_context(QCBORDecodeContext *decode_ctx,
                                          QCBOREncodeContext *encode_ctx,
                                          int32_t client_id)
@@ -111,12 +112,20 @@
     QCBORError qcbor_err;
     UsefulBufC out;
     int context_handle;
+    int32_t target_locality;
     bool retain_parent_context;
     bool allow_new_context_to_derive;
     bool create_certificate;
+    bool return_certificate;
+    bool allow_new_context_to_export;
+    bool export_cdi;
     DiceInputValues dice_inputs;
     int new_context_handle;
     int new_parent_context_handle;
+    uint8_t new_certificate_buf[DICE_CERT_SIZE];
+    uint8_t exported_cdi_buf[DICE_MAX_ENCODED_CDI_SIZE];
+    size_t new_certificate_actual_size = 0;
+    size_t exported_cdi_actual_size = 0;
 
     /* Decode DeriveContext command */
     QCBORDecode_EnterMap(decode_ctx, NULL);
@@ -142,6 +151,22 @@
         return dpe_err;
     }
 
+    QCBORDecode_GetByteStringInMapN(decode_ctx, DPE_DERIVE_CONTEXT_TARGET_LOCALITY,
+                                    &out);
+    if (out.len != sizeof(target_locality)) {
+        return DPE_INVALID_COMMAND;
+    }
+    memcpy(&target_locality, out.ptr, out.len);
+
+    QCBORDecode_GetBoolInMapN(decode_ctx, DPE_DERIVE_CONTEXT_RETURN_CERTIFICATE,
+                              &return_certificate);
+
+    QCBORDecode_GetBoolInMapN(decode_ctx, DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_EXPORT,
+                              &allow_new_context_to_export);
+
+    QCBORDecode_GetBoolInMapN(decode_ctx, DPE_DERIVE_CONTEXT_EXPORT_CDI,
+                              &export_cdi);
+
     QCBORDecode_ExitMap(decode_ctx);
 
     /* Exit top level array */
@@ -156,8 +181,18 @@
     dpe_err = derive_context_request(context_handle, retain_parent_context,
                                      allow_new_context_to_derive, create_certificate,
                                      &dice_inputs, client_id,
+                                     target_locality,
+                                     return_certificate,
+                                     allow_new_context_to_export,
+                                     export_cdi,
                                      &new_context_handle,
-                                     &new_parent_context_handle);
+                                     &new_parent_context_handle,
+                                     new_certificate_buf,
+                                     sizeof(new_certificate_buf),
+                                     &new_certificate_actual_size,
+                                     exported_cdi_buf,
+                                     sizeof(exported_cdi_buf),
+                                     &exported_cdi_actual_size);
     if (dpe_err != DPE_NO_ERROR) {
         return dpe_err;
     }
@@ -174,6 +209,19 @@
                                DPE_DERIVE_CONTEXT_PARENT_CONTEXT_HANDLE,
                                (UsefulBufC){ &new_parent_context_handle,
                                              sizeof(new_parent_context_handle) });
+
+    /* The certificate is already encoded into a CBOR array by the function
+     * add_encoded_layer_certificate. Add it as a byte string so that its
+     * decoding can be skipped and the CBOR returned to the caller.
+     */
+    QCBOREncode_AddBytesToMapN(encode_ctx, DPE_DERIVE_CONTEXT_NEW_CERTIFICATE,
+                               (UsefulBufC){ new_certificate_buf,
+                                             new_certificate_actual_size });
+
+    QCBOREncode_AddBytesToMapN(encode_ctx, DPE_DERIVE_CONTEXT_EXPORTED_CDI,
+                               (UsefulBufC){ exported_cdi_buf,
+                                             exported_cdi_actual_size });
+
     QCBOREncode_CloseMap(encode_ctx);
 
     QCBOREncode_CloseArray(encode_ctx);
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.c b/partitions/dice_protection_environment/dpe_context_mngr.c
index 4e05da3..3655c2f 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.c
+++ b/partitions/dice_protection_environment/dpe_context_mngr.c
@@ -63,12 +63,17 @@
 static void set_context_to_default(int i)
 {
     component_ctx_array[i].in_use = false;
-    component_ctx_array[i].is_leaf = false;
+    component_ctx_array[i].is_allowed_to_derive = true;
+    /* export CDI attribute is inherited and once disabled, a derived context
+     * and subsequent derivations cannot export CDI, hence enable by default
+     */
+    component_ctx_array[i].is_export_cdi_allowed = true;
     component_ctx_array[i].nonce = INVALID_NONCE_VALUE;
     component_ctx_array[i].parent_idx = INVALID_COMPONENT_IDX;
     component_ctx_array[i].linked_layer_idx = INVALID_LAYER_IDX;
     (void)memset(&component_ctx_array[i].data, 0, sizeof(struct component_context_data_t));
     //TODO: Question: how to initialise MHU Id mapping?
+    component_ctx_array[i].target_locality = 0;
     /* Allow component to be derived by default */
 }
 
@@ -76,6 +81,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;
     (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);
@@ -259,6 +265,39 @@
     return status;
 }
 
+static dpe_error_t get_encoded_cdi_to_export(struct layer_context_t *layer_ctx,
+                                             uint8_t *exported_cdi_buf,
+                                             size_t exported_cdi_buf_size,
+                                             size_t *exported_cdi_actual_size)
+{
+    uint8_t cdi_buf[DICE_CDI_SIZE];
+    size_t cdi_size;
+    psa_status_t status;
+    dpe_error_t err;
+
+    /* Get CDI value */
+    status = get_layer_cdi_value(layer_ctx,
+                                 cdi_buf,
+                                 sizeof(cdi_buf),
+                                 &cdi_size);
+    if (status != PSA_SUCCESS) {
+        return DPE_INTERNAL_ERROR;
+    }
+
+    /* Encode CDI value */
+    err = encode_cdi(cdi_buf,
+                     cdi_size,
+                     exported_cdi_buf,
+                     exported_cdi_buf_size,
+                     exported_cdi_actual_size);
+    if (err != DPE_NO_ERROR) {
+        return err;
+    }
+    layer_ctx->is_cdi_to_be_exported = true;
+
+    return DPE_NO_ERROR;
+}
+
 static dpe_error_t create_layer_certificate(uint16_t layer_idx)
 {
     uint16_t parent_layer_idx;
@@ -357,9 +396,10 @@
     layer_ctx_array[derived_ctx_layer].parent_layer_idx = parent_ctx_layer;
 }
 
-static inline bool is_input_client_id_valid(int32_t client_id)
+static inline bool is_input_client_id_valid(int32_t client_id, int32_t target_locality)
 {
-    //TODO: Waiting for implementation
+    // TODO: FIXME
+    //    return (client_id == target_locality);
     return true;
 }
 
@@ -467,13 +507,25 @@
                                    bool create_certificate,
                                    const DiceInputValues *dice_inputs,
                                    int32_t client_id,
+                                   int32_t target_locality,
+                                   bool return_certificate,
+                                   bool allow_new_context_to_export,
+                                   bool export_cdi,
                                    int *new_context_handle,
-                                   int *new_parent_context_handle)
+                                   int *new_parent_context_handle,
+                                   uint8_t *new_certificate_buf,
+                                   size_t new_certificate_buf_size,
+                                   size_t *new_certificate_actual_size,
+                                   uint8_t *exported_cdi_buf,
+                                   size_t exported_cdi_buf_size,
+                                   size_t *exported_cdi_actual_size)
 {
     dpe_error_t err;
     struct component_context_t *parent_ctx, *derived_ctx;
-    uint16_t parent_ctx_idx;
+    uint16_t parent_ctx_idx, linked_layer_idx;
     int free_component_idx;
+    struct layer_context_t *layer_ctx;
+    psa_status_t status;
 
     log_derive_context(input_ctx_handle, retain_parent_context,
                        allow_new_context_to_derive, create_certificate, dice_inputs,
@@ -490,6 +542,10 @@
     }
 #endif /* DPE_TEST_MODE */
 
+    if (export_cdi && !create_certificate) {
+        return DPE_INVALID_ARGUMENT;
+    }
+
     /* Validate dice inputs */
     if (!is_dice_input_valid(dice_inputs)) {
         return DPE_INVALID_ARGUMENT;
@@ -509,8 +565,13 @@
 
     parent_ctx = &component_ctx_array[parent_ctx_idx];
 
+    /* Check if parent context is allowed to derive */
+    if (!parent_ctx->is_allowed_to_derive) {
+        return DPE_INVALID_ARGUMENT;
+    }
+
     //TODO:  Question: how to get mhu id of incoming request?
-    if (!is_input_client_id_valid(client_id)) {
+    if (!is_input_client_id_valid(client_id, parent_ctx->target_locality)) {
         return DPE_INVALID_ARGUMENT;
     }
 
@@ -521,11 +582,25 @@
     }
 
     derived_ctx = &component_ctx_array[free_component_idx];
+    if (parent_ctx->is_export_cdi_allowed && allow_new_context_to_export) {
+        /* If parent context has export enabled and input allow_new_context_to_export
+         * is true, then allow context CDI to be exported for derived context
+         */
+        derived_ctx->is_export_cdi_allowed = true;
+    } else {
+        /* Export of new context CDI is NOT allowed */
+        derived_ctx->is_export_cdi_allowed = false;
+        if (export_cdi) {
+            return DPE_INVALID_ARGUMENT;
+        }
+    }
+
     /* Copy dice input to the new derived component context */
     err = copy_dice_input(derived_ctx, dice_inputs);
     if (err != DPE_NO_ERROR) {
         return err;
     }
+    derived_ctx->target_locality = target_locality;
 
     /* Update parent idx in new derived component context */
     derived_ctx->parent_idx = parent_ctx_idx;
@@ -551,7 +626,7 @@
         parent_ctx->nonce = INVALID_NONCE_VALUE;
     }
 
-    if (allow_new_context_to_derive) {
+    if (!export_cdi) {
         /* Return handle to derived context */
         *new_context_handle = SET_IDX(*new_context_handle, free_component_idx);
         err = renew_nonce(new_context_handle);
@@ -568,7 +643,38 @@
     }
 
     if (create_certificate) {
-        err = create_layer_certificate(derived_ctx->linked_layer_idx);
+        linked_layer_idx = derived_ctx->linked_layer_idx;
+        assert(linked_layer_idx < MAX_NUM_OF_LAYERS);
+        layer_ctx = &layer_ctx_array[linked_layer_idx];
+        layer_ctx->is_cdi_to_be_exported = export_cdi;
+
+        err = create_layer_certificate(linked_layer_idx);
+        if (err != DPE_NO_ERROR) {
+            return err;
+        }
+
+        if (return_certificate) {
+            if (new_certificate_buf_size < layer_ctx->data.cert_buf_len) {
+                return DPE_INVALID_ARGUMENT;
+            }
+
+            /* Encode and return generated layer certificate */
+            err = add_encoded_layer_certificate(layer_ctx->data.cert_buf,
+                                                layer_ctx->data.cert_buf_len,
+                                                new_certificate_buf,
+                                                new_certificate_buf_size,
+                                                new_certificate_actual_size);
+            if (err != DPE_NO_ERROR) {
+                return err;
+            }
+        }
+    }
+
+    if (export_cdi) {
+        err = get_encoded_cdi_to_export(layer_ctx,
+                                        exported_cdi_buf,
+                                        exported_cdi_buf_size,
+                                        exported_cdi_actual_size);
         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 df6cafd..625834f 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.h
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -19,6 +19,8 @@
 #endif
 
 #define DICE_CERT_SIZE  3072
+/* Below encoded CDI size accomodate both Attest and Seal CDI */
+#define DICE_MAX_ENCODED_CDI_SIZE ((2 * DICE_CDI_SIZE) + 16)
 
 #define INVALID_HANDLE 0xFFFFFFFF
 #define INVALID_COMPONENT_IDX 0xFFFF
@@ -58,10 +60,13 @@
 struct component_context_t {
     struct component_context_data_t data;   /* Component context data */
     bool in_use;                            /* Flag to indicate if element is used */
-    bool is_leaf;                           /* Is the component allowed to derive */
+    bool is_allowed_to_derive;              /* Is the component allowed to derive */
+    bool is_export_cdi_allowed;             /* Is CDI allowed to export */
     uint16_t nonce;                         /* Context handle nonce for the component */
     uint16_t parent_idx;                    /* Parent component's index */
     uint16_t linked_layer_idx;              /* Layer component is linked to */
+    int32_t  target_locality;               /* Identifies the locality to which the
+                                             * derived context will be bound */
     uint32_t expected_mhu_id;               /* Expected mhu to authorise derivation */
 };
 
@@ -90,6 +95,7 @@
     uint8_t attest_cdi_hash_input[DPE_HASH_ALG_SIZE];
     enum layer_state_t state;
     bool is_external_pub_key_provided;
+    bool is_cdi_to_be_exported;
 };
 
 /**
@@ -105,10 +111,10 @@
  * \brief Derives a component context and optionally creates certificate
  *        chain.
  *
- * \param[in]  input_context_handle        Input handle to parent component context
+ * \param[in]  input_context_handle        Input handle to parent component context.
  * \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
+ *                                         is calling DPE commands multiple times.
  * \param[in]  allow_new_context_to_derive Flag to indicate if derived context can
  *                                         derive further.
  * \param[in]  create_certificate          Flag to indicate if certificate needs
@@ -117,19 +123,47 @@
  * \param[in]  dice_inputs                 Pointer to dice_input buffer.
  * \param[in]  client_id                   Identifier of the client calling the
  *                                         service.
+ * \param[in]  target_locality             Identifier of the locality to which the
+ *                                         derived context should be bound to.
+ * \param[in]  return_certificate          Indicates whether to return the generated
+ *                                         certificate when create_certificate is true.
+ * \param[in]  allow_new_context_to_export Indicates whether the DPE permits export of
+ *                                         the CDI from the newly derived context.
+ * \param[in]  export_cdi                  Indicates whether to export derived CDI.
  * \param[out] new_context_handle          A new handle for derived context.
  * \param[out] new_parent_context_handle   A new handle for parent context.
+ * \param[out] new_certificate_buf         If create_certificate and return_certificate
+ *                                         are both true, this argument holds the new
+ *                                         certificate generated for the new context.
+ * \param[in]  new_certificate_buf_size    Size of the allocated buffer for
+ *                                         new certificate.
+ * \param[out] new_certificate_actual_size Actual size of the new certificate.
+ * \param[out] exported_cdi_buf            If export_cdi is true, this is the
+ *                                         exported CDI value.
+ * \param[in]  exported_cdi_buf_size       Size of the allocated buffer for
+ *                                         exported CDI.
+ * \param[out] exported_cdi_actual_size    Actual size of the exported CDI.
  *
  * \return Returns error code of type dpe_error_t
  */
-dpe_error_t derive_context_request(int input_context_handle,
+dpe_error_t derive_context_request(int input_ctx_handle,
                                    bool retain_parent_context,
                                    bool allow_new_context_to_derive,
                                    bool create_certificate,
                                    const DiceInputValues *dice_inputs,
                                    int32_t client_id,
+                                   int32_t target_locality,
+                                   bool return_certificate,
+                                   bool allow_new_context_to_export,
+                                   bool export_cdi,
                                    int *new_context_handle,
-                                   int *new_parent_context_handle);
+                                   int *new_parent_context_handle,
+                                   uint8_t *new_certificate_buf,
+                                   size_t new_certificate_buf_size,
+                                   size_t *new_certificate_actual_size,
+                                   uint8_t *exported_cdi_buf,
+                                   size_t exported_cdi_buf_size,
+                                   size_t *exported_cdi_actual_size);
 
 /**
  * \brief Destroys a component context and optionally depending on argument
diff --git a/partitions/dice_protection_environment/dpe_crypto_config.h b/partitions/dice_protection_environment/dpe_crypto_config.h
index a41052a..e89c23a 100644
--- a/partitions/dice_protection_environment/dpe_crypto_config.h
+++ b/partitions/dice_protection_environment/dpe_crypto_config.h
@@ -22,13 +22,14 @@
 #define DPE_CDI_KEY_TYPE  PSA_KEY_TYPE_DERIVE
 #define DPE_CDI_KEY_ALG   PSA_ALG_HKDF(PSA_ALG_SHA_256)
 #define DPE_CDI_KEY_BITS  256
-#define DPE_CDI_KEY_USAGE PSA_KEY_USAGE_DERIVE
+#define DPE_CDI_KEY_USAGE PSA_KEY_USAGE_DERIVE | PSA_KEY_USAGE_EXPORT
 
 /* Below labels as per
  * https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#attestation-cdi
  */
 #define DPE_ATTEST_CDI_LABEL "CDI_Attest"
 #define DPE_SEAL_CDI_LABEL   "CDI_Seal"
+#define DPE_ATTEST_EXPORTED_CDI_LABEL "Exported_CDI_Attest" /* Custom Label - yet to be specified */
 
 #define DPE_ATTEST_KEY_CURVE_TYPE PSA_ECC_FAMILY_SECP_R1
 #define DPE_ATTEST_KEY_TYPE       PSA_KEY_TYPE_ECC_KEY_PAIR(DPE_ATTEST_KEY_CURVE_TYPE)
diff --git a/partitions/dice_protection_environment/dpe_crypto_interface.c b/partitions/dice_protection_environment/dpe_crypto_interface.c
index 5f4e1a5..9a18d91 100644
--- a/partitions/dice_protection_environment/dpe_crypto_interface.c
+++ b/partitions/dice_protection_environment/dpe_crypto_interface.c
@@ -15,6 +15,7 @@
 #include "tfm_crypto_defs.h"
 
 static const char attest_cdi_label[] = DPE_ATTEST_CDI_LABEL;
+static const char exported_attest_cdi_label[] = DPE_ATTEST_EXPORTED_CDI_LABEL;
 static const char attest_key_pair_label[] = DPE_ATTEST_KEY_PAIR_LABEL;
 static const char id_label[] = DPE_ID_LABEL;
 static const uint8_t attest_key_salt[] = DPE_ATTEST_KEY_SALT;
@@ -92,13 +93,25 @@
 
     /* Perform CDI derivation */
     /* Parent layer CDI is the base key (input secret to key derivation) */
-    return perform_derivation(parent_layer_ctx->data.cdi_key_id,
-                              &derive_key_attr,
-                              (uint8_t *) &attest_cdi_label[0],
-                              sizeof(attest_cdi_label),
-                              layer_ctx->attest_cdi_hash_input,
-                              sizeof(layer_ctx->attest_cdi_hash_input),
-                              &layer_ctx->data.cdi_key_id);
+
+    if (layer_ctx->is_cdi_to_be_exported) {
+        return perform_derivation(parent_layer_ctx->data.cdi_key_id,
+                                &derive_key_attr,
+                                (uint8_t *) &exported_attest_cdi_label[0],
+                                sizeof(exported_attest_cdi_label),
+                                layer_ctx->attest_cdi_hash_input,
+                                sizeof(layer_ctx->attest_cdi_hash_input),
+                                &layer_ctx->data.cdi_key_id);
+
+    } else {
+        return perform_derivation(parent_layer_ctx->data.cdi_key_id,
+                                &derive_key_attr,
+                                (uint8_t *) &attest_cdi_label[0],
+                                sizeof(attest_cdi_label),
+                                layer_ctx->attest_cdi_hash_input,
+                                sizeof(layer_ctx->attest_cdi_hash_input),
+                                &layer_ctx->data.cdi_key_id);
+    }
 }
 
 psa_status_t derive_attestation_key(struct layer_context_t *layer_ctx)
@@ -266,3 +279,22 @@
     (void)psa_key_derivation_abort(&op);
     return status;
 }
+
+psa_status_t get_layer_cdi_value(const struct layer_context_t *layer_ctx,
+                                 uint8_t *cdi_buf,
+                                 size_t cdi_buf_size,
+                                 size_t *cdi_actual_size)
+{
+    psa_status_t status;
+
+    //TODO: Sealing CDI to be added later
+    status = psa_export_key(layer_ctx->data.cdi_key_id,
+                            cdi_buf,
+                            sizeof(cdi_buf),
+                            &cdi_actual_size);
+    if (status != PSA_SUCCESS) {
+        *cdi_actual_size = 0;
+    }
+
+    return status;
+}
diff --git a/partitions/dice_protection_environment/dpe_crypto_interface.h b/partitions/dice_protection_environment/dpe_crypto_interface.h
index 539dd35..880398d 100644
--- a/partitions/dice_protection_environment/dpe_crypto_interface.h
+++ b/partitions/dice_protection_environment/dpe_crypto_interface.h
@@ -88,6 +88,20 @@
 psa_status_t derive_cdi_id(psa_key_id_t attest_key_id, uint8_t *cdi_id,
                            size_t cdi_id_size);
 
+/**
+ * \brief Gets the layer's CDI value.
+ *
+ * \param[in]  layer_ctx       Pointer to current layer context.
+ * \param[out] cdi_buf         Pointer to output CDI Buffer.
+ * \param[in]  cdi_buf_size    Size of CDI buffer.
+ * \param[out] cdi_actual_size Actual size of the CDI value.
+ *
+ * \return Returns error code as specified in \ref psa_status_t
+ */
+psa_status_t get_layer_cdi_value(const struct layer_context_t *layer_ctx,
+                                 uint8_t *cdi_buf,
+                                 size_t cdi_buf_size,
+                                 size_t *cdi_actual_size);
 #ifdef __cplusplus
 }
 #endif
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 952bcd3..08d0ad2 100644
--- a/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
+++ b/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
@@ -53,8 +53,27 @@
  *                                         intermediate certificate. True only if
  *                                         it is the last component in the layer.
  * \param[in]  dice_inputs                 DICE input values.
+ * \param[in]  target_locality             Identifies the locality to which the
+ *                                         derived context will be bound. Could be
+ *                                         MHU id.
+ * \param[in]  return_certificate          Indicates whether to return the generated
+ *                                         certificate when create_certificate is true.
+ * \param[in]  allow_new_context_to_export Indicates whether the DPE permits export of
+ *                                         the CDI from the newly derived context.
+ * \param[in]  export_cdi                  Indicates whether to export derived CDI.
  * \param[out] new_context_handle          New handle for the derived context.
  * \param[out] new_parent_context_handle   New handle for the parent context.
+ * \param[out] new_certificate_buf         If create_certificate and return_certificate
+ *                                         are both true, this argument holds the new
+ *                                         certificate generated for the new context
+ * \param[in]  new_certificate_buf_size    Size of the allocated buffer for
+ *                                         new certificate.
+ * \param[out] new_certificate_actual_size Actual size of the new certificate.
+ * \param[out] exported_cdi_buf            If export_cdi is true, this is the
+ *                                         exported CDI value.
+ * \param[in]  exported_cdi_buf_size       Size of the allocated buffer for
+ *                                         exported cdi.
+ * \param[out] exported_cdi_actual_size    Actual size of the exported cdi.
  *
  * \return Returns error code of type dpe_error_t
  */
@@ -64,8 +83,18 @@
                    bool                   allow_new_context_to_derive,
                    bool                   create_certificate,
                    const DiceInputValues *dice_inputs,
+                   int32_t                target_locality,
+                   bool                   return_certificate,
+                   bool                   allow_new_context_to_export,
+                   bool                   export_cdi,
                    int                   *new_context_handle,
-                   int                   *new_parent_context_handle);
+                   int                   *new_parent_context_handle,
+                   uint8_t               *new_certificate_buf,
+                   size_t                 new_certificate_buf_size,
+                   size_t                *new_certificate_actual_size,
+                   uint8_t               *exported_cdi_buf,
+                   size_t                 exported_cdi_buf_size,
+                   size_t                *exported_cdi_actual_size);
 
 /**
  * \brief Destroys a DPE context.
diff --git a/partitions/dice_protection_environment/interface/include/dpe_client.h b/partitions/dice_protection_environment/interface/include/dpe_client.h
index eebef48..aa9d511 100644
--- a/partitions/dice_protection_environment/interface/include/dpe_client.h
+++ b/partitions/dice_protection_environment/interface/include/dpe_client.h
@@ -57,6 +57,10 @@
     DPE_DERIVE_CONTEXT_NEW_SESSION_IS_MIGRATABLE = 6,
     DPE_DERIVE_CONTEXT_INPUT_DATA = 7,
     DPE_DERIVE_CONTEXT_INTERNAL_INPUTS = 8,
+    DPE_DERIVE_CONTEXT_TARGET_LOCALITY = 9,
+    DPE_DERIVE_CONTEXT_RETURN_CERTIFICATE = 10,
+    DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_EXPORT = 11,
+    DPE_DERIVE_CONTEXT_EXPORT_CDI = 12,
 };
 
 enum dpe_destroy_context_input_labels_t {
@@ -68,6 +72,8 @@
     DPE_DERIVE_CONTEXT_NEW_CONTEXT_HANDLE = 1,
     DPE_DERIVE_CONTEXT_NEW_SESSION_RESPONDER_HANDSHAKE = 2,
     DPE_DERIVE_CONTEXT_PARENT_CONTEXT_HANDLE = 3,
+    DPE_DERIVE_CONTEXT_NEW_CERTIFICATE = 4,
+    DPE_DERIVE_CONTEXT_EXPORTED_CDI = 5,
 };
 
 enum dpe_certify_key_input_labels_t {
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 2c4757d..4d1b551 100644
--- a/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
+++ b/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
@@ -18,11 +18,19 @@
     bool allow_new_context_to_derive;
     bool create_certificate;
     const DiceInputValues *dice_inputs;
+    int32_t target_locality;
+    bool return_certificate;
+    bool allow_new_context_to_export;
+    bool export_cdi;
 };
 
 struct derive_context_output_t {
     int new_context_handle;
     int new_parent_context_handle;
+    const uint8_t *new_certificate;
+    size_t new_certificate_size;
+    const uint8_t *exported_cdi;
+    size_t exported_cdi_size;
 };
 
 struct destroy_context_input_t {
@@ -118,6 +126,15 @@
     QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CREATE_CERTIFICATE,
                               args->create_certificate);
     encode_dice_inputs(&encode_ctx, args->dice_inputs);
+    QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_TARGET_LOCALITY,
+                               (UsefulBufC){ &args->target_locality,
+                                             sizeof(args->target_locality) });
+    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_RETURN_CERTIFICATE,
+                              args->return_certificate);
+    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_EXPORT,
+                              args->allow_new_context_to_export);
+    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_EXPORT_CDI,
+                              args->export_cdi);
     QCBOREncode_CloseMap(&encode_ctx);
 
     QCBOREncode_CloseArray(&encode_ctx);
@@ -186,6 +203,18 @@
         }
         memcpy(&args->new_parent_context_handle, out.ptr, out.len);
 
+        QCBORDecode_GetByteStringInMapN(&decode_ctx,
+                                        DPE_DERIVE_CONTEXT_NEW_CERTIFICATE,
+                                        &out);
+        args->new_certificate = out.ptr;
+        args->new_certificate_size = out.len;
+
+        QCBORDecode_GetByteStringInMapN(&decode_ctx,
+                                        DPE_DERIVE_CONTEXT_EXPORTED_CDI,
+                                        &out);
+        args->exported_cdi = out.ptr;
+        args->exported_cdi_size = out.len;
+
         QCBORDecode_ExitMap(&decode_ctx);
     }
 
@@ -291,13 +320,24 @@
     return QCBORDecode_Finish(&decode_ctx);
 }
 
-dpe_error_t dpe_derive_context(int context_handle,
-                               bool retain_parent_context,
-                               bool allow_new_context_to_derive,
-                               bool create_certificate,
-                               const DiceInputValues *dice_inputs,
-                               int *new_context_handle,
-                               int *new_parent_context_handle)
+dpe_error_t
+dpe_derive_context(int                    context_handle,
+                   bool                   retain_parent_context,
+                   bool                   allow_new_context_to_derive,
+                   bool                   create_certificate,
+                   const DiceInputValues *dice_inputs,
+                   int32_t                target_locality,
+                   bool                   return_certificate,
+                   bool                   allow_new_context_to_export,
+                   bool                   export_cdi,
+                   int                   *new_context_handle,
+                   int                   *new_parent_context_handle,
+                   uint8_t               *new_certificate_buf,
+                   size_t                 new_certificate_buf_size,
+                   size_t                *new_certificate_actual_size,
+                   uint8_t               *exported_cdi_buf,
+                   size_t                 exported_cdi_buf_size,
+                   size_t                *exported_cdi_actual_size)
 {
     int32_t service_err;
     dpe_error_t dpe_err;
@@ -311,6 +351,10 @@
         allow_new_context_to_derive,
         create_certificate,
         dice_inputs,
+        target_locality,
+        return_certificate,
+        allow_new_context_to_export,
+        export_cdi,
     };
     struct derive_context_output_t out_args;
 
@@ -337,6 +381,20 @@
     *new_context_handle = out_args.new_context_handle;
     *new_parent_context_handle = out_args.new_parent_context_handle;
 
+    if (out_args.new_certificate_size > new_certificate_buf_size) {
+        return DPE_INVALID_ARGUMENT;
+    }
+    memcpy(new_certificate_buf, out_args.new_certificate,
+           out_args.new_certificate_size);
+    *new_certificate_actual_size = out_args.new_certificate_size;
+
+    if (out_args.exported_cdi_size > exported_cdi_buf_size) {
+        return DPE_INVALID_ARGUMENT;
+    }
+    memcpy(exported_cdi_buf, out_args.exported_cdi,
+           out_args.exported_cdi_size);
+    *exported_cdi_actual_size = out_args.exported_cdi_size;
+
     return DPE_NO_ERROR;
 }