Port custom allocator from minidriver, and add test case for it.

lib/internal.h

@@ -31,6 +31,8 @@
 #ifndef YKPIV_INTERNAL_H
 #define YKPIV_INTERNAL_H
 
+#include "ykpiv.h"
+
 #include <stdbool.h>
 
 #if BACKEND_PCSC
@@ -48,18 +50,6 @@
 #define DES_LEN_3DES  8*3
 #define CB_MGM_KEY DES_LEN_3DES
 
-  typedef void* (*ykpiv_pfn_alloc)(void* alloc_data, size_t size);
-  typedef void* (*ykpiv_pfn_realloc)(void* alloc_data, void* address, size_t size);
-  typedef void  (*ykpiv_pfn_free)(void* alloc_data, void* address);
-  typedef struct {
-    ykpiv_pfn_alloc   pfn_alloc;
-    ykpiv_pfn_realloc pfn_realloc;
-    ykpiv_pfn_free    pfn_free;
-    void *         alloc_data;
-  } ykpiv_allocator;
-
-extern ykpiv_allocator _mem_default_allocator;
-
 struct ykpiv_state {
   SCARDCONTEXT context;
   SCARDHANDLE card;

lib/tests/util.c

@@ -58,13 +58,24 @@ void setup(void) {
 
   res = ykpiv_init(&g_state, true);
   ck_assert_int_eq(res, YKPIV_OK);
+  ck_assert_ptr_nonnull(g_state);
 
   res = ykpiv_connect(g_state, NULL);
   ck_assert_int_eq(res, YKPIV_OK);
 }
 
 void teardown(void) {
-  ykpiv_done(g_state);
+  ykpiv_rc res;
+
+  // This is the expected case, if the allocator test ran, since it de-inits.
+  if (NULL == g_state)
+    return;
+
+  res = ykpiv_disconnect(g_state);
+  ck_assert_int_eq(res, YKPIV_OK);
+
+  res = ykpiv_done(g_state);
+  ck_assert_int_eq(res, YKPIV_OK);
 }
 
 START_TEST(test_devicemodel) {
@@ -208,6 +219,91 @@ START_TEST(test_reset) {
 }
 END_TEST
 
+
+struct t_alloc_data{
+  uint32_t count;
+} g_alloc_data;
+
+static void* _test_alloc(void *data, size_t cb) {
+  ck_assert_ptr_eq(data, &g_alloc_data);
+  ((struct t_alloc_data*)data)->count++;
+  return calloc(cb, 1);
+}
+
+static void * _test_realloc(void *data, void *p, size_t cb) {
+  ck_assert_ptr_eq(data, &g_alloc_data);
+  return realloc(p, cb);
+}
+
+static void _test_free(void *data, void *p) {
+  fflush(stderr);
+  ck_assert_ptr_eq(data, &g_alloc_data);
+  ((struct t_alloc_data*)data)->count--;
+  free(p);
+}
+
+ykpiv_allocator test_allocator_cbs = {
+  .pfn_alloc = _test_alloc,
+  .pfn_realloc = _test_realloc,
+  .pfn_free = _test_free,
+  .alloc_data = &g_alloc_data
+};
+
+uint8_t *alloc_auth_cert() {
+  ykpiv_rc res;
+  uint8_t *read_cert = NULL;
+  size_t read_cert_len = 0;
+
+  res = ykpiv_util_write_cert(g_state, YKPIV_KEY_AUTHENTICATION, (uint8_t*)g_cert, sizeof(g_cert));
+  ck_assert_int_eq(res, YKPIV_OK);
+
+  res = ykpiv_util_read_cert(g_state, YKPIV_KEY_AUTHENTICATION, &read_cert, &read_cert_len);
+  ck_assert_int_eq(res, YKPIV_OK);
+  ck_assert_ptr_nonnull(read_cert);
+  ck_assert_int_eq(read_cert_len, sizeof(g_cert));
+  ck_assert_mem_eq(g_cert, read_cert, sizeof(g_cert));
+  return read_cert;
+}
+
+START_TEST(test_allocator) {
+  ykpiv_rc res;
+  const ykpiv_allocator allocator;
+  uint8_t *cert1, *cert2;
+
+  res = ykpiv_done(g_state);
+  ck_assert_int_eq(res, YKPIV_OK);
+  g_state = NULL;
+
+  res = ykpiv_init_with_allocator(&g_state, false, &test_allocator_cbs);
+  ck_assert_int_eq(res, YKPIV_OK);
+  ck_assert_ptr_nonnull(g_state);
+
+  // Verify we can communicate with device and make some allocations
+  res = ykpiv_connect(g_state, NULL);
+  ck_assert_int_eq(res, YKPIV_OK);
+  test_authenticate(0);
+  cert1 = alloc_auth_cert();
+  cert2 = alloc_auth_cert();
+
+  // Verify allocations went through custom allocator, and still live
+  ck_assert_int_gt(g_alloc_data.count, 1);
+
+  // Free and shutdown everything
+  ykpiv_util_free(g_state, cert2);
+  ykpiv_util_free(g_state, cert1);
+  res = ykpiv_disconnect(g_state);
+  ck_assert_int_eq(res, YKPIV_OK);
+  res = ykpiv_done(g_state);
+  ck_assert_int_eq(res, YKPIV_OK);
+
+  // Verify equal number of frees as allocations
+  ck_assert_int_eq(g_alloc_data.count, 0);
+
+  // Clear g_state so teardown() is skipped
+  g_state = NULL;
+}
+END_TEST
+
 int confirm_destruction(void) {
   char verify[16];
 
@@ -242,16 +338,20 @@ Suite *test_suite(void) {
 #ifdef HW_TESTS
   tcase_add_unchecked_fixture(tc, setup, teardown);
 
-  // Reset first.  Tests run serially, and depend on a clean slate.
+  // Must be first: Reset device.  Tests run serially, and depend on a clean slate.
   tcase_add_test(tc, test_reset);
 
   // Authenticate after reset.
   tcase_add_test(tc, test_authenticate);
 
+  // Test util functionality
   tcase_add_test(tc, test_devicemodel);
   tcase_add_test(tc, test_get_set_cardid);
   tcase_add_test(tc, test_read_write_list_delete_cert);
   tcase_add_test(tc, test_generate_key);
+
+  // Must be last: tear down and re-test with custom memory allocator
+  tcase_add_test(tc, test_allocator);
 #endif
   suite_add_tcase(s, tc);
 
@@ -267,7 +367,7 @@ int main(void)
   s = test_suite();
   sr = srunner_create(s);
   srunner_set_fork_status(sr, CK_NOFORK);
-  srunner_run_all(sr, CK_NORMAL);
+  srunner_run_all(sr, CK_VERBOSE);
   number_failed = srunner_ntests_failed(sr);
   srunner_free(sr);
 

lib/util.c

@@ -129,23 +129,6 @@ prng_rc prng_generate(unsigned char *buffer, const size_t cb_req) {
   return rc;
 }
 
-/* Memory helper functions */
-
-static void* _alloc(ykpiv_state *state, size_t size) {
-  if (!state || !(state->allocator.pfn_alloc)) return NULL;
-  return state->allocator.pfn_alloc(state->allocator.alloc_data, size);
-}
-
-static void* _realloc(ykpiv_state *state, void *address, size_t size) {
-  if (!state || !(state->allocator.pfn_realloc)) return NULL;
-  return state->allocator.pfn_realloc(state->allocator.alloc_data, address, size);
-}
-
-static void _free(ykpiv_state *state, void *data) {
-  if (!data || !state || (!(state->allocator.pfn_free))) return;
-  state->allocator.pfn_free(state->allocator.alloc_data, data);
-}
-
 static size_t _obj_size_max(ykpiv_state *state) {
   return (state && state->isNEO) ? CB_OBJ_MAX_NEO : CB_OBJ_MAX;
 }
@@ -153,6 +136,9 @@ static size_t _obj_size_max(ykpiv_state *state) {
 #define MAX(a,b) (a) > (b) ? (a) : (b)
 #define MIN(a,b) (a) < (b) ? (a) : (b)
 
+void* _ykpiv_alloc(ykpiv_state *state, size_t size);
+void* _ykpiv_realloc(ykpiv_state *state, void *address, size_t size);
+void _ykpiv_free(ykpiv_state *state, void *data);
 int _ykpiv_set_length(unsigned char *buffer, size_t length);
 int _ykpiv_get_length(const unsigned char *buffer, size_t *len);
 ykpiv_rc _ykpiv_begin_transaction(ykpiv_state *state);
@@ -288,7 +274,7 @@ ykpiv_rc ykpiv_util_list_keys(ykpiv_state *state, uint8_t *key_count, ykpiv_key
   *data_len = 0;
 
   // allocate initial page of buffer
-  if (NULL == (pData = _alloc(state, CB_PAGE))) {
+  if (NULL == (pData = _ykpiv_alloc(state, CB_PAGE))) {
     res = YKPIV_MEMORY_ERROR;
     goto Cleanup;
   }
@@ -304,7 +290,7 @@ ykpiv_rc ykpiv_util_list_keys(ykpiv_state *state, uint8_t *key_count, ykpiv_key
       cbRealloc = (sizeof(ykpiv_key) + cbBuf - 1) > (cbData - offset) ? MAX((sizeof(ykpiv_key) + cbBuf - 1) - (cbData - offset), CB_PAGE) : 0;
 
       if (0 != cbRealloc) {
-        if (NULL == (pData = _realloc(state, pData, cbData + cbRealloc))) {
+        if (NULL == (pData = _ykpiv_realloc(state, pData, cbData + cbRealloc))) {
           res = YKPIV_MEMORY_ERROR;
           goto Cleanup;
         }
@@ -338,7 +324,7 @@ ykpiv_rc ykpiv_util_list_keys(ykpiv_state *state, uint8_t *key_count, ykpiv_key
 
 Cleanup:
 
-  if (pData) { _free(state, pData); }
+  if (pData) { _ykpiv_free(state, pData); }
 
   _ykpiv_end_transaction(state);
   return res;
@@ -348,7 +334,7 @@ ykpiv_rc ykpiv_util_free(ykpiv_state *state, void *data) {
   if (!data) return YKPIV_OK;
   if (!state || (!(state->allocator.pfn_free))) return YKPIV_GENERIC_ERROR;
 
-  _free(state, data);
+  _ykpiv_free(state, data);
 
   return YKPIV_OK;
 }
@@ -367,7 +353,7 @@ ykpiv_rc ykpiv_util_read_cert(ykpiv_state *state, uint8_t slot, uint8_t **data,
   *data_len = 0;
 
   if (YKPIV_OK == (res = _read_certificate(state, slot, buf, &cbBuf))) {
-    if (NULL == (*data = _alloc(state, cbBuf))) {
+    if (NULL == (*data = _ykpiv_alloc(state, cbBuf))) {
       res = YKPIV_MEMORY_ERROR;
       goto Cleanup;
     }
@@ -433,7 +419,7 @@ ykpiv_rc ykpiv_util_read_mscmap(ykpiv_state *state, ykpiv_container **containers
         goto Cleanup;
       }
 
-      if (NULL == (*containers = _alloc(state, len))) {
+      if (NULL == (*containers = _ykpiv_alloc(state, len))) {
         res = YKPIV_MEMORY_ERROR;
         goto Cleanup;
       }
@@ -522,7 +508,7 @@ ykpiv_rc ykpiv_util_read_msroots(ykpiv_state *state, uint8_t **data, size_t *dat
 
   // allocate first page
   cbData = _obj_size_max(state);
-  if (NULL == (pData = _alloc(state, cbData))) { res = YKPIV_MEMORY_ERROR; goto Cleanup; }
+  if (NULL == (pData = _ykpiv_alloc(state, cbData))) { res = YKPIV_MEMORY_ERROR; goto Cleanup; }
 
   for (object_id = YKPIV_OBJ_MSROOTS1; object_id <= YKPIV_OBJ_MSROOTS5; object_id++) {
     cbBuf = sizeof(buf);
@@ -558,7 +544,7 @@ ykpiv_rc ykpiv_util_read_msroots(ykpiv_state *state, uint8_t **data, size_t *dat
     cbRealloc = len > (cbData - offset) ? len - (cbData - offset) : 0;
 
     if (0 != cbRealloc) {
-      if (NULL == (pData = _realloc(state, pData, cbData + cbRealloc))) {
+      if (NULL == (pData = _ykpiv_realloc(state, pData, cbData + cbRealloc))) {
         res = YKPIV_MEMORY_ERROR;
         goto Cleanup;
       }
@@ -583,7 +569,7 @@ ykpiv_rc ykpiv_util_read_msroots(ykpiv_state *state, uint8_t **data, size_t *dat
 
 Cleanup:
 
-  if (pData) { _free(state, pData); }
+  if (pData) { _ykpiv_free(state, pData); }
 
   _ykpiv_end_transaction(state);
   return res;
@@ -777,7 +763,7 @@ ykpiv_rc ykpiv_util_generate_key(ykpiv_state *state, uint8_t slot, uint8_t algor
     data_ptr += _ykpiv_get_length(data_ptr, &len);
 
     cb_modulus = len;
-    if (NULL == (ptr_modulus = _alloc(state, cb_modulus))) {
+    if (NULL == (ptr_modulus = _ykpiv_alloc(state, cb_modulus))) {
       if (state->verbose) { fprintf(stderr, "Failed to allocate memory for modulus.\n"); }
       res = YKPIV_MEMORY_ERROR;
       goto Cleanup;
@@ -797,7 +783,7 @@ ykpiv_rc ykpiv_util_generate_key(ykpiv_state *state, uint8_t slot, uint8_t algor
     data_ptr += _ykpiv_get_length(data_ptr, &len);
 
     cb_exp = len;
-    if (NULL == (ptr_exp = _alloc(state, cb_exp))) {
+    if (NULL == (ptr_exp = _ykpiv_alloc(state, cb_exp))) {
       if (state->verbose) { fprintf(stderr, "Failed to allocate memory for public exponent.\n"); }
       res = YKPIV_MEMORY_ERROR;
       goto Cleanup;
@@ -838,7 +824,7 @@ ykpiv_rc ykpiv_util_generate_key(ykpiv_state *state, uint8_t slot, uint8_t algor
     }
 
     cb_point = len;
-    if (NULL == (ptr_point = _alloc(state, cb_point))) {
+    if (NULL == (ptr_point = _ykpiv_alloc(state, cb_point))) {
       if (state->verbose) { fprintf(stderr, "Failed to allocate memory for public point.\n"); }
       res = YKPIV_MEMORY_ERROR;
       goto Cleanup;
@@ -860,9 +846,9 @@ ykpiv_rc ykpiv_util_generate_key(ykpiv_state *state, uint8_t slot, uint8_t algor
 
 Cleanup:
 
-  if (ptr_modulus) { _free(state, modulus); }
+  if (ptr_modulus) { _ykpiv_free(state, modulus); }
-  if (ptr_exp) { _free(state, ptr_exp); }
+  if (ptr_exp) { _ykpiv_free(state, ptr_exp); }
-  if (ptr_point) { _free(state, ptr_exp); }
+  if (ptr_point) { _ykpiv_free(state, ptr_exp); }
 
   _ykpiv_end_transaction(state);
   return res;

lib/ykpiv.c

@@ -71,6 +71,23 @@ ykpiv_allocator _default_allocator = {
   .alloc_data = 0
 };
 
+/* Memory helper functions */
+
+void* _ykpiv_alloc(ykpiv_state *state, size_t size) {
+  if (!state || !(state->allocator.pfn_alloc)) return NULL;
+  return state->allocator.pfn_alloc(state->allocator.alloc_data, size);
+}
+
+void* _ykpiv_realloc(ykpiv_state *state, void *address, size_t size) {
+  if (!state || !(state->allocator.pfn_realloc)) return NULL;
+  return state->allocator.pfn_realloc(state->allocator.alloc_data, address, size);
+}
+
+void _ykpiv_free(ykpiv_state *state, void *data) {
+  if (!data || !state || (!(state->allocator.pfn_free))) return;
+  state->allocator.pfn_free(state->allocator.alloc_data, data);
+}
+
 static void dump_hex(const unsigned char *buf, unsigned int len) {
   unsigned int i;
   for (i = 0; i < len; i++) {
@@ -152,8 +169,9 @@ ykpiv_rc ykpiv_init(ykpiv_state **state, int verbose) {
 
 ykpiv_rc ykpiv_done(ykpiv_state *state) {
   ykpiv_disconnect(state);
-  free(state->pin);
+  if (state->pin)
-  free(state);
+    _ykpiv_free(state, state->pin);
+  _ykpiv_free(state, state);
   return YKPIV_OK;
 }
 
@@ -839,8 +857,8 @@ ykpiv_rc ykpiv_verify(ykpiv_state *state, const char *pin, int *tries) {
     return res;
   } else if(sw == SW_SUCCESS) {
     if (pin) {
-      free(state->pin);
+      _ykpiv_free(state, state->pin);
-      state->pin = malloc(len * sizeof(char) + 1);
+      state->pin = _ykpiv_alloc(state, len * sizeof(char) + 1);
       if (state->pin == NULL) {
         return YKPIV_MEMORY_ERROR;
       }
@@ -911,8 +929,8 @@ static ykpiv_rc change_pin_internal(ykpiv_state *state, int action, const char *
 ykpiv_rc ykpiv_change_pin(ykpiv_state *state, const char * current_pin, size_t current_pin_len, const char * new_pin, size_t new_pin_len, int *tries) {
   ykpiv_rc res = change_pin_internal(state, CHREF_ACT_CHANGE_PIN, current_pin, current_pin_len, new_pin, new_pin_len, tries);
   if (res == YKPIV_OK && new_pin != NULL) {
-    free(state->pin);
+    _ykpiv_free(state, state->pin);
-    state->pin = malloc(new_pin_len * sizeof(char) + 1);
+    state->pin = _ykpiv_alloc(state, new_pin_len * sizeof(char) + 1);
     if (state->pin == NULL) {
       return YKPIV_MEMORY_ERROR;
     }

lib/ykpiv.h

@@ -63,10 +63,21 @@ extern "C"
     YKPIV_RANGE_ERROR = -15 //i.e. value range error
   } ykpiv_rc;
 
+  typedef void* (*ykpiv_pfn_alloc)(void* alloc_data, size_t size);
+  typedef void* (*ykpiv_pfn_realloc)(void* alloc_data, void* address, size_t size);
+  typedef void  (*ykpiv_pfn_free)(void* alloc_data, void* address);
+  typedef struct ykpiv_allocator {
+    ykpiv_pfn_alloc   pfn_alloc;
+    ykpiv_pfn_realloc pfn_realloc;
+    ykpiv_pfn_free    pfn_free;
+    void *         alloc_data;
+  } ykpiv_allocator;
+
   const char *ykpiv_strerror(ykpiv_rc err);
   const char *ykpiv_strerror_name(ykpiv_rc err);
 
   ykpiv_rc ykpiv_init(ykpiv_state **state, int verbose);
+  ykpiv_rc ykpiv_init_with_allocator(ykpiv_state **state, int verbose, const ykpiv_allocator *allocator);
   ykpiv_rc ykpiv_done(ykpiv_state *state);
   ykpiv_rc ykpiv_connect(ykpiv_state *state, const char *wanted);
   ykpiv_rc ykpiv_list_readers(ykpiv_state *state, char *readers, size_t *len);