DPE: Refactor certificate creation to reduce memory usage

The layer_ctx_array eats up too much memory because a
DICE_CERT_SIZE (3KB) buffer is allocated per context.

These buffers are meant to create and store the certificates.
The certificates are created when create-certificate=true
and from that time onwards they are stored in encoded format.

The memory usage is reduced by:
 - Per context buffers are removed.
 - Certificates are created on the fly when a query arrives.
   The create certificate step is split into preparing and
   encoding the certificate.
 - Certificates are not stored anymore in an encoded format,
   just their raw content.
 - Certificates are always created in place within cmd_buf
   to avoid stack usage.

Change-Id: I6e8a4f2aecd4c935d41458d4ba5a0a1026b903ec
Signed-off-by: Tamas Ban <tamas.ban@arm.com>
diff --git a/partitions/dice_protection_environment/dpe_certificate.c b/partitions/dice_protection_environment/dpe_certificate.c
index 913b3f5..6dfb0f5 100644
--- a/partitions/dice_protection_environment/dpe_certificate.c
+++ b/partitions/dice_protection_environment/dpe_certificate.c
@@ -18,11 +18,6 @@
 #define ID_HEX_SIZE (2 * DICE_ID_SIZE)      /* Size of CDI encoded to ascii hex */
 #define LABEL_HEX_SIZE (2 * DPE_EXTERNAL_LABEL_MAX_SIZE)
 
-struct dpe_cert_encode_ctx {
-    QCBOREncodeContext           cbor_enc_ctx;
-    struct t_cose_sign1_sign_ctx signer_ctx;
-};
-
 static void convert_to_ascii_hex(const uint8_t *in,
                                  size_t in_size,
                                  char *out,
@@ -57,8 +52,8 @@
     }
 }
 
-static dpe_error_t certificate_encode_start(struct dpe_cert_encode_ctx *me,
-                                            const UsefulBuf out_buf,
+static dpe_error_t certificate_encode_start(QCBOREncodeContext *cbor_enc_ctx,
+                                            struct t_cose_sign1_sign_ctx *signer_ctx,
                                             psa_key_handle_t private_key)
 {
     enum t_cose_err_t t_cose_err;
@@ -66,44 +61,39 @@
     UsefulBufC attest_key_id = {NULL, 0};
 
     /* DPE Certificate is untagged COSE_Sign1 message */
-    t_cose_sign1_sign_init(&(me->signer_ctx), T_COSE_OPT_OMIT_CBOR_TAG, DPE_T_COSE_ALG);
+    t_cose_sign1_sign_init(signer_ctx, T_COSE_OPT_OMIT_CBOR_TAG, DPE_T_COSE_ALG);
 
     attest_key.crypto_lib = T_COSE_CRYPTO_LIB_PSA;
     attest_key.k.key_handle = private_key;
 
-    t_cose_sign1_set_signing_key(&(me->signer_ctx),
-                                 attest_key,
-                                 attest_key_id);
+    t_cose_sign1_set_signing_key(signer_ctx, attest_key, attest_key_id);
 
-    /* Spin up the CBOR encoder */
-    QCBOREncode_Init(&(me->cbor_enc_ctx), out_buf);
+    /* It is expected that the CBOR encoder is already initialized */
 
-    /* This will cause the cose headers to be encoded and written into
-     *  out_buf using me->cbor_enc_ctx
-     */
-    t_cose_err = t_cose_sign1_encode_parameters(&(me->signer_ctx),
-                                                &(me->cbor_enc_ctx));
+    /* This encodes and writes the cose headers to be into the CBOR context. */
+    t_cose_err = t_cose_sign1_encode_parameters(signer_ctx,
+                                                cbor_enc_ctx);
     if (t_cose_err) {
         return t_cose_err_to_dpe_err(t_cose_err);
     }
 
-    QCBOREncode_OpenMap(&(me->cbor_enc_ctx));
+    QCBOREncode_OpenMap(cbor_enc_ctx);
 
     return DPE_NO_ERROR;
 }
 
-static void add_key_usage_claim(struct dpe_cert_encode_ctx *me)
+static void add_key_usage_claim(QCBOREncodeContext *cbor_enc_ctx)
 {
     uint8_t key_usage = DPE_CERT_KEY_USAGE_CERT_SIGN;
 
     /* Encode key usage as byte string */
-    QCBOREncode_AddBytesToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddBytesToMapN(cbor_enc_ctx,
                                DPE_CERT_LABEL_KEY_USAGE,
                                (UsefulBufC){ &key_usage,
                                              sizeof(key_usage) });
 }
 
-static void add_label_claim(struct dpe_cert_encode_ctx *me,
+static void add_label_claim(QCBOREncodeContext *cbor_enc_ctx,
                             const uint8_t *label,
                             size_t label_size)
 {
@@ -117,22 +107,22 @@
                              sizeof(label_hex));
 
         /* Encode label as text string */
-        QCBOREncode_AddTextToMapN(&me->cbor_enc_ctx,
+        QCBOREncode_AddTextToMapN(cbor_enc_ctx,
                                   DPE_CERT_LABEL_EXTERNAL_LABEL,
                                   (UsefulBufC){ &label_hex[0],
                                                 label_size });
     }
 }
 
-static void add_cdi_export_claim(struct dpe_cert_encode_ctx *me,
+static void add_cdi_export_claim(QCBOREncodeContext *cbor_enc_ctx,
                                  struct layer_context_t *layer_ctx)
 {
-    QCBOREncode_AddBoolToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddBoolToMapN(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,
+static void add_subject_claim(QCBOREncodeContext *cbor_enc_ctx,
                               struct layer_context_t *layer_ctx)
 {
     char cdi_id_hex[ID_HEX_SIZE];
@@ -142,13 +132,13 @@
                          &cdi_id_hex[0],
                          sizeof(cdi_id_hex));
     /* Encode subject as text string */
-    QCBOREncode_AddTextToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddTextToMapN(cbor_enc_ctx,
                               DPE_CERT_LABEL_SUBJECT,
                               (UsefulBufC){ &cdi_id_hex[0],
                                             sizeof(cdi_id_hex) });
 }
 
-static void encode_issuer_claim(struct dpe_cert_encode_ctx *me,
+static void encode_issuer_claim(QCBOREncodeContext *cbor_enc_ctx,
                                 const uint8_t *issuer,
                                 size_t issuer_size)
 {
@@ -160,13 +150,13 @@
                          sizeof(cdi_id_hex));
 
     /* Encode issuer as text string */
-    QCBOREncode_AddTextToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddTextToMapN(cbor_enc_ctx,
                               DPE_CERT_LABEL_ISSUER,
                               (UsefulBufC){ &cdi_id_hex[0],
                                             sizeof(cdi_id_hex) });
 }
 
-static void encode_public_key(struct dpe_cert_encode_ctx *me,
+static void encode_public_key(QCBOREncodeContext *cbor_enc_ctx,
                               const uint8_t *pub_key,
                               size_t pub_key_size)
 {
@@ -176,26 +166,26 @@
     const int64_t cose_key_ec2_curve_value = DPE_T_COSE_KEY_EC2_CURVE_VAL;
     const int64_t cose_key_alg_value = DPE_T_COSE_KEY_ALG_VAL;
 
-    QCBOREncode_OpenMap(&me->cbor_enc_ctx);
+    QCBOREncode_OpenMap(cbor_enc_ctx);
 
     /* Add the key type as int */
-    QCBOREncode_AddInt64ToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddInt64ToMapN(cbor_enc_ctx,
                                DPE_CERT_LABEL_COSE_KEY_TYPE,
                                cose_key_type_value);
 
     /* Add the algorithm as int */
-    QCBOREncode_AddInt64ToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddInt64ToMapN(cbor_enc_ctx,
                                DPE_CERT_LABEL_COSE_KEY_ALG,
                                cose_key_alg_value);
 
     /* Add the key operation as [+ (tstr/int)] */
-    QCBOREncode_OpenArrayInMapN(&me->cbor_enc_ctx, DPE_CERT_LABEL_COSE_KEY_OPS);
-    QCBOREncode_AddInt64(&me->cbor_enc_ctx,
+    QCBOREncode_OpenArrayInMapN(cbor_enc_ctx, DPE_CERT_LABEL_COSE_KEY_OPS);
+    QCBOREncode_AddInt64(cbor_enc_ctx,
                                cose_key_ops_value);
-    QCBOREncode_CloseArray(&me->cbor_enc_ctx);
+    QCBOREncode_CloseArray(cbor_enc_ctx);
 
     /* Add the curve */
-    QCBOREncode_AddInt64ToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddInt64ToMapN(cbor_enc_ctx,
                                DPE_CERT_LABEL_COSE_KEY_EC2_CURVE,
                                cose_key_ec2_curve_value);
 
@@ -217,21 +207,21 @@
      * compressed or uncompressed. The uncompressed form is indicated by 0x04
      * and the compressed form is indicated by either 0x02 or 0x03.
      */
-    QCBOREncode_AddBytesToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddBytesToMapN(cbor_enc_ctx,
                                DPE_CERT_LABEL_COSE_KEY_EC2_X,
                                (UsefulBufC){ &pub_key[1],
                                              pub_key_size / 2 });
 
-    QCBOREncode_AddBytesToMapN(&me->cbor_enc_ctx,
+    QCBOREncode_AddBytesToMapN(cbor_enc_ctx,
                                DPE_CERT_LABEL_COSE_KEY_EC2_Y,
                                (UsefulBufC){ &pub_key[1 + (pub_key_size / 2)],
                                              pub_key_size / 2 });
 
-    QCBOREncode_CloseMap(&me->cbor_enc_ctx);
+    QCBOREncode_CloseMap(cbor_enc_ctx);
 
 }
 
-static void add_public_key_claim(struct dpe_cert_encode_ctx *me,
+static void add_public_key_claim(QCBOREncodeContext *cbor_enc_ctx,
                                  const uint8_t *pub_key,
                                  size_t pub_key_size)
 {
@@ -239,54 +229,61 @@
 
     /* Cose key is encoded as a map. This map is wrapped into the a
      * byte string and it is further encoded as a map */
-    QCBOREncode_BstrWrapInMapN(&me->cbor_enc_ctx, DPE_CERT_LABEL_SUBJECT_PUBLIC_KEY);
-    encode_public_key(me, pub_key, pub_key_size);
-    QCBOREncode_CloseBstrWrap2(&me->cbor_enc_ctx, true, &wrapped);
+    QCBOREncode_BstrWrapInMapN(cbor_enc_ctx, DPE_CERT_LABEL_SUBJECT_PUBLIC_KEY);
+    encode_public_key(cbor_enc_ctx, pub_key, pub_key_size);
+    QCBOREncode_CloseBstrWrap2(cbor_enc_ctx, true, &wrapped);
     assert(wrapped.len <= DICE_MAX_ENCODED_PUBLIC_KEY_SIZE);
 }
 
-static void add_public_key_to_certificate_chain(struct dpe_cert_encode_ctx *me,
+static void add_public_key_to_certificate_chain(QCBOREncodeContext *cbor_enc_ctx,
                                                 const uint8_t *pub_key,
                                                 size_t pub_key_size)
 {
     UsefulBufC wrapped;
 
     /* Cose key is encoded as a map wrapped into a byte string */
-    QCBOREncode_BstrWrap(&me->cbor_enc_ctx);
-    encode_public_key(me, pub_key, pub_key_size);
-    QCBOREncode_CloseBstrWrap2(&me->cbor_enc_ctx, true, &wrapped);
+    QCBOREncode_BstrWrap(cbor_enc_ctx);
+    encode_public_key(cbor_enc_ctx, pub_key, pub_key_size);
+    QCBOREncode_CloseBstrWrap2(cbor_enc_ctx, true, &wrapped);
     assert(wrapped.len <= DICE_MAX_ENCODED_PUBLIC_KEY_SIZE);
 }
 
-static dpe_error_t certificate_encode_finish(struct dpe_cert_encode_ctx *me,
+static dpe_error_t certificate_encode_finish(QCBOREncodeContext *cbor_enc_ctx,
+                                             struct t_cose_sign1_sign_ctx *signer_ctx,
+                                             bool finish_cbor_encoding,
                                              UsefulBufC *completed_cert)
 {
     QCBORError qcbor_result;
     enum t_cose_err_t cose_return_value;
 
-    QCBOREncode_CloseMap(&(me->cbor_enc_ctx));
+    QCBOREncode_CloseMap(cbor_enc_ctx);
 
     /* -- Finish up the COSE_Sign1. This is where the signing happens -- */
-    cose_return_value = t_cose_sign1_encode_signature(&(me->signer_ctx),
-                                                      &(me->cbor_enc_ctx));
+    cose_return_value = t_cose_sign1_encode_signature(signer_ctx,
+                                                      cbor_enc_ctx);
     if (cose_return_value) {
         /* Main errors are invoking the hash or signature */
         return t_cose_err_to_dpe_err(cose_return_value);
     }
 
-    /* Finally close off the CBOR formatting and get the pointer and length
-     * of the resulting COSE_Sign1
+    /* If only a single certificate is created then encoding can be finished.
+     * Otherwise, when multiple certifcate is encoded in a raw
+     * (GetCertificateChain) then encoding will be finished
+     * by close_certificate_chain().
      */
-    qcbor_result = QCBOREncode_Finish(&(me->cbor_enc_ctx), completed_cert);
-    if (qcbor_result == QCBOR_ERR_BUFFER_TOO_SMALL) {
-        return DPE_INSUFFICIENT_MEMORY;
-
-    } else if (qcbor_result != QCBOR_SUCCESS) {
-        /* likely from array not closed, too many closes, ... */
-        return DPE_INTERNAL_ERROR;
-
-    } else {
-        return DPE_NO_ERROR;
+    if (finish_cbor_encoding) {
+       /* Finally close off the CBOR formatting and get the pointer and length
+        * of the resulting COSE_Sign1.
+        */
+        qcbor_result = QCBOREncode_Finish(cbor_enc_ctx, completed_cert);
+        if (qcbor_result == QCBOR_ERR_BUFFER_TOO_SMALL) {
+            return DPE_INSUFFICIENT_MEMORY;
+        } else if (qcbor_result != QCBOR_SUCCESS) {
+            /* likely from array not closed, too many closes, ... */
+            return DPE_INTERNAL_ERROR;
+        } else {
+            return DPE_NO_ERROR;
+        }
     }
 }
 
@@ -348,7 +345,7 @@
 }
 
 static void encode_layer_sw_components_array(uint16_t layer_idx,
-                                             struct dpe_cert_encode_ctx *me)
+                                             QCBOREncodeContext *cbor_enc_ctx)
 {
     int i, cnt;
     struct component_context_t *component_ctx;
@@ -361,20 +358,20 @@
 
             if (cnt == 1) {
                 /* Open array which stores SW components claims. */
-                QCBOREncode_OpenArrayInMapN(&me->cbor_enc_ctx,
+                QCBOREncode_OpenArrayInMapN(cbor_enc_ctx,
                                             DPE_CERT_LABEL_SW_COMPONENTS);
             }
-            encode_sw_component_measurements(&me->cbor_enc_ctx, component_ctx);
+            encode_sw_component_measurements(cbor_enc_ctx, component_ctx);
         }
     }
 
     if (cnt != 0) {
         /* Close array which stores SW components claims. */
-        QCBOREncode_CloseArray(&me->cbor_enc_ctx);
+        QCBOREncode_CloseArray(cbor_enc_ctx);
     }
 }
 
-static dpe_error_t add_issuer_claim(struct dpe_cert_encode_ctx *me,
+static dpe_error_t add_issuer_claim(QCBOREncodeContext *cbor_enc_ctx,
                                     uint16_t layer_idx,
                                     psa_key_id_t root_attest_key_id,
                                     const struct layer_context_t *parent_layer_ctx)
@@ -388,11 +385,11 @@
             return DPE_INTERNAL_ERROR;
         }
 
-        encode_issuer_claim(me,
+        encode_issuer_claim(cbor_enc_ctx,
                             rot_cdi_id,
                             sizeof(rot_cdi_id));
     } else {
-        encode_issuer_claim(me,
+        encode_issuer_claim(cbor_enc_ctx,
                             parent_layer_ctx->data.cdi_id,
                             sizeof(parent_layer_ctx->data.cdi_id));
     }
@@ -400,18 +397,30 @@
     return DPE_NO_ERROR;
 }
 
-dpe_error_t encode_layer_certificate(uint16_t layer_idx,
-                                     struct layer_context_t *layer_ctx,
-                                     const struct layer_context_t *parent_layer_ctx)
+static dpe_error_t encode_layer_certificate_internal(uint16_t layer_idx,
+                                                     QCBOREncodeContext *cbor_enc_ctx,
+                                                     bool finish_cbor_encoding,
+                                                     size_t *cert_actual_size)
 {
+    struct t_cose_sign1_sign_ctx signer_ctx;
+    struct layer_context_t *layer_ctx, *parent_layer_ctx;
+    uint16_t parent_layer_idx;
     dpe_error_t err;
-    struct dpe_cert_encode_ctx dpe_cert_ctx;
-    UsefulBuf cert;
     UsefulBufC completed_cert;
     psa_key_id_t attest_key_id;
 
+    /* Valid options: true & !NULL OR false & NULL */
+    assert(finish_cbor_encoding ^ (cert_actual_size == NULL));
+
+    assert(layer_idx < MAX_NUM_OF_LAYERS);
+    layer_ctx = get_layer_ctx_ptr(layer_idx);
+
+    parent_layer_idx = layer_ctx->parent_layer_idx;
+    assert(parent_layer_idx < MAX_NUM_OF_LAYERS);
+    parent_layer_ctx = get_layer_ctx_ptr(parent_layer_idx);
+
     /* The RoT layer certificate is signed by the provisioned attestation key,
-     * all other layers are signed by the parent layer's key.
+     * all other layers are signed by the parent layer's attestation key.
      */
     if (layer_idx == DPE_ROT_LAYER_IDX) {
         attest_key_id = dpe_plat_get_root_attest_key_id();
@@ -419,14 +428,11 @@
         attest_key_id = parent_layer_ctx->data.attest_key_id;
     }
 
-    /* Get started creating the certificate/token. This sets up the CBOR and
+    /* Get started creating the certificate. This sets up the CBOR and
      * COSE contexts which causes the COSE headers to be constructed.
      */
-    cert.ptr = &layer_ctx->data.cert_buf[0];
-    cert.len = sizeof(layer_ctx->data.cert_buf);
-
-    err = certificate_encode_start(&dpe_cert_ctx,
-                                   cert,
+    err = certificate_encode_start(cbor_enc_ctx,
+                                   &signer_ctx,
                                    attest_key_id);
     if (err != DPE_NO_ERROR) {
         return err;
@@ -434,13 +440,13 @@
 
     /* Add all the required claims */
     /* Add issuer/authority claim */
-    err = add_issuer_claim(&dpe_cert_ctx, layer_idx, attest_key_id, parent_layer_ctx);
+    err = add_issuer_claim(cbor_enc_ctx, layer_idx, attest_key_id, parent_layer_ctx);
     if (err != DPE_NO_ERROR) {
         return err;
     }
 
     /* Add subject claim */
-    add_subject_claim(&dpe_cert_ctx, layer_ctx);
+    add_subject_claim(cbor_enc_ctx, layer_ctx);
 
     /* Encode all firmware measurements for the components linked to this layer */
     //TODO:
@@ -449,40 +455,59 @@
      * an array is created for all the components' measurements and within the
      * array, there are multiple maps, one for each SW component
      */
-    encode_layer_sw_components_array(layer_idx, &dpe_cert_ctx);
+    encode_layer_sw_components_array(layer_idx, cbor_enc_ctx);
 
     /* Add label claim */
-    add_label_claim(&dpe_cert_ctx,
+    add_label_claim(cbor_enc_ctx,
                     &layer_ctx->data.external_key_deriv_label[0],
                     layer_ctx->data.external_key_deriv_label_len);
 
     /* Add public key claim */
-    add_public_key_claim(&dpe_cert_ctx,
+    add_public_key_claim(cbor_enc_ctx,
                          &layer_ctx->data.attest_pub_key[0],
                          layer_ctx->data.attest_pub_key_len);
 
     /* Add key usage claim */
-    add_key_usage_claim(&dpe_cert_ctx);
+    add_key_usage_claim(cbor_enc_ctx);
 
     /* Add CDI exported claim */
     if (layer_ctx->is_cdi_to_be_exported) {
-        add_cdi_export_claim(&dpe_cert_ctx, layer_ctx);
+        add_cdi_export_claim(cbor_enc_ctx, layer_ctx);
     }
 
-    /* Finish up creating the token. This is where the actual signature
-     * is generated. This finishes up the CBOR encoding too.
+    /* Finish up creating the certificate. This is where the actual signature
+     * is generated.
      */
-    err = certificate_encode_finish(&dpe_cert_ctx, &completed_cert);
+    err = certificate_encode_finish(cbor_enc_ctx, &signer_ctx,
+                                    finish_cbor_encoding, &completed_cert);
     if (err != DPE_NO_ERROR) {
         return err;
     }
 
-    /* Update the final size of the token/certificate */
-    layer_ctx->data.cert_buf_len = completed_cert.len;
+    /* Update the final size of the certificate if requested */
+    if (cert_actual_size != NULL) {
+        *cert_actual_size = completed_cert.len;
+    }
 
     return err;
 }
 
+dpe_error_t encode_layer_certificate(uint16_t layer_idx,
+                                     uint8_t *cert_buf,
+                                     size_t cert_buf_size,
+                                     size_t *cert_actual_size)
+{
+    QCBOREncodeContext cbor_enc_ctx;
+
+    QCBOREncode_Init(&cbor_enc_ctx,
+                     (UsefulBuf){ cert_buf,
+                                  cert_buf_size });
+
+    /* Only a single certificate is encoded */
+    return encode_layer_certificate_internal(layer_idx, &cbor_enc_ctx,
+                                             true, cert_actual_size);
+}
+
 dpe_error_t store_layer_certificate(struct layer_context_t *layer_ctx)
 {
     //TODO:
@@ -490,35 +515,26 @@
     return DPE_NO_ERROR;
 }
 
-static void open_certificate_chain(struct dpe_cert_encode_ctx *me,
+static void open_certificate_chain(QCBOREncodeContext *cbor_enc_ctx,
                                    uint8_t *cert_chain_buf,
                                    size_t cert_chain_buf_size)
 {
     /* Set up encoding context with output buffer. */
-    QCBOREncode_Init(&me->cbor_enc_ctx,
+    QCBOREncode_Init(cbor_enc_ctx,
                      (UsefulBuf){ &cert_chain_buf[0],
                                   cert_chain_buf_size });
-    QCBOREncode_OpenArray(&me->cbor_enc_ctx);
+    QCBOREncode_OpenArray(cbor_enc_ctx);
 }
 
-static void add_certificate_to_chain(struct dpe_cert_encode_ctx *me,
-                                     struct layer_context_t *layer_ctx)
-{
-    /* Add already encoded layers certificate to the chain */
-    QCBOREncode_AddEncoded(&me->cbor_enc_ctx,
-                           (UsefulBufC){ &layer_ctx->data.cert_buf[0],
-                                         layer_ctx->data.cert_buf_len });
-}
-
-static dpe_error_t close_certificate_chain(struct dpe_cert_encode_ctx *me,
+static dpe_error_t close_certificate_chain(QCBOREncodeContext *cbor_enc_ctx,
                                            size_t *cert_chain_actual_size)
 {
     QCBORError encode_error;
     UsefulBufC completed_cert_chain;
 
-    QCBOREncode_CloseArray(&me->cbor_enc_ctx);
+    QCBOREncode_CloseArray(cbor_enc_ctx);
 
-    encode_error = QCBOREncode_Finish(&me->cbor_enc_ctx,
+    encode_error = QCBOREncode_Finish(cbor_enc_ctx,
                                       &completed_cert_chain);
 
     /* Check for any encoding errors. */
@@ -533,7 +549,7 @@
     return DPE_NO_ERROR;
 }
 
-static dpe_error_t add_root_attestation_public_key(struct dpe_cert_encode_ctx *me)
+static dpe_error_t add_root_attestation_public_key(QCBOREncodeContext *cbor_enc_ctx)
 {
     psa_status_t status;
     psa_key_id_t attest_key_id;
@@ -548,7 +564,7 @@
         return DPE_INTERNAL_ERROR;
     }
 
-    add_public_key_to_certificate_chain(me, &attest_pub_key[0], attest_pub_key_len);
+    add_public_key_to_certificate_chain(cbor_enc_ctx, &attest_pub_key[0], attest_pub_key_len);
 
     return DPE_NO_ERROR;
 }
@@ -559,18 +575,18 @@
                                   size_t *cert_chain_actual_size)
 {
     struct layer_context_t *layer_ctx;
-    struct dpe_cert_encode_ctx dpe_cert_chain_ctx;
+    QCBOREncodeContext cbor_enc_ctx;
     dpe_error_t err;
     int i;
     uint16_t layer_chain[MAX_NUM_OF_LAYERS];
     uint16_t layer_cnt = 0;
 
-    open_certificate_chain(&dpe_cert_chain_ctx,
+    open_certificate_chain(&cbor_enc_ctx,
                            cert_chain_buf,
                            cert_chain_buf_size);
 
     /* Add DICE/Root public key (IAK public key) as the first entry of array */
-    err = add_root_attestation_public_key(&dpe_cert_chain_ctx);
+    err = add_root_attestation_public_key(&cbor_enc_ctx);
     if (err != DPE_NO_ERROR) {
         return err;
     }
@@ -596,13 +612,16 @@
 
     /* Add certificate from RoT to leaf layer order */
     while (i >= DPE_ROT_LAYER_IDX) {
-        layer_ctx = get_layer_ctx_ptr(layer_chain[i]);
-        assert(layer_ctx != NULL);
-        add_certificate_to_chain(&dpe_cert_chain_ctx, layer_ctx);
+        /* Might multiple certificate is encoded */
+        err = encode_layer_certificate_internal(layer_chain[i], &cbor_enc_ctx,
+                                                false, NULL);
+        if (err != DPE_NO_ERROR) {
+            return err;
+        }
         i--;
     }
 
-    return close_certificate_chain(&dpe_cert_chain_ctx,
+    return close_certificate_chain(&cbor_enc_ctx,
                                    cert_chain_actual_size);
 }
 
@@ -642,60 +661,3 @@
 
     return DPE_NO_ERROR;
 }
-
-void clear_certificate_chain(uint16_t layer_idx,
-                             struct layer_context_t *layer_ctx)
-{
-    uint16_t layer_cnt = 0;
-
-    // Q - Confirm - Clear all the certificate info till threshold layer?
-    while ((layer_idx >= DPE_DESTROY_CONTEXT_THRESHOLD_LAYER_IDX) &&
-           (layer_cnt < MAX_NUM_OF_LAYERS)) {
-
-        layer_cnt++;
-        memset(&layer_ctx->data.cert_buf[0], 0, layer_ctx->data.cert_buf_len);
-
-        if (layer_idx == DPE_DESTROY_CONTEXT_THRESHOLD_LAYER_IDX) {
-            /* Cannot clear data from this point onwards */
-            break;
-        }
-
-        layer_ctx = get_layer_ctx_ptr(layer_idx);
-        assert(layer_ctx->parent_layer_idx < layer_idx);
-        /* Move to the parent layer */
-        layer_idx = layer_ctx->parent_layer_idx;
-    }
-}
-
-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 4cb0ef5..95cbbf2 100644
--- a/partitions/dice_protection_environment/dpe_certificate.h
+++ b/partitions/dice_protection_environment/dpe_certificate.h
@@ -22,15 +22,17 @@
 /**
  * \brief Encodes and signs the certificate for a layer
  *
- * \param[in] layer_idx         Index of the current layer context.
- * \param[in] layer_ctx         Pointer to current layer context.
- * \param[in] parent_layer_ctx  Pointer to parent layer context.
- *
+ * \param[in]  layer_idx         Index of the current layer context.
+ * \param[out] cert_buf          Pointer to the output cert buffer.
+ * \param[in]  cert_buf_size     Size of the output cert buffer.
+ * \param[out] cert_actual_size  Actual size of the final certificate.
+ * *
  * \return Returns error code of type dpe_error_t
  */
 dpe_error_t encode_layer_certificate(uint16_t layer_idx,
-                                     struct layer_context_t *layer_ctx,
-                                     const struct layer_context_t *parent_layer_ctx);
+                                     uint8_t *cert_buf,
+                                     size_t cert_buf_size,
+                                     size_t *cert_actual_size);
 
 /**
  * \brief Stores signed certificate for a layer
@@ -73,33 +75,6 @@
                        size_t encoded_cdi_buf_size,
                        size_t *encoded_cdi_actual_size);
 
-/**
- * \brief Clears the certificate chain.
- *
- * \param[in] layer_idx  Index of the current layer context.
- * \param[in] layer_ctx  Pointer to current layer context.
- *
- */
-void clear_certificate_chain(uint16_t layer_idx,
-                             struct layer_context_t *layer_ctx);
-
-/**
- * \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 143e12c..ad5f811 100644
--- a/partitions/dice_protection_environment/dpe_cmd_decode.c
+++ b/partitions/dice_protection_environment/dpe_cmd_decode.c
@@ -308,7 +308,7 @@
                                              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
+     * encode_layer_certificate(). Add it as a byte string so that its
      * decoding can be skipped and the CBOR returned to the caller.
      */
     CHECK_OVERFLOW_TO_TEMP_BUF;
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.c b/partitions/dice_protection_environment/dpe_context_mngr.c
index ac74c33..39abeea 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.c
+++ b/partitions/dice_protection_environment/dpe_context_mngr.c
@@ -289,11 +289,10 @@
     return DPE_NO_ERROR;
 }
 
-static dpe_error_t create_layer_certificate(uint16_t layer_idx)
+static dpe_error_t prepare_layer_certificate(uint16_t layer_idx)
 {
     uint16_t parent_layer_idx;
     psa_status_t status;
-    dpe_error_t err;
     struct layer_context_t *layer_ctx, *parent_layer_ctx;
 
     assert(layer_idx < MAX_NUM_OF_LAYERS);
@@ -341,14 +340,7 @@
         return DPE_INTERNAL_ERROR;
     }
 
-    err = encode_layer_certificate(layer_idx,
-                                   layer_ctx,
-                                   parent_layer_ctx);
-    if (err != DPE_NO_ERROR) {
-        return err;
-    }
-
-    return store_layer_certificate(layer_ctx);
+    return DPE_NO_ERROR;
 }
 
 static uint16_t open_new_layer(void)
@@ -664,22 +656,17 @@
 
         /* Finalise the layer */
         layer_ctx->state = LAYER_STATE_FINALISED;
-        err = create_layer_certificate(linked_layer_idx);
+        err = prepare_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);
+            err = encode_layer_certificate(linked_layer_idx,
+                                           new_certificate_buf,
+                                           new_certificate_buf_size,
+                                           new_certificate_actual_size);
             if (err != DPE_NO_ERROR) {
                 return err;
             }
@@ -703,8 +690,8 @@
     log_dpe_layer_metadata(layer_ctx, linked_layer_idx);
     if (create_certificate) {
         log_intermediate_certificate(linked_layer_idx,
-                                     &layer_ctx->data.cert_buf[0],
-                                     layer_ctx->data.cert_buf_len);
+                                     new_certificate_buf,
+                                     *new_certificate_actual_size);
     }
 
     return DPE_NO_ERROR;
@@ -859,7 +846,15 @@
      * derive context command
      */
     /* Create leaf certificate */
-    err = create_layer_certificate(input_layer_idx);
+    err = prepare_layer_certificate(input_layer_idx);
+    if (err != DPE_NO_ERROR) {
+        return err;
+    }
+
+    err = encode_layer_certificate(input_layer_idx,
+                                   certificate_buf,
+                                   certificate_buf_size,
+                                   certificate_actual_size);
     if (err != DPE_NO_ERROR) {
         return err;
     }
@@ -878,15 +873,6 @@
            parent_layer_ctx->data.attest_pub_key_len);
     *derived_public_key_actual_size = parent_layer_ctx->data.attest_pub_key_len;
 
-    /* Get certificate */
-    if (certificate_buf_size < layer_ctx->data.cert_buf_len) {
-        return DPE_INVALID_ARGUMENT;
-    }
-    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, if requested */
     if (retain_context) {
         *new_context_handle = input_ctx_handle;
@@ -908,8 +894,8 @@
            layer_ctx->data.attest_pub_key_len);
     log_certify_key_output_handle(*new_context_handle);
     log_intermediate_certificate(input_layer_idx,
-                                 &layer_ctx->data.cert_buf[0],
-                                 layer_ctx->data.cert_buf_len);
+                                 certificate_buf,
+                                 *certificate_actual_size);
 
     return DPE_NO_ERROR;
 }
@@ -962,8 +948,10 @@
         component_ctx_array[input_ctx_idx].nonce = GET_NONCE(*new_context_handle);
 
         if (clear_from_context) {
-            /* Clear all the accumulated certificate information so far */
-            clear_certificate_chain(input_layer_idx, layer_ctx);
+        //TODO: Reimplement the clear_from_context functionality after memory
+        //      optimization; Certificates are not ready made and they are not
+        //      stored in the layer context anymore. They are created on-the-fly
+        //      when requested. Add a test as well.
         }
 
     } else {
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.h b/partitions/dice_protection_environment/dpe_context_mngr.h
index 0a35b7c..d861d43 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.h
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -83,8 +83,6 @@
     size_t attest_pub_key_len;
     uint8_t external_key_deriv_label[DPE_EXTERNAL_LABEL_MAX_SIZE];
     size_t external_key_deriv_label_len;
-    uint8_t cert_buf[DICE_CERT_SIZE];
-    size_t cert_buf_len;
 };
 
 enum layer_state_t {