added server context init flags

wolfSSL · Mar 22, 2024 · e651636 · e651636
1 parent f956767
commit e651636
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/src/wh_server.c b/src/wh_server.c
@@ -65,6 +65,7 @@ int wh_Server_Init(whServerContext* server, whServerConfig* config)
         (void)wh_Server_Cleanup(server);
         return WH_ERROR_ABORTED;
     }
+    server->flags.wcInitFlag = true;
 
 #if defined(WOLF_CRYPTO_CB)
     server->crypto->devId = config->devId;
@@ -81,20 +82,20 @@ int wh_Server_Init(whServerContext* server, whServerConfig* config)
 #else
     server->crypto->devId = INVALID_DEVID;
 #endif
+    server->flags.wcDevIdInitFlag = true;
 
     rc = wc_InitRng_ex(server->crypto->rng, NULL, server->crypto->devId);
     if (rc != 0) {
         (void)wh_Server_Cleanup(server);
         return WH_ERROR_ABORTED;
     }
+    server->flags.wcRngInitFlag = true;
 
     rc = wh_Nvm_Init(server->nvm, config->nvm_config);
     if (rc != 0) {
         (void)wh_Server_Cleanup(server);
         return WH_ERROR_ABORTED;
     }
-    server->nvm->cb      = config->nvm_config->cb;
-    server->nvm->context = config->nvm_config->context;
 
     rc = wh_CommServer_Init(server->comm, config->comm_config);
     if (rc != 0) {
@@ -112,10 +113,24 @@ int wh_Server_Cleanup(whServerContext* server)
 
     (void)wh_CommServer_Cleanup(server->comm);
     (void)wh_Nvm_Cleanup(server->nvm);
-    (void)wc_FreeRng(server->crypto->rng);
-    (void)wolfCrypt_Cleanup();
+
+#if defined(WOLF_CRYPTO_CB)
+    if (server->flags.wcDevIdInitFlag &&
+        server->crypto->devId != INVALID_DEVID) {
+        (void)wc_CryptoCb_UnRegisterDevice(server->crypto->devId);
+    }
+#endif
+
+    if (server->flags.wcRngInitFlag) {
+        (void)wc_FreeRng(server->crypto->rng);
+    }
+
+    if (server->flags.wcInitFlag) {
+        (void)wolfCrypt_Cleanup();
+    }
 
     memset(server, 0, sizeof(*server));
+
     return 0;
 }
 

diff --git a/wolfhsm/wh_server.h b/wolfhsm/wh_server.h
@@ -7,6 +7,7 @@
  */
 
 #include <stdint.h>
+#include <stdbool.h>
 
 #include "wolfhsm/wh_common.h"
 #include "wolfhsm/wh_comm.h"
@@ -30,8 +31,15 @@ typedef struct {
     WC_RNG rng[1];
 } crypto_context;
 
+typedef struct {
+    bool wcInitFlag: 1;
+    bool wcRngInitFlag: 1;
+    bool wcDevIdInitFlag: 1;
+} whServerFlags;
+
 /* Context structure to maintain the state of an HSM server */
 typedef struct whServerContext_t {
+    whServerFlags flags;
     whCommServer comm[1];
     whNvmContext nvm[1];
     crypto_context crypto[1];