DPE: Add context manager code

It manages handles, component contexts and layer contexts.

Signed-off-by: Maulik Patel <maulik.patel@arm.com>
Change-Id: I280a5096792f7c286695e9d567e09e00850ac101
diff --git a/partitions/dice_protection_environment/CMakeLists.txt b/partitions/dice_protection_environment/CMakeLists.txt
index 447ee8a..b5eb637 100644
--- a/partitions/dice_protection_environment/CMakeLists.txt
+++ b/partitions/dice_protection_environment/CMakeLists.txt
@@ -20,6 +20,7 @@
 target_sources(tfm_app_rot_partition_dpe
     PRIVATE
         dpe_cmd_decode.c
+        dpe_context_mngr.c
         dpe_impl.c
         dpe_log.c
         dpe_req_mngr.c
@@ -45,6 +46,11 @@
         ${CMAKE_BINARY_DIR}/generated/secure_fw/partitions/dice_protection_environment
 )
 
+target_compile_definitions(tfm_app_rot_partition_dpe
+    PRIVATE
+        $<$<BOOL:${TFM_S_REG_TEST}>:TFM_S_REG_TEST>
+)
+
 target_link_libraries(tfm_app_rot_partition_dpe
     PRIVATE
         tfm_sprt
diff --git a/partitions/dice_protection_environment/dpe_cmd_decode.c b/partitions/dice_protection_environment/dpe_cmd_decode.c
index 1f14225..0c4215f 100644
--- a/partitions/dice_protection_environment/dpe_cmd_decode.c
+++ b/partitions/dice_protection_environment/dpe_cmd_decode.c
@@ -10,6 +10,7 @@
 #include <string.h>
 
 #include "dpe_client.h"
+#include "dpe_context_mngr.h"
 #include "dpe_impl.h"
 #include "qcbor/qcbor_encode.h"
 #include "qcbor/qcbor_decode.h"
@@ -103,7 +104,8 @@
 }
 
 static dpe_error_t decode_derive_child(QCBORDecodeContext *decode_ctx,
-                                       QCBOREncodeContext *encode_ctx)
+                                       QCBOREncodeContext *encode_ctx,
+                                       int32_t client_id)
 {
     dpe_error_t dpe_err;
     QCBORError qcbor_err;
@@ -151,10 +153,11 @@
         return DPE_INVALID_COMMAND;
     }
 
-    dpe_err = dpe_derive_child_impl(context_handle, retain_parent_context,
-                                    allow_child_to_derive, create_certificate,
-                                    &dice_inputs, &new_child_context_handle,
-                                    &new_parent_context_handle);
+    dpe_err = derive_child_request(context_handle, retain_parent_context,
+                                   allow_child_to_derive, create_certificate,
+                                   &dice_inputs, client_id,
+                                   &new_child_context_handle,
+                                   &new_parent_context_handle);
     if (dpe_err != DPE_NO_ERROR) {
         return dpe_err;
     }
@@ -305,7 +308,7 @@
     if (qcbor_err == QCBOR_SUCCESS) {
         switch (command_id) {
         case DPE_DERIVE_CHILD:
-            dpe_err = decode_derive_child(&decode_ctx, &encode_ctx);
+            dpe_err = decode_derive_child(&decode_ctx, &encode_ctx, client_id);
             break;
         case DPE_CERTIFY_KEY:
             dpe_err = decode_certify_key(&decode_ctx, &encode_ctx);
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.c b/partitions/dice_protection_environment/dpe_context_mngr.c
new file mode 100644
index 0000000..d2bc2ab
--- /dev/null
+++ b/partitions/dice_protection_environment/dpe_context_mngr.c
@@ -0,0 +1,412 @@
+/*
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#include "dpe_context_mngr.h"
+#include <assert.h>
+#include <string.h>
+#include "dice_protection_environment.h"
+#include "dpe_log.h"
+#include "psa/crypto.h"
+
+static struct component_context_t component_ctx_array[MAX_NUM_OF_COMPONENTS];
+static struct layer_context_t layer_ctx_array[MAX_NUM_OF_LAYERS];
+
+static int get_free_component_context_index(void)
+{
+    int i;
+
+    for (i = 0; i < MAX_NUM_OF_COMPONENTS; i++) {
+        if (!component_ctx_array[i].in_use) {
+            break;
+        }
+    }
+
+    if (i >= MAX_NUM_OF_COMPONENTS) {
+        /* No free index left in the array -- all used up! */
+        return -1;
+    }
+
+    return i;
+}
+
+static inline void invalidate_handle(int *handle)
+{
+    *handle = INVALID_HANDLE;
+}
+
+static dpe_error_t renew_nonce(int *handle)
+{
+    uint16_t nonce;
+
+    psa_status_t status = psa_generate_random((uint8_t *)&nonce, sizeof(nonce));
+    if (status != PSA_SUCCESS) {
+        return DPE_INTERNAL_ERROR;
+    }
+    *handle = SET_NONCE(*handle, nonce);
+
+    return DPE_NO_ERROR;
+}
+
+static dpe_error_t generate_new_handle(int *out_handle)
+{
+    /* Find the free component array element */
+    int free_component_idx = get_free_component_context_index();
+    if (free_component_idx < 0) {
+        return DPE_INSUFFICIENT_MEMORY;
+    }
+
+    *out_handle = SET_IDX(*out_handle, free_component_idx);
+
+    return renew_nonce(out_handle);
+}
+
+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].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?
+    /* Allow component to be derived by default */
+}
+
+static void invalidate_layer(int i)
+{
+    layer_ctx_array[i].state = LAYER_STATE_CLOSED;
+    layer_ctx_array[i].parent_layer_idx = INVALID_LAYER_IDX;
+    (void)memset(&layer_ctx_array[i].data, 0, sizeof(struct layer_context_data_t));
+}
+
+void initialise_all_dpe_contexts(void)
+{
+    int i;
+
+    for (i = 0; i < MAX_NUM_OF_COMPONENTS; i++) {
+        set_context_to_default(i);
+    }
+
+    for (i = 0; i < MAX_NUM_OF_LAYERS; i++) {
+        invalidate_layer(i);
+    }
+}
+
+static dpe_error_t copy_dice_input(struct component_context_t *dest_ctx,
+                                   const DiceInputValues *dice_inputs)
+{
+    size_t hash_len;
+    psa_status_t status;
+
+    memcpy(&dest_ctx->data.measurement_value, dice_inputs->code_hash,
+           DICE_HASH_SIZE);
+    memcpy(&dest_ctx->data.measurement_descriptor,
+           dice_inputs->code_descriptor,
+           dice_inputs->code_descriptor_size);
+
+    dest_ctx->data.measurement_descriptor_size =
+                                      dice_inputs->code_descriptor_size;
+
+    memcpy(&dest_ctx->data.signer_id, dice_inputs->authority_hash, DICE_HASH_SIZE);
+    memcpy(&dest_ctx->data.signer_id_descriptor,
+           dice_inputs->authority_descriptor,
+           dice_inputs->authority_descriptor_size);
+
+    dest_ctx->data.signer_id_descriptor_size =
+                                         dice_inputs->authority_descriptor_size;
+
+    if (dice_inputs->config_type == kDiceConfigTypeInline) {
+        /* Copy config_value */
+        memcpy(&dest_ctx->data.config_value, dice_inputs->config_value,
+               DICE_INLINE_CONFIG_SIZE);
+
+    } else {
+        /* Copy config descriptor */
+        memcpy(&dest_ctx->data.config_descriptor, dice_inputs->config_descriptor,
+                dice_inputs->config_descriptor_size);
+        dest_ctx->data.config_descriptor_size = dice_inputs->config_descriptor_size;
+
+        /* Calculate config value as hash of input config descriptor */
+        status = psa_hash_compute(PSA_ALG_SHA_256,
+                                  dice_inputs->config_descriptor,
+                                  dice_inputs->config_descriptor_size,
+                                  dest_ctx->data.config_value,
+                                  sizeof(dest_ctx->data.config_value),
+                                  &hash_len);
+
+        if (status != PSA_SUCCESS) {
+            return DPE_INTERNAL_ERROR;
+        }
+    }
+
+    dest_ctx->data.mode = dice_inputs->mode;
+    memcpy(&dest_ctx->data.hidden, dice_inputs->hidden, DICE_HIDDEN_SIZE);
+
+    return DPE_NO_ERROR;
+}
+
+static bool is_dice_input_valid(const DiceInputValues *dice_inputs)
+{
+    if ((dice_inputs->code_descriptor_size > DICE_CODE_DESCRIPTOR_MAX_SIZE) ||
+        (dice_inputs->authority_descriptor_size > DICE_AUTHORITY_DESCRIPTOR_MAX_SIZE) ||
+        (dice_inputs->config_descriptor_size > DICE_CONFIG_DESCRIPTOR_MAX_SIZE)) {
+        return false;
+    }
+
+    return true;
+}
+
+static bool is_input_handle_valid(int input_context_handle)
+{
+    uint16_t idx = GET_IDX(input_context_handle);
+    uint16_t nonce = GET_NONCE(input_context_handle);
+
+    /* Validate input handle id and nonce */
+    if ((idx >= MAX_NUM_OF_COMPONENTS) || (nonce == INVALID_NONCE_VALUE)) {
+        return false;
+    }
+
+    if (nonce == component_ctx_array[idx].nonce) {
+        return true;
+    }
+
+    return false;
+}
+
+static dpe_error_t derive_child_create_certificate(uint16_t curr_idx)
+{
+    //TODO: Implementation pending
+    /* Finalise the layer */
+    layer_ctx_array[curr_idx].state = LAYER_STATE_FINALISED;
+    return DPE_NO_ERROR;
+}
+
+static uint16_t open_new_layer(void)
+{
+    int i;
+
+    for (i = 0; i < MAX_NUM_OF_LAYERS; i++) {
+        if (layer_ctx_array[i].state == LAYER_STATE_CLOSED) {
+            layer_ctx_array[i].state = LAYER_STATE_OPEN;
+            return i;
+        }
+    }
+
+    return INVALID_LAYER_IDX;
+}
+
+static inline void link_layer(uint16_t child_layer, uint16_t parent_layer)
+{
+    layer_ctx_array[child_layer].parent_layer_idx = parent_layer;
+}
+
+dpe_error_t derive_rot_context(const DiceInputValues *dice_inputs,
+                               int *new_child_ctx_handle,
+                               int *new_parent_ctx_handle)
+{
+    int status;
+    struct component_context_t *child_comp_ctx, *new_child_ctx;
+    uint16_t new_layer_idx;
+
+    log_derive_rot_context(dice_inputs);
+
+    /* Validate dice inputs */
+    if (!is_dice_input_valid(dice_inputs)) {
+        return DPE_INVALID_ARGUMENT;
+    }
+
+    child_comp_ctx = &component_ctx_array[0];
+    status = copy_dice_input(child_comp_ctx, dice_inputs);
+    if (status != DPE_NO_ERROR) {
+        return status;
+    }
+
+    child_comp_ctx->in_use = true;
+    /* Link context to RoT Layer */
+    child_comp_ctx->linked_layer_idx = DPE_ROT_LAYER_IDX;
+    /* Parent is same as child for RoT context */
+    child_comp_ctx->parent_idx = 0;
+    /* Parent not deriving any more children */
+    invalidate_handle(new_parent_ctx_handle);
+
+    //TODO: Update expected_mhu_id of derived child
+    /* Create certificate for RoT layer */
+    status = derive_child_create_certificate(DPE_ROT_LAYER_IDX);
+    if (status != DPE_NO_ERROR) {
+        return status;
+    }
+
+    /* Generate new handle for child for subsequent requests */
+    if (generate_new_handle(new_child_ctx_handle) != DPE_NO_ERROR) {
+        return DPE_INTERNAL_ERROR;
+    }
+
+    /* Update the component context array element as pointed by newly generated handle */
+    new_child_ctx = &component_ctx_array[GET_IDX(*new_child_ctx_handle)];
+    new_child_ctx->nonce = GET_NONCE(*new_child_ctx_handle);
+    new_child_ctx->in_use = true;
+    /* New child's parent is current RoT component which is evaluated as 0 */
+    new_child_ctx->parent_idx = 0;
+
+    /* Open new layer since RoT layer is finalised and
+     * link the new child to this new layer
+     */
+    new_layer_idx = open_new_layer();
+    new_child_ctx->linked_layer_idx = new_layer_idx;
+
+    /* Link this new layer to the RoT Layer */
+    link_layer(new_layer_idx, DPE_ROT_LAYER_IDX);
+
+    return DPE_NO_ERROR;
+}
+
+static inline bool is_input_client_id_valid(int32_t client_id)
+{
+    //TODO: Waiting for implementation
+    return true;
+}
+
+static void assign_layer_to_context(struct component_context_t *new_ctx)
+{
+    uint16_t new_layer_idx, parent_layer_idx;
+
+    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 child is finalised; open a new layer */
+        new_layer_idx = open_new_layer();
+        /* 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;
+    }
+}
+
+dpe_error_t derive_child_request(int input_ctx_handle,
+                                 bool retain_parent_context,
+                                 bool allow_child_to_derive,
+                                 bool create_certificate,
+                                 const DiceInputValues *dice_inputs,
+                                 int32_t client_id,
+                                 int *new_child_ctx_handle,
+                                 int *new_parent_ctx_handle)
+{
+    dpe_error_t status;
+    struct component_context_t *child_ctx, *parent_ctx, *new_ctx;
+    uint16_t input_child_idx, input_parent_idx;
+
+#ifdef TFM_S_REG_TEST
+    if (layer_ctx_array[DPE_ROT_LAYER_IDX].state != LAYER_STATE_FINALISED) {
+        return derive_rot_context(dice_inputs,
+                                  new_child_ctx_handle,
+                                  new_parent_ctx_handle);
+    }
+#endif /* TFM_S_REG_TEST */
+
+    log_derive_child(input_ctx_handle, retain_parent_context,
+                     allow_child_to_derive, create_certificate, dice_inputs,
+                     client_id);
+
+    /* Validate dice inputs */
+    if (!is_dice_input_valid(dice_inputs)) {
+        return DPE_INVALID_ARGUMENT;
+    }
+
+    /* Validate input handle */
+    if (!is_input_handle_valid(input_ctx_handle)) {
+        return DPE_INVALID_ARGUMENT;
+    }
+    /* Get child component index from the input handle */
+    input_child_idx = GET_IDX(input_ctx_handle);
+    /* Get parent index of input referenced child component */
+    input_parent_idx = component_ctx_array[input_child_idx].parent_idx;
+
+    /* Below check is for safety only; It should not happen
+     * input_child_idx is already checked above in is_input_handle_valid()
+     */
+    assert(input_parent_idx < MAX_NUM_OF_COMPONENTS);
+
+    child_ctx = &component_ctx_array[input_child_idx];
+    parent_ctx = &component_ctx_array[input_parent_idx];
+
+    //TODO:  Question: how to get mhu id of incoming request?
+    if (!is_input_client_id_valid(client_id)) {
+        return DPE_INVALID_ARGUMENT;
+    }
+
+    /* Copy dice input to the child component context */
+    status = copy_dice_input(child_ctx, dice_inputs);
+    if (status != DPE_NO_ERROR) {
+        return status;
+    }
+
+    if (create_certificate) {
+        status = derive_child_create_certificate(child_ctx->linked_layer_idx);
+        if (status != DPE_NO_ERROR) {
+            return status;
+        }
+    }
+
+    if (allow_child_to_derive) {
+        /* Generate new handle for child for subsequent requests */
+        if (generate_new_handle(new_child_ctx_handle) != DPE_NO_ERROR) {
+            return DPE_INTERNAL_ERROR;
+        }
+        /* Update the component context array element as pointed by newly generated handle */
+        new_ctx = &component_ctx_array[GET_IDX(*new_child_ctx_handle)];
+        /* Update nonce in new child component context */
+        new_ctx->nonce = GET_NONCE(*new_child_ctx_handle);
+        /* Update parent idx in new child component context */
+        new_ctx->parent_idx = input_child_idx;
+        /* Mark new child component index as in use */
+        new_ctx->in_use = true;
+        assign_layer_to_context(new_ctx);
+
+    } else {
+        /* Child not deriving any children */
+        /* Tag this component as a leaf */
+        child_ctx->is_leaf = true;
+        invalidate_handle(new_child_ctx_handle);
+        /* Renew nonce of child context so it cannot be used again */
+        child_ctx->nonce = INVALID_NONCE_VALUE;
+    }
+
+    if (retain_parent_context) {
+        /* Parent deriving multiple children */
+        /* Generate new handle for child for the same parent for subsequent requests */
+        if (generate_new_handle(new_parent_ctx_handle) != DPE_NO_ERROR) {
+            return DPE_INTERNAL_ERROR;
+        }
+        /* Update the component context array element as pointed by newly generated handle */
+        new_ctx = &component_ctx_array[GET_IDX(*new_parent_ctx_handle)];
+        /* Update nonce in new child component context */
+        new_ctx->nonce = GET_NONCE(*new_parent_ctx_handle);
+        /* Update parent idx in new child component context */
+        new_ctx->parent_idx = input_parent_idx;
+        /* Mark new child component index as in use */
+        new_ctx->in_use = true;
+        assign_layer_to_context(new_ctx);
+
+    } else {
+        /* Parent not deriving any more children */
+        /* No need to return parent handle */
+        invalidate_handle(new_parent_ctx_handle);
+        /* Renew nonce of parent context so it cannot be used again */
+        parent_ctx->nonce = INVALID_NONCE_VALUE;
+    }
+
+    return DPE_NO_ERROR;
+}
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.h b/partitions/dice_protection_environment/dpe_context_mngr.h
new file mode 100644
index 0000000..8335263
--- /dev/null
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -0,0 +1,135 @@
+/*
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#ifndef __DPE_CONTEXT_MNGR_H__
+#define __DPE_CONTEXT_MNGR_H__
+
+#include <stddef.h>
+#include <stdint.h>
+#include <stdbool.h>
+#include "dice_protection_environment.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define DICE_WRAPPING_KEY_SIZE  32
+#define DICE_CERT_SIZE  1024
+
+#define INVALID_HANDLE 0xFFFFFFFF
+#define INVALID_COMPONENT_IDX 0xFFFF
+#define INVALID_NONCE_VALUE  0xFFFF
+#define MAX_NUM_OF_COMPONENTS 30
+#define DPE_ROT_LAYER_IDX 0
+#define MAX_NUM_OF_LAYERS 10
+#define INVALID_LAYER_IDX 65535
+
+/* Most significant 16 bits represent nonce & remaining 16 bits represent component index */
+#define GET_IDX(handle) ((handle) & 0xffff)
+#define GET_NONCE(handle) ((handle >> 16) & 0xffff)
+
+#define SET_IDX(handle, idx) ((handle & 0xffff0000) | idx)
+#define SET_NONCE(handle, nonce) ((handle & 0x00ffff) | (nonce << 16))
+
+struct component_context_data_t {
+    uint8_t        measurement_value[DICE_HASH_SIZE];
+    uint8_t        measurement_descriptor[DICE_CODE_DESCRIPTOR_MAX_SIZE];
+    size_t         measurement_descriptor_size;
+    uint8_t        signer_id[DICE_HASH_SIZE];
+    uint8_t        signer_id_descriptor[DICE_AUTHORITY_DESCRIPTOR_MAX_SIZE];
+    size_t         signer_id_descriptor_size;
+    uint8_t        config_value[DICE_INLINE_CONFIG_SIZE];
+    uint8_t        config_descriptor[DICE_CONFIG_DESCRIPTOR_MAX_SIZE];
+    size_t         config_descriptor_size;
+    DiceMode       mode;
+    uint8_t        hidden[DICE_HIDDEN_SIZE];
+};
+
+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 */
+    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 */
+    uint32_t expected_mhu_id;               /* Expected mhu to authorise derivation */
+};
+
+struct layer_context_data_t {
+    uint8_t cdi_attest[DICE_CDI_SIZE];
+    uint8_t cdi_seal[DICE_CDI_SIZE];
+    uint8_t wrapping_key[DICE_WRAPPING_KEY_SIZE];
+    uint8_t cert_buf[DICE_CERT_SIZE];
+    size_t cert_buf_size;
+};
+
+enum layer_state_t {
+    LAYER_STATE_CLOSED = 0,
+    LAYER_STATE_OPEN,
+    LAYER_STATE_FINALISED
+};
+
+struct layer_context_t {
+    struct layer_context_data_t data;
+    uint16_t parent_layer_idx;
+    enum layer_state_t state;
+};
+
+/**
+ * \brief Derives a root of trust component context and creates certificate.
+ *
+ * \param[in]  dice_inputs               Pointer to dice_input buffer.
+ * \param[out] new_child_context_handle  A new handle for child context.
+ * \param[out] new_parent_context_handle A new handle for parent context.
+ *
+ * \return Returns error code of type dpe_error_t
+ */
+dpe_error_t derive_rot_context(const DiceInputValues *dice_inputs,
+                               int *new_child_ctx_handle,
+                               int *new_parent_ctx_handle);
+
+/**
+ * \brief Derives a child component context and optionally creates certificate
+ *        chain.
+ *
+ * \param[in]  input_context_handle      Input handle to child 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
+ * \param[in]  allow_child_to_derive     Flag to indicate if requested child can
+ *                                       derive further.
+ * \param[in]  create_certificate        Flag to indicate if certificate needs
+ *                                       to be created. TRUE only if it is the
+ *                                       last component in the layer.
+ * \param[in]  dice_inputs               Pointer to dice_input buffer.
+ * \param[in]  client_id                 Identifier of the client calling the
+ *                                       service.
+ * \param[out] new_child_context_handle  A new handle for child context.
+ * \param[out] new_parent_context_handle A new handle for parent context.
+ *
+ * \return Returns error code of type dpe_error_t
+ */
+dpe_error_t derive_child_request(int input_context_handle,
+                                 bool retain_parent_context,
+                                 bool allow_child_to_derive,
+                                 bool create_certificate,
+                                 const DiceInputValues *dice_inputs,
+                                 int32_t client_id,
+                                 int *new_child_context_handle,
+                                 int *new_parent_context_handle);
+
+/**
+ * \brief Initialise all DPE Layer and component contexts
+ *
+ */
+void initialise_all_dpe_contexts(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* __DPE_CONTEXT_MNGR_H__ */
diff --git a/partitions/dice_protection_environment/dpe_impl.c b/partitions/dice_protection_environment/dpe_impl.c
index 8167a4f..c80cf6c 100644
--- a/partitions/dice_protection_environment/dpe_impl.c
+++ b/partitions/dice_protection_environment/dpe_impl.c
@@ -11,23 +11,6 @@
 
 #include "dpe_log.h"
 
-dpe_error_t dpe_derive_child_impl(int context_handle,
-                                  bool retain_parent_context,
-                                  bool allow_child_to_derive,
-                                  bool create_certificate,
-                                  const DiceInputValues *dice_inputs,
-                                  int *child_context_handle,
-                                  int *new_context_handle)
-{
-    log_derive_child(context_handle, retain_parent_context,
-                     allow_child_to_derive, create_certificate, dice_inputs);
-
-    *child_context_handle = 123;
-    *new_context_handle = 456;
-
-    return DPE_NO_ERROR;
-}
-
 dpe_error_t dpe_certify_key_impl(int context_handle,
                                  bool retain_context,
                                  const uint8_t *public_key,
diff --git a/partitions/dice_protection_environment/dpe_impl.h b/partitions/dice_protection_environment/dpe_impl.h
index 5a5fe0a..c1e4de4 100644
--- a/partitions/dice_protection_environment/dpe_impl.h
+++ b/partitions/dice_protection_environment/dpe_impl.h
@@ -28,15 +28,6 @@
  */
 #define DPE_CERTIFICATE_CHAIN_MAX_SIZE 2048
 
-/* Internal DPE service implementation of dpe_derive_child() */
-dpe_error_t dpe_derive_child_impl(int context_handle,
-                                  bool retain_parent_context,
-                                  bool allow_child_to_derive,
-                                  bool create_certificate,
-                                  const DiceInputValues *dice_inputs,
-                                  int *child_context_handle,
-                                  int *new_context_handle);
-
 /* Internal DPE service implementation of dpe_certify_key() */
 dpe_error_t dpe_certify_key_impl(int context_handle,
                                  bool retain_context,
diff --git a/partitions/dice_protection_environment/dpe_log.c b/partitions/dice_protection_environment/dpe_log.c
index a290891..5a37706 100644
--- a/partitions/dice_protection_environment/dpe_log.c
+++ b/partitions/dice_protection_environment/dpe_log.c
@@ -6,6 +6,7 @@
  */
 
 #include "dpe_log.h"
+#include "dpe_context_mngr.h"
 
 #if (TFM_PARTITION_LOG_LEVEL >= TFM_PARTITION_LOG_LEVEL_DEBUG)
 
@@ -50,18 +51,27 @@
     print_byte_array(input->hidden, sizeof(input->hidden));
 }
 
+void log_derive_rot_context(const DiceInputValues *dice_inputs)
+{
+    LOG_DBGFMT("DPE DeriveRoTContext:\r\n");
+    log_dice_inputs(dice_inputs);
+}
+
 void log_derive_child(int context_handle,
                       bool retain_parent_context,
                       bool allow_child_to_derive,
                       bool create_certificate,
-                      const DiceInputValues *dice_inputs)
+                      const DiceInputValues *dice_inputs,
+                      int32_t client_id)
 {
     LOG_DBGFMT("DPE DeriveChild:\r\n");
-    LOG_DBGFMT(" - context_handle = %d\r\n", context_handle);
+    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(" - retain_parent_context = %d\r\n", retain_parent_context);
     LOG_DBGFMT(" - allow_child_to_derive = %d\r\n", allow_child_to_derive);
     LOG_DBGFMT(" - create_certificate = %d\r\n", create_certificate);
     log_dice_inputs(dice_inputs);
+    LOG_DBGFMT(" - client_id = %d\r\n", client_id);
 }
 
 void log_certify_key(int context_handle,
@@ -72,7 +82,8 @@
                      size_t label_size)
 {
     LOG_DBGFMT("DPE CertifyKey:\r\n");
-    LOG_DBGFMT(" - context_handle = %d\r\n", context_handle);
+    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(" - retain_context = %d\r\n", retain_context);
     LOG_DBGFMT(" - public_key =");
     print_byte_array(public_key, public_key_size);
diff --git a/partitions/dice_protection_environment/dpe_log.h b/partitions/dice_protection_environment/dpe_log.h
index d30b2a3..97eb752 100644
--- a/partitions/dice_protection_environment/dpe_log.h
+++ b/partitions/dice_protection_environment/dpe_log.h
@@ -18,13 +18,19 @@
 #if (TFM_PARTITION_LOG_LEVEL >= TFM_PARTITION_LOG_LEVEL_DEBUG)
 
 /**
+ * \brief Log the derive rot context command parameters.
+ */
+void log_derive_rot_context(const DiceInputValues *dice_inputs);
+
+/**
  * \brief Log the derive child command parameters.
  */
 void log_derive_child(int context_handle,
                       bool retain_parent_context,
                       bool allow_child_to_derive,
                       bool create_certificate,
-                      const DiceInputValues *dice_inputs);
+                      const DiceInputValues *dice_inputs,
+                      int32_t client_id);
 
 /**
  * \brief Log the certify key command parameters.
@@ -38,6 +44,7 @@
 
 #else /* TFM_PARTITION_LOG_LEVEL */
 
+#define log_derive_rot_context(...)
 #define log_derive_child(...)
 #define log_certify_key(...)
 
diff --git a/partitions/dice_protection_environment/dpe_req_mngr.c b/partitions/dice_protection_environment/dpe_req_mngr.c
index c1cbe06..06d6f1e 100644
--- a/partitions/dice_protection_environment/dpe_req_mngr.c
+++ b/partitions/dice_protection_environment/dpe_req_mngr.c
@@ -8,6 +8,7 @@
 #include <string.h>
 
 #include "dpe_cmd_decode.h"
+#include "dpe_context_mngr.h"
 #include "psa/service.h"
 
 #define MIN(x, y) (((x) < (y)) ? (x) : (y))
@@ -18,6 +19,8 @@
 
 psa_status_t tfm_dpe_init(void)
 {
+    initialise_all_dpe_contexts();
+
     return PSA_SUCCESS;
 }
 
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 cc34a38..8aefec2 100644
--- a/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
+++ b/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
@@ -18,14 +18,21 @@
 extern "C" {
 #endif
 
+/* Additional defines for max size limit */
+#define DICE_AUTHORITY_DESCRIPTOR_MAX_SIZE  64
+#define DICE_CONFIG_DESCRIPTOR_MAX_SIZE     64
+/* The theoretical maximum image version is: "255.255.65535\0" */
+#define DICE_CODE_DESCRIPTOR_MAX_SIZE 14
+
 typedef int32_t dpe_error_t;
 
-#define DPE_NO_ERROR               ((dpe_error_t)0)
-#define DPE_INTERNAL_ERROR         ((dpe_error_t)1)
-#define DPE_INVALID_COMMAND        ((dpe_error_t)2)
-#define DPE_INVALID_ARGUMENT       ((dpe_error_t)3)
-#define DPE_ARGUMENT_NOT_SUPPORTED ((dpe_error_t)4)
-#define DPE_SESSION_EXHAUSTED      ((dpe_error_t)5)
+#define DPE_NO_ERROR                  ((dpe_error_t)0)
+#define DPE_INTERNAL_ERROR            ((dpe_error_t)1)
+#define DPE_INVALID_COMMAND           ((dpe_error_t)2)
+#define DPE_INVALID_ARGUMENT          ((dpe_error_t)3)
+#define DPE_ARGUMENT_NOT_SUPPORTED    ((dpe_error_t)4)
+#define DPE_SESSION_EXHAUSTED         ((dpe_error_t)5)
+#define DPE_INSUFFICIENT_MEMORY       ((dpe_error_t)128)
 
 /**
  * \brief Performs the DICE computation to derive a child context and optionally