diff --git a/fs/ggml/ggml.go b/fs/ggml/ggml.go index 4b4fcb16..7d83b93a 100644 --- a/fs/ggml/ggml.go +++ b/fs/ggml/ggml.go @@ -520,10 +520,49 @@ func (t Tensor) Elements() uint64 { 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 { 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 { return TensorType(t.Kind).String() } diff --git a/fs/ggml/gguf.go b/fs/ggml/gguf.go index 33560182..07c26182 100644 --- a/fs/ggml/gguf.go +++ b/fs/ggml/gguf.go @@ -14,9 +14,14 @@ import ( "strings" "github.com/ollama/ollama/fs" + fsgguf "github.com/ollama/ollama/fs/gguf" "golang.org/x/sync/errgroup" ) +const ( + maxGGUFWriteConcurrency = 2 +) + type containerGGUF struct { ByteOrder binary.ByteOrder @@ -202,6 +207,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { if err != nil { 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) 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.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 llm.kv["general.parameter_count"] = llm.parameters alignment := llm.kv.Uint("general.alignment", 32) + if alignment == 0 { + return fmt.Errorf("invalid GGUF alignment: 0") + } offset, err := rs.Seek(0, io.SeekCurrent) if err != nil { @@ -243,6 +261,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { } padding := ggufPadding(offset, int64(alignment)) + if padding > int64(maxInt64())-offset { + return fmt.Errorf("GGUF tensor offset overflow") + } llm.tensorOffset = uint64(offset + padding) // get file size to validate tensor bounds @@ -256,7 +277,15 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { } 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) { 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) } - 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) } } @@ -298,11 +330,22 @@ func readGGUFV1String(llm *gguf, r io.Reader) (string, error) { if err := binary.Read(r, llm.ByteOrder, &length); err != nil { 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 - if _, err := io.CopyN(&b, r, int64(length)); err != nil { + if _, err := io.CopyN(&b, r, int64(size)); err != nil { return "", err } + if b.Bytes()[b.Len()-1] != 0 { + return "", fmt.Errorf("invalid GGUF v1 string terminator") + } // gguf v1 strings are null-terminated 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 } 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 } - size := int(llm.ByteOrder.Uint64(buf)) + size, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string") + if err != nil { + return err + } for size > 0 { n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))]) if err != nil { @@ -356,7 +404,10 @@ func readGGUFString(llm *gguf, r io.Reader) (string, error) { 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) { buf = make([]byte, length) } else { @@ -394,7 +445,9 @@ func readGGUFStringsData(llm *gguf, r io.Reader, a *array[string]) (any, error) a.values[i] = e } 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) } -func newArray[T any](size, maxSize int) *array[T] { - a := array[T]{size: size} - if maxSize < 0 || size <= maxSize { - a.values = make([]T, size) +func newArray[T any](size uint64, maxSize int) (*array[T], error) { + if size > uint64(maxInt()) { + return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, maxInt()) } - 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) { @@ -434,40 +495,32 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) { switch t { case ggufTypeUint8: - a := newArray[uint8](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint8](llm, r, n) case ggufTypeInt8: - a := newArray[int8](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int8](llm, r, n) case ggufTypeUint16: - a := newArray[uint16](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint16](llm, r, n) case ggufTypeInt16: - a := newArray[int16](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int16](llm, r, n) case ggufTypeUint32: - a := newArray[uint32](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint32](llm, r, n) case ggufTypeInt32: - a := newArray[int32](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int32](llm, r, n) case ggufTypeUint64: - a := newArray[uint64](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint64](llm, r, n) case ggufTypeInt64: - a := newArray[int64](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int64](llm, r, n) case ggufTypeFloat32: - a := newArray[float32](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[float32](llm, r, n) case ggufTypeFloat64: - a := newArray[float64](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[float64](llm, r, n) case ggufTypeBool: - a := newArray[bool](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[bool](llm, r, n) 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 { 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) { for i := range a.size { 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 } +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 func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error { 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)) var g errgroup.Group - g.SetLimit(runtime.GOMAXPROCS(0)) - // TODO consider reducing if tensors size * gomaxprocs is larger than free memory + g.SetLimit(min(runtime.GOMAXPROCS(0), maxGGUFWriteConcurrency)) for _, t := range ts { w := io.NewOffsetWriter(f, offset+int64(t.Offset)) g.Go(func() error { diff --git a/fs/ggml/gguf_test.go b/fs/ggml/gguf_test.go index 4d114567..5e0db169 100644 --- a/fs/ggml/gguf_test.go +++ b/fs/ggml/gguf_test.go @@ -2,12 +2,14 @@ package ggml import ( "bytes" + "encoding/binary" "math/rand/v2" "os" "strings" "testing" "github.com/google/go-cmp/cmp" + fsgguf "github.com/ollama/ollama/fs/gguf" ) 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) + } +} diff --git a/fs/gguf/gguf.go b/fs/gguf/gguf.go index bbb9bb41..4468a314 100644 --- a/fs/gguf/gguf.go +++ b/fs/gguf/gguf.go @@ -13,6 +13,13 @@ import ( "strings" ) +const ( + MaxStringLength = 16 << 20 + MaxArraySize = 64 << 20 + + MaxTensorDims = 4 +) + const ( typeUint8 uint32 = iota typeInt8 @@ -78,6 +85,9 @@ func Open(path string) (f *File, err error) { offset := f.reader.offset 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 return nil } @@ -100,6 +110,9 @@ func (f *File) readTensor() (TensorInfo, error) { if err != nil { return TensorInfo{}, err } + if dims > MaxTensorDims { + return TensorInfo{}, fmt.Errorf("%w tensor dimensions %d exceeds maximum %d", ErrUnsupported, dims, MaxTensorDims) + } shape := make([]uint64, dims) for i := range dims { @@ -119,12 +132,16 @@ func (f *File) readTensor() (TensorInfo, error) { return TensorInfo{}, err } - return TensorInfo{ + ti := TensorInfo{ Name: name, Offset: offset, Shape: shape, 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) { @@ -191,11 +208,16 @@ func readString(f *File) (string, error) { return "", err } - if int(n) > len(f.bts) { - f.bts = make([]byte, n) + length, err := checkedLength(n, "string", MaxStringLength) + 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 { return "", err } @@ -246,8 +268,13 @@ func readArray(f *File) (any, error) { } func readArrayData[T any](f *File, n uint64) (s []T, err error) { - s = make([]T, n) - for i := range n { + size, err := checkedLength(n, "array size", MaxArraySize) + if err != nil { + return nil, err + } + + s = make([]T, size) + for i := range size { e, err := read[T](f) if err != nil { 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) { - s = make([]string, n) - for i := range n { + size, err := checkedLength(n, "array size", MaxArraySize) + if err != nil { + return nil, err + } + + s = make([]string, size) + for i := range size { e, err := readString(f) if err != nil { return nil, err @@ -273,6 +305,24 @@ func readArrayString(f *File, n uint64) (s []string, err error) { 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 { f.keyValues.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) { 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) } // fast forward through tensor info if we haven't already - _ = f.tensors.rest() - return t, io.NewSectionReader(f.file, f.offset+int64(t.Offset), t.NumBytes()), nil + f.tensors.rest() + 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 } diff --git a/fs/gguf/gguf_internal_test.go b/fs/gguf/gguf_internal_test.go new file mode 100644 index 00000000..3bf0dcf9 --- /dev/null +++ b/fs/gguf/gguf_internal_test.go @@ -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() +} diff --git a/fs/gguf/lazy.go b/fs/gguf/lazy.go index 16ab9909..c2c2423d 100644 --- a/fs/gguf/lazy.go +++ b/fs/gguf/lazy.go @@ -2,8 +2,8 @@ package gguf import ( "encoding/binary" + "fmt" "iter" - "log/slog" ) type lazy[T any] struct { @@ -11,6 +11,7 @@ type lazy[T any] struct { next func() (T, bool) stop func() values []T + err error // successFunc is called when all values have been successfully read. 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 { 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.next, it.stop = iter.Pull(func(yield func(T) bool) { for i := range it.count { t, err := fn() 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 } @@ -38,7 +42,10 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) { } 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] { return func(yield func(int, T) bool) { - for i := range int(g.count) { - if i < len(g.values) { - if !yield(i, g.values[i]) { + for i := range g.count { + n := int(i) + if n < len(g.values) { + if !yield(n, g.values[n]) { break } } else { @@ -68,7 +76,7 @@ func (g *lazy[T]) All() iter.Seq2[int, T] { break } - if !yield(i, t) { + if !yield(n, t) { break } } @@ -87,3 +95,7 @@ func (g *lazy[T]) rest() (collected bool) { return collected } + +func (g *lazy[T]) Err() error { + return g.err +} diff --git a/fs/gguf/tensor.go b/fs/gguf/tensor.go index 194c1d73..2b17ce03 100644 --- a/fs/gguf/tensor.go +++ b/fs/gguf/tensor.go @@ -24,11 +24,57 @@ func (ti TensorInfo) NumValues() int64 { 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. func (ti TensorInfo) NumBytes() int64 { 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 { return slog.GroupValue( slog.String("name", ti.Name), diff --git a/server/create.go b/server/create.go index fb07e725..7d3ce33b 100644 --- a/server/create.go +++ b/server/create.go @@ -11,6 +11,7 @@ import ( "io/fs" "log/slog" "maps" + "math" "net" "net/http" "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] if ok { - val := v.(string) - return val + val, ok := v.(string) + 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] if ok { - val := v.(float64) - return val + val, ok := v.(float64) + 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 != "" { config.ModelFamilies = []string{config.ModelFamily} } - config.BaseName = strFromInfo("base_name") - config.FileType = strFromInfo("quantization_level") - config.ModelType = strFromInfo("parameter_size") - config.ContextLen = int(vFromInfo("context_length")) - config.EmbedLen = int(vFromInfo("embedding_length")) + if config.BaseName, err = strFromInfo("base_name"); err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } + 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 { @@ -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 { for fn := range files { if strings.HasSuffix(fn, ".safetensors") { diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 1069afdc..b110659d 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -17,6 +17,7 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/fs/ggml" + fsgguf "github.com/ollama/ollama/fs/gguf" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/model/parsers" 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 { + if _, err := checkedModelListGGUFLength(count, "array size", fsgguf.MaxArraySize); err != nil { + return err + } + var size uint64 switch arrayType { case modelListGGUFTypeUint8, modelListGGUFTypeInt8, modelListGGUFTypeBool: @@ -727,6 +732,10 @@ func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uin default: 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)) } @@ -736,15 +745,16 @@ func readModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version ui return "", err } - if length == 0 { - return "", nil + n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength) + if err != nil { + return "", err } - bts := make([]byte, length) + bts := make([]byte, n) if _, err := io.ReadFull(r, bts); err != nil { 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] } 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 { 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 { @@ -766,6 +781,24 @@ func discardModelListGGUFBytes(r io.Reader, n int64) error { 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 { for _, capability := range values { capabilities = appendModelListCapability(capabilities, capability) diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go index 1f40c7b3..047ecd81 100644 --- a/server/model_list_cache_test.go +++ b/server/model_list_cache_test.go @@ -1,15 +1,20 @@ package server import ( + "bytes" "context" + "encoding/binary" "errors" "net/http" + "os" "slices" + "strings" "testing" "github.com/gin-gonic/gin" "github.com/ollama/ollama/api" + fsgguf "github.com/ollama/ollama/fs/gguf" "github.com/ollama/ollama/manifest" "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) { t.Helper() _, 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()) } } + +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) + } +} diff --git a/server/routes_create_test.go b/server/routes_create_test.go index 1a6205b7..01705d27 100644 --- a/server/routes_create_test.go +++ b/server/routes_create_test.go @@ -102,6 +102,58 @@ func createRequest(t *testing.T, fn func(*gin.Context), body any) *httptest.Resp 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 { t.Helper()