DPE: Add optional parameters tests

Signed-off-by: Maulik Patel <maulik.patel@arm.com>
Change-Id: I433396ad8225ddbb1898c68b24d465e9967c8fa8
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.h b/partitions/dice_protection_environment/dpe_context_mngr.h
index b34fb9a..dbbc23d 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.h
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -27,9 +27,7 @@
 #define INVALID_LAYER_IDX 65535
 #define DPE_ROT_LAYER_IDX 0
 
-/* Below configuration defines are platform dependant */
-#define DICE_CERT_SIZE  3072
-#define DICE_CERT_CHAIN_SIZE  5600
+/* Below configuration defines are platform dependent */
 #define MAX_NUM_OF_COMPONENTS 20
 #define MAX_NUM_OF_LAYERS 4
 #define DPE_PLATFORM_LAYER_IDX 1
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 cddb9ed..d062196 100644
--- a/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
+++ b/partitions/dice_protection_environment/interface/include/dice_protection_environment.h
@@ -34,9 +34,12 @@
 #define DPE_INSUFFICIENT_MEMORY       ((dpe_error_t)128)
 #define DPE_ERR_CBOR_FORMATTING       ((dpe_error_t)129)
 
+/* Below custom configuration defines are platform dependent */
 #define DPE_CERT_ID_INVALID 0
 #define DPE_ROT_CERT_ID 0x100
 #define DPE_CERT_ID_SAME_AS_PARENT 0xFFFFFFFF
+#define DICE_CERT_SIZE  3072
+#define DICE_CERT_CHAIN_SIZE  5600
 
 /**
  * \brief Performs the DICE computation to derive a new context and optionally
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 c929f61..576d05d 100644
--- a/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
+++ b/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
@@ -13,8 +13,6 @@
 #include "qcbor/qcbor_decode.h"
 #include "qcbor/qcbor_spiffy_decode.h"
 
-#define DPE_CERT_CHAIN_BUF_SIZE 4096
-
 static void encode_dice_inputs(QCBOREncodeContext *encode_ctx,
                                const DiceInputValues *input)
 {
@@ -500,7 +498,7 @@
     dpe_error_t dpe_err;
     QCBORError qcbor_err;
     UsefulBufC encoded_buf;
-    UsefulBuf_MAKE_STACK_UB(cmd_buf, 2500);
+    UsefulBuf_MAKE_STACK_UB(cmd_buf, DICE_CERT_SIZE);
 
     const struct certify_key_input_t in_args = {
         context_handle,
@@ -576,7 +574,7 @@
     dpe_error_t dpe_err;
     QCBORError qcbor_err;
     UsefulBufC encoded_buf;
-    UsefulBuf_MAKE_STACK_UB(cmd_buf, DPE_CERT_CHAIN_BUF_SIZE);
+    UsefulBuf_MAKE_STACK_UB(cmd_buf, DICE_CERT_CHAIN_SIZE);
 
     const struct get_certificate_chain_input_t in_args = {
         context_handle,
diff --git a/partitions/dice_protection_environment/test/dpe_certify_key_test.c b/partitions/dice_protection_environment/test/dpe_certify_key_test.c
index b0c3fb6..6fb7c81 100644
--- a/partitions/dice_protection_environment/test/dpe_certify_key_test.c
+++ b/partitions/dice_protection_environment/test/dpe_certify_key_test.c
@@ -10,6 +10,22 @@
 #include "dpe_test.h"
 #include "dpe_test_data.h"
 
+#define CALL_CERTIFY_KEY_WITH_TEST_PARAM() \
+        dpe_certify_key_with_test_param(out_ctx_handle, /* input_ctx_handle */  \
+              true,                             /* retain_context */            \
+              NULL,                             /* public_key */                \
+              0,                                /* public_key_size */           \
+              NULL,                             /* label */                     \
+              0,                                /* label_size */                \
+              certificate_buf,                  /* certificate_buf */           \
+              sizeof(certificate_buf),          /* certificate_buf_size */      \
+              &certificate_actual_size,         /* certificate_actual_size */   \
+              derived_public_key_buf,           /* derived_public_key_buf */    \
+              sizeof(derived_public_key_buf),   /* derived_public_key_buf_size */           \
+              &derived_public_key_actual_size,  /* derived_public_key_buf_actual_size */    \
+              &new_context_handle,              /* new_context_handle */        \
+              &test_params);                    /* test_params */
+
 extern struct dpe_derive_context_test_data_t
               derive_context_test_dataset_1[DERIVE_CONTEXT_TEST_DATA1_SIZE];
 extern struct dpe_derive_context_test_data_t
@@ -25,7 +41,7 @@
     int in_handle, out_ctx_handle, out_parent_handle, new_context_handle;
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
     int saved_handles_cnt, i, j;
-    uint8_t certificate_buf[3072];
+    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;
@@ -342,6 +358,13 @@
         return;
     }
 
+    /* Destroy other derived contexts for subsequent test */
+    dpe_err = dpe_destroy_context(out_ctx_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
+
     /* Save the last handle for the subsequent test */
     retained_rot_ctx_handle = out_parent_handle;
 
@@ -485,6 +508,7 @@
     uint8_t derived_public_key_buf[10];
     size_t derived_public_key_actual_size;
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
+    struct dpe_certify_key_test_params_t test_params = {0};
 
     dpe_err = dpe_derive_context(retained_rot_ctx_handle,       /* input_ctx_handle */
                                  DPE_PLATFORM_CERT_ID,          /* cert_id */
@@ -509,21 +533,9 @@
         return;
     }
 
-    /* Call test encode function with corrupt_encoded_cbor = true */
-    dpe_err = dpe_certify_key_with_test_param(out_ctx_handle,   /* input_ctx_handle */
-                              true,                             /* retain_context */
-                              NULL,                             /* public_key */
-                              0,                                /* public_key_size */
-                              NULL,                             /* label */
-                              0,                                /* label_size */
-                              certificate_buf,                  /* certificate_buf */
-                              sizeof(certificate_buf),          /* certificate_buf_size */
-                              &certificate_actual_size,         /* certificate_actual_size */
-                              derived_public_key_buf,           /* derived_public_key_buf */
-                              sizeof(derived_public_key_buf),   /* derived_public_key_buf_size */
-                              &derived_public_key_actual_size,  /* derived_public_key_buf_actual_size */
-                              &new_context_handle,              /* new_context_handle */
-                              true);                            /* corrupt_encoded_cbor */
+    /* Call test encode function with is_encoded_cbor_corrupt = true */
+    test_params.is_encoded_cbor_corrupt = true;
+    dpe_err = CALL_CERTIFY_KEY_WITH_TEST_PARAM();
     if (dpe_err != DPE_INVALID_COMMAND) {
         TEST_FAIL("DPE CertifyKey test: Invalid CBOR construct should return invalid command");
         return;
@@ -534,3 +546,88 @@
 
     ret->val = TEST_PASSED;
 }
+
+void certify_key_without_optional_args_test(struct test_result_t *ret)
+{
+    dpe_error_t dpe_err;
+    int out_ctx_handle, out_parent_handle, new_context_handle;
+    uint8_t certificate_buf[2500];
+    size_t certificate_actual_size;
+    uint8_t derived_public_key_buf[DPE_ATTEST_PUB_KEY_SIZE];
+    size_t derived_public_key_actual_size;
+    DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
+    struct dpe_certify_key_test_params_t test_params = {0};
+
+    dpe_err = dpe_derive_context(retained_rot_ctx_handle,       /* input_ctx_handle */
+                                 DPE_PLATFORM_CERT_ID,          /* cert_id */
+                                 true,                          /* retain_parent_context */
+                                 true,                          /* allow_new_context_to_derive */
+                                 true,                          /* create_certificate */
+                                 &dice_inputs,                  /* dice_inputs */
+                                 0,                             /* target_locality */
+                                 false,                         /* return_certificate */
+                                 true,                          /* allow_new_context_to_export */
+                                 false,                         /* export_cdi */
+                                 &out_ctx_handle,               /* new_context_handle */
+                                 &out_parent_handle,            /* new_parent_context_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 (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext call failed");
+        return;
+    }
+    retained_rot_ctx_handle = out_parent_handle;
+
+    test_params.is_public_key_missing = true;
+    dpe_err = CALL_CERTIFY_KEY_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE CertifyKey test: Without optional parameter should not fail");
+        return;
+    }
+    // TODO: If public key is omitted, key pair should be derived from context and label.
+    //       Validate it.
+    if (derived_public_key_actual_size == 0) {
+        TEST_FAIL("DPE CertifyKey test: Without optional parameter should not fail");
+        return;
+    }
+
+    test_params.is_public_key_missing = false;
+    test_params.is_label_missing = true;
+    out_ctx_handle = new_context_handle;
+    derived_public_key_actual_size = 0;
+    dpe_err = CALL_CERTIFY_KEY_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE CertifyKey test: Without optional parameter should not fail");
+        return;
+    }
+    // TODO: If label is omitted, empty label should be used for key derivation. Validate it.
+    if (derived_public_key_actual_size == 0) {
+        TEST_FAIL("DPE CertifyKey test: Without optional parameter should not fail");
+        return;
+    }
+
+    test_params.is_label_missing = false;
+    test_params.is_retain_context_missing = true;
+    out_ctx_handle = new_context_handle;
+    /* This test will create undestroyable context as default value of
+     * retain_context is false
+     */
+    dpe_err = CALL_CERTIFY_KEY_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE CertifyKey test: Without optional parameter should not fail");
+        return;
+    }
+    /* Default value of retain_context = false, hence it should NOT
+     * return valid context handle
+     */
+    if (new_context_handle != INVALID_HANDLE) {
+        TEST_FAIL("DPE CertifyKey test: Without optional parameter should not fail");
+        return;
+    }
+
+    ret->val = TEST_PASSED;
+}
diff --git a/partitions/dice_protection_environment/test/dpe_derive_context_test.c b/partitions/dice_protection_environment/test/dpe_derive_context_test.c
index f4bc7b5..76a3711 100644
--- a/partitions/dice_protection_environment/test/dpe_derive_context_test.c
+++ b/partitions/dice_protection_environment/test/dpe_derive_context_test.c
@@ -12,23 +12,23 @@
 
 #define CALL_DERIVE_CONTEXT_WITH_TEST_PARAM() \
         dpe_derive_context_with_test_param(retained_rot_ctx_handle, /* input_ctx_handle */  \
-            DPE_PLATFORM_CERT_ID,     /* cert_id */                                          \
+            DPE_PLATFORM_CERT_ID,    /* cert_id */                                          \
             true,                    /* retain_parent_context */                            \
             true,                    /* allow_new_context_to_derive */                      \
             false,                   /* create_certificate */                               \
             &dice_inputs,            /* dice_inputs */                                      \
             0,                       /* target_locality */                                  \
-            false,                   /* return_certificate */                               \
+            return_certificate,      /* return_certificate */                               \
             true,                    /* allow_new_context_to_export */                      \
             false,                   /* export_cdi */                                       \
             &out_ctx_handle,         /* new_context_handle */                               \
             &out_parent_handle,      /* new_parent_context_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 */                         \
+            certificate_buf,         /* new_certificate_buf */                              \
+            sizeof(certificate_buf), /* new_certificate_buf_size */                         \
+            &certificate_actual_size,/* new_certificate_actual_size */                      \
+            exported_cdi_buf,        /* exported_cdi_buf */                                 \
+            sizeof(exported_cdi_buf),/* exported_cdi_buf_size */                            \
+            &exported_cdi_actual_size,/* exported_cdi_actual_size */                        \
             &test_params);           /* test_parameters */
 
 extern int retained_rot_ctx_handle;
@@ -362,6 +362,11 @@
     dpe_error_t dpe_err;
     int out_ctx_handle;
     int out_parent_handle;
+    uint8_t certificate_buf[DICE_CERT_SIZE];
+    size_t certificate_actual_size;
+    uint8_t exported_cdi_buf[DICE_MAX_ENCODED_CDI_SIZE];
+    size_t exported_cdi_actual_size;
+    bool return_certificate = false;
 
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
     struct dpe_derive_context_test_params_t test_params = {0};
@@ -418,10 +423,15 @@
     dpe_error_t dpe_err;
     int out_ctx_handle;
     int out_parent_handle;
+    uint8_t certificate_buf[DICE_CERT_SIZE];
+    size_t certificate_actual_size;
+    uint8_t exported_cdi_buf[DICE_MAX_ENCODED_CDI_SIZE];
+    size_t exported_cdi_actual_size;
+    bool return_certificate = false;
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
     struct dpe_derive_context_test_params_t test_params = {0};
 
-    test_params.corrupt_encoded_cbor = true;
+    test_params.is_encoded_cbor_corrupt = true;
     dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
     if (dpe_err != DPE_INVALID_COMMAND) {
         TEST_FAIL("DPE DeriveContext test: Invalid CBOR construct "
@@ -816,3 +826,115 @@
 
     ret->val = TEST_PASSED;
 }
+
+void derive_context_without_optional_args_test(struct test_result_t *ret)
+{
+    dpe_error_t dpe_err;
+    int out_ctx_handle;
+    int out_parent_handle;
+    uint8_t certificate_buf[DICE_CERT_SIZE];
+    size_t certificate_actual_size;
+    uint8_t exported_cdi_buf[DICE_MAX_ENCODED_CDI_SIZE];
+    size_t exported_cdi_actual_size;
+    bool return_certificate = false;
+    DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
+    struct dpe_derive_context_test_params_t test_params = {0};
+
+    test_params.is_allow_new_context_to_derive_missing = true;
+    dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+
+    /* Default value of allow_new_context_to_derive = true, hence it should
+     * return valid context handle
+     */
+    if (out_ctx_handle == INVALID_HANDLE) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+
+    retained_rot_ctx_handle = out_parent_handle;
+    test_params.is_allow_new_context_to_derive_missing = false;
+    test_params.is_create_certificate_missing = true;
+    return_certificate = true;
+    certificate_actual_size = 0;
+    dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+    /* Default value of create_certificate = true, hence it should return
+     * valid certificate
+     */
+    if (certificate_actual_size == 0) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+
+    retained_rot_ctx_handle = out_parent_handle;
+    test_params.is_create_certificate_missing = false;
+    test_params.is_return_certificate_missing = true;
+    certificate_actual_size = 0;
+    dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+    /* Default value of return_certificate = false, hence it should NOT
+     * return valid certificate
+     */
+    if (certificate_actual_size != 0) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+
+    retained_rot_ctx_handle = out_parent_handle;
+    test_params.is_return_certificate_missing = false;
+    test_params.is_allow_new_context_to_export_missing = true;
+    dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+    //TODO: Side effect validation as below
+    // Will need to call DeriveContext again and check if CDI cannot be exported,
+    // but it also depends on few other arguments which will make this test case complex.
+
+    retained_rot_ctx_handle = out_parent_handle;
+    test_params.is_allow_new_context_to_export_missing = false;
+    test_params.is_export_cdi_missing = true;
+    exported_cdi_actual_size = 0;
+    dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+    /* Default value of export_cdi = false, hence it should NOT return CDI */
+    if (exported_cdi_actual_size != 0) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+
+    retained_rot_ctx_handle = out_parent_handle;
+    test_params.is_export_cdi_missing = false;
+    test_params.is_retain_parent_context_missing = true;
+    /* This test will create undestroyable context as default value of
+     * retain_parent_context is false
+     */
+    dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+    /* Default value of retain_parent_context = false, hence it should NOT
+     * return valid parent handle
+     */
+    if (out_parent_handle != INVALID_HANDLE) {
+        TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
+        return;
+    }
+
+    ret->val = TEST_PASSED;
+}
diff --git a/partitions/dice_protection_environment/test/dpe_test.h b/partitions/dice_protection_environment/test/dpe_test.h
index 066ccfc..ab7df86 100644
--- a/partitions/dice_protection_environment/test/dpe_test.h
+++ b/partitions/dice_protection_environment/test/dpe_test.h
@@ -158,6 +158,13 @@
 void derive_context_with_parent_leaf_component_test(struct test_result_t *ret);
 
 /**
+ * \brief Test the DPE DeriveContext without optional arguments.
+ *
+ * \param[out] ret  Test result
+ */
+void derive_context_without_optional_args_test(struct test_result_t *ret);
+
+/**
  * \brief Test the DPE CertifyKey functionality.
  *        Note: this test will create context(s) which cannot be destroyed.
  *
@@ -214,6 +221,13 @@
  */
 void certify_key_invalid_cbor_encoded_input_test(struct test_result_t *ret);
 
+/**
+ * \brief Test the DPE CertifyKey without optional arguments.
+ *
+ * \param[out] ret  Test result
+ */
+void certify_key_without_optional_args_test(struct test_result_t *ret);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/partitions/dice_protection_environment/test/dpe_test_cmd_encode.c b/partitions/dice_protection_environment/test/dpe_test_cmd_encode.c
index cae3084..cb7787a 100644
--- a/partitions/dice_protection_environment/test/dpe_test_cmd_encode.c
+++ b/partitions/dice_protection_environment/test/dpe_test_cmd_encode.c
@@ -89,32 +89,44 @@
                                (UsefulBufC){ &args->context_handle,
                                              sizeof(args->context_handle) });
     if (!test_params->is_cert_id_missing) {
-        QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CERT_ID,
-                                   (UsefulBufC){ &args->cert_id,
-                                                 sizeof(args->cert_id) });
+        QCBOREncode_AddUInt64ToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CERT_ID,
+                                    args->cert_id);
     }
-    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_RETAIN_PARENT_CONTEXT,
-                              args->retain_parent_context);
-    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_DERIVE,
-                              args->allow_new_context_to_derive);
-    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CREATE_CERTIFICATE,
-                              args->create_certificate);
+
+    if (!test_params->is_retain_parent_context_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_RETAIN_PARENT_CONTEXT,
+                                args->retain_parent_context);
+    }
+    if (!test_params->is_allow_new_context_to_derive_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_DERIVE,
+                                  args->allow_new_context_to_derive);
+    }
+    if (!test_params->is_create_certificate_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_CREATE_CERTIFICATE,
+                                  args->create_certificate);
+    }
     if (!test_params->is_input_dice_data_missing) {
         encode_dice_inputs(&encode_ctx, args->dice_inputs, test_params);
     }
     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);
+    if (!test_params->is_return_certificate_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_RETURN_CERTIFICATE,
+                                  args->return_certificate);
+    }
+    if (!test_params->is_allow_new_context_to_export_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_ALLOW_NEW_CONTEXT_TO_EXPORT,
+                                  args->allow_new_context_to_export);
+    }
+    if (!test_params->is_export_cdi_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_DERIVE_CONTEXT_EXPORT_CDI,
+                                  args->export_cdi);
+    }
 
     QCBOREncode_CloseMap(&encode_ctx);
 
-    if (test_params->corrupt_encoded_cbor) {
+    if (test_params->is_encoded_cbor_corrupt) {
         /* Deliberately corrupt CBOR map metadata and construct */
         *((uint8_t *)encode_ctx.OutBuf.UB.ptr + 1) = 0xff;
     }
@@ -180,11 +192,10 @@
     return QCBORDecode_Finish(&decode_ctx);
 }
 
-
 static QCBORError encode_certify_key(const struct certify_key_input_t *args,
                                      UsefulBuf buf,
                                      UsefulBufC *encoded_buf,
-                                     bool corrupt_encoded_cbor)
+                                     struct dpe_certify_key_test_params_t *test_params)
 {
     QCBOREncodeContext encode_ctx;
 
@@ -198,16 +209,22 @@
     QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_CERTIFY_KEY_CONTEXT_HANDLE,
                                (UsefulBufC){ &args->context_handle,
                                              sizeof(args->context_handle) });
-    QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_CERTIFY_KEY_RETAIN_CONTEXT,
-                              args->retain_context);
-    QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_CERTIFY_KEY_PUBLIC_KEY,
-                               (UsefulBufC){ args->public_key,
-                                             args->public_key_size });
-    QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_CERTIFY_KEY_LABEL,
-                               (UsefulBufC){ args->label, args->label_size} );
+    if (!test_params->is_retain_context_missing) {
+        QCBOREncode_AddBoolToMapN(&encode_ctx, DPE_CERTIFY_KEY_RETAIN_CONTEXT,
+                                  args->retain_context);
+    }
+    if (!test_params->is_public_key_missing) {
+        QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_CERTIFY_KEY_PUBLIC_KEY,
+                                   (UsefulBufC){ args->public_key,
+                                                 args->public_key_size });
+    }
+    if (!test_params->is_label_missing) {
+        QCBOREncode_AddBytesToMapN(&encode_ctx, DPE_CERTIFY_KEY_LABEL,
+                                   (UsefulBufC){ args->label, args->label_size} );
+    }
     QCBOREncode_CloseMap(&encode_ctx);
 
-    if (corrupt_encoded_cbor) {
+    if (test_params->is_encoded_cbor_corrupt) {
         /* Deliberately corrupt CBOR map metadata and construct */
         *((uint8_t *)encode_ctx.OutBuf.UB.ptr + 1) = 0xff;
     }
@@ -290,7 +307,7 @@
     dpe_error_t dpe_err;
     QCBORError qcbor_err;
     UsefulBufC encoded_buf;
-    UsefulBuf_MAKE_STACK_UB(cmd_buf, 612);
+    UsefulBuf_MAKE_STACK_UB(cmd_buf, DICE_CERT_SIZE);
 
     const struct derive_context_input_t in_args = {
         context_handle,
@@ -348,6 +365,8 @@
         memcpy(new_certificate_buf, out_args.new_certificate,
                out_args.new_certificate_size);
         *new_certificate_actual_size = out_args.new_certificate_size;
+    } else {
+        *new_certificate_actual_size = 0;
     }
 
     if (export_cdi) {
@@ -357,6 +376,8 @@
         memcpy(exported_cdi_buf, out_args.exported_cdi,
                out_args.exported_cdi_size);
         *exported_cdi_actual_size = out_args.exported_cdi_size;
+    } else {
+        *exported_cdi_actual_size = 0;
     }
 
     return DPE_NO_ERROR;
@@ -375,13 +396,13 @@
                             size_t derived_public_key_buf_size,
                             size_t *derived_public_key_actual_size,
                             int *new_context_handle,
-                            bool corrupt_encoded_cbor)
+                            struct dpe_certify_key_test_params_t *test_params)
 {
     int32_t service_err;
     dpe_error_t dpe_err;
     QCBORError qcbor_err;
     UsefulBufC encoded_buf;
-    UsefulBuf_MAKE_STACK_UB(cmd_buf, 2500);
+    UsefulBuf_MAKE_STACK_UB(cmd_buf, DICE_CERT_SIZE);
 
     const struct certify_key_input_t in_args = {
         context_handle,
@@ -401,7 +422,7 @@
         return DPE_INVALID_ARGUMENT;
     }
 
-    qcbor_err = encode_certify_key(&in_args, cmd_buf, &encoded_buf, corrupt_encoded_cbor);
+    qcbor_err = encode_certify_key(&in_args, cmd_buf, &encoded_buf, test_params);
     if (qcbor_err != QCBOR_SUCCESS) {
         return DPE_INTERNAL_ERROR;
     }
diff --git a/partitions/dice_protection_environment/test/dpe_test_data.h b/partitions/dice_protection_environment/test/dpe_test_data.h
index 2efbde6..a8163a7 100644
--- a/partitions/dice_protection_environment/test/dpe_test_data.h
+++ b/partitions/dice_protection_environment/test/dpe_test_data.h
@@ -53,9 +53,22 @@
     bool is_config_descriptor_missing;
     bool is_authority_hash_missing;
     bool is_mode_missing;
-    bool corrupt_encoded_cbor;
+    bool is_encoded_cbor_corrupt;
     bool is_input_dice_data_missing;
     bool is_cert_id_missing;
+    bool is_retain_parent_context_missing;
+    bool is_allow_new_context_to_derive_missing;
+    bool is_create_certificate_missing;
+    bool is_return_certificate_missing;
+    bool is_allow_new_context_to_export_missing;
+    bool is_export_cdi_missing;
+};
+
+struct dpe_certify_key_test_params_t {
+    bool is_encoded_cbor_corrupt;
+    bool is_retain_context_missing;
+    bool is_public_key_missing;
+    bool is_label_missing;
 };
 
 dpe_error_t
@@ -80,19 +93,19 @@
                    struct dpe_derive_context_test_params_t *test_params);
 
 dpe_error_t dpe_certify_key_with_test_param(int context_handle,
-                    bool                        retain_context,
-                    const uint8_t              *public_key,
-                    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                    *derived_public_key_buf,
-                    size_t                      derived_public_key_buf_size,
-                    size_t                     *derived_public_key_actual_size,
-                    int                        *new_context_handle,
-                    bool                        corrupt_encoded_cbor);
+                    bool                    retain_context,
+                    const uint8_t          *public_key,
+                    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                *derived_public_key_buf,
+                    size_t                  derived_public_key_buf_size,
+                    size_t                 *derived_public_key_actual_size,
+                    int                    *new_context_handle,
+                    struct dpe_certify_key_test_params_t *test_params);
 
 #ifdef __cplusplus
 }
diff --git a/partitions/dice_protection_environment/test/secure/dpe_s_interface_testsuite.c b/partitions/dice_protection_environment/test/secure/dpe_s_interface_testsuite.c
index af25907..a8d9c0b 100644
--- a/partitions/dice_protection_environment/test/secure/dpe_s_interface_testsuite.c
+++ b/partitions/dice_protection_environment/test/secure/dpe_s_interface_testsuite.c
@@ -58,8 +58,12 @@
      "DPE CertifyKey - invalid public key buffer size"},
     {&certify_key_invalid_cbor_encoded_input_test, "DPE_S_TEST_1024",
      "DPE CertifyKey - invalid cbor encoded input"},
-    {&derive_context_with_parent_leaf_component_test, "DPE_S_TEST_1025",
+    {&certify_key_without_optional_args_test, "DPE_S_TEST_1025",
+     "DPE CertifyKey - without optional arguments"},
+    {&derive_context_with_parent_leaf_component_test, "DPE_S_TEST_1026",
      "DPE DeriveContext - Leaf component"},
+    {&derive_context_without_optional_args_test, "DPE_S_TEST_1027",
+     "DPE DeriveContext - without optional arguments"},
 };
 
 void register_testsuite_extra_s_interface(struct test_suite_t *p_test_suite)