DPE: Fix all test failures

- For reg test, increase number of layers to 6.
- Separate certificate IDs for all the tests where undestroyable
  contexts are created.
- For now, until allocated buffer sizes are passed to the service, return valid
  context handle if DPE service has successfully derived the context.

Signed-off-by: Maulik Patel <maulik.patel@arm.com>
Change-Id: I5766b0df21d79dff569789aedee94bd50c2b9d9a
diff --git a/partitions/dice_protection_environment/dpe_context_mngr.h b/partitions/dice_protection_environment/dpe_context_mngr.h
index dbbc23d..0a35b7c 100644
--- a/partitions/dice_protection_environment/dpe_context_mngr.h
+++ b/partitions/dice_protection_environment/dpe_context_mngr.h
@@ -29,9 +29,13 @@
 
 /* Below configuration defines are platform dependent */
 #define MAX_NUM_OF_COMPONENTS 20
-#define MAX_NUM_OF_LAYERS 4
 #define DPE_PLATFORM_LAYER_IDX 1
 #define DPE_SECURE_WORLD_AND_HYPERVISOR_LAYER_IDX 2
+#ifdef DPE_TEST_MODE
+#define MAX_NUM_OF_LAYERS 6
+#else
+#define MAX_NUM_OF_LAYERS 4
+#endif /* DPE_TEST_MODE */
 
 /* Below threshold defines the threshold below which a context cannot be destroyed */
 #define DPE_DESTROY_CONTEXT_THRESHOLD_LAYER_IDX DPE_SECURE_WORLD_AND_HYPERVISOR_LAYER_IDX
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 247cb9e..a4dca0a 100644
--- a/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
+++ b/partitions/dice_protection_environment/interface/src/dpe_cmd_encode.c
@@ -538,6 +538,12 @@
     }
 
     /* Copy returned values into caller's memory */
+    /* Output buffer sizes are checked by client side API implementation and
+     * not by DPE service. Until allocated buffer sizes are passed to the service,
+     * and checked if their size is sufficient, return output handle
+     * from the command
+     */
+    *new_context_handle = out_args.new_context_handle;
     if (out_args.certificate_chain_size > certificate_chain_buf_size) {
         return DPE_INVALID_ARGUMENT;
     }
@@ -554,8 +560,6 @@
         *derived_public_key_actual_size = out_args.derived_public_key_size;
     }
 
-    *new_context_handle = out_args.new_context_handle;
-
     return DPE_NO_ERROR;
 }
 
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 fa86b10..14cb248 100644
--- a/partitions/dice_protection_environment/test/dpe_certify_key_test.c
+++ b/partitions/dice_protection_environment/test/dpe_certify_key_test.c
@@ -247,7 +247,7 @@
     int out_parent_handle;
 
     dpe_err = dpe_derive_context(retained_rot_ctx_handle,       /* context_handle */
-                                 DPE_PLATFORM_CERT_ID,          /* cert_id */
+                                 DPE_UNDESTROYABLE_CTX_CERT_ID_3, /* cert_id */
                                  true,                          /* retain_parent_context */
                                  true,                          /* allow_new_context_to_derive */
                                  true,                          /* create_certificate */
@@ -434,6 +434,15 @@
     /* Save the last handle for the subsequent test */
     retained_rot_ctx_handle = out_parent_handle;
 
+    /* Since certificate buffer size is checked by client side API implementation,
+     * it derives a valid DPE context within the service, so destroy that context
+     */
+    dpe_err = dpe_destroy_context(new_context_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
+
     ret->val = TEST_PASSED;
 }
 
@@ -491,6 +500,15 @@
     /* Save the last handle for the subsequent test */
     retained_rot_ctx_handle = out_parent_handle;
 
+    /* Since public key buffer size is checked by client side API implementation,
+     * it derives a valid DPE context within the service, so destroy that context
+     */
+    dpe_err = dpe_destroy_context(new_context_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
+
     ret->val = TEST_PASSED;
 }
 
@@ -536,6 +554,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;
 
@@ -554,7 +579,7 @@
     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 */
+                                 DPE_UNDESTROYABLE_CTX_CERT_ID_4, /* cert_id */
                                  true,                          /* retain_parent_context */
                                  true,                          /* allow_new_context_to_derive */
                                  true,                          /* create_certificate */
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 76a3711..d8daf0d 100644
--- a/partitions/dice_protection_environment/test/dpe_derive_context_test.c
+++ b/partitions/dice_protection_environment/test/dpe_derive_context_test.c
@@ -12,7 +12,7 @@
 
 #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 */                                          \
+            cert_id,                 /* cert_id */                                          \
             true,                    /* retain_parent_context */                            \
             true,                    /* allow_new_context_to_derive */                      \
             false,                   /* create_certificate */                               \
@@ -367,6 +367,7 @@
     uint8_t exported_cdi_buf[DICE_MAX_ENCODED_CDI_SIZE];
     size_t exported_cdi_actual_size;
     bool return_certificate = false;
+    uint32_t cert_id = DPE_PLATFORM_CERT_ID;
 
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
     struct dpe_derive_context_test_params_t test_params = {0};
@@ -430,6 +431,7 @@
     bool return_certificate = false;
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
     struct dpe_derive_context_test_params_t test_params = {0};
+    uint32_t cert_id = DPE_PLATFORM_CERT_ID;
 
     test_params.is_encoded_cbor_corrupt = true;
     dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
@@ -504,7 +506,7 @@
      * hence use invalid cert id.
      */
     dpe_err = dpe_derive_context(retained_rot_ctx_handle,       /* input_ctx_handle */
-                                 DPE_PLATFORM_CERT_ID,          /* cert_id */
+                                 DPE_UNDESTROYABLE_CTX_CERT_ID_1, /* cert_id */
                                  true,                          /* retain_parent_context */
                                  true,                          /* allow_new_context_to_derive */
                                  true,                          /* create_certificate */
@@ -682,7 +684,7 @@
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
 
     dpe_err = dpe_derive_context(retained_rot_ctx_handle,       /* input_ctx_handle */
-                                 DPE_PLATFORM_CERT_ID,          /* cert_id */
+                                 DPE_UNDESTROYABLE_CTX_CERT_ID_2, /* cert_id */
                                  true,                          /* retain_parent_context */
                                  true,                          /* allow_new_context_to_derive */
                                  true,                          /* create_certificate */
@@ -839,6 +841,7 @@
     bool return_certificate = false;
     DiceInputValues dice_inputs = DEFAULT_DICE_INPUT;
     struct dpe_derive_context_test_params_t test_params = {0};
+    uint32_t cert_id = DPE_PLATFORM_CERT_ID;
 
     test_params.is_allow_new_context_to_derive_missing = true;
     dpe_err = CALL_DERIVE_CONTEXT_WITH_TEST_PARAM();
@@ -854,6 +857,11 @@
         TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
         return;
     }
+    dpe_err = dpe_destroy_context(out_ctx_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
 
     retained_rot_ctx_handle = out_parent_handle;
     test_params.is_allow_new_context_to_derive_missing = false;
@@ -872,6 +880,11 @@
         TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
         return;
     }
+    dpe_err = dpe_destroy_context(out_ctx_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
 
     retained_rot_ctx_handle = out_parent_handle;
     test_params.is_create_certificate_missing = false;
@@ -889,6 +902,11 @@
         TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
         return;
     }
+    dpe_err = dpe_destroy_context(out_ctx_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
 
     retained_rot_ctx_handle = out_parent_handle;
     test_params.is_return_certificate_missing = false;
@@ -901,6 +919,11 @@
     //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.
+    dpe_err = dpe_destroy_context(out_ctx_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
 
     retained_rot_ctx_handle = out_parent_handle;
     test_params.is_allow_new_context_to_export_missing = false;
@@ -916,6 +939,11 @@
         TEST_FAIL("DPE DeriveContext test: Without optional parameter should not fail");
         return;
     }
+    dpe_err = dpe_destroy_context(out_ctx_handle, false);
+    if (dpe_err != DPE_NO_ERROR) {
+        TEST_FAIL("DPE DestroyContext call failed");
+        return;
+    }
 
     retained_rot_ctx_handle = out_parent_handle;
     test_params.is_export_cdi_missing = false;
@@ -923,6 +951,7 @@
     /* This test will create undestroyable context as default value of
      * retain_parent_context is false
      */
+    cert_id = DPE_UNDESTROYABLE_CTX_CERT_ID_5;
     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");
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 cb7787a..4aba44f 100644
--- a/partitions/dice_protection_environment/test/dpe_test_cmd_encode.c
+++ b/partitions/dice_protection_environment/test/dpe_test_cmd_encode.c
@@ -442,6 +442,13 @@
     }
 
     /* Copy returned values into caller's memory */
+    /* Output buffer sizes are checked by client side API implementation and
+     * not by DPE service. Until allocated buffer sizes are passed to the service,
+     * and checked if their size is sufficient, return output handle
+     * from the command
+     */
+    *new_context_handle = out_args.new_context_handle;
+
     if (out_args.certificate_chain_size > certificate_chain_buf_size) {
         return DPE_INVALID_ARGUMENT;
     }
@@ -458,9 +465,5 @@
         *derived_public_key_actual_size = out_args.derived_public_key_size;
     }
 
-    if (retain_context) {
-        *new_context_handle = out_args.new_context_handle;
-    }
-
     return DPE_NO_ERROR;
 }
diff --git a/partitions/dice_protection_environment/test/dpe_test_data.h b/partitions/dice_protection_environment/test/dpe_test_data.h
index a8163a7..b30ae24 100644
--- a/partitions/dice_protection_environment/test/dpe_test_data.h
+++ b/partitions/dice_protection_environment/test/dpe_test_data.h
@@ -16,6 +16,14 @@
 #define INVALID_COMPONENT_IDX 0xFFFF
 
 #define DPE_PLATFORM_CERT_ID 0x200
+/* Certificate IDs used for tests where layers are finalized and undestroyable
+ * contexts are created (i.e. no valid handle is returned for the context)
+ */
+#define DPE_UNDESTROYABLE_CTX_CERT_ID_1 0x901
+#define DPE_UNDESTROYABLE_CTX_CERT_ID_2 0x902
+#define DPE_UNDESTROYABLE_CTX_CERT_ID_3 0x903
+#define DPE_UNDESTROYABLE_CTX_CERT_ID_4 0x904
+#define DPE_UNDESTROYABLE_CTX_CERT_ID_5 0x905
 
 #define DERIVE_CONTEXT_TEST_DATA1_SIZE 3
 #define DERIVE_CONTEXT_TEST_DATA2_SIZE 1