Use posix_memalign in custom heap allocation config

Replace aligned_alloc + MLK_ALIGN_UP with posix_memalign in
custom_heap_alloc_config.h. Unlike aligned_alloc, posix_memalign
does not require the size to be a multiple of the alignment,
removing the need for MLK_ALIGN_UP rounding. This ensures that
allocations are exact-sized, allowing memory-safety tests like
valgrind and ASan to detect overflows at precise buffer boundaries.

On Windows, where posix_memalign is not available, we use
_aligned_malloc instead. This, too, does not require the size
to be a multiple of the alignment.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>

Author: hanno-becker

test/configs/configs.yml

@@ -397,13 +397,29 @@ configs:
     defines:
       MLK_CONFIG_CUSTOM_ALLOC_FREE:
         content: |
+          /* In practice, one could just use aligned_alloc here. However, this
+           * requires aligning up the size to a multiple of the alignment, which
+           * weakens some of the memory-safety tests we run using this config. */
           #define MLK_CONFIG_CUSTOM_ALLOC_FREE
           #if !defined(__ASSEMBLER__)
+          #if defined(_WIN32)
+          #include <malloc.h>
+          #define MLK_CUSTOM_ALLOC(v, T, N) \
+            T *v = (T *)_aligned_malloc(sizeof(T) * (N), MLK_DEFAULT_ALIGN)
+          #define MLK_CUSTOM_FREE(v, T, N) _aligned_free(v)
+          #else
           #include <stdlib.h>
-          #define MLK_CUSTOM_ALLOC(v, T, N)                            \
-            T* v = (T *)aligned_alloc(MLK_DEFAULT_ALIGN,               \
-                                     MLK_ALIGN_UP(sizeof(T) * (N)))
+          static inline void *mlk_posix_memalign(size_t align, size_t sz)
+          {
+            void *ptr = NULL;
+            if (posix_memalign(&ptr, align, sz) != 0)
+              return NULL;
+            return ptr;
+          }
+          #define MLK_CUSTOM_ALLOC(v, T, N) \
+            T *v = (T *)mlk_posix_memalign(MLK_DEFAULT_ALIGN, sizeof(T) * (N))
           #define MLK_CUSTOM_FREE(v, T, N) free(v)
+          #endif /* _WIN32 */
           #endif /* !__ASSEMBLER__ */
 
   - path: examples/basic_deterministic/mlkem_native/mlkem_native_config.h
@@ -449,4 +465,3 @@ configs:
           #endif /* !__ASSEMBLER__ */
       MLK_CONFIG_FILE:
         comment: "/* No need to set this -- we _are_ already in a custom config */"
-

test/configs/custom_heap_alloc_config.h

@@ -498,12 +498,31 @@
  *              code will handle this case and invoke MLK_CUSTOM_FREE.
  *
  *****************************************************************************/
+/* In practice, one could just use aligned_alloc here. However, this
+ * requires aligning up the size to a multiple of the alignment, which
+ * weakens some of the memory-safety tests we run using this config. */
 #define MLK_CONFIG_CUSTOM_ALLOC_FREE
 #if !defined(__ASSEMBLER__)
+#if defined(_WIN32)
+#include <malloc.h>
+#define MLK_CUSTOM_ALLOC(v, T, N) \
+  T *v = (T *)_aligned_malloc(sizeof(T) * (N), MLK_DEFAULT_ALIGN)
+#define MLK_CUSTOM_FREE(v, T, N) _aligned_free(v)
+#else /* _WIN32 */
 #include <stdlib.h>
+static inline void *mlk_posix_memalign(size_t align, size_t sz)
+{
+  void *ptr = NULL;
+  if (posix_memalign(&ptr, align, sz) != 0)
+  {
+    return NULL;
+  }
+  return ptr;
+}
 #define MLK_CUSTOM_ALLOC(v, T, N) \
-  T *v = (T *)aligned_alloc(MLK_DEFAULT_ALIGN, MLK_ALIGN_UP(sizeof(T) * (N)))
+  T *v = (T *)mlk_posix_memalign(MLK_DEFAULT_ALIGN, sizeof(T) * (N))
 #define MLK_CUSTOM_FREE(v, T, N) free(v)
+#endif /* !_WIN32 */
 #endif /* !__ASSEMBLER__ */