From da9e3a84722771f9454ee3428f6a16cfbe57d2ae Mon Sep 17 00:00:00 2001
From: "Ching-Hsin,Lee" <chinglee@amazon.com>
Date: Thu, 16 May 2024 04:13:14 +0000
Subject: [PATCH] Fix for CBMC

---
 source/core_pkcs11.c                   | 19 ++---
 test/wrapper_utest/core_pkcs11_utest.c | 97 +++++++++++++++++++++++++-
 2 files changed, 107 insertions(+), 9 deletions(-)

diff --git a/source/core_pkcs11.c b/source/core_pkcs11.c
index 5dd1dd17..eb595601 100644
--- a/source/core_pkcs11.c
+++ b/source/core_pkcs11.c
@@ -164,21 +164,24 @@ CK_RV xInitializePkcs11Token( void )
     CK_FLAGS xTokenFlags = 0;
     CK_TOKEN_INFO_PTR pxTokenInfo = NULL;
 
-    xResult = C_GetFunctionList( &pxFunctionList );
+    xResult = xInitializePKCS11();
 
-    if( ( pxFunctionList == NULL ) || ( pxFunctionList->C_GetTokenInfo == NULL ) || ( pxFunctionList->C_InitToken == NULL ) )
+    if( ( xResult == CKR_OK ) || ( xResult == CKR_CRYPTOKI_ALREADY_INITIALIZED ) )
     {
-        xResult = CKR_FUNCTION_FAILED;
+        xResult = xGetSlotList( &pxSlotId, &xSlotCount );
     }
 
     if( xResult == CKR_OK )
     {
-        xResult = xInitializePKCS11();
-    }
+        xResult = C_GetFunctionList( &pxFunctionList );
 
-    if( ( xResult == CKR_OK ) || ( xResult == CKR_CRYPTOKI_ALREADY_INITIALIZED ) )
-    {
-        xResult = xGetSlotList( &pxSlotId, &xSlotCount );
+        if( xResult == CKR_OK )
+        {
+            if( ( pxFunctionList == NULL ) || ( pxFunctionList->C_GetTokenInfo == NULL ) || ( pxFunctionList->C_InitToken == NULL ) )
+            {
+                xResult = CKR_FUNCTION_FAILED;
+            }
+        }
     }
 
     if( xResult == CKR_OK )
diff --git a/test/wrapper_utest/core_pkcs11_utest.c b/test/wrapper_utest/core_pkcs11_utest.c
index aeefc446..a2cd04fc 100644
--- a/test/wrapper_utest/core_pkcs11_utest.c
+++ b/test/wrapper_utest/core_pkcs11_utest.c
@@ -176,6 +176,58 @@ static CK_RV prvSetFunctionList( CK_FUNCTION_LIST_PTR_PTR ppxPtr )
     return CKR_OK;
 }
 
+/*!
+ * @brief Create a stub for the PKCS #11 function list.
+ *
+ * Fails on the fourth call in order to create coverage for a nested branch.
+ *
+ */
+static CK_RV prvSetFunctionList2( CK_FUNCTION_LIST_PTR_PTR ppxPtr )
+{
+    static uint32_t ulCalls = 0;
+    CK_RV xResult = CKR_OK;
+
+    ulCalls++;
+
+    if( ulCalls == 3 )
+    {
+        xResult = CKR_ARGUMENTS_BAD;
+        *ppxPtr = NULL;
+    }
+    else
+    {
+        *ppxPtr = &prvP11FunctionList;
+    }
+
+    return xResult;
+}
+
+/*!
+ * @brief Create a stub for the PKCS #11 function list.
+ *
+ * Fails on the fourth call in order to create coverage for a nested branch.
+ *
+ */
+static CK_RV prvSetFunctionList3( CK_FUNCTION_LIST_PTR_PTR ppxPtr )
+{
+    static uint32_t ulCalls = 0;
+    CK_RV xResult = CKR_OK;
+
+    ulCalls++;
+
+    if( ulCalls == 3 )
+    {
+        xResult = CKR_OK;
+        *ppxPtr = NULL;
+    }
+    else
+    {
+        *ppxPtr = &prvP11FunctionList;
+    }
+
+    return xResult;
+}
+
 /*!
  * @brief Return empty function list
  *
@@ -552,6 +604,23 @@ void test_IotPkcs11_xInitializePkcs11TokenAlreadyInit( void )
     TEST_ASSERT_EQUAL( CKR_OK, xResult );
 }
 
+/*!
+ * @brief xInitializePkcs11Token xInitializePKCS11 return error.
+ *
+ */
+void test_IotPkcs11_xInitializePkcs11TokenInitFailed( void )
+{
+    CK_RV xResult = CKR_OK;
+
+    C_GetFunctionList_IgnoreAndReturn( CKR_OK );
+    C_GetFunctionList_Stub( ( void * ) &prvSetFunctionList );
+    C_Initialize_IgnoreAndReturn( CKR_GENERAL_ERROR );
+
+    xResult = xInitializePkcs11Token();
+
+    TEST_ASSERT_EQUAL( CKR_GENERAL_ERROR, xResult );
+}
+
 /*!
  * @brief xInitializePkcs11Token C_GetTokenInfo failure due to memory constraint.
  *
@@ -617,7 +686,33 @@ void test_IotPkcs11_xInitializePkcs11TokenBadFunctionList( void )
 {
     CK_RV xResult = CKR_OK;
 
-    C_GetFunctionList_IgnoreAndReturn( CKR_ARGUMENTS_BAD );
+    C_GetFunctionList_IgnoreAndReturn( CKR_OK );
+    C_GetFunctionList_Stub( ( void * ) &prvSetFunctionList2 );
+    C_Initialize_IgnoreAndReturn( CKR_OK );
+    pvPkcs11Malloc_Stub( pvPkcs11MallocCb );
+    vPkcs11Free_Stub( vPkcs11FreeCb );
+    C_GetSlotList_Stub( ( void * ) xGet1Item );
+
+    xResult = xInitializePkcs11Token();
+
+    TEST_ASSERT_EQUAL( CKR_ARGUMENTS_BAD, xResult );
+}
+
+/*!
+ * @brief xInitializePkcs11Token failure due to bad C_GetFunctionList.
+ *
+ */
+void test_IotPkcs11_xInitializePkcs11TokenEmptyFunctionList( void )
+{
+    CK_RV xResult = CKR_OK;
+
+    C_GetFunctionList_IgnoreAndReturn( CKR_OK );
+    C_GetFunctionList_Stub( ( void * ) &prvSetFunctionList3 );
+    C_Initialize_IgnoreAndReturn( CKR_OK );
+    pvPkcs11Malloc_Stub( pvPkcs11MallocCb );
+    vPkcs11Free_Stub( vPkcs11FreeCb );
+    C_GetSlotList_Stub( ( void * ) xGet1Item );
+
     xResult = xInitializePkcs11Token();
 
     TEST_ASSERT_EQUAL( CKR_FUNCTION_FAILED, xResult );