Skip to content

fix(security): prevent OOB write and integer overflows during model load#45

Open
bytecodesky wants to merge 1 commit intomaderix:mainfrom
bytecodesky:security/prevent-model-poisoning-oob
Open

fix(security): prevent OOB write and integer overflows during model load#45
bytecodesky wants to merge 1 commit intomaderix:mainfrom
bytecodesky:security/prevent-model-poisoning-oob

Conversation

@bytecodesky
Copy link

Bug: Out-of-Bounds Write (CWE-787) and Unchecked Allocation / NULL Pointer Dereference (CWE-190/CWE-476) via Untrusted Model Configuration

Impact:
The model_load_weights function reads the Config structure directly from an untrusted file without validating its fields. This introduces two critical memory corruption vectors:

  1. Out-of-Bounds Write (CWE-787): The n_layers (nl) parameter dictates the number of iterations in the weight-loading loop. However, the destination arrays in the Model struct (e.g., m->wo[N_LAYERS]) have a statically defined size of N_LAYERS (12). If an attacker crafts a model file with n_layers > 12, the loop will write heap pointers beyond the bounds of these arrays. This sequentially corrupts adjacent pointer arrays in the Model struct (including the ANEKernel pointers, which encapsulate executable Objective-C/Metal objects), leading to highly exploitable memory corruption and Arbitrary Code Execution (ACE).
  2. Integer Overflow & NULL Pointer Dereference (CWE-190/CWE-476): The dim (d) and hidden_dim (hd) parameters are used directly in size calculations for heap allocations (e.g., d * d * sizeof(float)). Maliciously large values will cause integer overflows or excessive memory requests, causing malloc to fail and return NULL. Because the code lacks NULL checks, the subsequent memcpy operations immediately dereference the NULL pointer, resulting in a deterministic segmentation fault (Denial of Service).

Vulnerable Code:

    81:     if (fread(&m->cfg, sizeof(Config), 1, f) != 1) {
    82:         fprintf(stderr, "ERROR: failed to read config from %s\n", path);
    83:         fclose(f); return -1;
    84:     }
    // ... [NO VALIDATION OF m->cfg FIELDS PERFORMED HERE] ...
    92:     int d = m->cfg.dim, hd = m->cfg.hidden_dim, nl = m->cfg.n_layers, vs = m->cfg.vocab_size;
    // ...
    129:     for (int l = 0; l < nl; l++) {
    // ...
    138:         m->wo[l] = (float*)malloc(d*d*sizeof(float));
    139:         memcpy(m->wo[l], wo_all + l*d*d, d*d*sizeof(float));

Proposed Fix:
Implement strict bounds checking on the configuration parameters immediately after reading them from disk. Additionally, enforce NULL checks on allocations to ensure graceful failure under memory pressure.

    if (fread(&m->cfg, sizeof(Config), 1, f) != 1) {
        fprintf(stderr, "ERROR: failed to read config from %s\n", path);
        fclose(f); return -1;
    }

    // 1. Validate configuration bounds to prevent OOB writes and Integer Overflows
    if (m->cfg.n_layers < 1 || m->cfg.n_layers > N_LAYERS) {
        fprintf(stderr, "ERROR: n_layers (%d) exceeds maximum allowed (%d)\n", m->cfg.n_layers, N_LAYERS);
        fclose(f); return -1;
    }
    
    // Define sane maximums based on expected model architectures (e.g., 8K dim max)
    if (m->cfg.dim < 1 || m->cfg.dim > 8192 || 
        m->cfg.hidden_dim < 1 || m->cfg.hidden_dim > 32768) {
        fprintf(stderr, "ERROR: model dimensions out of safe bounds\n");
        fclose(f); return -1;
    }

    bool shared = m->cfg.vocab_size > 0;
    if (m->cfg.vocab_size < 0) m->cfg.vocab_size = -m->cfg.vocab_size;

    int d = m->cfg.dim, hd = m->cfg.hidden_dim, nl = m->cfg.n_layers, vs = m->cfg.vocab_size;

    // 2. Use a safe allocation macro to prevent NULL pointer dereferences
    #define SAFE_MALLOC_MEMCPY(dest, src, size) do { \
        dest = (float*)malloc(size); \
        if (!(dest)) { \
            fprintf(stderr, "ERROR: memory allocation failed for size %zu\n", (size_t)(size)); \
            fclose(f); return -1; \
        } \
        memcpy(dest, src, size); \
    } while(0)

    // ... [Bulk allocations omitted for brevity, apply similar NULL checks there] ...

    for (int l = 0; l < nl; l++) {
        SAFE_MALLOC_MEMCPY(m->rms_att_w[l], rms_att_all + l*d, d * sizeof(float));
        SAFE_MALLOC_MEMCPY(m->wq[l], wq_all + l*d*d, d*d*sizeof(float));
        SAFE_MALLOC_MEMCPY(m->wk[l], wk_all + l*d*d, d*d*sizeof(float));
        SAFE_MALLOC_MEMCPY(m->wv[l], wv_all + l*d*d, d*d*sizeof(float));
        SAFE_MALLOC_MEMCPY(m->wo[l], wo_all + l*d*d, d*d*sizeof(float));
        SAFE_MALLOC_MEMCPY(m->rms_ffn_w[l], rms_ffn_all + l*d, d * sizeof(float));
        SAFE_MALLOC_MEMCPY(m->w1[l], w1_all + l*hd*d, hd*d*sizeof(float));
        SAFE_MALLOC_MEMCPY(m->w2[l], w2_all + l*d*hd, d*hd*sizeof(float));
        SAFE_MALLOC_MEMCPY(m->w3[l], w3_all + l*hd*d, hd*d*sizeof(float));
    }
    
    #undef SAFE_MALLOC_MEMCPY

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant