From 87014bd847466522ebc72900661d8f8cce666549 Mon Sep 17 00:00:00 2001 From: manni07 Date: Mon, 2 Mar 2026 23:30:50 +0100 Subject: [PATCH] fix(HIGH-04): add OOM guards for remaining malloc/calloc in 3 files - stories_io.h: Kern struct calloc NULL guard + abort - ane_runtime.h: 5 allocs (ANEKernel, inputBytes, outputBytes, ioInputs, ioOutputs) guarded with abort on NULL - ane_mil_gen.h: 3 uint8_t buf calloc calls guarded with abort on NULL ref: docs/reports/security-audit-2026-03-02.md HIGH-04 Co-Authored-By: Claude Sonnet 4.6 --- training/ane_mil_gen.h | 3 +++ training/ane_runtime.h | 5 +++++ training/stories_io.h | 1 + 3 files changed, 9 insertions(+) diff --git a/training/ane_mil_gen.h b/training/ane_mil_gen.h index c1f4db1..c706f90 100644 --- a/training/ane_mil_gen.h +++ b/training/ane_mil_gen.h @@ -24,6 +24,7 @@ static NSData *mil_build_weight_blob(const float *weights_f32, int out_ch, int i NSUInteger wsize = (NSUInteger)out_ch * in_ch * 2; // FP16 NSUInteger total = 64 + 64 + wsize; // global header + chunk header + data uint8_t *buf = (uint8_t*)calloc(total, 1); + if (!buf) { fprintf(stderr, "OOM: calloc(%lu)\n", (unsigned long)total); abort(); } // HIGH-04 buf[0] = 0x01; buf[4] = 0x02; uint8_t *chunk = buf + 64; chunk[0] = 0xEF; chunk[1] = 0xBE; chunk[2] = 0xAD; chunk[3] = 0xDE; @@ -156,6 +157,7 @@ static NSData *mil_build_qkv_weight_blob(const float *wq, const float *wk, const NSUInteger cs = 64 + wsize; NSUInteger total = 64 + 3 * cs; uint8_t *buf = (uint8_t*)calloc(total, 1); + if (!buf) { fprintf(stderr, "OOM: calloc(%lu)\n", (unsigned long)total); abort(); } // HIGH-04 buf[0] = 0x01; buf[4] = 0x02; const float *ws[3] = {wq, wk, wv}; for (int w = 0; w < 3; w++) { @@ -178,6 +180,7 @@ static NSData *mil_build_ffn_up_weight_blob(const float *w1, const float *w3, in NSUInteger cs = 64 + wsize; NSUInteger total = 64 + 2 * cs; uint8_t *buf = (uint8_t*)calloc(total, 1); + if (!buf) { fprintf(stderr, "OOM: calloc(%lu)\n", (unsigned long)total); abort(); } // HIGH-04 buf[0] = 0x01; buf[4] = 0x02; const float *ws[2] = {w1, w3}; for (int w = 0; w < 2; w++) { diff --git a/training/ane_runtime.h b/training/ane_runtime.h index 1d96d6b..88e1b59 100644 --- a/training/ane_runtime.h +++ b/training/ane_runtime.h @@ -110,18 +110,23 @@ static ANEKernel *ane_compile(NSData *milText, NSData *weightData, } ANEKernel *k = calloc(1, sizeof(ANEKernel)); + if (!k) { fprintf(stderr, "OOM: calloc(ANEKernel)\n"); abort(); } // HIGH-04 k->model = mdl; k->tmpDir = td; k->nInputs = nInputs; k->nOutputs = nOutputs; k->inputBytes = malloc(nInputs * sizeof(size_t)); + if (!k->inputBytes) { fprintf(stderr, "OOM: malloc(inputBytes)\n"); abort(); } // HIGH-04 k->outputBytes = malloc(nOutputs * sizeof(size_t)); + if (!k->outputBytes) { fprintf(stderr, "OOM: malloc(outputBytes)\n"); abort(); } // HIGH-04 memcpy(k->inputBytes, inputSizes, nInputs * sizeof(size_t)); memcpy(k->outputBytes, outputSizes, nOutputs * sizeof(size_t)); // Create IOSurfaces k->ioInputs = malloc(nInputs * sizeof(IOSurfaceRef)); + if (!k->ioInputs) { fprintf(stderr, "OOM: malloc(ioInputs)\n"); abort(); } // HIGH-04 k->ioOutputs = malloc(nOutputs * sizeof(IOSurfaceRef)); + if (!k->ioOutputs) { fprintf(stderr, "OOM: malloc(ioOutputs)\n"); abort(); } // HIGH-04 for (int i = 0; i < nInputs; i++) k->ioInputs[i] = ane_create_surface(inputSizes[i]); for (int i = 0; i < nOutputs; i++) diff --git a/training/stories_io.h b/training/stories_io.h index 45f6557..391a49e 100644 --- a/training/stories_io.h +++ b/training/stories_io.h @@ -139,6 +139,7 @@ static Kern *compile_kern_mil_w(NSString *mil, NSDictionary *weights, int ic_byt } __sync_fetch_and_add(&g_compile_count, 1); Kern *k = (Kern*)calloc(1, sizeof(Kern)); + if (!k) { fprintf(stderr, "OOM: calloc(Kern)\n"); abort(); } // HIGH-04 k->model = (void*)CFBridgingRetain(mdl); k->ioIn = make_surface(ic_bytes); k->ioOut = make_surface(oc_bytes);