From a291117ed2be9c69ffe89794805c816a5c25d815 Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Tue, 2 Apr 2024 08:38:48 -0600
Subject: [PATCH 1/6] added custom user-callback feature

---
 src/wh_client.c             | 40 ++++++++++++++++++++-
 src/wh_message_custom.c     | 33 +++++++++++++++++
 src/wh_server.c             | 29 ++++-----------
 src/wh_server_custom.c      | 72 +++++++++++++++++++++++++++++++++++++
 test/Makefile               |  2 ++
 test/wh_config.h            |  2 +-
 test/wh_test_clientserver.c | 23 ++++++++++++
 wolfhsm/wh_client.h         |  5 +++
 wolfhsm/wh_error.h          |  3 ++
 wolfhsm/wh_message.h        |  5 +--
 wolfhsm/wh_message_custom.h | 25 +++++++++++++
 wolfhsm/wh_server_custom.h  | 19 ++++++++++
 12 files changed, 231 insertions(+), 27 deletions(-)
 create mode 100644 src/wh_message_custom.c
 create mode 100644 src/wh_server_custom.c
 create mode 100644 wolfhsm/wh_message_custom.h
 create mode 100644 wolfhsm/wh_server_custom.h

diff --git a/src/wh_client.c b/src/wh_client.c
index 2c8e3ea0..9e267d71 100644
--- a/src/wh_client.c
+++ b/src/wh_client.c
@@ -23,6 +23,7 @@
 /* Message definitions */
 #include "wolfhsm/wh_message.h"
 #include "wolfhsm/wh_message_comm.h"
+#include "wolfhsm/wh_message_custom.h"
 
 #include "wolfhsm/wh_client.h"
 
@@ -34,7 +35,7 @@ int wh_Client_Init(whClientContext* c, const whClientConfig* config)
     }
 
     memset(c, 0, sizeof(*c));
-    
+
     if (    ((rc = wh_CommClient_Init(c->comm, config->comm)) == 0) &&
             ((rc = wolfCrypt_Init()) == 0) &&
             ((rc = wc_CryptoCb_RegisterDevice(WOLFHSM_DEV_ID, wolfHSM_CryptoCb, c)) == 0) &&
@@ -201,3 +202,40 @@ int wh_Client_Echo(whClientContext* c, uint16_t snd_len, const void* snd_data,
     return rc;
 }
 
+
+int wh_Client_CustomRequest(whClientContext* c, uint16_t action)
+{
+    whMessageCustom_Request req = {0};
+
+    if (NULL == c || action >= WH_MESSAGE_ACTION_MAX) {
+        return WH_ERROR_BADARGS;
+    }
+
+    req.id = action;
+
+    return wh_Client_SendRequest(c, WH_MESSAGE_GROUP_CUSTOM, action,
+                                 sizeof(req), &req);
+}
+
+
+int wh_Client_CustomResponse(whClientContext* c, int32_t *out_rc)
+{
+    whMessageCustom_Response resp = {0};
+    uint16_t resp_group = 0;
+    uint16_t resp_action = 0;
+    uint16_t resp_size = 0;
+    int32_t rc = 0;
+
+    if (NULL == c || out_rc == NULL) {
+        return WH_ERROR_BADARGS;
+    }
+
+    rc = wh_Client_RecvResponse(c, &resp_group, &resp_action, &resp_size, &resp);
+    if (rc != WH_ERROR_OK) {
+        return rc;
+    }
+
+    *out_rc = resp.rc;
+
+    return WH_ERROR_OK;
+}
\ No newline at end of file
diff --git a/src/wh_message_custom.c b/src/wh_message_custom.c
new file mode 100644
index 00000000..91689a8a
--- /dev/null
+++ b/src/wh_message_custom.c
@@ -0,0 +1,33 @@
+#include <stddef.h>
+
+#include "wolfhsm/wh_message_custom.h"
+#include "wolfhsm/wh_error.h"
+#include "wolfhsm/wh_comm.h"
+
+
+int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
+                                      const whMessageCustom_Request* src,
+                                      whMessageCustom_Request*       dst)
+{
+    if ((src == NULL) || (dst == NULL)) {
+        return WH_ERROR_BADARGS;
+    }
+
+    dst->id = wh_Translate16(magic, src->id);
+
+    return WH_ERROR_OK;
+}
+
+int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
+                                       const whMessageCustom_Response* src,
+                                       whMessageCustom_Response*       dst)
+{
+    if ((src == NULL) || (dst == NULL)) {
+        return WH_ERROR_BADARGS;
+    }
+
+    dst->rc  = wh_Translate32(magic, src->rc);
+    dst->err = wh_Translate32(magic, src->err);
+
+    return WH_ERROR_OK;
+}
\ No newline at end of file
diff --git a/src/wh_server.c b/src/wh_server.c
index e64725d8..ed642456 100644
--- a/src/wh_server.c
+++ b/src/wh_server.c
@@ -18,8 +18,9 @@
 
 /* Server API's */
 #include "wolfhsm/wh_server.h"
-#include "wolfhsm/wh_server_crypto.h"
 #include "wolfhsm/wh_server_internal.h"
+#include "wolfhsm/wh_server_crypto.h"
+#include "wolfhsm/wh_server_custom.h"
 #if defined(WOLFHSM_SHE_EXTENSION)
 #include "wolfhsm/wh_server_she.h"
 #endif
@@ -44,11 +45,6 @@ static int _wh_Server_HandleSheRequest(whServerContext* server,
         uint16_t req_size, const void* req_packet,
         uint16_t *out_resp_size, void* resp_packet);
 #endif
-static int _wh_Server_HandleCustomRequest(whServerContext* server,
-        uint16_t magic, uint16_t action, uint16_t seq,
-        uint16_t req_size, const void* req_packet,
-        uint16_t *out_resp_size, void* resp_packet);
-
 
 int wh_Server_Init(whServerContext* server, whServerConfig* config)
 {
@@ -215,20 +211,7 @@ static int _wh_Server_HandleSheRequest(whServerContext* server,
 }
 #endif
 
-static int _wh_Server_HandleCustomRequest(whServerContext* server,
-        uint16_t magic, uint16_t action, uint16_t seq,
-        uint16_t req_size, const void* req_packet,
-        uint16_t *out_resp_size, void* resp_packet)
-{
-    int rc = 0;
-    switch (action) {
-    /* TODO: Add custom/user callback message handling here */
-    default:
-        /* Unknown request. Respond with empty packet */
-        *out_resp_size = 0;
-    }
-    return rc;
-}
+
 
 int wh_Server_HandleRequestMessage(whServerContext* server)
 {
@@ -266,11 +249,11 @@ int wh_Server_HandleRequestMessage(whServerContext* server)
             rc = _wh_Server_HandleKeyRequest(server, magic, action, seq,
                     size, data, &size, data);
         break;
-        
+
         case WH_MESSAGE_GROUP_CRYPTO:
             rc = wh_Server_HandleCryptoRequest(server, action, data, &size);
         break;
-        
+
         case WH_MESSAGE_GROUP_PKCS11:
             rc = _wh_Server_HandlePkcs11Request(server, magic, action, seq,
                     size, data, &size, data);
@@ -283,7 +266,7 @@ int wh_Server_HandleRequestMessage(whServerContext* server)
 #endif
 
         case WH_MESSAGE_GROUP_CUSTOM:
-            rc = _wh_Server_HandleCustomRequest(server, magic, action, seq,
+            rc = wh_Server_HandleCustomRequest(server, magic, action, seq,
                     size, data, &size, data);
             break;
 
diff --git a/src/wh_server_custom.c b/src/wh_server_custom.c
new file mode 100644
index 00000000..3e2806d6
--- /dev/null
+++ b/src/wh_server_custom.c
@@ -0,0 +1,72 @@
+#include "wolfhsm/wh_server_custom.h"
+#include "wolfhsm/wh_error.h"
+#include "wolfhsm/wh_message.h"
+#include "wolfhsm/wh_message_custom.h"
+
+
+/* Statically allocated callback table holding user callbacks */
+static whServerCustomCb customHandlerTable[WH_MESSAGE_ACTION_MAX] = {NULL};
+
+
+int wh_Server_RegisterCustomCb(uint16_t action, whServerCustomCb handler)
+{
+    if (NULL == handler || action >= WH_MESSAGE_ACTION_MAX) {
+        return WH_ERROR_BADARGS;
+    }
+
+    customHandlerTable[action] = handler;
+
+    return WH_ERROR_OK;
+}
+
+
+int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
+                                  uint16_t action, uint16_t seq,
+                                  uint16_t req_size, const void* req_packet,
+                                  uint16_t* out_resp_size, void* resp_packet)
+{
+    int                      rc   = 0;
+    whMessageCustom_Request  req  = {0};
+    whMessageCustom_Response resp = {0};
+
+
+    if (action >= WH_MESSAGE_ACTION_MAX) {
+        /* Invalid callback index  */
+        return WH_ERROR_BADARGS;
+    }
+
+    if (req_size != sizeof(whMessageCustom_Request)) {
+        /* Request is malformed */
+        return WH_ERROR_ABORTED;
+    }
+
+    /* Translate the request */
+    if ((rc = wh_MessageCustom_TranslateRequest(magic, req_packet, &req)) !=
+        WH_ERROR_OK) {
+        return rc;
+    }
+
+    if (customHandlerTable[action] != NULL) {
+        /* Invoke the registered callback, storing the return value in the
+         * reponse  */
+        resp.rc = customHandlerTable[action](server, &req, &resp);
+        /* TODO: propagate wolfHSM error codes (requires modifiying caller
+         * function)*/
+        resp.err = WH_ERROR_OK;
+    }
+    else {
+        /* No callback was registered, populate response errors, but we must
+         * return success to ensure the "error" response is sent  */
+        /* TODO: what should we set resp.rc to? */
+        resp.err = WH_ERROR_NO_HANDLER;
+    }
+
+    /* Translate the response and set output size */
+    if ((rc = wh_MessageCustom_TranslateResponse(magic, &resp, resp_packet)) !=
+        WH_ERROR_OK) {
+        return rc;
+    }
+    *out_resp_size = sizeof(resp);
+
+    return WH_ERROR_OK;
+}
diff --git a/test/Makefile b/test/Makefile
index c64a9283..6a45c258 100644
--- a/test/Makefile
+++ b/test/Makefile
@@ -80,11 +80,13 @@ SRC_C += \
             $(WOLFHSM_DIR)/src/wh_flash_unit.c \
             $(WOLFHSM_DIR)/src/wh_message_comm.c \
             $(WOLFHSM_DIR)/src/wh_message_nvm.c \
+            $(WOLFHSM_DIR)/src/wh_message_custom.c \
             $(WOLFHSM_DIR)/src/wh_nvm.c \
             $(WOLFHSM_DIR)/src/wh_nvm_flash.c \
             $(WOLFHSM_DIR)/src/wh_server.c \
             $(WOLFHSM_DIR)/src/wh_server_nvm.c \
             $(WOLFHSM_DIR)/src/wh_server_crypto.c \
+            $(WOLFHSM_DIR)/src/wh_server_custom.c \
             $(WOLFHSM_DIR)/src/wh_client_cryptocb.c \
             $(WOLFHSM_DIR)/src/wh_transport_mem.c \
             $(WOLFHSM_DIR)/src/wh_flash_ramsim.c
diff --git a/test/wh_config.h b/test/wh_config.h
index 8f76b552..0bff3b6b 100644
--- a/test/wh_config.h
+++ b/test/wh_config.h
@@ -2,6 +2,6 @@
 #define WH_CONFIG_H_
 
 #define WH_CFG_TEST_VERBOSE
-#define WH_CFG_TEST_POSIX
+//#define WH_CFG_TEST_POSIX
 
 #endif /* WH_CONFIG_H_*/
diff --git a/test/wh_test_clientserver.c b/test/wh_test_clientserver.c
index 28010353..3520eab0 100644
--- a/test/wh_test_clientserver.c
+++ b/test/wh_test_clientserver.c
@@ -17,6 +17,8 @@
 #include "wolfhsm/wh_flash_ramsim.h"
 
 #include "wolfhsm/wh_server.h"
+#include "wolfhsm/wh_server_custom.h"
+#include "wolfhsm/wh_message.h"
 #include "wolfhsm/wh_client.h"
 
 #if defined(WH_CFG_TEST_POSIX)
@@ -31,6 +33,16 @@
 #define REPEAT_COUNT 10
 #define ONE_MS 1000
 
+
+/* Dummy callback that returns the registered callback ID */
+static int _customServerCb(whServerContext*               server,
+                           const whMessageCustom_Request* req,
+                           whMessageCustom_Response*      resp)
+{
+    return req->id;
+}
+
+
 int whTest_ClientServerSequential(void)
 {
     int ret = 0;
@@ -430,6 +442,17 @@ int whTest_ClientServerSequential(void)
                 ret, server_rc, avail_size, avail_objects, reclaim_size, reclaim_objects);
     #endif
 
+    /* Test custom registered callbacks */
+    for (counter = 0; counter < WH_MESSAGE_ACTION_MAX; counter++) {
+        int rc = 0;
+        WH_TEST_RETURN_ON_FAIL(wh_Server_RegisterCustomCb(counter, _customServerCb));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, counter));
+        WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &rc));
+        WH_TEST_ASSERT_RETURN(rc == counter);
+    }
+
+
     WH_TEST_RETURN_ON_FAIL(wh_Server_Cleanup(server));
     WH_TEST_RETURN_ON_FAIL(wh_Client_Cleanup(client));
 
diff --git a/wolfhsm/wh_client.h b/wolfhsm/wh_client.h
index 3b175b7a..a0f08848 100644
--- a/wolfhsm/wh_client.h
+++ b/wolfhsm/wh_client.h
@@ -177,4 +177,9 @@ int wh_Client_NvmReadDma(whClientContext* c,
         whNvmId id, whNvmSize offset, whNvmSize data_len, uint8_t* data,
         int32_t *out_rc);
 
+
+/* Client custom-callback support */
+int wh_Client_CustomRequest(whClientContext* c, uint16_t action); /* TODO: user data */
+int wh_Client_CustomResponse(whClientContext* c, int32_t *out_rc);
+
 #endif /* WOLFHSM_WH_CLIENT_H_ */
diff --git a/wolfhsm/wh_error.h b/wolfhsm/wh_error.h
index 5ba1b9dc..d49ac946 100644
--- a/wolfhsm/wh_error.h
+++ b/wolfhsm/wh_error.h
@@ -24,6 +24,9 @@ enum {
     WH_ERROR_NOTBLANK       = -413, /* Area is no blank */
     WH_ERROR_NOTFOUND       = -414, /* Matching object not found */
     WH_ERROR_NOSPACE        = -415, /* No available space */
+
+    /* Custom-callback status returns */
+    WH_ERROR_NO_HANDLER     = -420, /* No handler registered for action */
 };
 
 #endif /* WOLFHSM_WH_ERROR_H_ */
diff --git a/wolfhsm/wh_message.h b/wolfhsm/wh_message.h
index dfe3e57e..3b794d41 100644
--- a/wolfhsm/wh_message.h
+++ b/wolfhsm/wh_message.h
@@ -23,8 +23,9 @@ enum {
     WH_MESSAGE_GROUP_SHE            = 0x0700, /* SHE protocol */
     WH_MESSAGE_GROUP_CUSTOM         = 0x1000, /* User-specified features */
 
-    WH_MESSAGE_ACTION_MASK         = 0x00FF, /* 255 subtypes per group*/
-    WH_MESSAGE_ACTION_NONE         = 0x0000, /* No action. Invalid. */
+    WH_MESSAGE_ACTION_MASK         = 0x00FF,  /* 255 subtypes per group*/
+    WH_MESSAGE_ACTION_NONE         = 0x0000,  /* No action. Invalid. */
+    WH_MESSAGE_ACTION_MAX          = 0x00FF,  /* Max action value */
 };
 
 /* Construct the message kind based on group and action */
diff --git a/wolfhsm/wh_message_custom.h b/wolfhsm/wh_message_custom.h
new file mode 100644
index 00000000..71399af7
--- /dev/null
+++ b/wolfhsm/wh_message_custom.h
@@ -0,0 +1,25 @@
+#ifndef WH_MESSAGE_CUSTOM_H_
+#define WH_MESSAGE_CUSTOM_H_
+
+#include <stdint.h>
+
+typedef struct {
+    uint16_t id;
+    /* TODO: pass I/O pointers and sizes? Mimic NVM DMA request? */
+} whMessageCustom_Request;
+
+typedef struct {
+    int32_t rc;  /* Return code from custom callback */
+    int32_t err; /* wolfHSM-specific error */
+} whMessageCustom_Response;
+
+
+int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
+                                      const whMessageCustom_Request* src,
+                                      whMessageCustom_Request*       dst);
+
+int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
+                                       const whMessageCustom_Response* src,
+                                       whMessageCustom_Response*       dst);
+
+#endif /* WH_MESSAGE_CUSTOM_H_*/
\ No newline at end of file
diff --git a/wolfhsm/wh_server_custom.h b/wolfhsm/wh_server_custom.h
new file mode 100644
index 00000000..4e2ddea1
--- /dev/null
+++ b/wolfhsm/wh_server_custom.h
@@ -0,0 +1,19 @@
+#ifndef WH_SERVER_CUSTOM_H_
+#define WH_SERVER_CUSTOM_H_
+
+#include "wolfhsm/wh_server.h"
+#include "wolfhsm/wh_message_custom.h"
+
+
+typedef int (*whServerCustomCb)(
+    whServerContext*               server, /* TODO: should this be const? */
+    const whMessageCustom_Request* req, whMessageCustom_Response* resp);
+
+int wh_Server_RegisterCustomCb(uint16_t actionId, whServerCustomCb cb);
+
+int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
+                                  uint16_t action, uint16_t seq,
+                                  uint16_t req_size, const void* req_packet,
+                                  uint16_t* out_resp_size, void* resp_packet);
+
+#endif /* WH_SERVER_CUSTOM_H_ */
\ No newline at end of file

From 3f6296783870b2932c3f779edcdb88c22c47ce38 Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Tue, 2 Apr 2024 13:30:51 -0600
Subject: [PATCH 2/6] custom request/callbacks now support arbitrary
 input/output data

---
 src/wh_client.c             | 40 +++++++++---------
 src/wh_message_custom.c     | 51 ++++++++++++++++++++++-
 test/wh_config.h            |  2 +-
 test/wh_test_clientserver.c | 81 ++++++++++++++++++++++++++++++++-----
 wolfhsm/wh_client.h         |  5 ++-
 wolfhsm/wh_message_custom.h | 62 ++++++++++++++++++++++++++--
 6 files changed, 203 insertions(+), 38 deletions(-)

diff --git a/src/wh_client.c b/src/wh_client.c
index 9e267d71..21b054c6 100644
--- a/src/wh_client.c
+++ b/src/wh_client.c
@@ -202,40 +202,42 @@ int wh_Client_Echo(whClientContext* c, uint16_t snd_len, const void* snd_data,
     return rc;
 }
 
-
-int wh_Client_CustomRequest(whClientContext* c, uint16_t action)
+int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req)
 {
-    whMessageCustom_Request req = {0};
-
-    if (NULL == c || action >= WH_MESSAGE_ACTION_MAX) {
+    if (NULL == c || req == NULL || req->id >= WH_MESSAGE_ACTION_MAX) {
         return WH_ERROR_BADARGS;
     }
 
-    req.id = action;
-
-    return wh_Client_SendRequest(c, WH_MESSAGE_GROUP_CUSTOM, action,
-                                 sizeof(req), &req);
+    return wh_Client_SendRequest(c, WH_MESSAGE_GROUP_CUSTOM, req->id,
+                                 sizeof(*req), req);
 }
 
-
-int wh_Client_CustomResponse(whClientContext* c, int32_t *out_rc)
+int wh_Client_CustomResponse(whClientContext*          c,
+                             whMessageCustom_Response* outResp)
 {
-    whMessageCustom_Response resp = {0};
-    uint16_t resp_group = 0;
-    uint16_t resp_action = 0;
-    uint16_t resp_size = 0;
-    int32_t rc = 0;
+    whMessageCustom_Response resp;
+    uint16_t                 resp_group  = 0;
+    uint16_t                 resp_action = 0;
+    uint16_t                 resp_size   = 0;
+    int32_t                  rc          = 0;
 
-    if (NULL == c || out_rc == NULL) {
+    if (NULL == c || outResp == NULL) {
         return WH_ERROR_BADARGS;
     }
 
-    rc = wh_Client_RecvResponse(c, &resp_group, &resp_action, &resp_size, &resp);
+    rc =
+        wh_Client_RecvResponse(c, &resp_group, &resp_action, &resp_size, &resp);
     if (rc != WH_ERROR_OK) {
         return rc;
     }
 
-    *out_rc = resp.rc;
+    if (resp_size != sizeof(resp) || resp_group != WH_MESSAGE_GROUP_CUSTOM ||
+        resp_action >= WH_MESSAGE_ACTION_MAX) {
+        /* message invalid */
+        return WH_ERROR_ABORTED;
+    }
+
+    memcpy(outResp, &resp, sizeof(resp));
 
     return WH_ERROR_OK;
 }
\ No newline at end of file
diff --git a/src/wh_message_custom.c b/src/wh_message_custom.c
index 91689a8a..4591adfd 100644
--- a/src/wh_message_custom.c
+++ b/src/wh_message_custom.c
@@ -1,10 +1,49 @@
 #include <stddef.h>
+#include <string.h>
 
 #include "wolfhsm/wh_message_custom.h"
 #include "wolfhsm/wh_error.h"
 #include "wolfhsm/wh_comm.h"
 
 
+static void _translateCustomData(uint16_t magic, uint32_t translatedType,
+                                 const whMessageCustom_Data* src,
+                                 whMessageCustom_Data*       dst)
+{
+    if (translatedType < WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START) {
+        switch (translatedType) {
+            case WH_MESSAGE_CUSTOM_TYPE_DMA32: {
+                dst->dma32.client_addr =
+                    wh_Translate32(magic, src->dma32.client_addr);
+                dst->dma32.client_sz =
+                    wh_Translate32(magic, src->dma32.client_sz);
+                dst->dma32.server_addr =
+                    wh_Translate32(magic, src->dma32.server_addr);
+                dst->dma32.server_sz =
+                    wh_Translate32(magic, src->dma32.server_sz);
+            } break;
+            case WH_MESSAGE_CUSTOM_TYPE_DMA64: {
+                dst->dma64.client_addr =
+                    wh_Translate64(magic, src->dma64.client_addr);
+                dst->dma64.client_sz =
+                    wh_Translate64(magic, src->dma64.client_sz);
+                dst->dma64.server_addr =
+                    wh_Translate64(magic, src->dma64.server_addr);
+                dst->dma64.server_sz =
+                    wh_Translate64(magic, src->dma64.server_sz);
+            } break;
+            default: {
+                /* reserved message types - no translation for now */
+            } break;
+        }
+    }
+    else {
+        memcpy(dst->buffer.data, src->buffer.data,
+               sizeof(dst->buffer.data));
+    }
+}
+
+
 int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
                                       const whMessageCustom_Request* src,
                                       whMessageCustom_Request*       dst)
@@ -12,12 +51,14 @@ int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
     if ((src == NULL) || (dst == NULL)) {
         return WH_ERROR_BADARGS;
     }
-
-    dst->id = wh_Translate16(magic, src->id);
+    dst->id   = wh_Translate16(magic, src->id);
+    dst->type = wh_Translate32(magic, src->type);
+    _translateCustomData(magic, dst->type, &src->data, &dst->data);
 
     return WH_ERROR_OK;
 }
 
+
 int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
                                        const whMessageCustom_Response* src,
                                        whMessageCustom_Response*       dst)
@@ -29,5 +70,11 @@ int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
     dst->rc  = wh_Translate32(magic, src->rc);
     dst->err = wh_Translate32(magic, src->err);
 
+    /* TODO: should we continue to translate responses for err != 0?
+     * Probably still should...*/
+    dst->id   = wh_Translate16(magic, src->id);
+    dst->type = wh_Translate32(magic, src->type);
+    _translateCustomData(magic, dst->type, &src->data, &dst->data);
+
     return WH_ERROR_OK;
 }
\ No newline at end of file
diff --git a/test/wh_config.h b/test/wh_config.h
index 0bff3b6b..8f76b552 100644
--- a/test/wh_config.h
+++ b/test/wh_config.h
@@ -2,6 +2,6 @@
 #define WH_CONFIG_H_
 
 #define WH_CFG_TEST_VERBOSE
-//#define WH_CFG_TEST_POSIX
+#define WH_CFG_TEST_POSIX
 
 #endif /* WH_CONFIG_H_*/
diff --git a/test/wh_test_clientserver.c b/test/wh_test_clientserver.c
index 3520eab0..87afcc20 100644
--- a/test/wh_test_clientserver.c
+++ b/test/wh_test_clientserver.c
@@ -34,14 +34,83 @@
 #define ONE_MS 1000
 
 
-/* Dummy callback that returns the registered callback ID */
+/* Dummy callback that copies server data to client data */
 static int _customServerCb(whServerContext*               server,
                            const whMessageCustom_Request* req,
                            whMessageCustom_Response*      resp)
 {
+    uint8_t* serverPtr = NULL;
+    uint8_t* clientPtr = NULL;
+    size_t copySz = 0;
+
+    if (req->type == WH_MESSAGE_CUSTOM_TYPE_DMA64) {
+        clientPtr = (uint8_t*)((uintptr_t)req->data.dma64.client_addr);
+        serverPtr = (uint8_t*)((uintptr_t)req->data.dma64.server_addr);
+        resp->data.dma64.client_sz = req->data.dma64.server_sz;
+        copySz = req->data.dma64.server_sz;
+    }
+    else if (req->type == WH_MESSAGE_CUSTOM_TYPE_DMA32) {
+        clientPtr = (uint8_t*)((uintptr_t)req->data.dma32.client_addr);
+        serverPtr = (uint8_t*)((uintptr_t)req->data.dma32.server_addr);
+        resp->data.dma32.client_sz = req->data.dma32.server_sz;
+        copySz = req->data.dma32.server_sz;
+    }
+
+    memcpy(clientPtr, serverPtr, copySz);
+
     return req->id;
 }
 
+/* Helper function to test client server callbacks. Client and server must be
+ * already initialized */
+static int _testCallbacks(whServerContext* server, whClientContext* client)
+{
+    size_t                   counter;
+    whMessageCustom_Request  req  = {0};
+    whMessageCustom_Response resp = {0};
+
+    const char input[] = "The answer to the ultimate question of life, the "
+                         "universe and everything is 42";
+    char       output[sizeof(input)] = {0};
+
+    for (counter = 0; counter < WH_MESSAGE_ACTION_MAX; counter++) {
+        WH_TEST_RETURN_ON_FAIL(
+            wh_Server_RegisterCustomCb(counter, _customServerCb));
+
+        /* prepare the request */
+        req.id = counter;
+        /* 64-bit host system */
+        if (sizeof(uintptr_t) == sizeof(uint64_t)) {
+            req.type                   = WH_MESSAGE_CUSTOM_TYPE_DMA64;
+            req.data.dma64.server_addr = (uint64_t)((uintptr_t)input);
+            req.data.dma64.server_sz   = sizeof(input);
+            req.data.dma64.client_addr = (uint64_t)((uintptr_t)output);
+            req.data.dma64.client_sz   = 0;
+        }
+        /* 32-bit host system */
+        else if (sizeof(uintptr_t) == sizeof(uint32_t)) {
+            req.type                   = WH_MESSAGE_CUSTOM_TYPE_DMA32;
+            req.data.dma32.server_addr = (uint32_t)((uintptr_t)&input);
+            req.data.dma32.server_sz   = sizeof(input);
+            req.data.dma32.client_addr = (uint32_t)((uintptr_t)output);
+            req.data.dma32.client_sz   = 0;
+        }
+
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, &req));
+        WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
+        WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_OK);
+        WH_TEST_ASSERT_RETURN(resp.rc == counter);
+        WH_TEST_ASSERT_RETURN(0 == memcmp(output, input, sizeof(input)));
+
+        memset(output, 0, sizeof(output));
+        memset(&req, 0, sizeof(req));
+        memset(&resp, 0, sizeof(resp));
+    }
+
+    return WH_ERROR_OK;
+}
+
 
 int whTest_ClientServerSequential(void)
 {
@@ -443,15 +512,7 @@ int whTest_ClientServerSequential(void)
     #endif
 
     /* Test custom registered callbacks */
-    for (counter = 0; counter < WH_MESSAGE_ACTION_MAX; counter++) {
-        int rc = 0;
-        WH_TEST_RETURN_ON_FAIL(wh_Server_RegisterCustomCb(counter, _customServerCb));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, counter));
-        WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &rc));
-        WH_TEST_ASSERT_RETURN(rc == counter);
-    }
-
+    WH_TEST_RETURN_ON_FAIL(_testCallbacks(server, client));
 
     WH_TEST_RETURN_ON_FAIL(wh_Server_Cleanup(server));
     WH_TEST_RETURN_ON_FAIL(wh_Client_Cleanup(client));
diff --git a/wolfhsm/wh_client.h b/wolfhsm/wh_client.h
index a0f08848..c524f0ed 100644
--- a/wolfhsm/wh_client.h
+++ b/wolfhsm/wh_client.h
@@ -26,6 +26,7 @@
 
 /* Component includes */
 #include "wolfhsm/wh_comm.h"
+#include "wolfhsm/wh_message_custom.h"
 
 /* Client context */
 struct whClientContext_t {
@@ -179,7 +180,7 @@ int wh_Client_NvmReadDma(whClientContext* c,
 
 
 /* Client custom-callback support */
-int wh_Client_CustomRequest(whClientContext* c, uint16_t action); /* TODO: user data */
-int wh_Client_CustomResponse(whClientContext* c, int32_t *out_rc);
+int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req);
+int wh_Client_CustomResponse(whClientContext* c, whMessageCustom_Response *resp);
 
 #endif /* WOLFHSM_WH_CLIENT_H_ */
diff --git a/wolfhsm/wh_message_custom.h b/wolfhsm/wh_message_custom.h
index 71399af7..365aa82a 100644
--- a/wolfhsm/wh_message_custom.h
+++ b/wolfhsm/wh_message_custom.h
@@ -3,21 +3,75 @@
 
 #include <stdint.h>
 
+#define WH_MESSAGE_CUSTOM_BUF_SIZE (256)
+
+/* Type indicator for custom request/response messages. Indicates how
+ * to interpret whMessageCustomData */
+typedef enum {
+    /* message types reserved for internal usage*/
+    WH_MESSAGE_CUSTOM_TYPE_DMA32      = 0,
+    WH_MESSAGE_CUSTOM_TYPE_DMA64      = 1,
+    WH_MESSAGE_CUSTOM_TYPE_RESERVED_2 = 2,
+    WH_MESSAGE_CUSTOM_TYPE_RESERVED_3 = 3,
+    WH_MESSAGE_CUSTOM_TYPE_RESERVED_4 = 4,
+    WH_MESSAGE_CUSTOM_TYPE_RESERVED_5 = 5,
+    WH_MESSAGE_CUSTOM_TYPE_RESERVED_6 = 6,
+    WH_MESSAGE_CUSTOM_TYPE_RESERVED_7 = 7,
+    /* User-defined types start from here, up to UINT32_MAX */
+    WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START = 8,
+} whMessageCustom_Type;
+
+
+/* union providing some helpful abstractions for passing pointers in/out of
+ * custom callbacks on top of a raw data buffer */
+typedef union {
+    /* pointer/size pairs for 32-bit systems */
+    struct {
+        uint32_t client_addr;
+        uint32_t client_sz;
+        uint32_t server_addr;
+        uint32_t server_sz;
+    } dma32;
+    /* pointer/size pairs for 64-bit systems */
+    struct {
+        uint64_t client_addr;
+        uint64_t client_sz;
+        uint64_t server_addr;
+        uint64_t server_sz;
+    } dma64;
+    /* raw data buffer for user-defined schema */
+    struct {
+        uint8_t data[WH_MESSAGE_CUSTOM_BUF_SIZE];
+    } buffer;
+} whMessageCustom_Data;
+
+/* request message to the custom server callback */
 typedef struct {
-    uint16_t id;
-    /* TODO: pass I/O pointers and sizes? Mimic NVM DMA request? */
+    uint16_t             id;   /* indentifier of registered callback  */
+    uint32_t             type; /* whMessageCustom_Type */
+    whMessageCustom_Data data;
 } whMessageCustom_Request;
 
+/* response message from the custom server callback */
 typedef struct {
-    int32_t rc;  /* Return code from custom callback */
-    int32_t err; /* wolfHSM-specific error */
+    uint16_t id;   /* indentifier of registered callback  */
+    uint32_t type; /* whMessageCustom_Type */
+    int32_t  rc;   /* Return code from custom callback. Invalid if err != 0 */
+    int32_t  err;  /* wolfHSM-specific error. If err != 0, rc is invalid */
+    whMessageCustom_Data data;
 } whMessageCustom_Response;
 
 
+/* Translates a custom request message. The whMessageCustom_Request.data field
+ * will not be translated for whMessageCustom_Request.type values greater than
+ * WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START */
 int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
                                       const whMessageCustom_Request* src,
                                       whMessageCustom_Request*       dst);
 
+/* Translates a custom response message. The whMessageCustom_Request.data field
+ * will not be translated for whMessageCustom_Request.type values greater than
+ * WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START */
 int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
                                        const whMessageCustom_Response* src,
                                        whMessageCustom_Response*       dst);

From a5055f335a4676e71e1add09e6c154dec133f782 Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Tue, 2 Apr 2024 13:57:56 -0600
Subject: [PATCH 3/6] added some pointer arg checks, added negative test
 ensuring that requesting an unregistered callbacks return error

---
 src/wh_server_custom.c      |  6 ++++++
 test/wh_test_clientserver.c | 17 +++++++++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/wh_server_custom.c b/src/wh_server_custom.c
index 3e2806d6..d141390a 100644
--- a/src/wh_server_custom.c
+++ b/src/wh_server_custom.c
@@ -29,9 +29,14 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
     whMessageCustom_Request  req  = {0};
     whMessageCustom_Response resp = {0};
 
+    if (NULL == server || NULL == req_packet || NULL == resp_packet ||
+        out_resp_size == NULL) {
+        return WH_ERROR_BADARGS;
+    }
 
     if (action >= WH_MESSAGE_ACTION_MAX) {
         /* Invalid callback index  */
+        /* TODO: is this the appropriate error to return? */
         return WH_ERROR_BADARGS;
     }
 
@@ -66,6 +71,7 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
         WH_ERROR_OK) {
         return rc;
     }
+
     *out_resp_size = sizeof(resp);
 
     return WH_ERROR_OK;
diff --git a/test/wh_test_clientserver.c b/test/wh_test_clientserver.c
index 87afcc20..7bd9ecff 100644
--- a/test/wh_test_clientserver.c
+++ b/test/wh_test_clientserver.c
@@ -74,21 +74,28 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
     char       output[sizeof(input)] = {0};
 
     for (counter = 0; counter < WH_MESSAGE_ACTION_MAX; counter++) {
+        /* First, test that an unregistered callback returns error */
+        req.id = counter;
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, &req));
+        WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
+        WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_NO_HANDLER);
+
+        /* Register a custom callback */
         WH_TEST_RETURN_ON_FAIL(
             wh_Server_RegisterCustomCb(counter, _customServerCb));
 
-        /* prepare the request */
-        req.id = counter;
-        /* 64-bit host system */
+        /* prepare the rest of the request */
         if (sizeof(uintptr_t) == sizeof(uint64_t)) {
+            /* 64-bit host system */
             req.type                   = WH_MESSAGE_CUSTOM_TYPE_DMA64;
             req.data.dma64.server_addr = (uint64_t)((uintptr_t)input);
             req.data.dma64.server_sz   = sizeof(input);
             req.data.dma64.client_addr = (uint64_t)((uintptr_t)output);
             req.data.dma64.client_sz   = 0;
         }
-        /* 32-bit host system */
         else if (sizeof(uintptr_t) == sizeof(uint32_t)) {
+            /* 32-bit host system */
             req.type                   = WH_MESSAGE_CUSTOM_TYPE_DMA32;
             req.data.dma32.server_addr = (uint32_t)((uintptr_t)&input);
             req.data.dma32.server_sz   = sizeof(input);
@@ -108,6 +115,8 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
         memset(&resp, 0, sizeof(resp));
     }
 
+    /* TODO: query registered callback list and ensure it is populated correctly */
+
     return WH_ERROR_OK;
 }
 

From 12c810eee75adf4083b0dc8b71d1d7de2edfa94c Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Tue, 2 Apr 2024 14:34:17 -0600
Subject: [PATCH 4/6] added client custom callback query capability

---
 src/wh_client.c             | 15 +++++++++++++++
 src/wh_message_custom.c     |  4 ++++
 src/wh_server_custom.c      | 15 ++++++++-------
 test/wh_test_clientserver.c | 17 ++++++++++++++---
 wolfhsm/wh_client.h         |  1 +
 wolfhsm/wh_message_custom.h |  6 +++---
 6 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/src/wh_client.c b/src/wh_client.c
index 21b054c6..9fa71d22 100644
--- a/src/wh_client.c
+++ b/src/wh_client.c
@@ -240,4 +240,19 @@ int wh_Client_CustomResponse(whClientContext*          c,
     memcpy(outResp, &resp, sizeof(resp));
 
     return WH_ERROR_OK;
+}
+
+int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id)
+{
+    whMessageCustom_Request req = {0};
+
+    if (c == NULL || id >= WH_MESSAGE_ACTION_MAX) {
+        return WH_ERROR_BADARGS;
+    }
+
+    req.id = id;
+    req.type = WH_MESSAGE_CUSTOM_TYPE_QUERY;
+
+    return wh_Client_SendRequest(c, WH_MESSAGE_GROUP_CUSTOM, req.id,
+                                 sizeof(req), &req);
 }
\ No newline at end of file
diff --git a/src/wh_message_custom.c b/src/wh_message_custom.c
index 4591adfd..8c6a2870 100644
--- a/src/wh_message_custom.c
+++ b/src/wh_message_custom.c
@@ -12,6 +12,9 @@ static void _translateCustomData(uint16_t magic, uint32_t translatedType,
 {
     if (translatedType < WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START) {
         switch (translatedType) {
+            case WH_MESSAGE_CUSTOM_TYPE_QUERY: {
+                /* right now, no further translations required */
+            } break;
             case WH_MESSAGE_CUSTOM_TYPE_DMA32: {
                 dst->dma32.client_addr =
                     wh_Translate32(magic, src->dma32.client_addr);
@@ -51,6 +54,7 @@ int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
     if ((src == NULL) || (dst == NULL)) {
         return WH_ERROR_BADARGS;
     }
+
     dst->id   = wh_Translate16(magic, src->id);
     dst->type = wh_Translate32(magic, src->type);
     _translateCustomData(magic, dst->type, &src->data, &dst->data);
diff --git a/src/wh_server_custom.c b/src/wh_server_custom.c
index d141390a..54cbe9cd 100644
--- a/src/wh_server_custom.c
+++ b/src/wh_server_custom.c
@@ -52,17 +52,18 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
     }
 
     if (customHandlerTable[action] != NULL) {
-        /* Invoke the registered callback, storing the return value in the
-         * reponse  */
-        resp.rc = customHandlerTable[action](server, &req, &resp);
-        /* TODO: propagate wolfHSM error codes (requires modifiying caller
-         * function)*/
+        /* If this isn't a query to check if the callback exists, invoke the
+         * registered callback, storing the return value in the reponse  */
+        if (req.type != WH_MESSAGE_CUSTOM_TYPE_QUERY) {
+            resp.rc = customHandlerTable[action](server, &req, &resp);
+        }
+        /* TODO: propagate other wolfHSM error codes (requires modifiying caller
+         * function) once generic server code supports it */
         resp.err = WH_ERROR_OK;
     }
     else {
-        /* No callback was registered, populate response errors, but we must
+        /* No callback was registered, populate response error. We must
          * return success to ensure the "error" response is sent  */
-        /* TODO: what should we set resp.rc to? */
         resp.err = WH_ERROR_NO_HANDLER;
     }
 
diff --git a/test/wh_test_clientserver.c b/test/wh_test_clientserver.c
index 7bd9ecff..8d3904d9 100644
--- a/test/wh_test_clientserver.c
+++ b/test/wh_test_clientserver.c
@@ -74,8 +74,15 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
     char       output[sizeof(input)] = {0};
 
     for (counter = 0; counter < WH_MESSAGE_ACTION_MAX; counter++) {
-        /* First, test that an unregistered callback returns error */
         req.id = counter;
+
+        /* Check that the callback shows as unregistered */
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequestCheckRegistered(client, req.id));
+        WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
+        WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_NO_HANDLER);
+
+        /* Test that calling an unregistered callback returns error */
         WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, &req));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
         WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
@@ -85,6 +92,12 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
         WH_TEST_RETURN_ON_FAIL(
             wh_Server_RegisterCustomCb(counter, _customServerCb));
 
+        /* Check that the callback now shows as registered */
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequestCheckRegistered(client, req.id));
+        WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
+        WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_OK);
+
         /* prepare the rest of the request */
         if (sizeof(uintptr_t) == sizeof(uint64_t)) {
             /* 64-bit host system */
@@ -115,8 +128,6 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
         memset(&resp, 0, sizeof(resp));
     }
 
-    /* TODO: query registered callback list and ensure it is populated correctly */
-
     return WH_ERROR_OK;
 }
 
diff --git a/wolfhsm/wh_client.h b/wolfhsm/wh_client.h
index c524f0ed..a8430a7b 100644
--- a/wolfhsm/wh_client.h
+++ b/wolfhsm/wh_client.h
@@ -180,6 +180,7 @@ int wh_Client_NvmReadDma(whClientContext* c,
 
 
 /* Client custom-callback support */
+int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id);
 int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req);
 int wh_Client_CustomResponse(whClientContext* c, whMessageCustom_Response *resp);
 
diff --git a/wolfhsm/wh_message_custom.h b/wolfhsm/wh_message_custom.h
index 365aa82a..b2eb498a 100644
--- a/wolfhsm/wh_message_custom.h
+++ b/wolfhsm/wh_message_custom.h
@@ -9,9 +9,9 @@
  * to interpret whMessageCustomData */
 typedef enum {
     /* message types reserved for internal usage*/
-    WH_MESSAGE_CUSTOM_TYPE_DMA32      = 0,
-    WH_MESSAGE_CUSTOM_TYPE_DMA64      = 1,
-    WH_MESSAGE_CUSTOM_TYPE_RESERVED_2 = 2,
+    WH_MESSAGE_CUSTOM_TYPE_QUERY      = 0,
+    WH_MESSAGE_CUSTOM_TYPE_DMA32      = 1,
+    WH_MESSAGE_CUSTOM_TYPE_DMA64      = 2,
     WH_MESSAGE_CUSTOM_TYPE_RESERVED_3 = 3,
     WH_MESSAGE_CUSTOM_TYPE_RESERVED_4 = 4,
     WH_MESSAGE_CUSTOM_TYPE_RESERVED_5 = 5,

From 590975893f8d03cd0cb1f96397e8eee7e39eaf51 Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Wed, 3 Apr 2024 16:55:25 -0600
Subject: [PATCH 5/6] Addressed review comments: - made custom callback table
 part of server context - reduced number of callbacks to 8 - deleted
 wh_server_custom.h and moved contents to wh_server.h - Added
 wh_Client_CustomResponseCheckRegistered() helper

---
 src/wh_client.c             | 37 ++++++++++++++++++++++++++++++++++---
 src/wh_message_custom.c     |  3 ++-
 src/wh_server.c             |  1 -
 src/wh_server_custom.c      | 25 +++++++++++++------------
 test/wh_test_clientserver.c | 19 +++++++++++--------
 wolfhsm/wh_client.h         |  8 +++++++-
 wolfhsm/wh_common.h         |  3 +++
 wolfhsm/wh_server.h         | 24 ++++++++++++++++++++++++
 wolfhsm/wh_server_custom.h  | 19 -------------------
 9 files changed, 94 insertions(+), 45 deletions(-)
 delete mode 100644 wolfhsm/wh_server_custom.h

diff --git a/src/wh_client.c b/src/wh_client.c
index 9fa71d22..1ba38a36 100644
--- a/src/wh_client.c
+++ b/src/wh_client.c
@@ -204,7 +204,7 @@ int wh_Client_Echo(whClientContext* c, uint16_t snd_len, const void* snd_data,
 
 int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req)
 {
-    if (NULL == c || req == NULL || req->id >= WH_MESSAGE_ACTION_MAX) {
+    if (NULL == c || req == NULL || req->id >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
         return WH_ERROR_BADARGS;
     }
 
@@ -232,7 +232,7 @@ int wh_Client_CustomResponse(whClientContext*          c,
     }
 
     if (resp_size != sizeof(resp) || resp_group != WH_MESSAGE_GROUP_CUSTOM ||
-        resp_action >= WH_MESSAGE_ACTION_MAX) {
+        resp_action >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
         /* message invalid */
         return WH_ERROR_ABORTED;
     }
@@ -246,7 +246,7 @@ int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id)
 {
     whMessageCustom_Request req = {0};
 
-    if (c == NULL || id >= WH_MESSAGE_ACTION_MAX) {
+    if (c == NULL || id >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
         return WH_ERROR_BADARGS;
     }
 
@@ -255,4 +255,35 @@ int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id)
 
     return wh_Client_SendRequest(c, WH_MESSAGE_GROUP_CUSTOM, req.id,
                                  sizeof(req), &req);
+}
+
+
+int wh_Client_CustomResponseCheckRegistered(whClientContext* c, uint16_t* outId, int* responseError)
+{
+    int rc = 0;
+    whMessageCustom_Response resp = {0};
+
+    if (c == NULL || outId == NULL || responseError == NULL) {
+        return WH_ERROR_BADARGS;
+    }
+
+    rc = wh_Client_CustomResponse(c, &resp);
+    if (rc != WH_ERROR_OK) {
+        return rc;
+    }
+
+    if (resp.type != WH_MESSAGE_CUSTOM_TYPE_QUERY) {
+        /* message invalid */
+        return WH_ERROR_ABORTED;
+    }
+
+    if (resp.err != WH_ERROR_OK && resp.err != WH_ERROR_NO_HANDLER) {
+        /* error codes that aren't related to the query should be fatal */
+        return WH_ERROR_ABORTED;
+    }
+
+    *outId = resp.id;
+    *responseError = resp.err;
+
+    return WH_ERROR_OK;
 }
\ No newline at end of file
diff --git a/src/wh_message_custom.c b/src/wh_message_custom.c
index 8c6a2870..b8ed5963 100644
--- a/src/wh_message_custom.c
+++ b/src/wh_message_custom.c
@@ -41,7 +41,8 @@ static void _translateCustomData(uint16_t magic, uint32_t translatedType,
         }
     }
     else {
-        memcpy(dst->buffer.data, src->buffer.data,
+        /* use memmove in case data is translated "in place" */
+        memmove(dst->buffer.data, src->buffer.data,
                sizeof(dst->buffer.data));
     }
 }
diff --git a/src/wh_server.c b/src/wh_server.c
index ed642456..d302df2e 100644
--- a/src/wh_server.c
+++ b/src/wh_server.c
@@ -20,7 +20,6 @@
 #include "wolfhsm/wh_server.h"
 #include "wolfhsm/wh_server_internal.h"
 #include "wolfhsm/wh_server_crypto.h"
-#include "wolfhsm/wh_server_custom.h"
 #if defined(WOLFHSM_SHE_EXTENSION)
 #include "wolfhsm/wh_server_she.h"
 #endif
diff --git a/src/wh_server_custom.c b/src/wh_server_custom.c
index 54cbe9cd..c588a298 100644
--- a/src/wh_server_custom.c
+++ b/src/wh_server_custom.c
@@ -1,20 +1,18 @@
-#include "wolfhsm/wh_server_custom.h"
+#include "wolfhsm/wh_server.h"
+#include "wolfhsm/wh_common.h"
 #include "wolfhsm/wh_error.h"
 #include "wolfhsm/wh_message.h"
 #include "wolfhsm/wh_message_custom.h"
 
 
-/* Statically allocated callback table holding user callbacks */
-static whServerCustomCb customHandlerTable[WH_MESSAGE_ACTION_MAX] = {NULL};
-
-
-int wh_Server_RegisterCustomCb(uint16_t action, whServerCustomCb handler)
+int wh_Server_RegisterCustomCb(whServerContext* server, uint16_t action, whServerCustomCb handler)
 {
-    if (NULL == handler || action >= WH_MESSAGE_ACTION_MAX) {
+    if (NULL == server || NULL == handler ||
+        action >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
         return WH_ERROR_BADARGS;
     }
 
-    customHandlerTable[action] = handler;
+    server->customHandlerTable[action] = handler;
 
     return WH_ERROR_OK;
 }
@@ -34,7 +32,7 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
         return WH_ERROR_BADARGS;
     }
 
-    if (action >= WH_MESSAGE_ACTION_MAX) {
+    if (action >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
         /* Invalid callback index  */
         /* TODO: is this the appropriate error to return? */
         return WH_ERROR_BADARGS;
@@ -51,11 +49,11 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
         return rc;
     }
 
-    if (customHandlerTable[action] != NULL) {
+    if (server->customHandlerTable[action] != NULL) {
         /* If this isn't a query to check if the callback exists, invoke the
          * registered callback, storing the return value in the reponse  */
         if (req.type != WH_MESSAGE_CUSTOM_TYPE_QUERY) {
-            resp.rc = customHandlerTable[action](server, &req, &resp);
+            resp.rc = server->customHandlerTable[action](server, &req, &resp);
         }
         /* TODO: propagate other wolfHSM error codes (requires modifiying caller
          * function) once generic server code supports it */
@@ -67,7 +65,10 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
         resp.err = WH_ERROR_NO_HANDLER;
     }
 
-    /* Translate the response and set output size */
+    /* tag response with requested callback ID for client-side bookkeeping*/
+    resp.id = req.id;
+
+    /* Translate the response */
     if ((rc = wh_MessageCustom_TranslateResponse(magic, &resp, resp_packet)) !=
         WH_ERROR_OK) {
         return rc;
diff --git a/test/wh_test_clientserver.c b/test/wh_test_clientserver.c
index 8d3904d9..83d37d32 100644
--- a/test/wh_test_clientserver.c
+++ b/test/wh_test_clientserver.c
@@ -17,7 +17,6 @@
 #include "wolfhsm/wh_flash_ramsim.h"
 
 #include "wolfhsm/wh_server.h"
-#include "wolfhsm/wh_server_custom.h"
 #include "wolfhsm/wh_message.h"
 #include "wolfhsm/wh_client.h"
 
@@ -34,7 +33,7 @@
 #define ONE_MS 1000
 
 
-/* Dummy callback that copies server data to client data */
+/* Dummy callback that loopback-copies client data */
 static int _customServerCb(whServerContext*               server,
                            const whMessageCustom_Request* req,
                            whMessageCustom_Response*      resp)
@@ -68,19 +67,22 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
     size_t                   counter;
     whMessageCustom_Request  req  = {0};
     whMessageCustom_Response resp = {0};
+    uint16_t                 outId   = 0;
+    int                      respErr = 0;
 
     const char input[] = "The answer to the ultimate question of life, the "
                          "universe and everything is 42";
     char       output[sizeof(input)] = {0};
 
-    for (counter = 0; counter < WH_MESSAGE_ACTION_MAX; counter++) {
+    for (counter = 0; counter < WH_CUSTOM_RQST_NUM_CALLBACKS; counter++) {
         req.id = counter;
 
         /* Check that the callback shows as unregistered */
         WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequestCheckRegistered(client, req.id));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
-        WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_NO_HANDLER);
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponseCheckRegistered(client, &outId, &respErr));
+        WH_TEST_ASSERT_RETURN(outId == req.id);
+        WH_TEST_ASSERT_RETURN(respErr == WH_ERROR_NO_HANDLER);
 
         /* Test that calling an unregistered callback returns error */
         WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, &req));
@@ -90,13 +92,14 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
 
         /* Register a custom callback */
         WH_TEST_RETURN_ON_FAIL(
-            wh_Server_RegisterCustomCb(counter, _customServerCb));
+            wh_Server_RegisterCustomCb(server, counter, _customServerCb));
 
         /* Check that the callback now shows as registered */
         WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequestCheckRegistered(client, req.id));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
-        WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_OK);
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponseCheckRegistered(client, &outId, &respErr));
+        WH_TEST_ASSERT_RETURN(outId == req.id);
+        WH_TEST_ASSERT_RETURN(respErr == WH_ERROR_OK);
 
         /* prepare the rest of the request */
         if (sizeof(uintptr_t) == sizeof(uint64_t)) {
diff --git a/wolfhsm/wh_client.h b/wolfhsm/wh_client.h
index a8430a7b..b19d176a 100644
--- a/wolfhsm/wh_client.h
+++ b/wolfhsm/wh_client.h
@@ -180,8 +180,14 @@ int wh_Client_NvmReadDma(whClientContext* c,
 
 
 /* Client custom-callback support */
-int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id);
 int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req);
 int wh_Client_CustomResponse(whClientContext* c, whMessageCustom_Response *resp);
+/* Instructs server to query if a callback is registered */
+int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id);
+/* Processes a server response to callback query. OutId is set to the ID of the
+ * received query. ResponseError is set to WH_ERROR_OK if the callback is
+ * registered, and WH_ERROR_NO_HANDLER if not */
+int wh_Client_CustomResponseCheckRegistered(whClientContext* c, uint16_t* outId, int* responseError);
+
 
 #endif /* WOLFHSM_WH_CLIENT_H_ */
diff --git a/wolfhsm/wh_common.h b/wolfhsm/wh_common.h
index eda1adfd..cec5ab69 100644
--- a/wolfhsm/wh_common.h
+++ b/wolfhsm/wh_common.h
@@ -98,4 +98,7 @@ typedef struct {
 
 extern const whManifest_ex manifests[WOLFHSM_NUM_MANIFESTS];
 
+/* Custom request shared defs */
+#define WH_CUSTOM_RQST_NUM_CALLBACKS 8
+
 #endif /* WOLFHSM_WH_COMMON_H_ */
diff --git a/wolfhsm/wh_server.h b/wolfhsm/wh_server.h
index 226cc3ce..3f8656c0 100644
--- a/wolfhsm/wh_server.h
+++ b/wolfhsm/wh_server.h
@@ -12,6 +12,7 @@
 #include "wolfhsm/wh_common.h"
 #include "wolfhsm/wh_comm.h"
 #include "wolfhsm/wh_nvm.h"
+#include "wolfhsm/wh_message_custom.h"
 
 #include "wolfssl/wolfcrypt/settings.h"
 #include "wolfssl/wolfcrypt/random.h"
@@ -37,6 +38,17 @@ typedef struct {
     bool wcDevIdInitFlag: 1;
 } whServerFlags;
 
+/* Forward declaration of the server structure so its elements can reference
+ * itself  (e.g. server argument to custom callback) */
+struct whServerContext_t;
+
+/* Type definition for a custom server callback  */
+typedef int (*whServerCustomCb)(
+    struct whServerContext_t* server,   /* points to dispatching server ctx */
+    const whMessageCustom_Request* req, /* request from client to callback */
+    whMessageCustom_Response*      resp /* response from callback to client */
+);
+
 /* Context structure to maintain the state of an HSM server */
 typedef struct whServerContext_t {
     whServerFlags flags;
@@ -44,6 +56,7 @@ typedef struct whServerContext_t {
     whNvmContext nvm[1];
     crypto_context crypto[1];
     CacheSlot cache[WOLFHSM_NUM_RAMKEYS];
+    whServerCustomCb customHandlerTable[WH_CUSTOM_RQST_NUM_CALLBACKS];
 } whServerContext;
 
 typedef struct whServerConfig_t {
@@ -67,4 +80,15 @@ int wh_Server_HandleRequestMessage(whServerContext* server);
  */
 int wh_Server_Cleanup(whServerContext* server);
 
+/* Registers a custom callback with the server
+*/
+int wh_Server_RegisterCustomCb(whServerContext* server, uint16_t actionId, whServerCustomCb cb);
+
+/* Receive and handle an incoming custom callback request
+*/
+int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
+                                  uint16_t action, uint16_t seq,
+                                  uint16_t req_size, const void* req_packet,
+                                  uint16_t* out_resp_size, void* resp_packet);
+
 #endif /* WOLFHSM_WH_SERVER_H_ */
diff --git a/wolfhsm/wh_server_custom.h b/wolfhsm/wh_server_custom.h
deleted file mode 100644
index 4e2ddea1..00000000
--- a/wolfhsm/wh_server_custom.h
+++ /dev/null
@@ -1,19 +0,0 @@
-#ifndef WH_SERVER_CUSTOM_H_
-#define WH_SERVER_CUSTOM_H_
-
-#include "wolfhsm/wh_server.h"
-#include "wolfhsm/wh_message_custom.h"
-
-
-typedef int (*whServerCustomCb)(
-    whServerContext*               server, /* TODO: should this be const? */
-    const whMessageCustom_Request* req, whMessageCustom_Response* resp);
-
-int wh_Server_RegisterCustomCb(uint16_t actionId, whServerCustomCb cb);
-
-int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
-                                  uint16_t action, uint16_t seq,
-                                  uint16_t req_size, const void* req_packet,
-                                  uint16_t* out_resp_size, void* resp_packet);
-
-#endif /* WH_SERVER_CUSTOM_H_ */
\ No newline at end of file

From 102930731b69dec343eb2a69fcbed4933217966b Mon Sep 17 00:00:00 2001
From: Brett Nicholas <7547222+bigbrett@users.noreply.github.com>
Date: Wed, 3 Apr 2024 20:00:39 -0600
Subject: [PATCH 6/6] Review feedback: - Renamed custom structs and functions
 to customcb - renamed files - added blocking client helper - changed
 request/response id from uint16_t to uint32_t

---
 src/wh_client.c                               | 52 ++++++++----
 ...message_custom.c => wh_message_customcb.c} | 33 ++++----
 src/wh_server.c                               |  2 +-
 ...h_server_custom.c => wh_server_customcb.c} | 33 ++++----
 test/Makefile                                 |  4 +-
 test/wh_test_clientserver.c                   | 34 ++++----
 wolfhsm/wh_client.h                           | 10 +--
 wolfhsm/wh_common.h                           |  2 +-
 wolfhsm/wh_message_custom.h                   | 79 -------------------
 wolfhsm/wh_message_customcb.h                 | 79 +++++++++++++++++++
 wolfhsm/wh_server.h                           | 10 +--
 11 files changed, 180 insertions(+), 158 deletions(-)
 rename src/{wh_message_custom.c => wh_message_customcb.c} (66%)
 rename src/{wh_server_custom.c => wh_server_customcb.c} (64%)
 delete mode 100644 wolfhsm/wh_message_custom.h
 create mode 100644 wolfhsm/wh_message_customcb.h

diff --git a/src/wh_client.c b/src/wh_client.c
index 1ba38a36..9a87124d 100644
--- a/src/wh_client.c
+++ b/src/wh_client.c
@@ -23,7 +23,7 @@
 /* Message definitions */
 #include "wolfhsm/wh_message.h"
 #include "wolfhsm/wh_message_comm.h"
-#include "wolfhsm/wh_message_custom.h"
+#include "wolfhsm/wh_message_customcb.h"
 
 #include "wolfhsm/wh_client.h"
 
@@ -202,9 +202,9 @@ int wh_Client_Echo(whClientContext* c, uint16_t snd_len, const void* snd_data,
     return rc;
 }
 
-int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req)
+int wh_Client_CustomCbRequest(whClientContext* c, const whMessageCustomCb_Request* req)
 {
-    if (NULL == c || req == NULL || req->id >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
+    if (NULL == c || req == NULL || req->id >= WH_CUSTOM_CB_NUM_CALLBACKS) {
         return WH_ERROR_BADARGS;
     }
 
@@ -212,10 +212,10 @@ int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* r
                                  sizeof(*req), req);
 }
 
-int wh_Client_CustomResponse(whClientContext*          c,
-                             whMessageCustom_Response* outResp)
+int wh_Client_CustomCbResponse(whClientContext*          c,
+                             whMessageCustomCb_Response* outResp)
 {
-    whMessageCustom_Response resp;
+    whMessageCustomCb_Response resp;
     uint16_t                 resp_group  = 0;
     uint16_t                 resp_action = 0;
     uint16_t                 resp_size   = 0;
@@ -232,7 +232,7 @@ int wh_Client_CustomResponse(whClientContext*          c,
     }
 
     if (resp_size != sizeof(resp) || resp_group != WH_MESSAGE_GROUP_CUSTOM ||
-        resp_action >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
+        resp_action >= WH_CUSTOM_CB_NUM_CALLBACKS) {
         /* message invalid */
         return WH_ERROR_ABORTED;
     }
@@ -242,37 +242,37 @@ int wh_Client_CustomResponse(whClientContext*          c,
     return WH_ERROR_OK;
 }
 
-int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id)
+int wh_Client_CustomCheckRegisteredRequest(whClientContext* c, uint32_t id)
 {
-    whMessageCustom_Request req = {0};
+    whMessageCustomCb_Request req = {0};
 
-    if (c == NULL || id >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
+    if (c == NULL || id >= WH_CUSTOM_CB_NUM_CALLBACKS) {
         return WH_ERROR_BADARGS;
     }
 
     req.id = id;
-    req.type = WH_MESSAGE_CUSTOM_TYPE_QUERY;
+    req.type = WH_MESSAGE_CUSTOM_CB_TYPE_QUERY;
 
     return wh_Client_SendRequest(c, WH_MESSAGE_GROUP_CUSTOM, req.id,
                                  sizeof(req), &req);
 }
 
 
-int wh_Client_CustomResponseCheckRegistered(whClientContext* c, uint16_t* outId, int* responseError)
+int wh_Client_CustomCbCheckRegisteredResponse(whClientContext* c, uint16_t* outId, int* responseError)
 {
     int rc = 0;
-    whMessageCustom_Response resp = {0};
+    whMessageCustomCb_Response resp = {0};
 
     if (c == NULL || outId == NULL || responseError == NULL) {
         return WH_ERROR_BADARGS;
     }
 
-    rc = wh_Client_CustomResponse(c, &resp);
+    rc = wh_Client_CustomCbResponse(c, &resp);
     if (rc != WH_ERROR_OK) {
         return rc;
     }
 
-    if (resp.type != WH_MESSAGE_CUSTOM_TYPE_QUERY) {
+    if (resp.type != WH_MESSAGE_CUSTOM_CB_TYPE_QUERY) {
         /* message invalid */
         return WH_ERROR_ABORTED;
     }
@@ -286,4 +286,26 @@ int wh_Client_CustomResponseCheckRegistered(whClientContext* c, uint16_t* outId,
     *responseError = resp.err;
 
     return WH_ERROR_OK;
+}
+
+
+int wh_Client_CustomCbCheckRegistered(whClientContext* c, uint16_t id, int* responseError)
+{
+    int rc = 0;
+
+    if (NULL == c || NULL == responseError || id >= WH_CUSTOM_CB_NUM_CALLBACKS) {
+        return WH_ERROR_BADARGS;
+    }
+
+    do {
+        rc = wh_Client_CustomCheckRegisteredRequest(c, id);
+    } while (rc == WH_ERROR_NOTREADY);
+
+    if (rc == WH_ERROR_OK) {
+        do {
+            rc = wh_Client_CustomCbCheckRegisteredResponse(c, &id, responseError);
+        } while (rc == WH_ERROR_NOTREADY);
+    }
+
+    return rc;
 }
\ No newline at end of file
diff --git a/src/wh_message_custom.c b/src/wh_message_customcb.c
similarity index 66%
rename from src/wh_message_custom.c
rename to src/wh_message_customcb.c
index b8ed5963..1512d4be 100644
--- a/src/wh_message_custom.c
+++ b/src/wh_message_customcb.c
@@ -1,21 +1,21 @@
 #include <stddef.h>
 #include <string.h>
 
-#include "wolfhsm/wh_message_custom.h"
+#include "wolfhsm/wh_message_customcb.h"
 #include "wolfhsm/wh_error.h"
 #include "wolfhsm/wh_comm.h"
 
 
 static void _translateCustomData(uint16_t magic, uint32_t translatedType,
-                                 const whMessageCustom_Data* src,
-                                 whMessageCustom_Data*       dst)
+                                 const whMessageCustomCb_Data* src,
+                                 whMessageCustomCb_Data*       dst)
 {
-    if (translatedType < WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START) {
+    if (translatedType < WH_MESSAGE_CUSTOM_CB_TYPE_USER_DEFINED_START) {
         switch (translatedType) {
-            case WH_MESSAGE_CUSTOM_TYPE_QUERY: {
+            case WH_MESSAGE_CUSTOM_CB_TYPE_QUERY: {
                 /* right now, no further translations required */
             } break;
-            case WH_MESSAGE_CUSTOM_TYPE_DMA32: {
+            case WH_MESSAGE_CUSTOM_CB_TYPE_DMA32: {
                 dst->dma32.client_addr =
                     wh_Translate32(magic, src->dma32.client_addr);
                 dst->dma32.client_sz =
@@ -25,7 +25,7 @@ static void _translateCustomData(uint16_t magic, uint32_t translatedType,
                 dst->dma32.server_sz =
                     wh_Translate32(magic, src->dma32.server_sz);
             } break;
-            case WH_MESSAGE_CUSTOM_TYPE_DMA64: {
+            case WH_MESSAGE_CUSTOM_CB_TYPE_DMA64: {
                 dst->dma64.client_addr =
                     wh_Translate64(magic, src->dma64.client_addr);
                 dst->dma64.client_sz =
@@ -42,21 +42,20 @@ static void _translateCustomData(uint16_t magic, uint32_t translatedType,
     }
     else {
         /* use memmove in case data is translated "in place" */
-        memmove(dst->buffer.data, src->buffer.data,
-               sizeof(dst->buffer.data));
+        memmove(dst->buffer.data, src->buffer.data, sizeof(dst->buffer.data));
     }
 }
 
 
-int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
-                                      const whMessageCustom_Request* src,
-                                      whMessageCustom_Request*       dst)
+int wh_MessageCustomCb_TranslateRequest(uint16_t                         magic,
+                                        const whMessageCustomCb_Request* src,
+                                        whMessageCustomCb_Request*       dst)
 {
     if ((src == NULL) || (dst == NULL)) {
         return WH_ERROR_BADARGS;
     }
 
-    dst->id   = wh_Translate16(magic, src->id);
+    dst->id   = wh_Translate32(magic, src->id);
     dst->type = wh_Translate32(magic, src->type);
     _translateCustomData(magic, dst->type, &src->data, &dst->data);
 
@@ -64,9 +63,9 @@ int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
 }
 
 
-int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
-                                       const whMessageCustom_Response* src,
-                                       whMessageCustom_Response*       dst)
+int wh_MessageCustomCb_TranslateResponse(uint16_t magic,
+                                         const whMessageCustomCb_Response* src,
+                                         whMessageCustomCb_Response*       dst)
 {
     if ((src == NULL) || (dst == NULL)) {
         return WH_ERROR_BADARGS;
@@ -77,7 +76,7 @@ int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
 
     /* TODO: should we continue to translate responses for err != 0?
      * Probably still should...*/
-    dst->id   = wh_Translate16(magic, src->id);
+    dst->id   = wh_Translate32(magic, src->id);
     dst->type = wh_Translate32(magic, src->type);
     _translateCustomData(magic, dst->type, &src->data, &dst->data);
 
diff --git a/src/wh_server.c b/src/wh_server.c
index d302df2e..edd9ee19 100644
--- a/src/wh_server.c
+++ b/src/wh_server.c
@@ -265,7 +265,7 @@ int wh_Server_HandleRequestMessage(whServerContext* server)
 #endif
 
         case WH_MESSAGE_GROUP_CUSTOM:
-            rc = wh_Server_HandleCustomRequest(server, magic, action, seq,
+            rc = wh_Server_HandleCustomCbRequest(server, magic, action, seq,
                     size, data, &size, data);
             break;
 
diff --git a/src/wh_server_custom.c b/src/wh_server_customcb.c
similarity index 64%
rename from src/wh_server_custom.c
rename to src/wh_server_customcb.c
index c588a298..5685b7bb 100644
--- a/src/wh_server_custom.c
+++ b/src/wh_server_customcb.c
@@ -2,13 +2,14 @@
 #include "wolfhsm/wh_common.h"
 #include "wolfhsm/wh_error.h"
 #include "wolfhsm/wh_message.h"
-#include "wolfhsm/wh_message_custom.h"
+#include "wolfhsm/wh_message_customcb.h"
 
 
-int wh_Server_RegisterCustomCb(whServerContext* server, uint16_t action, whServerCustomCb handler)
+int wh_Server_RegisterCustomCb(whServerContext* server, uint16_t action,
+                               whServerCustomCb handler)
 {
     if (NULL == server || NULL == handler ||
-        action >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
+        action >= WH_CUSTOM_CB_NUM_CALLBACKS) {
         return WH_ERROR_BADARGS;
     }
 
@@ -18,33 +19,33 @@ int wh_Server_RegisterCustomCb(whServerContext* server, uint16_t action, whServe
 }
 
 
-int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
-                                  uint16_t action, uint16_t seq,
-                                  uint16_t req_size, const void* req_packet,
-                                  uint16_t* out_resp_size, void* resp_packet)
+int wh_Server_HandleCustomCbRequest(whServerContext* server, uint16_t magic,
+                                    uint16_t action, uint16_t seq,
+                                    uint16_t req_size, const void* req_packet,
+                                    uint16_t* out_resp_size, void* resp_packet)
 {
-    int                      rc   = 0;
-    whMessageCustom_Request  req  = {0};
-    whMessageCustom_Response resp = {0};
+    int                        rc   = 0;
+    whMessageCustomCb_Request  req  = {0};
+    whMessageCustomCb_Response resp = {0};
 
     if (NULL == server || NULL == req_packet || NULL == resp_packet ||
         out_resp_size == NULL) {
         return WH_ERROR_BADARGS;
     }
 
-    if (action >= WH_CUSTOM_RQST_NUM_CALLBACKS) {
+    if (action >= WH_CUSTOM_CB_NUM_CALLBACKS) {
         /* Invalid callback index  */
         /* TODO: is this the appropriate error to return? */
         return WH_ERROR_BADARGS;
     }
 
-    if (req_size != sizeof(whMessageCustom_Request)) {
+    if (req_size != sizeof(whMessageCustomCb_Request)) {
         /* Request is malformed */
         return WH_ERROR_ABORTED;
     }
 
     /* Translate the request */
-    if ((rc = wh_MessageCustom_TranslateRequest(magic, req_packet, &req)) !=
+    if ((rc = wh_MessageCustomCb_TranslateRequest(magic, req_packet, &req)) !=
         WH_ERROR_OK) {
         return rc;
     }
@@ -52,7 +53,7 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
     if (server->customHandlerTable[action] != NULL) {
         /* If this isn't a query to check if the callback exists, invoke the
          * registered callback, storing the return value in the reponse  */
-        if (req.type != WH_MESSAGE_CUSTOM_TYPE_QUERY) {
+        if (req.type != WH_MESSAGE_CUSTOM_CB_TYPE_QUERY) {
             resp.rc = server->customHandlerTable[action](server, &req, &resp);
         }
         /* TODO: propagate other wolfHSM error codes (requires modifiying caller
@@ -69,8 +70,8 @@ int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
     resp.id = req.id;
 
     /* Translate the response */
-    if ((rc = wh_MessageCustom_TranslateResponse(magic, &resp, resp_packet)) !=
-        WH_ERROR_OK) {
+    if ((rc = wh_MessageCustomCb_TranslateResponse(
+             magic, &resp, resp_packet)) != WH_ERROR_OK) {
         return rc;
     }
 
diff --git a/test/Makefile b/test/Makefile
index 6a45c258..489a6e52 100644
--- a/test/Makefile
+++ b/test/Makefile
@@ -80,13 +80,13 @@ SRC_C += \
             $(WOLFHSM_DIR)/src/wh_flash_unit.c \
             $(WOLFHSM_DIR)/src/wh_message_comm.c \
             $(WOLFHSM_DIR)/src/wh_message_nvm.c \
-            $(WOLFHSM_DIR)/src/wh_message_custom.c \
+            $(WOLFHSM_DIR)/src/wh_message_customcb.c \
             $(WOLFHSM_DIR)/src/wh_nvm.c \
             $(WOLFHSM_DIR)/src/wh_nvm_flash.c \
             $(WOLFHSM_DIR)/src/wh_server.c \
             $(WOLFHSM_DIR)/src/wh_server_nvm.c \
             $(WOLFHSM_DIR)/src/wh_server_crypto.c \
-            $(WOLFHSM_DIR)/src/wh_server_custom.c \
+            $(WOLFHSM_DIR)/src/wh_server_customcb.c \
             $(WOLFHSM_DIR)/src/wh_client_cryptocb.c \
             $(WOLFHSM_DIR)/src/wh_transport_mem.c \
             $(WOLFHSM_DIR)/src/wh_flash_ramsim.c
diff --git a/test/wh_test_clientserver.c b/test/wh_test_clientserver.c
index 83d37d32..560d1a46 100644
--- a/test/wh_test_clientserver.c
+++ b/test/wh_test_clientserver.c
@@ -35,20 +35,20 @@
 
 /* Dummy callback that loopback-copies client data */
 static int _customServerCb(whServerContext*               server,
-                           const whMessageCustom_Request* req,
-                           whMessageCustom_Response*      resp)
+                           const whMessageCustomCb_Request* req,
+                           whMessageCustomCb_Response*      resp)
 {
     uint8_t* serverPtr = NULL;
     uint8_t* clientPtr = NULL;
     size_t copySz = 0;
 
-    if (req->type == WH_MESSAGE_CUSTOM_TYPE_DMA64) {
+    if (req->type == WH_MESSAGE_CUSTOM_CB_TYPE_DMA64) {
         clientPtr = (uint8_t*)((uintptr_t)req->data.dma64.client_addr);
         serverPtr = (uint8_t*)((uintptr_t)req->data.dma64.server_addr);
         resp->data.dma64.client_sz = req->data.dma64.server_sz;
         copySz = req->data.dma64.server_sz;
     }
-    else if (req->type == WH_MESSAGE_CUSTOM_TYPE_DMA32) {
+    else if (req->type == WH_MESSAGE_CUSTOM_CB_TYPE_DMA32) {
         clientPtr = (uint8_t*)((uintptr_t)req->data.dma32.client_addr);
         serverPtr = (uint8_t*)((uintptr_t)req->data.dma32.server_addr);
         resp->data.dma32.client_sz = req->data.dma32.server_sz;
@@ -65,8 +65,8 @@ static int _customServerCb(whServerContext*               server,
 static int _testCallbacks(whServerContext* server, whClientContext* client)
 {
     size_t                   counter;
-    whMessageCustom_Request  req  = {0};
-    whMessageCustom_Response resp = {0};
+    whMessageCustomCb_Request  req  = {0};
+    whMessageCustomCb_Response resp = {0};
     uint16_t                 outId   = 0;
     int                      respErr = 0;
 
@@ -74,20 +74,20 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
                          "universe and everything is 42";
     char       output[sizeof(input)] = {0};
 
-    for (counter = 0; counter < WH_CUSTOM_RQST_NUM_CALLBACKS; counter++) {
+    for (counter = 0; counter < WH_CUSTOM_CB_NUM_CALLBACKS; counter++) {
         req.id = counter;
 
         /* Check that the callback shows as unregistered */
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequestCheckRegistered(client, req.id));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCheckRegisteredRequest(client, req.id));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponseCheckRegistered(client, &outId, &respErr));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCbCheckRegisteredResponse(client, &outId, &respErr));
         WH_TEST_ASSERT_RETURN(outId == req.id);
         WH_TEST_ASSERT_RETURN(respErr == WH_ERROR_NO_HANDLER);
 
         /* Test that calling an unregistered callback returns error */
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, &req));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCbRequest(client, &req));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCbResponse(client, &resp));
         WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_NO_HANDLER);
 
         /* Register a custom callback */
@@ -95,16 +95,16 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
             wh_Server_RegisterCustomCb(server, counter, _customServerCb));
 
         /* Check that the callback now shows as registered */
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequestCheckRegistered(client, req.id));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCheckRegisteredRequest(client, req.id));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponseCheckRegistered(client, &outId, &respErr));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCbCheckRegisteredResponse(client, &outId, &respErr));
         WH_TEST_ASSERT_RETURN(outId == req.id);
         WH_TEST_ASSERT_RETURN(respErr == WH_ERROR_OK);
 
         /* prepare the rest of the request */
         if (sizeof(uintptr_t) == sizeof(uint64_t)) {
             /* 64-bit host system */
-            req.type                   = WH_MESSAGE_CUSTOM_TYPE_DMA64;
+            req.type                   = WH_MESSAGE_CUSTOM_CB_TYPE_DMA64;
             req.data.dma64.server_addr = (uint64_t)((uintptr_t)input);
             req.data.dma64.server_sz   = sizeof(input);
             req.data.dma64.client_addr = (uint64_t)((uintptr_t)output);
@@ -112,16 +112,16 @@ static int _testCallbacks(whServerContext* server, whClientContext* client)
         }
         else if (sizeof(uintptr_t) == sizeof(uint32_t)) {
             /* 32-bit host system */
-            req.type                   = WH_MESSAGE_CUSTOM_TYPE_DMA32;
+            req.type                   = WH_MESSAGE_CUSTOM_CB_TYPE_DMA32;
             req.data.dma32.server_addr = (uint32_t)((uintptr_t)&input);
             req.data.dma32.server_sz   = sizeof(input);
             req.data.dma32.client_addr = (uint32_t)((uintptr_t)output);
             req.data.dma32.client_sz   = 0;
         }
 
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomRequest(client, &req));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCbRequest(client, &req));
         WH_TEST_RETURN_ON_FAIL(wh_Server_HandleRequestMessage(server));
-        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomResponse(client, &resp));
+        WH_TEST_RETURN_ON_FAIL(wh_Client_CustomCbResponse(client, &resp));
         WH_TEST_ASSERT_RETURN(resp.err == WH_ERROR_OK);
         WH_TEST_ASSERT_RETURN(resp.rc == counter);
         WH_TEST_ASSERT_RETURN(0 == memcmp(output, input, sizeof(input)));
diff --git a/wolfhsm/wh_client.h b/wolfhsm/wh_client.h
index b19d176a..cf7beedc 100644
--- a/wolfhsm/wh_client.h
+++ b/wolfhsm/wh_client.h
@@ -26,7 +26,7 @@
 
 /* Component includes */
 #include "wolfhsm/wh_comm.h"
-#include "wolfhsm/wh_message_custom.h"
+#include "wolfhsm/wh_message_customcb.h"
 
 /* Client context */
 struct whClientContext_t {
@@ -180,14 +180,14 @@ int wh_Client_NvmReadDma(whClientContext* c,
 
 
 /* Client custom-callback support */
-int wh_Client_CustomRequest(whClientContext* c, const whMessageCustom_Request* req);
-int wh_Client_CustomResponse(whClientContext* c, whMessageCustom_Response *resp);
+int wh_Client_CustomCbRequest(whClientContext* c, const whMessageCustomCb_Request* req);
+int wh_Client_CustomCbResponse(whClientContext* c, whMessageCustomCb_Response *resp);
 /* Instructs server to query if a callback is registered */
-int wh_Client_CustomRequestCheckRegistered(whClientContext* c, uint32_t id);
+int wh_Client_CustomCheckRegisteredRequest(whClientContext* c, uint32_t id);
 /* Processes a server response to callback query. OutId is set to the ID of the
  * received query. ResponseError is set to WH_ERROR_OK if the callback is
  * registered, and WH_ERROR_NO_HANDLER if not */
-int wh_Client_CustomResponseCheckRegistered(whClientContext* c, uint16_t* outId, int* responseError);
+int wh_Client_CustomCbCheckRegisteredResponse(whClientContext* c, uint16_t* outId, int* responseError);
 
 
 #endif /* WOLFHSM_WH_CLIENT_H_ */
diff --git a/wolfhsm/wh_common.h b/wolfhsm/wh_common.h
index cec5ab69..f023c26e 100644
--- a/wolfhsm/wh_common.h
+++ b/wolfhsm/wh_common.h
@@ -99,6 +99,6 @@ typedef struct {
 extern const whManifest_ex manifests[WOLFHSM_NUM_MANIFESTS];
 
 /* Custom request shared defs */
-#define WH_CUSTOM_RQST_NUM_CALLBACKS 8
+#define WH_CUSTOM_CB_NUM_CALLBACKS 8
 
 #endif /* WOLFHSM_WH_COMMON_H_ */
diff --git a/wolfhsm/wh_message_custom.h b/wolfhsm/wh_message_custom.h
deleted file mode 100644
index b2eb498a..00000000
--- a/wolfhsm/wh_message_custom.h
+++ /dev/null
@@ -1,79 +0,0 @@
-#ifndef WH_MESSAGE_CUSTOM_H_
-#define WH_MESSAGE_CUSTOM_H_
-
-#include <stdint.h>
-
-#define WH_MESSAGE_CUSTOM_BUF_SIZE (256)
-
-/* Type indicator for custom request/response messages. Indicates how
- * to interpret whMessageCustomData */
-typedef enum {
-    /* message types reserved for internal usage*/
-    WH_MESSAGE_CUSTOM_TYPE_QUERY      = 0,
-    WH_MESSAGE_CUSTOM_TYPE_DMA32      = 1,
-    WH_MESSAGE_CUSTOM_TYPE_DMA64      = 2,
-    WH_MESSAGE_CUSTOM_TYPE_RESERVED_3 = 3,
-    WH_MESSAGE_CUSTOM_TYPE_RESERVED_4 = 4,
-    WH_MESSAGE_CUSTOM_TYPE_RESERVED_5 = 5,
-    WH_MESSAGE_CUSTOM_TYPE_RESERVED_6 = 6,
-    WH_MESSAGE_CUSTOM_TYPE_RESERVED_7 = 7,
-    /* User-defined types start from here, up to UINT32_MAX */
-    WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START = 8,
-} whMessageCustom_Type;
-
-
-/* union providing some helpful abstractions for passing pointers in/out of
- * custom callbacks on top of a raw data buffer */
-typedef union {
-    /* pointer/size pairs for 32-bit systems */
-    struct {
-        uint32_t client_addr;
-        uint32_t client_sz;
-        uint32_t server_addr;
-        uint32_t server_sz;
-    } dma32;
-    /* pointer/size pairs for 64-bit systems */
-    struct {
-        uint64_t client_addr;
-        uint64_t client_sz;
-        uint64_t server_addr;
-        uint64_t server_sz;
-    } dma64;
-    /* raw data buffer for user-defined schema */
-    struct {
-        uint8_t data[WH_MESSAGE_CUSTOM_BUF_SIZE];
-    } buffer;
-} whMessageCustom_Data;
-
-/* request message to the custom server callback */
-typedef struct {
-    uint16_t             id;   /* indentifier of registered callback  */
-    uint32_t             type; /* whMessageCustom_Type */
-    whMessageCustom_Data data;
-} whMessageCustom_Request;
-
-/* response message from the custom server callback */
-typedef struct {
-    uint16_t id;   /* indentifier of registered callback  */
-    uint32_t type; /* whMessageCustom_Type */
-    int32_t  rc;   /* Return code from custom callback. Invalid if err != 0 */
-    int32_t  err;  /* wolfHSM-specific error. If err != 0, rc is invalid */
-    whMessageCustom_Data data;
-} whMessageCustom_Response;
-
-
-/* Translates a custom request message. The whMessageCustom_Request.data field
- * will not be translated for whMessageCustom_Request.type values greater than
- * WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START */
-int wh_MessageCustom_TranslateRequest(uint16_t                       magic,
-                                      const whMessageCustom_Request* src,
-                                      whMessageCustom_Request*       dst);
-
-/* Translates a custom response message. The whMessageCustom_Request.data field
- * will not be translated for whMessageCustom_Request.type values greater than
- * WH_MESSAGE_CUSTOM_TYPE_USER_DEFINED_START */
-int wh_MessageCustom_TranslateResponse(uint16_t                        magic,
-                                       const whMessageCustom_Response* src,
-                                       whMessageCustom_Response*       dst);
-
-#endif /* WH_MESSAGE_CUSTOM_H_*/
\ No newline at end of file
diff --git a/wolfhsm/wh_message_customcb.h b/wolfhsm/wh_message_customcb.h
new file mode 100644
index 00000000..c4b07213
--- /dev/null
+++ b/wolfhsm/wh_message_customcb.h
@@ -0,0 +1,79 @@
+#ifndef WH_MESSAGE_CUSTOM_CB_H_
+#define WH_MESSAGE_CUSTOM_CB_H_
+
+#include <stdint.h>
+
+#define WH_MESSAGE_CUSTOM_CB_BUF_SIZE (256)
+
+/* Type indicator for custom request/response messages. Indicates how
+ * to interpret whMessageCustomData */
+typedef enum {
+    /* message types reserved for internal usage*/
+    WH_MESSAGE_CUSTOM_CB_TYPE_QUERY      = 0,
+    WH_MESSAGE_CUSTOM_CB_TYPE_DMA32      = 1,
+    WH_MESSAGE_CUSTOM_CB_TYPE_DMA64      = 2,
+    WH_MESSAGE_CUSTOM_CB_TYPE_RESERVED_3 = 3,
+    WH_MESSAGE_CUSTOM_CB_TYPE_RESERVED_4 = 4,
+    WH_MESSAGE_CUSTOM_CB_TYPE_RESERVED_5 = 5,
+    WH_MESSAGE_CUSTOM_CB_TYPE_RESERVED_6 = 6,
+    WH_MESSAGE_CUSTOM_CB_TYPE_RESERVED_7 = 7,
+    /* User-defined types start from here, up to UINT32_MAX */
+    WH_MESSAGE_CUSTOM_CB_TYPE_USER_DEFINED_START = 8,
+} whMessageCustomCb_Type;
+
+
+/* union providing some helpful abstractions for passing pointers in/out of
+ * custom callbacks on top of a raw data buffer */
+typedef union {
+    /* pointer/size pairs for 32-bit systems */
+    struct {
+        uint32_t client_addr;
+        uint32_t client_sz;
+        uint32_t server_addr;
+        uint32_t server_sz;
+    } dma32;
+    /* pointer/size pairs for 64-bit systems */
+    struct {
+        uint64_t client_addr;
+        uint64_t client_sz;
+        uint64_t server_addr;
+        uint64_t server_sz;
+    } dma64;
+    /* raw data buffer for user-defined schema */
+    struct {
+        uint8_t data[WH_MESSAGE_CUSTOM_CB_BUF_SIZE];
+    } buffer;
+} whMessageCustomCb_Data;
+
+/* request message to the custom server callback */
+typedef struct {
+    uint32_t               id;   /* indentifier of registered callback  */
+    uint32_t               type; /* whMessageCustomCb_Type */
+    whMessageCustomCb_Data data;
+} whMessageCustomCb_Request;
+
+/* response message from the custom server callback */
+typedef struct {
+    uint32_t id;   /* indentifier of registered callback  */
+    uint32_t type; /* whMessageCustomCb_Type */
+    int32_t  rc;   /* Return code from custom callback. Invalid if err != 0 */
+    int32_t  err;  /* wolfHSM-specific error. If err != 0, rc is invalid */
+    whMessageCustomCb_Data data;
+} whMessageCustomCb_Response;
+
+
+/* Translates a custom request message. The whMessageCustomCb_Request.data field
+ * will not be translated for whMessageCustomCb_Request.type values greater than
+ * WH_MESSAGE_CUSTOM_CB_TYPE_USER_DEFINED_START */
+int wh_MessageCustomCb_TranslateRequest(uint16_t                         magic,
+                                        const whMessageCustomCb_Request* src,
+                                        whMessageCustomCb_Request*       dst);
+
+/* Translates a custom response message. The whMessageCustomCb_Request.data
+ * field will not be translated for whMessageCustomCb_Request.type values
+ * greater than WH_MESSAGE_CUSTOM_CB_TYPE_USER_DEFINED_START */
+int wh_MessageCustomCb_TranslateResponse(uint16_t magic,
+                                         const whMessageCustomCb_Response* src,
+                                         whMessageCustomCb_Response*       dst);
+
+#endif /* WH_MESSAGE_CUSTOM_CB_H_*/
\ No newline at end of file
diff --git a/wolfhsm/wh_server.h b/wolfhsm/wh_server.h
index 3f8656c0..8944b767 100644
--- a/wolfhsm/wh_server.h
+++ b/wolfhsm/wh_server.h
@@ -12,7 +12,7 @@
 #include "wolfhsm/wh_common.h"
 #include "wolfhsm/wh_comm.h"
 #include "wolfhsm/wh_nvm.h"
-#include "wolfhsm/wh_message_custom.h"
+#include "wolfhsm/wh_message_customcb.h"
 
 #include "wolfssl/wolfcrypt/settings.h"
 #include "wolfssl/wolfcrypt/random.h"
@@ -45,8 +45,8 @@ struct whServerContext_t;
 /* Type definition for a custom server callback  */
 typedef int (*whServerCustomCb)(
     struct whServerContext_t* server,   /* points to dispatching server ctx */
-    const whMessageCustom_Request* req, /* request from client to callback */
-    whMessageCustom_Response*      resp /* response from callback to client */
+    const whMessageCustomCb_Request* req, /* request from client to callback */
+    whMessageCustomCb_Response*      resp /* response from callback to client */
 );
 
 /* Context structure to maintain the state of an HSM server */
@@ -56,7 +56,7 @@ typedef struct whServerContext_t {
     whNvmContext nvm[1];
     crypto_context crypto[1];
     CacheSlot cache[WOLFHSM_NUM_RAMKEYS];
-    whServerCustomCb customHandlerTable[WH_CUSTOM_RQST_NUM_CALLBACKS];
+    whServerCustomCb customHandlerTable[WH_CUSTOM_CB_NUM_CALLBACKS];
 } whServerContext;
 
 typedef struct whServerConfig_t {
@@ -86,7 +86,7 @@ int wh_Server_RegisterCustomCb(whServerContext* server, uint16_t actionId, whSer
 
 /* Receive and handle an incoming custom callback request
 */
-int wh_Server_HandleCustomRequest(whServerContext* server, uint16_t magic,
+int wh_Server_HandleCustomCbRequest(whServerContext* server, uint16_t magic,
                                   uint16_t action, uint16_t seq,
                                   uint16_t req_size, const void* req_packet,
                                   uint16_t* out_resp_size, void* resp_packet);