diff --git a/partitions/dice_protection_environment/dpe_cmd_decode.c b/partitions/dice_protection_environment/dpe_cmd_decode.c
index 935078d..5112063 100644
--- a/partitions/dice_protection_environment/dpe_cmd_decode.c
+++ b/partitions/dice_protection_environment/dpe_cmd_decode.c
@@ -287,8 +287,8 @@
     size_t public_key_size;
     const uint8_t *label;
     size_t label_size;
-    uint8_t certificate_chain_buf[DICE_CERT_CHAIN_SIZE];
-    size_t certificate_chain_actual_size;
+    uint8_t certificate_buf[DICE_CERT_SIZE];
+    size_t certificate_actual_size;
     uint8_t derived_public_key_buf[DPE_ATTEST_PUB_KEY_SIZE];
     size_t derived_public_key_actual_size;
     int new_context_handle;
@@ -328,9 +328,9 @@
 
     dpe_err = certify_key_request(context_handle, retain_context, public_key,
                                   public_key_size, label, label_size,
-                                  certificate_chain_buf,
-                                  sizeof(certificate_chain_buf),
-                                  &certificate_chain_actual_size,
+                                  certificate_buf,
+                                  sizeof(certificate_buf),
+                                  &certificate_actual_size,
                                   derived_public_key_buf,
                                   sizeof(derived_public_key_buf),
                                   &derived_public_key_actual_size,
@@ -349,9 +349,9 @@
      * key implementation. 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_CERTIFY_KEY_CERTIFICATE_CHAIN,
-                               (UsefulBufC){ certificate_chain_buf,
-                                             certificate_chain_actual_size });
+    QCBOREncode_AddBytesToMapN(encode_ctx, DPE_CERTIFY_KEY_CERTIFICATE,
+                               (UsefulBufC){ certificate_buf,
+                                             certificate_actual_size });
 
     QCBOREncode_AddBytesToMapN(encode_ctx, DPE_CERTIFY_KEY_DERIVED_PUBLIC_KEY,
                                (UsefulBufC){ derived_public_key_buf,
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.c b/partitions/dice_protection_environment/dpe_context_mngr.c
index fc542a0..e279b39 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.c
+++ b/partitions/dice_protection_environment/dpe_context_mngr.c
@@ -307,8 +307,6 @@
 
     assert(layer_idx < MAX_NUM_OF_LAYERS);
     layer_ctx = &layer_ctx_array[layer_idx];
-   /* Finalise the layer */
-    layer_ctx->state = LAYER_STATE_FINALISED;
     parent_layer_idx = layer_ctx->parent_layer_idx;
     assert(parent_layer_idx < MAX_NUM_OF_LAYERS);
     parent_layer_ctx = &layer_ctx_array[parent_layer_idx];
@@ -648,6 +646,8 @@
         layer_ctx = &layer_ctx_array[linked_layer_idx];
         layer_ctx->is_cdi_to_be_exported = export_cdi;
 
+        /* Finalise the layer */
+        layer_ctx->state = LAYER_STATE_FINALISED;
         err = create_layer_certificate(linked_layer_idx);
         if (err != DPE_NO_ERROR) {
             return err;
@@ -765,9 +765,9 @@
                                 size_t public_key_size,
                                 const uint8_t *label,
                                 size_t label_size,
-                                uint8_t *certificate_chain_buf,
-                                size_t certificate_chain_buf_size,
-                                size_t *certificate_chain_actual_size,
+                                uint8_t *certificate_buf,
+                                size_t certificate_buf_size,
+                                size_t *certificate_actual_size,
                                 uint8_t *derived_public_key_buf,
                                 size_t derived_public_key_buf_size,
                                 size_t *derived_public_key_actual_size,
@@ -831,7 +831,7 @@
     /* Correct layer should already be assigned in last call of
      * derive context command
      */
-    /* Finalise the current layer & create leaf certificate */
+    /* Create leaf certificate */
     err = create_layer_certificate(input_layer_idx);
     if (err != DPE_NO_ERROR) {
         return err;
@@ -851,16 +851,14 @@
            parent_layer_ctx->data.attest_pub_key_len);
     *derived_public_key_actual_size = parent_layer_ctx->data.attest_pub_key_len;
 
-    /* Get certificate chain */
-    err = get_certificate_chain(input_layer_idx,
-                                certificate_chain_buf,
-                                certificate_chain_buf_size,
-                                certificate_chain_actual_size);
-    if (err != DPE_NO_ERROR) {
-        return err;
+    /* Get certificate */
+    if (certificate_buf_size < layer_ctx->data.cert_buf_len) {
+        return DPE_INVALID_ARGUMENT;
     }
-
-    log_certificate_chain(certificate_chain_buf, *certificate_chain_actual_size);
+    memcpy(certificate_buf,
+           &layer_ctx->data.cert_buf[0],
+           layer_ctx->data.cert_buf_len);
+    *certificate_actual_size = layer_ctx->data.cert_buf_len;
 
     /* Renew handle for the same context */
     *new_context_handle = input_ctx_handle;
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.h b/partitions/dice_protection_environment/dpe_context_mngr.h
index 0222640..cf16a7d 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.h
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -203,8 +203,8 @@
 struct layer_context_t* get_layer_ctx_ptr(uint16_t layer_idx);
 
 /**
- * \brief Generates a leaf certificate and returns all the certificate chain
- *        leading to it. This command functionality depends on whether:
+ * \brief Certifies the attestation key and generates a leaf certificate.
+ *        This command functionality depends on whether:
  *        - last layer is finalised
  *        - public key is supplied to the command
  *        - label is supplied to the command
@@ -221,20 +221,19 @@
  *  |               |            | see Note F | no label       |
  *  +---------------+------------+------------+----------------+
  *
- *  A - Opens a new layer (if not opened), creates a leaf certificate which
- *      includes supplied key and generates certificate chain.
+ *  A - Opens a new layer (if not opened), and creates a leaf certificate which
+ *      includes supplied key.
  *  B - Creates certificate for current (existing) layer, which includes supplied
- *      key and generates certificate chain.
+ *      key.
  *  C - Opens a new layer (if not opened), performs derivation which includes
- *      supplied label, creates leaf certificate (including supplied label as a
- *      claim) and generates certificate chain.
+ *      supplied label, and creates leaf certificate (including supplied label
+ *      as a claim).
  *  D - Opens a new layer (if not opened), performs standard derivation,
- *      creates a leaf certificate and generates certificate chain.
- *  E - Performs derivation (which includes supplied label) for current/existing layer,
- *      creates certificate which includes supplied label as a claim, and generates
- *      certificate chain.
- *  F - Performs standard derivation for current/existing layer, creates certificate
- *      and generates certificate chain.
+ *      and creates a leaf certificate.
+ *  E - Performs derivation (which includes supplied label) for current/existing layer
+ *      and creates certificate which includes supplied label as a claim.
+ *  F - Performs standard derivation for current/existing layer, and creates
+ *      certificate.
  *
  * \param[in]  input_ctx_handle                Input handle to component context.
  * \param[in]  retain_context                  Flag to indicate if context needs
@@ -249,12 +248,11 @@
  *                                             already provided, this argument is
  *                                             ignored.
  * \param[in]  label_size                      Size of the input label.
- * \param[out] certificate_chain_buf           Pointer to the buffer where
- *                                             certificate chain will be stored.
- * \param[in]  certificate_chain_buf_size      Size of the allocated buffer for
- *                                             certificate chain.
- * \param[out] certificate_chain_actual_size   Actual size of the certificate
- *                                             chain.
+ * \param[out] certificate_buf                 Pointer to the buffer where
+ *                                             the certificate will be stored.
+ * \param[in]  certificate_buf_size            Size of the allocated buffer for
+ *                                             the certificate.
+ * \param[out] certificate_actual_size         Actual size of the certificate.
  * \param[out] derived_public_key_buf          Pointer to the buffer where
  *                                             derived public key will be stored.
  * \param[in]  derived_public_key_buf_size     Size of the allocated buffer for
@@ -271,9 +269,9 @@
                                 size_t public_key_size,
                                 const uint8_t *label,
                                 size_t label_size,
-                                uint8_t *certificate_chain_buf,
-                                size_t certificate_chain_buf_size,
-                                size_t *certificate_chain_actual_size,
+                                uint8_t *certificate_buf,
+                                size_t certificate_buf_size,
+                                size_t *certificate_actual_size,
                                 uint8_t *derived_public_key_buf,
                                 size_t derived_public_key_buf_size,
                                 size_t *derived_public_key_actual_size,
diff --git a/partitions/dice_protection_environment/interface/include/dpe_client.h b/partitions/dice_protection_environment/interface/include/dpe_client.h
index 79f7b05..eed9598 100644
--- a/partitions/dice_protection_environment/interface/include/dpe_client.h
+++ b/partitions/dice_protection_environment/interface/include/dpe_client.h
@@ -83,7 +83,7 @@
 };
 
 enum dpe_certify_key_output_labels_t {
-    DPE_CERTIFY_KEY_CERTIFICATE_CHAIN = 1,
+    DPE_CERTIFY_KEY_CERTIFICATE = 1,
     DPE_CERTIFY_KEY_DERIVED_PUBLIC_KEY = 2,
     DPE_CERTIFY_KEY_NEW_CONTEXT_HANDLE = 3,
 };
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 2d79b09..48f622c 100644
--- a/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
+++ b/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
@@ -306,7 +306,7 @@
         QCBORDecode_EnterMap(&decode_ctx, NULL);
 
         QCBORDecode_GetByteStringInMapN(&decode_ctx,
-                                        DPE_CERTIFY_KEY_CERTIFICATE_CHAIN,
+                                        DPE_CERTIFY_KEY_CERTIFICATE,
                                         &out);
         args->certificate_chain = out.ptr;
         args->certificate_chain_size = out.len;