create: harden GGUF create flows (#17062)

* create: harden GGUF create flows

* lint
This commit is contained in:
Daniel Hiltgen
2026-07-06 16:20:20 -07:00
committed by GitHub
parent 87b64213b4
commit 67b6a1c2d4
11 changed files with 883 additions and 76 deletions
+39
View File
@@ -520,10 +520,49 @@ func (t Tensor) Elements() uint64 {
return count return count
} }
func (t Tensor) elements() (uint64, bool) {
var count uint64 = 1
for _, n := range t.Shape {
if n != 0 && count > ^uint64(0)/n {
return 0, false
}
count *= n
}
return count, true
}
func (t Tensor) Size() uint64 { func (t Tensor) Size() uint64 {
return t.Elements() * t.typeSize() / t.blockSize() return t.Elements() * t.typeSize() / t.blockSize()
} }
func (t Tensor) size() (uint64, bool) {
elements, ok := t.elements()
if !ok {
return 0, false
}
typeSize := t.typeSize()
blockSize := t.blockSize()
if typeSize == 0 || blockSize == 0 {
return 0, false
}
rowSize := uint64(1)
if len(t.Shape) > 0 {
rowSize = t.Shape[0]
}
if rowSize%blockSize != 0 {
return 0, false
}
blocks := elements / blockSize
if blocks > ^uint64(0)/typeSize {
return 0, false
}
return blocks * typeSize, true
}
func (t Tensor) Type() string { func (t Tensor) Type() string {
return TensorType(t.Kind).String() return TensorType(t.Kind).String()
} }
+131 -38
View File
@@ -14,9 +14,14 @@ import (
"strings" "strings"
"github.com/ollama/ollama/fs" "github.com/ollama/ollama/fs"
fsgguf "github.com/ollama/ollama/fs/gguf"
"golang.org/x/sync/errgroup" "golang.org/x/sync/errgroup"
) )
const (
maxGGUFWriteConcurrency = 2
)
type containerGGUF struct { type containerGGUF struct {
ByteOrder binary.ByteOrder ByteOrder binary.ByteOrder
@@ -202,6 +207,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
if err != nil { if err != nil {
return fmt.Errorf("failed to read tensor dimensions: %w", err) return fmt.Errorf("failed to read tensor dimensions: %w", err)
} }
if dims > fsgguf.MaxTensorDims {
return fmt.Errorf("tensor %q dimensions %d exceeds maximum %d", name, dims, fsgguf.MaxTensorDims)
}
shape := make([]uint64, dims) shape := make([]uint64, dims)
for i := 0; uint32(i) < dims; i++ { for i := 0; uint32(i) < dims; i++ {
@@ -229,13 +237,23 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
} }
llm.tensors = append(llm.tensors, &tensor) llm.tensors = append(llm.tensors, &tensor)
llm.parameters += tensor.Elements() elements, ok := tensor.elements()
if !ok {
return fmt.Errorf("tensor %q elements overflow", tensor.Name)
}
if llm.parameters > maxUint64()-elements {
return fmt.Errorf("parameter count overflow")
}
llm.parameters += elements
} }
// patch KV with parameter count // patch KV with parameter count
llm.kv["general.parameter_count"] = llm.parameters llm.kv["general.parameter_count"] = llm.parameters
alignment := llm.kv.Uint("general.alignment", 32) alignment := llm.kv.Uint("general.alignment", 32)
if alignment == 0 {
return fmt.Errorf("invalid GGUF alignment: 0")
}
offset, err := rs.Seek(0, io.SeekCurrent) offset, err := rs.Seek(0, io.SeekCurrent)
if err != nil { if err != nil {
@@ -243,6 +261,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
} }
padding := ggufPadding(offset, int64(alignment)) padding := ggufPadding(offset, int64(alignment))
if padding > int64(maxInt64())-offset {
return fmt.Errorf("GGUF tensor offset overflow")
}
llm.tensorOffset = uint64(offset + padding) llm.tensorOffset = uint64(offset + padding)
// get file size to validate tensor bounds // get file size to validate tensor bounds
@@ -256,7 +277,15 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
} }
for _, tensor := range llm.tensors { for _, tensor := range llm.tensors {
tensorEnd := llm.tensorOffset + tensor.Offset + tensor.Size() tensorSize, ok := tensor.size()
if !ok {
return fmt.Errorf("tensor %q size overflow", tensor.Name)
}
tensorEnd, ok := checkedAdd(llm.tensorOffset, tensor.Offset, tensorSize)
if !ok {
return fmt.Errorf("tensor %q offset+size overflows", tensor.Name)
}
if tensorEnd > uint64(fileSize) { if tensorEnd > uint64(fileSize) {
return fmt.Errorf("tensor %q offset+size (%d) exceeds file size (%d)", tensor.Name, tensorEnd, fileSize) return fmt.Errorf("tensor %q offset+size (%d) exceeds file size (%d)", tensor.Name, tensorEnd, fileSize)
} }
@@ -271,7 +300,10 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error {
return fmt.Errorf("failed to seek to init padding: %w", err) return fmt.Errorf("failed to seek to init padding: %w", err)
} }
if _, err := rs.Seek(int64(tensor.Size()), io.SeekCurrent); err != nil { if tensorSize > maxInt64() {
return fmt.Errorf("tensor %q size %d exceeds maximum %d", tensor.Name, tensorSize, maxInt64())
}
if _, err := rs.Seek(int64(tensorSize), io.SeekCurrent); err != nil {
return fmt.Errorf("failed to seek to tensor: %w", err) return fmt.Errorf("failed to seek to tensor: %w", err)
} }
} }
@@ -298,11 +330,22 @@ func readGGUFV1String(llm *gguf, r io.Reader) (string, error) {
if err := binary.Read(r, llm.ByteOrder, &length); err != nil { if err := binary.Read(r, llm.ByteOrder, &length); err != nil {
return "", err return "", err
} }
if length == 0 {
return "", fmt.Errorf("invalid GGUF v1 string length: 0")
}
size, err := checkGGUFLength(length, "string")
if err != nil {
return "", err
}
var b bytes.Buffer var b bytes.Buffer
if _, err := io.CopyN(&b, r, int64(length)); err != nil { if _, err := io.CopyN(&b, r, int64(size)); err != nil {
return "", err return "", err
} }
if b.Bytes()[b.Len()-1] != 0 {
return "", fmt.Errorf("invalid GGUF v1 string terminator")
}
// gguf v1 strings are null-terminated // gguf v1 strings are null-terminated
b.Truncate(b.Len() - 1) b.Truncate(b.Len() - 1)
@@ -320,7 +363,9 @@ func readGGUFV1StringsData(llm *gguf, r io.Reader, a *array[string]) (any, error
a.values[i] = e a.values[i] = e
} else { } else {
_ = discardGGUFString(llm, r) if err := discardGGUFString(llm, r); err != nil {
return nil, err
}
} }
} }
@@ -334,7 +379,10 @@ func discardGGUFString(llm *gguf, r io.Reader) error {
return err return err
} }
size := int(llm.ByteOrder.Uint64(buf)) size, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string")
if err != nil {
return err
}
for size > 0 { for size > 0 {
n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))]) n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))])
if err != nil { if err != nil {
@@ -356,7 +404,10 @@ func readGGUFString(llm *gguf, r io.Reader) (string, error) {
return "", err return "", err
} }
length := int(llm.ByteOrder.Uint64(buf)) length, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string")
if err != nil {
return "", err
}
if length > len(llm.scratch) { if length > len(llm.scratch) {
buf = make([]byte, length) buf = make([]byte, length)
} else { } else {
@@ -394,7 +445,9 @@ func readGGUFStringsData(llm *gguf, r io.Reader, a *array[string]) (any, error)
a.values[i] = e a.values[i] = e
} else { } else {
discardGGUFString(llm, r) if err := discardGGUFString(llm, r); err != nil {
return nil, err
}
} }
} }
@@ -413,12 +466,20 @@ func (a *array[T]) MarshalJSON() ([]byte, error) {
return json.Marshal(a.values) return json.Marshal(a.values)
} }
func newArray[T any](size, maxSize int) *array[T] { func newArray[T any](size uint64, maxSize int) (*array[T], error) {
a := array[T]{size: size} if size > uint64(maxInt()) {
if maxSize < 0 || size <= maxSize { return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, maxInt())
a.values = make([]T, size)
} }
return &a if size > fsgguf.MaxArraySize {
return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, fsgguf.MaxArraySize)
}
n := int(size)
a := array[T]{size: n}
if maxSize < 0 || n <= maxSize {
a.values = make([]T, n)
}
return &a, nil
} }
func readGGUFArray(llm *gguf, r io.Reader) (any, error) { func readGGUFArray(llm *gguf, r io.Reader) (any, error) {
@@ -434,40 +495,32 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) {
switch t { switch t {
case ggufTypeUint8: case ggufTypeUint8:
a := newArray[uint8](int(n), llm.maxArraySize) return readGGUFArrayOf[uint8](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeInt8: case ggufTypeInt8:
a := newArray[int8](int(n), llm.maxArraySize) return readGGUFArrayOf[int8](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeUint16: case ggufTypeUint16:
a := newArray[uint16](int(n), llm.maxArraySize) return readGGUFArrayOf[uint16](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeInt16: case ggufTypeInt16:
a := newArray[int16](int(n), llm.maxArraySize) return readGGUFArrayOf[int16](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeUint32: case ggufTypeUint32:
a := newArray[uint32](int(n), llm.maxArraySize) return readGGUFArrayOf[uint32](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeInt32: case ggufTypeInt32:
a := newArray[int32](int(n), llm.maxArraySize) return readGGUFArrayOf[int32](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeUint64: case ggufTypeUint64:
a := newArray[uint64](int(n), llm.maxArraySize) return readGGUFArrayOf[uint64](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeInt64: case ggufTypeInt64:
a := newArray[int64](int(n), llm.maxArraySize) return readGGUFArrayOf[int64](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeFloat32: case ggufTypeFloat32:
a := newArray[float32](int(n), llm.maxArraySize) return readGGUFArrayOf[float32](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeFloat64: case ggufTypeFloat64:
a := newArray[float64](int(n), llm.maxArraySize) return readGGUFArrayOf[float64](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeBool: case ggufTypeBool:
a := newArray[bool](int(n), llm.maxArraySize) return readGGUFArrayOf[bool](llm, r, n)
return readGGUFArrayData(llm, r, a)
case ggufTypeString: case ggufTypeString:
a := newArray[string](int(n), llm.maxArraySize) a, err := newArray[string](n, llm.maxArraySize)
if err != nil {
return nil, err
}
if llm.Version == 1 { if llm.Version == 1 {
return readGGUFV1StringsData(llm, r, a) return readGGUFV1StringsData(llm, r, a)
} }
@@ -478,6 +531,14 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) {
} }
} }
func readGGUFArrayOf[T any](llm *gguf, r io.Reader, n uint64) (any, error) {
a, err := newArray[T](n, llm.maxArraySize)
if err != nil {
return nil, err
}
return readGGUFArrayData(llm, r, a)
}
func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error) { func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error) {
for i := range a.size { for i := range a.size {
e, err := readGGUF[T](llm, r) e, err := readGGUF[T](llm, r)
@@ -493,6 +554,39 @@ func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error)
return a, nil return a, nil
} }
func checkGGUFLength(n uint64, kind string) (int, error) {
if n > uint64(maxInt()) {
return 0, fmt.Errorf("GGUF %s length %d exceeds maximum %d", kind, n, maxInt())
}
if n > fsgguf.MaxStringLength {
return 0, fmt.Errorf("GGUF %s length %d exceeds maximum %d", kind, n, fsgguf.MaxStringLength)
}
return int(n), nil
}
func maxInt() int {
return int(^uint(0) >> 1)
}
func maxInt64() uint64 {
return 1<<63 - 1
}
func maxUint64() uint64 {
return ^uint64(0)
}
func checkedAdd(ns ...uint64) (uint64, bool) {
var sum uint64
for _, n := range ns {
if sum > maxUint64()-n {
return 0, false
}
sum += n
}
return sum, true
}
// writeGGUFArray writes a slice s of type E to the write with a gguf type of t // writeGGUFArray writes a slice s of type E to the write with a gguf type of t
func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error { func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil { if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil {
@@ -580,8 +674,7 @@ func WriteGGUF(f *os.File, kv fs.Config, ts []*Tensor) error {
offset += ggufPadding(offset, int64(alignment)) offset += ggufPadding(offset, int64(alignment))
var g errgroup.Group var g errgroup.Group
g.SetLimit(runtime.GOMAXPROCS(0)) g.SetLimit(min(runtime.GOMAXPROCS(0), maxGGUFWriteConcurrency))
// TODO consider reducing if tensors size * gomaxprocs is larger than free memory
for _, t := range ts { for _, t := range ts {
w := io.NewOffsetWriter(f, offset+int64(t.Offset)) w := io.NewOffsetWriter(f, offset+int64(t.Offset))
g.Go(func() error { g.Go(func() error {
+140
View File
@@ -2,12 +2,14 @@ package ggml
import ( import (
"bytes" "bytes"
"encoding/binary"
"math/rand/v2" "math/rand/v2"
"os" "os"
"strings" "strings"
"testing" "testing"
"github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp"
fsgguf "github.com/ollama/ollama/fs/gguf"
) )
func TestWriteGGUF(t *testing.T) { func TestWriteGGUF(t *testing.T) {
@@ -127,3 +129,141 @@ func TestWriteGGUF(t *testing.T) {
} }
}) })
} }
func TestDecodeRejectsMalformedGGUFMetadata(t *testing.T) {
testCases := []struct {
name string
data func() []byte
err string
}{
{
name: "v1_zero_length_string",
data: func() []byte {
var b bytes.Buffer
writeRaw(t, &b, []byte("GGUF"))
writeRaw(t, &b, uint32(1))
writeRaw(t, &b, uint32(0)) // tensors
writeRaw(t, &b, uint32(1)) // key-values
writeRaw(t, &b, uint64(0))
return b.Bytes()
},
err: "invalid GGUF v1 string length",
},
{
name: "oversized_string",
data: func() []byte {
var b bytes.Buffer
writeRawGGUFHeader(t, &b, 0, 1)
writeRaw(t, &b, uint64(fsgguf.MaxStringLength+1))
return b.Bytes()
},
err: "string length",
},
{
name: "oversized_collected_array",
data: func() []byte {
var b bytes.Buffer
writeRawGGUFHeader(t, &b, 0, 1)
writeRawGGUFString(t, &b, "tokenizer.ggml.tokens")
writeRaw(t, &b, ggufTypeArray)
writeRaw(t, &b, ggufTypeString)
writeRaw(t, &b, uint64(fsgguf.MaxArraySize+1))
return b.Bytes()
},
err: "array size",
},
{
name: "zero_alignment",
data: func() []byte {
var b bytes.Buffer
writeRawGGUFHeader(t, &b, 0, 1)
writeRawGGUFString(t, &b, "general.alignment")
writeRaw(t, &b, ggufTypeUint32)
writeRaw(t, &b, uint32(0))
return b.Bytes()
},
err: "invalid GGUF alignment",
},
{
name: "tensor_elements_overflow",
data: func() []byte {
var b bytes.Buffer
writeRawGGUFHeader(t, &b, 1, 0)
writeRawGGUFString(t, &b, "bad.weight")
writeRaw(t, &b, uint32(2))
writeRaw(t, &b, ^uint64(0))
writeRaw(t, &b, uint64(2))
writeRaw(t, &b, uint32(TensorTypeF32))
writeRaw(t, &b, uint64(0))
return b.Bytes()
},
err: "elements overflow",
},
{
name: "too_many_tensor_dimensions",
data: func() []byte {
var b bytes.Buffer
writeRawGGUFHeader(t, &b, 1, 0)
writeRawGGUFString(t, &b, "bad.weight")
writeRaw(t, &b, uint32(fsgguf.MaxTensorDims+1))
return b.Bytes()
},
err: "dimensions",
},
{
name: "tensor_row_not_multiple_of_block_size",
data: func() []byte {
var b bytes.Buffer
writeRawGGUFHeader(t, &b, 1, 0)
writeRawGGUFString(t, &b, "bad.weight")
writeRaw(t, &b, uint32(1))
writeRaw(t, &b, uint64(31))
writeRaw(t, &b, uint32(TensorTypeQ4_0))
writeRaw(t, &b, uint64(0))
return b.Bytes()
},
err: "size overflow",
},
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
defer func() {
if r := recover(); r != nil {
t.Fatalf("Decode panicked: %v", r)
}
}()
_, err := Decode(bytes.NewReader(tt.data()), -1)
if err == nil {
t.Fatal("Decode unexpectedly succeeded")
}
if !strings.Contains(err.Error(), tt.err) {
t.Fatalf("Decode error = %q, want containing %q", err, tt.err)
}
})
}
}
func writeRawGGUFHeader(t *testing.T, b *bytes.Buffer, tensors, kv uint64) {
t.Helper()
writeRaw(t, b, []byte("GGUF"))
writeRaw(t, b, uint32(3))
writeRaw(t, b, tensors)
writeRaw(t, b, kv)
}
func writeRawGGUFString(t *testing.T, b *bytes.Buffer, s string) {
t.Helper()
writeRaw(t, b, uint64(len(s)))
writeRaw(t, b, []byte(s))
}
func writeRaw(t *testing.T, b *bytes.Buffer, v any) {
t.Helper()
if err := binary.Write(b, binary.LittleEndian, v); err != nil {
t.Fatal(err)
}
}
+108 -12
View File
@@ -13,6 +13,13 @@ import (
"strings" "strings"
) )
const (
MaxStringLength = 16 << 20
MaxArraySize = 64 << 20
MaxTensorDims = 4
)
const ( const (
typeUint8 uint32 = iota typeUint8 uint32 = iota
typeInt8 typeInt8
@@ -78,6 +85,9 @@ func Open(path string) (f *File, err error) {
offset := f.reader.offset offset := f.reader.offset
alignment := cmp.Or(f.KeyValue("general.alignment").Int(), 32) alignment := cmp.Or(f.KeyValue("general.alignment").Int(), 32)
if alignment <= 0 {
return fmt.Errorf("%w alignment %d", ErrUnsupported, alignment)
}
f.offset = offset + (alignment-offset%alignment)%alignment f.offset = offset + (alignment-offset%alignment)%alignment
return nil return nil
} }
@@ -100,6 +110,9 @@ func (f *File) readTensor() (TensorInfo, error) {
if err != nil { if err != nil {
return TensorInfo{}, err return TensorInfo{}, err
} }
if dims > MaxTensorDims {
return TensorInfo{}, fmt.Errorf("%w tensor dimensions %d exceeds maximum %d", ErrUnsupported, dims, MaxTensorDims)
}
shape := make([]uint64, dims) shape := make([]uint64, dims)
for i := range dims { for i := range dims {
@@ -119,12 +132,16 @@ func (f *File) readTensor() (TensorInfo, error) {
return TensorInfo{}, err return TensorInfo{}, err
} }
return TensorInfo{ ti := TensorInfo{
Name: name, Name: name,
Offset: offset, Offset: offset,
Shape: shape, Shape: shape,
Type: TensorType(type_), Type: TensorType(type_),
}, nil }
if _, ok := ti.numBytes(); !ok {
return TensorInfo{}, fmt.Errorf("%w tensor %q size overflows", ErrUnsupported, ti.Name)
}
return ti, nil
} }
func (f *File) readKeyValue() (KeyValue, error) { func (f *File) readKeyValue() (KeyValue, error) {
@@ -191,11 +208,16 @@ func readString(f *File) (string, error) {
return "", err return "", err
} }
if int(n) > len(f.bts) { length, err := checkedLength(n, "string", MaxStringLength)
f.bts = make([]byte, n) if err != nil {
return "", err
} }
bts := f.bts[:n] if length > len(f.bts) {
f.bts = make([]byte, length)
}
bts := f.bts[:length]
if _, err := io.ReadFull(f.reader, bts); err != nil { if _, err := io.ReadFull(f.reader, bts); err != nil {
return "", err return "", err
} }
@@ -246,8 +268,13 @@ func readArray(f *File) (any, error) {
} }
func readArrayData[T any](f *File, n uint64) (s []T, err error) { func readArrayData[T any](f *File, n uint64) (s []T, err error) {
s = make([]T, n) size, err := checkedLength(n, "array size", MaxArraySize)
for i := range n { if err != nil {
return nil, err
}
s = make([]T, size)
for i := range size {
e, err := read[T](f) e, err := read[T](f)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -260,8 +287,13 @@ func readArrayData[T any](f *File, n uint64) (s []T, err error) {
} }
func readArrayString(f *File, n uint64) (s []string, err error) { func readArrayString(f *File, n uint64) (s []string, err error) {
s = make([]string, n) size, err := checkedLength(n, "array size", MaxArraySize)
for i := range n { if err != nil {
return nil, err
}
s = make([]string, size)
for i := range size {
e, err := readString(f) e, err := readString(f)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -273,6 +305,24 @@ func readArrayString(f *File, n uint64) (s []string, err error) {
return s, nil return s, nil
} }
func checkedLength(n uint64, kind string, max uint64) (int, error) {
if n > uint64(maxInt()) {
return 0, fmt.Errorf("%s %d exceeds maximum %d", kind, n, maxInt())
}
if n > max {
return 0, fmt.Errorf("%s %d exceeds maximum %d", kind, n, max)
}
return int(n), nil
}
func maxInt() int {
return int(^uint(0) >> 1)
}
func maxInt64() uint64 {
return 1<<63 - 1
}
func (f *File) Close() error { func (f *File) Close() error {
f.keyValues.stop() f.keyValues.stop()
f.tensors.stop() f.tensors.stop()
@@ -337,11 +387,57 @@ func (f *File) TensorInfos() iter.Seq2[int, TensorInfo] {
func (f *File) TensorReader(name string) (TensorInfo, io.Reader, error) { func (f *File) TensorReader(name string) (TensorInfo, io.Reader, error) {
t := f.TensorInfo(name) t := f.TensorInfo(name)
if t.NumBytes() == 0 { if err := f.Err(); err != nil {
return TensorInfo{}, nil, err
}
if t.Name == "" {
return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name)
}
numBytes, ok := t.numBytes()
if !ok {
return TensorInfo{}, nil, fmt.Errorf("%w tensor %q size overflows", ErrUnsupported, t.Name)
}
if numBytes == 0 {
return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name) return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name)
} }
// fast forward through tensor info if we haven't already // fast forward through tensor info if we haven't already
_ = f.tensors.rest() f.tensors.rest()
return t, io.NewSectionReader(f.file, f.offset+int64(t.Offset), t.NumBytes()), nil if err := f.Err(); err != nil {
return TensorInfo{}, nil, err
}
if t.Offset > maxInt64() {
return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset %d exceeds maximum %d", ErrUnsupported, t.Name, t.Offset, maxInt64())
}
offset := f.offset + int64(t.Offset)
if offset < f.offset {
return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset overflows", ErrUnsupported, t.Name)
}
fileInfo, err := f.file.Stat()
if err != nil {
return TensorInfo{}, nil, err
}
if numBytes > fileInfo.Size()-offset {
return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset+size exceeds file size", ErrUnsupported, t.Name)
}
return t, io.NewSectionReader(f.file, offset, numBytes), nil
}
func (f *File) Err() error {
// Key/value and tensor metadata are read lazily, so parse errors can surface
// after Open succeeds.
if f.keyValues != nil {
if err := f.keyValues.Err(); err != nil {
return err
}
}
if f.tensors != nil {
if err := f.tensors.Err(); err != nil {
return err
}
}
return nil
} }
+158
View File
@@ -0,0 +1,158 @@
package gguf
import (
"bytes"
"encoding/binary"
"os"
"strings"
"testing"
)
func TestReadStringRejectsOversizedLength(t *testing.T) {
var b bytes.Buffer
writeInternalRaw(t, &b, uint64(MaxStringLength+1))
_, err := readString(testFile(b.Bytes()))
if err == nil {
t.Fatal("readString unexpectedly succeeded")
}
if !strings.Contains(err.Error(), "string") {
t.Fatalf("readString error = %q, want string length error", err)
}
}
func TestReadArrayRejectsOversizedCollectedArray(t *testing.T) {
var b bytes.Buffer
writeInternalRaw(t, &b, typeString)
writeInternalRaw(t, &b, uint64(MaxArraySize+1))
_, err := readArray(testFile(b.Bytes()))
if err == nil {
t.Fatal("readArray unexpectedly succeeded")
}
if !strings.Contains(err.Error(), "array size") {
t.Fatalf("readArray error = %q, want array size error", err)
}
}
func TestTensorInfoRejectsRowNotMultipleOfBlockSize(t *testing.T) {
ti := TensorInfo{
Name: "bad.weight",
Shape: []uint64{31},
Type: TensorTypeQ4_0,
}
if _, ok := ti.numBytes(); ok {
t.Fatal("numBytes unexpectedly succeeded")
}
}
func TestTensorReaderReturnsLazyParseError(t *testing.T) {
var b bytes.Buffer
writeInternalRaw(t, &b, []byte("GGUF"))
writeInternalRaw(t, &b, uint32(3))
writeInternalRaw(t, &b, uint64(1)) // tensors
writeInternalRaw(t, &b, uint64(0)) // key-values
writeInternalString(t, &b, "bad.weight")
writeInternalRaw(t, &b, uint32(MaxTensorDims+1))
p := writeTempFile(t, b.Bytes())
f, err := Open(p)
if err != nil {
t.Fatal(err)
}
defer f.Close()
_, _, err = f.TensorReader("bad.weight")
if err == nil {
t.Fatal("TensorReader unexpectedly succeeded")
}
if !strings.Contains(err.Error(), "dimensions") {
t.Fatalf("TensorReader error = %q, want dimensions error", err)
}
}
func TestTensorReaderRejectsInvalidOffset(t *testing.T) {
var b bytes.Buffer
writeInternalRaw(t, &b, []byte("GGUF"))
writeInternalRaw(t, &b, uint32(3))
writeInternalRaw(t, &b, uint64(1)) // tensors
writeInternalRaw(t, &b, uint64(0)) // key-values
writeInternalString(t, &b, "bad.weight")
writeInternalRaw(t, &b, uint32(1)) // dimensions
writeInternalRaw(t, &b, uint64(1))
writeInternalRaw(t, &b, uint32(TensorTypeF32))
writeInternalRaw(t, &b, ^uint64(0))
p := writeTempFile(t, b.Bytes())
f, err := Open(p)
if err != nil {
t.Fatal(err)
}
defer f.Close()
_, _, err = f.TensorReader("bad.weight")
if err == nil {
t.Fatal("TensorReader unexpectedly succeeded")
}
if !strings.Contains(err.Error(), "offset") {
t.Fatalf("TensorReader error = %q, want offset error", err)
}
}
func TestTensorReaderRejectsMissingTensor(t *testing.T) {
var b bytes.Buffer
writeInternalRaw(t, &b, []byte("GGUF"))
writeInternalRaw(t, &b, uint32(3))
writeInternalRaw(t, &b, uint64(0)) // tensors
writeInternalRaw(t, &b, uint64(0)) // key-values
p := writeTempFile(t, b.Bytes())
f, err := Open(p)
if err != nil {
t.Fatal(err)
}
defer f.Close()
_, _, err = f.TensorReader("missing.weight")
if err == nil {
t.Fatal("TensorReader unexpectedly succeeded")
}
if !strings.Contains(err.Error(), "not found") {
t.Fatalf("TensorReader error = %q, want not found error", err)
}
}
func testFile(data []byte) *File {
return &File{
reader: newBufferedReader(bytes.NewReader(data), 32<<10),
bts: make([]byte, 4096),
}
}
func writeInternalRaw(t *testing.T, b *bytes.Buffer, v any) {
t.Helper()
if err := binary.Write(b, binary.LittleEndian, v); err != nil {
t.Fatal(err)
}
}
func writeInternalString(t *testing.T, b *bytes.Buffer, s string) {
t.Helper()
writeInternalRaw(t, b, uint64(len(s)))
writeInternalRaw(t, b, []byte(s))
}
func writeTempFile(t *testing.T, data []byte) string {
t.Helper()
f, err := os.CreateTemp(t.TempDir(), "")
if err != nil {
t.Fatal(err)
}
defer f.Close()
if _, err := f.Write(data); err != nil {
t.Fatal(err)
}
return f.Name()
}
+19 -7
View File
@@ -2,8 +2,8 @@ package gguf
import ( import (
"encoding/binary" "encoding/binary"
"fmt"
"iter" "iter"
"log/slog"
) )
type lazy[T any] struct { type lazy[T any] struct {
@@ -11,6 +11,7 @@ type lazy[T any] struct {
next func() (T, bool) next func() (T, bool)
stop func() stop func()
values []T values []T
err error
// successFunc is called when all values have been successfully read. // successFunc is called when all values have been successfully read.
successFunc func() error successFunc func() error
@@ -21,13 +22,16 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) {
if err := binary.Read(f.reader, binary.LittleEndian, &it.count); err != nil { if err := binary.Read(f.reader, binary.LittleEndian, &it.count); err != nil {
return nil, err return nil, err
} }
if it.count > uint64(maxInt()) {
return nil, fmt.Errorf("GGUF item count %d exceeds maximum %d", it.count, maxInt())
}
it.values = make([]T, 0) it.values = make([]T, 0)
it.next, it.stop = iter.Pull(func(yield func(T) bool) { it.next, it.stop = iter.Pull(func(yield func(T) bool) {
for i := range it.count { for i := range it.count {
t, err := fn() t, err := fn()
if err != nil { if err != nil {
slog.Error("error reading tensor", "index", i, "error", err) it.err = fmt.Errorf("error reading GGUF item %d: %w", i, err)
return return
} }
@@ -38,7 +42,10 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) {
} }
if it.successFunc != nil { if it.successFunc != nil {
it.successFunc() if err := it.successFunc(); err != nil {
it.err = err
return
}
} }
}) })
@@ -57,9 +64,10 @@ func (g *lazy[T]) Values() iter.Seq[T] {
func (g *lazy[T]) All() iter.Seq2[int, T] { func (g *lazy[T]) All() iter.Seq2[int, T] {
return func(yield func(int, T) bool) { return func(yield func(int, T) bool) {
for i := range int(g.count) { for i := range g.count {
if i < len(g.values) { n := int(i)
if !yield(i, g.values[i]) { if n < len(g.values) {
if !yield(n, g.values[n]) {
break break
} }
} else { } else {
@@ -68,7 +76,7 @@ func (g *lazy[T]) All() iter.Seq2[int, T] {
break break
} }
if !yield(i, t) { if !yield(n, t) {
break break
} }
} }
@@ -87,3 +95,7 @@ func (g *lazy[T]) rest() (collected bool) {
return collected return collected
} }
func (g *lazy[T]) Err() error {
return g.err
}
+46
View File
@@ -24,11 +24,57 @@ func (ti TensorInfo) NumValues() int64 {
return numItems return numItems
} }
func (ti TensorInfo) numValues() (int64, bool) {
var numItems int64 = 1
for _, dim := range ti.Shape {
if dim > maxInt64() {
return 0, false
}
n := int64(dim)
if n != 0 && numItems > int64(maxInt64())/n {
return 0, false
}
numItems *= n
}
return numItems, true
}
// NumBytes returns the number of bytes in the tensor. // NumBytes returns the number of bytes in the tensor.
func (ti TensorInfo) NumBytes() int64 { func (ti TensorInfo) NumBytes() int64 {
return int64(float64(ti.NumValues()) * ti.Type.NumBytes()) return int64(float64(ti.NumValues()) * ti.Type.NumBytes())
} }
func (ti TensorInfo) numBytes() (int64, bool) {
numValues, ok := ti.numValues()
if !ok {
return 0, false
}
typeSize := ti.Type.typeSize()
blockSize := ti.Type.blockSize()
if typeSize == 0 || blockSize == 0 {
return 0, false
}
rowSize := int64(1)
if len(ti.Shape) > 0 {
if ti.Shape[0] > maxInt64() {
return 0, false
}
rowSize = int64(ti.Shape[0])
}
if rowSize%blockSize != 0 {
return 0, false
}
blocks := numValues / blockSize
if blocks > int64(maxInt64())/typeSize {
return 0, false
}
return blocks * typeSize, true
}
func (ti TensorInfo) LogValue() slog.Value { func (ti TensorInfo) LogValue() slog.Value {
return slog.GroupValue( return slog.GroupValue(
slog.String("name", ti.Name), slog.String("name", ti.Name),
+50 -14
View File
@@ -11,6 +11,7 @@ import (
"io/fs" "io/fs"
"log/slog" "log/slog"
"maps" "maps"
"math"
"net" "net"
"net/http" "net/http"
"net/url" "net/url"
@@ -268,34 +269,65 @@ func (s *Server) CreateHandler(c *gin.Context) {
} }
} }
strFromInfo := func(k string) string { strFromInfo := func(k string) (string, error) {
v, ok := r.Info[k] v, ok := r.Info[k]
if ok { if ok {
val := v.(string) val, ok := v.(string)
return val if !ok {
return "", fmt.Errorf("info field %q must be a string", k)
}
return val, nil
} }
return "" return "", nil
} }
vFromInfo := func(k string) float64 { intFromInfo := func(k string) (int, error) {
v, ok := r.Info[k] v, ok := r.Info[k]
if ok { if ok {
val := v.(float64) val, ok := v.(float64)
return val if !ok {
return 0, fmt.Errorf("info field %q must be a number", k)
}
if val < 0 || math.Trunc(val) != val || val > float64(maxCreateInfoInt()) {
return 0, fmt.Errorf("info field %q must be a non-negative integer", k)
}
return int(val), nil
} }
return 0 return 0, nil
} }
config.ModelFamily = strFromInfo("model_family") if config.ModelFamily, err = strFromInfo("model_family"); err != nil {
ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
return
}
if config.ModelFamily != "" { if config.ModelFamily != "" {
config.ModelFamilies = []string{config.ModelFamily} config.ModelFamilies = []string{config.ModelFamily}
} }
config.BaseName = strFromInfo("base_name") if config.BaseName, err = strFromInfo("base_name"); err != nil {
config.FileType = strFromInfo("quantization_level") ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
config.ModelType = strFromInfo("parameter_size") return
config.ContextLen = int(vFromInfo("context_length")) }
config.EmbedLen = int(vFromInfo("embedding_length")) if config.FileType, err = strFromInfo("quantization_level"); err != nil {
ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
return
}
if config.ModelType, err = strFromInfo("parameter_size"); err != nil {
ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
return
}
contextLen, err := intFromInfo("context_length")
if err != nil {
ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
return
}
config.ContextLen = contextLen
embedLen, err := intFromInfo("embedding_length")
if err != nil {
ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest}
return
}
config.EmbedLen = embedLen
} }
if err := createModel(r, name, baseLayers, config, fn); err != nil { if err := createModel(r, name, baseLayers, config, fn); err != nil {
@@ -440,6 +472,10 @@ func convertModelFromFilesWithMediaType(files map[string]string, baseLayers []*l
} }
} }
func maxCreateInfoInt() int {
return int(^uint(0) >> 1)
}
func detectModelTypeFromFiles(files map[string]string) string { func detectModelTypeFromFiles(files map[string]string) string {
for fn := range files { for fn := range files {
if strings.HasSuffix(fn, ".safetensors") { if strings.HasSuffix(fn, ".safetensors") {
+38 -5
View File
@@ -17,6 +17,7 @@ import (
"github.com/ollama/ollama/api" "github.com/ollama/ollama/api"
"github.com/ollama/ollama/fs/ggml" "github.com/ollama/ollama/fs/ggml"
fsgguf "github.com/ollama/ollama/fs/gguf"
"github.com/ollama/ollama/manifest" "github.com/ollama/ollama/manifest"
"github.com/ollama/ollama/model/parsers" "github.com/ollama/ollama/model/parsers"
ollamatemplate "github.com/ollama/ollama/template" ollamatemplate "github.com/ollama/ollama/template"
@@ -707,6 +708,10 @@ func skipModelListGGUFValue(r io.Reader, byteOrder binary.ByteOrder, version uin
} }
func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uint32, arrayType uint32, count uint64) error { func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uint32, arrayType uint32, count uint64) error {
if _, err := checkedModelListGGUFLength(count, "array size", fsgguf.MaxArraySize); err != nil {
return err
}
var size uint64 var size uint64
switch arrayType { switch arrayType {
case modelListGGUFTypeUint8, modelListGGUFTypeInt8, modelListGGUFTypeBool: case modelListGGUFTypeUint8, modelListGGUFTypeInt8, modelListGGUFTypeBool:
@@ -727,6 +732,10 @@ func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uin
default: default:
return fmt.Errorf("unsupported gguf array type %d", arrayType) return fmt.Errorf("unsupported gguf array type %d", arrayType)
} }
if count > uint64(maxModelListGGUFInt64())/size {
return fmt.Errorf("gguf array byte length %d*%d exceeds maximum %d", count, size, maxModelListGGUFInt64())
}
return discardModelListGGUFBytes(r, int64(count*size)) return discardModelListGGUFBytes(r, int64(count*size))
} }
@@ -736,15 +745,16 @@ func readModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version ui
return "", err return "", err
} }
if length == 0 { n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength)
return "", nil if err != nil {
return "", err
} }
bts := make([]byte, length) bts := make([]byte, n)
if _, err := io.ReadFull(r, bts); err != nil { if _, err := io.ReadFull(r, bts); err != nil {
return "", err return "", err
} }
if version == 1 && bts[len(bts)-1] == 0 { if version == 1 && len(bts) > 0 && bts[len(bts)-1] == 0 {
bts = bts[:len(bts)-1] bts = bts[:len(bts)-1]
} }
return string(bts), nil return string(bts), nil
@@ -755,7 +765,12 @@ func skipModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version ui
if err := binary.Read(r, byteOrder, &length); err != nil { if err := binary.Read(r, byteOrder, &length); err != nil {
return err return err
} }
return discardModelListGGUFBytes(r, int64(length))
n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength)
if err != nil {
return err
}
return discardModelListGGUFBytes(r, int64(n))
} }
func discardModelListGGUFBytes(r io.Reader, n int64) error { func discardModelListGGUFBytes(r io.Reader, n int64) error {
@@ -766,6 +781,24 @@ func discardModelListGGUFBytes(r io.Reader, n int64) error {
return err return err
} }
func checkedModelListGGUFLength(n uint64, kind string, max uint64) (int, error) {
if n > uint64(maxModelListGGUFInt()) {
return 0, fmt.Errorf("gguf %s %d exceeds maximum %d", kind, n, maxModelListGGUFInt())
}
if n > max {
return 0, fmt.Errorf("gguf %s %d exceeds maximum %d", kind, n, max)
}
return int(n), nil
}
func maxModelListGGUFInt() int {
return int(^uint(0) >> 1)
}
func maxModelListGGUFInt64() int64 {
return 1<<63 - 1
}
func appendModelListCapabilities(capabilities []model.Capability, values ...model.Capability) []model.Capability { func appendModelListCapabilities(capabilities []model.Capability, values ...model.Capability) []model.Capability {
for _, capability := range values { for _, capability := range values {
capabilities = appendModelListCapability(capabilities, capability) capabilities = appendModelListCapability(capabilities, capability)
+102
View File
@@ -1,15 +1,20 @@
package server package server
import ( import (
"bytes"
"context" "context"
"encoding/binary"
"errors" "errors"
"net/http" "net/http"
"os"
"slices" "slices"
"strings"
"testing" "testing"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/ollama/ollama/api" "github.com/ollama/ollama/api"
fsgguf "github.com/ollama/ollama/fs/gguf"
"github.com/ollama/ollama/manifest" "github.com/ollama/ollama/manifest"
"github.com/ollama/ollama/types/model" "github.com/ollama/ollama/types/model"
) )
@@ -218,6 +223,67 @@ func TestModelListCacheSyncDropsStaleEntryOnRefreshFailure(t *testing.T) {
} }
} }
func TestReadModelListGGUFRejectsMalformedMetadata(t *testing.T) {
cases := []struct {
name string
data []byte
want string
}{
{
name: "oversized key string",
data: modelListGGUFTestFile(func(b *bytes.Buffer) {
writeModelListGGUFHeader(t, b, 1)
writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1)
}),
want: "string",
},
{
name: "oversized skipped string",
data: modelListGGUFTestFile(func(b *bytes.Buffer) {
writeModelListGGUFHeader(t, b, 1)
writeModelListGGUFString(t, b, "unused")
writeModelListGGUFUint32(t, b, modelListGGUFTypeString)
writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1)
}),
want: "string",
},
{
name: "oversized skipped array",
data: modelListGGUFTestFile(func(b *bytes.Buffer) {
writeModelListGGUFHeader(t, b, 1)
writeModelListGGUFString(t, b, "unused")
writeModelListGGUFUint32(t, b, modelListGGUFTypeArray)
writeModelListGGUFUint32(t, b, modelListGGUFTypeUint8)
writeModelListGGUFUint64(t, b, fsgguf.MaxArraySize+1)
}),
want: "array size",
},
}
for _, tt := range cases {
t.Run(tt.name, func(t *testing.T) {
defer func() {
if r := recover(); r != nil {
t.Fatalf("readModelListGGUF panicked: %v", r)
}
}()
path := t.TempDir() + "/model.gguf"
if err := os.WriteFile(path, tt.data, 0o600); err != nil {
t.Fatal(err)
}
_, err := readModelListGGUF(path)
if err == nil {
t.Fatal("expected error")
}
if !strings.Contains(err.Error(), tt.want) {
t.Fatalf("error = %v, want substring %q", err, tt.want)
}
})
}
}
func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl string) { func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl string) {
t.Helper() t.Helper()
_, digest := createBinFile(t, kv, nil) _, digest := createBinFile(t, kv, nil)
@@ -237,3 +303,39 @@ func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl str
t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String()) t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String())
} }
} }
func modelListGGUFTestFile(fn func(*bytes.Buffer)) []byte {
var b bytes.Buffer
fn(&b)
return b.Bytes()
}
func writeModelListGGUFHeader(t *testing.T, b *bytes.Buffer, numKV uint64) {
t.Helper()
writeModelListGGUFUint32(t, b, modelListGGUFMagicLE)
writeModelListGGUFUint32(t, b, 3)
writeModelListGGUFUint64(t, b, 0)
writeModelListGGUFUint64(t, b, numKV)
}
func writeModelListGGUFString(t *testing.T, b *bytes.Buffer, s string) {
t.Helper()
writeModelListGGUFUint64(t, b, uint64(len(s)))
if _, err := b.WriteString(s); err != nil {
t.Fatal(err)
}
}
func writeModelListGGUFUint32(t *testing.T, b *bytes.Buffer, v uint32) {
t.Helper()
if err := binary.Write(b, binary.LittleEndian, v); err != nil {
t.Fatal(err)
}
}
func writeModelListGGUFUint64(t *testing.T, b *bytes.Buffer, v uint64) {
t.Helper()
if err := binary.Write(b, binary.LittleEndian, v); err != nil {
t.Fatal(err)
}
}
+52
View File
@@ -102,6 +102,58 @@ func createRequest(t *testing.T, fn func(*gin.Context), body any) *httptest.Resp
return w.ResponseRecorder return w.ResponseRecorder
} }
func TestCreateHandlerRejectsInvalidInfoTypes(t *testing.T) {
gin.SetMode(gin.TestMode)
testCases := []struct {
name string
info map[string]any
err string
}{
{
name: "string_field",
info: map[string]any{
"model_family": float64(1),
},
err: "model_family",
},
{
name: "fractional_integer_field",
info: map[string]any{
"context_length": 1.5,
},
err: "context_length",
},
}
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
var s Server
_, digest := createBinFile(t, nil, nil)
w := createRequest(t, s.CreateHandler, api.CreateRequest{
Model: "test-create-invalid-info",
Files: map[string]string{
"test.gguf": digest,
},
Info: tt.info,
Stream: &stream,
})
if w.Code != http.StatusBadRequest {
t.Fatalf("expected status code 400, got %d", w.Code)
}
var resp map[string]string
if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
t.Fatal(err)
}
if got := resp["error"]; !strings.Contains(got, tt.err) {
t.Fatalf("expected %s error, got %q", tt.err, got)
}
})
}
}
func readCreatedModelConfig(t *testing.T, name string) model.ConfigV2 { func readCreatedModelConfig(t *testing.T, name string) model.ConfigV2 {
t.Helper() t.Helper()