Skip to content

Validate safetensors header fields before using them as sizes - #713

Open
professor-moody wants to merge 1 commit into
ztxz16:masterfrom
professor-moody:fix-safetensors-header-reconciliation
Open

professor-moody wants to merge 1 commit into
ztxz16:masterfrom
professor-moody:fix-safetensors-header-reconciliation

Conversation

@professor-moody

Copy link
Copy Markdown

Fixes #712.

Three values from the .safetensors header were used as allocation sizes or read lengths without being checked against anything else. Loading a crafted file writes past the end of heap allocations.

SafeTensors (src/model.cpp:1430-1435). The 8-byte header length went straight into new char[configBytes + 5]. That addition wraps for values near UINT64_MAX, so the allocation becomes 0-4 bytes while the following fread still asks for configBytes. Now bounded against the real file size. The fopen result and the first fread return were also unchecked, so a short or unreadable file left configBytes uninitialised.

SafeTensorItem (src/model.cpp:900-912). data_offsets[1] was indexed without checking that two elements were parsed, and there was no check that the end offset is not before the beginning.

SafeTensorItem::CreateBuffer (src/model.cpp:1383-1385). The tensor buffer is allocated from the shape product and then filled by fread with the data_offsets span. Those are independent header fields and can disagree. The else branch has the mirror problem, where ConvertDataType copies len * unitSize out of a buffer allocated with bytes.

All three are now checked at parse time, where the fields are first read, rather than at each point of use.

Testing

Against the crafted files from #712, on a build with -fsanitize=address and _FORTIFY_SOURCE disabled so the writes are visible rather than caught by glibc:

input before after
valid model loads loads, unchanged
shape [32,8], data_offsets span 65536 heap overflow, WRITE of size 65536 rejected, "data_offsets span does not match shape and dtype"
data_offsets: [0] heap overflow, READ of size 8 rejected, "data_offsets must have exactly 2 elements"
data_offsets: [] SEGV rejected, same
header length 2^64-1, 4 KB body heap overflow, WRITE of size 4088 rejected, "header length exceeds the file size"
header length 2^64-1, no body heap overflow, WRITE of size 1 rejected, same

Seven crafted inputs in total are rejected, each by the specific check intended for it. A model directory whose header fields agree still loads to inference.

I have the crafted files if they would be useful for a regression test.

The safetensors header is untrusted input, and three values from it were used as
allocation sizes or read lengths without being checked against anything else.

- SafeTensors: the 8-byte header length was passed straight to
  new char[configBytes + 5]. That addition wraps for values near UINT64_MAX,
  giving a 0-4 byte allocation while the following fread still asks for
  configBytes. Bound it against the actual file size, and check the fopen and
  fread results, which were also unchecked.

- SafeTensorItem: data_offsets was indexed at [1] without checking that two
  elements were parsed, and there was no check that the end offset is not before
  the beginning.

- SafeTensorItem: the tensor buffer is allocated from the shape product but
  filled by fread with the data_offsets span. Those are independent header
  fields, so they can disagree. Reconcile them once, at parse time, against the
  size implied by shape and dtype.

The shape product and the shape-times-dtype product are also checked for
overflow.

Valid models are unaffected: the new checks only reject headers whose own fields
contradict each other.

Fixes ztxz16#712
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.

Heap buffer overflows loading a crafted .safetensors model

1 participant