security: reject malformed model tensors at the untrusted-mirror boundary

Colibri loads model directories and safetensors from mirrors it does not
control, so the file's declared shapes and byte spans are attacker-influenced
input. Three memory-safety holes on that boundary, independently confirmed
(incl. a from-scratch adversarial audit that re-derived the same two) and
present in the shipped v1.0.0:

- st.h st_read_f32: numel came from the shape, nbytes from the offsets, with
  no cross-check. A crafted tensor whose shape inflates numel past nbytes made
  the BF16/F16 loop read past the malloc'd raw buffer and the F32 memcpy write
  past the caller's config-sized destination (heap OOB read + write). Now
  enforce numel*esz == nbytes before any copy.
- st.h header parse: the shape product could overflow int64 to a small/negative
  numel that would then pass the cross-check. Guard each multiply.
- glm.c qt_resolve_fmt (new, replaces the three duplicated "?1:?2:3" fmt sites
  in qt_from_disk and both expert_load arms): the old inference SILENTLY fell
  to int2 for any unrecognized weight byte count, so a too-short weight became
  a valid int2 whose matmul read O*I nibbles past the buffer; and an oversized
  scale array overflowed the per-row t->s. Now the weight bytes must match a
  known int8/int4/int2 layout and the scale array must match the expected
  per-row (O) or grouped (O*ng) cardinality, else refuse.
- glm.c config/generation_config slurp: unbounded ftell -> malloc(n+1) gave a
  hostile file a load-time OOM, and on malloc failure b[got]=0 was a NULL
  deref. Cap at 256 MB and NULL-check.

Verified: TF token-exactness unchanged on every quant format (full-precision
32/32, int4 11/32, int2 1/32, mix 5/32 -- byte-identical to the pre-change
binary); fmt=4 grouped path preserved (the scale check is by construction the
same condition detect_group_size already imposed); a hand-crafted hostile
safetensors is refused cleanly; ASan+UBSan clean on legit and hostile loads
(only the pre-existing intentional startup leaks remain).

These are the C trust-boundary items of #368, landed as a minimal standalone
fix; the server-side and build items of that PR follow via its rebase.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
JustVugg
2026-07-19 13:05:56 +02:00
parent fae4b2cc3e
commit 72e36772f5
2 changed files with 58 additions and 16 deletions
+22 -1
View File
@@ -200,7 +200,19 @@ static void st_init(shards *S, const char *snap_dir) {
if (a0 < 0 || b0 < a0 || data_start + b0 > fsz) {
fprintf(stderr, "%s: tensor '%s' data_offsets [%lld,%lld] out of file bounds (%lld)\n",
files[fi], name, (long long)a0, (long long)b0, (long long)fsz); exit(1); }
int64_t numel = 1; for (int k = 0; k < shp->len; k++) numel *= (int64_t)shp->kids[k]->num;
/* SEC: lo shape viene da un file non fidato (mirror). Senza il guard
* di overflow, uno shape tipo [65535,65535,65535,...] fa avvolgere
* numel a un valore piccolo/negativo che poi passerebbe il cross-check
* numel*esz==nbytes in st_read_f32, riaprendo l'OOB. */
int64_t numel = 1; int bad_shape = 0;
for (int k = 0; k < shp->len; k++) {
int64_t d = (int64_t)shp->kids[k]->num;
if (d < 0 || (d != 0 && numel > INT64_MAX / d)) { bad_shape = 1; break; }
numel *= d;
}
if (bad_shape) {
fprintf(stderr, "%s: tensor '%s' shape overflows int64 — refusing (hostile or corrupt file)\n",
files[fi], name); exit(1); }
if (S->n == S->cap) { S->cap *= 2; S->t = realloc(S->t, S->cap*sizeof(st_tensor)); }
st_tensor *t = &S->t[S->n++];
t->name = strdup(name); t->fd = fd; t->off = data_start + a0;
@@ -249,6 +261,15 @@ static void st_prefetch(shards *S, const char *name) {
static int64_t st_read_f32(shards *S, const char *name, float *out, int drop) {
st_tensor *t = st_find(S, name);
if (!t) { fprintf(stderr, "missing tensor: %s\n", name); exit(1); }
/* SEC: numel viene dallo shape, nbytes dagli offset — due campi indipendenti
* del file. Se non concordano, la memcpy F32 (nbytes) o i loop BF16/F16
* (numel elementi da un raw di soli nbytes) sforano il buffer del chiamante,
* che e' dimensionato sul config, non sul file. Il chiamante che alloca su
* st_numel resta coerente; questo blocca l'ingresso ostile a monte. */
int esz = (t->dtype == 2) ? 4 : 2;
if (t->numel < 0 || t->numel > t->nbytes / esz || t->numel * (int64_t)esz != t->nbytes) {
fprintf(stderr, "%s: tensor '%s' shape/bytes mismatch (numel %lld, %lld bytes, dtype %d) — refusing (hostile or corrupt file)\n",
name, name, (long long)t->numel, (long long)t->nbytes, t->dtype); exit(1); }
void *raw = malloc(t->nbytes);
if (!raw) { fprintf(stderr, "malloc %lld bytes for tensor %s failed\n", (long long)t->nbytes, name); exit(1); }
st_pread_full(t->fd, raw, t->nbytes, t->off, "pread data");