fix(HIGH-01): add token-count guard and embed_lookup bounds check

- train_large.m: early exit if n_tokens < SEQ+1 prevents arithmetic
  underflow and downstream index corruption
- stories_cpu_ops.h: clamp tok >= VOCAB to 0 in embed_lookup() to
  prevent heap-buffer-overflow

ref: docs/reports/security-audit-2026-03-02.md HIGH-01
This commit is contained in:
manni07 2026-03-02 23:12:25 +01:00
parent a00ce3afad
commit 236e495a01
2 changed files with 839 additions and 834 deletions

View File

@ -112,6 +112,7 @@ static float cross_entropy_loss(float *dlogits, const float *logits, const uint1
static void embed_lookup(float *x, const float *embed, const uint16_t *tokens, int dim, int seq) { static void embed_lookup(float *x, const float *embed, const uint16_t *tokens, int dim, int seq) {
for (int t = 0; t < seq; t++) { for (int t = 0; t < seq; t++) {
int tok = tokens[t]; int tok = tokens[t];
if (tok >= VOCAB) { tok = 0; } // HIGH-01: clamp invalid token -> position 0
for (int d = 0; d < dim; d++) { for (int d = 0; d < dim; d++) {
x[d*seq + t] = embed[tok*dim + d]; x[d*seq + t] = embed[tok*dim + d];
} }

View File

@ -296,6 +296,10 @@ int main(int argc, char *argv[]) {
uint16_t *token_data = (uint16_t*)mmap(NULL, data_len, PROT_READ, MAP_PRIVATE, data_fd, 0); uint16_t *token_data = (uint16_t*)mmap(NULL, data_len, PROT_READ, MAP_PRIVATE, data_fd, 0);
if (token_data == MAP_FAILED) { printf("mmap failed\n"); return 1; } if (token_data == MAP_FAILED) { printf("mmap failed\n"); return 1; }
size_t n_tokens = data_len / 2; size_t n_tokens = data_len / 2;
if (n_tokens < (size_t)SEQ + 1) {
fprintf(stderr, "Token file too small: %zu tokens, need >%d\n", n_tokens, SEQ + 1);
return 1;
}
printf("Token data: %zu tokens (%.1f MB)\n", n_tokens, data_len/1e6); printf("Token data: %zu tokens (%.1f MB)\n", n_tokens, data_len/1e6);
// Gradient buffers shared across layers (reused each step) // Gradient buffers shared across layers (reused each step)