diff --git a/arrow/ipc/file_reader.go b/arrow/ipc/file_reader.go index 605b768e7..aa4e37706 100644 --- a/arrow/ipc/file_reader.go +++ b/arrow/ipc/file_reader.go @@ -22,6 +22,7 @@ import ( "errors" "fmt" "io" + "math" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" @@ -58,7 +59,39 @@ const footerSizeLen = 4 var minimumOffsetSize = int64(len(Magic)*2 + footerSizeLen) type basicReaderImpl struct { - r ReadAtSeeker + r ReadAtSeeker + maxMetadataSize int64 + maxBodySize int64 +} + +func validateFileBlock(offset int64, meta int32, body, fileSize, maxMetadataSize, maxBodySize int64) error { + if offset < 0 { + return fmt.Errorf("arrow/ipc: invalid file block offset %d", offset) + } + if meta < 4 { + return fmt.Errorf("arrow/ipc: invalid file block metadata length %d", meta) + } + if body < 0 { + return fmt.Errorf("arrow/ipc: invalid file block body length %d", body) + } + if maxMetadataSize > 0 && int64(meta) > maxMetadataSize { + return fmt.Errorf("arrow/ipc: file block metadata length %d exceeds limit %d", meta, maxMetadataSize) + } + if maxBodySize > 0 && body > maxBodySize { + return fmt.Errorf("arrow/ipc: file block body length %d exceeds limit %d", body, maxBodySize) + } + maxInt := int64(^uint(0) >> 1) + if body > maxInt { + return fmt.Errorf("arrow/ipc: file block body length %d is not addressable", body) + } + if body > math.MaxInt64-int64(meta) { + return fmt.Errorf("arrow/ipc: file block length overflows: metadata=%d body=%d", meta, body) + } + blockLen := int64(meta) + body + if offset > fileSize || blockLen > fileSize-offset { + return fmt.Errorf("arrow/ipc: file block at offset %d with length %d exceeds file size %d", offset, blockLen, fileSize) + } + return nil } func (r *basicReaderImpl) getBytes(offset, len int64) ([]byte, error) { @@ -82,6 +115,9 @@ func (r *basicReaderImpl) block(mem memory.Allocator, f *footerBlock, i int) (da if !f.data.RecordBatches(&blk, i) { return fileBlock{}, fmt.Errorf("arrow/ipc: could not extract file block %d", i) } + if err := validateFileBlock(blk.Offset(), blk.MetaDataLength(), blk.BodyLength(), f.offset, r.maxMetadataSize, r.maxBodySize); err != nil { + return fileBlock{}, err + } return fileBlock{ offset: blk.Offset(), @@ -97,6 +133,9 @@ func (r *basicReaderImpl) dict(mem memory.Allocator, f *footerBlock, i int) (dat if !f.data.Dictionaries(&blk, i) { return fileBlock{}, fmt.Errorf("arrow/ipc: could not extract dictionary block %d", i) } + if err := validateFileBlock(blk.Offset(), blk.MetaDataLength(), blk.BodyLength(), f.offset, r.maxMetadataSize, r.maxBodySize); err != nil { + return fileBlock{}, err + } return fileBlock{ offset: blk.Offset(), @@ -108,7 +147,9 @@ func (r *basicReaderImpl) dict(mem memory.Allocator, f *footerBlock, i int) (dat } type mappedReaderImpl struct { - data []byte + data []byte + maxMetadataSize int64 + maxBodySize int64 } func (r *mappedReaderImpl) getBytes(offset, length int64) ([]byte, error) { @@ -126,6 +167,9 @@ func (r *mappedReaderImpl) block(_ memory.Allocator, f *footerBlock, i int) (dat if !f.data.RecordBatches(&blk, i) { return mappedFileBlock{}, fmt.Errorf("arrow/ipc: could not extract file block %d", i) } + if err := validateFileBlock(blk.Offset(), blk.MetaDataLength(), blk.BodyLength(), int64(len(r.data)), r.maxMetadataSize, r.maxBodySize); err != nil { + return mappedFileBlock{}, err + } return mappedFileBlock{ offset: blk.Offset(), @@ -140,6 +184,9 @@ func (r *mappedReaderImpl) dict(_ memory.Allocator, f *footerBlock, i int) (data if !f.data.Dictionaries(&blk, i) { return mappedFileBlock{}, fmt.Errorf("arrow/ipc: could not extract dictionary block %d", i) } + if err := validateFileBlock(blk.Offset(), blk.MetaDataLength(), blk.BodyLength(), int64(len(r.data)), r.maxMetadataSize, r.maxBodySize); err != nil { + return mappedFileBlock{}, err + } return mappedFileBlock{ offset: blk.Offset(), @@ -181,7 +228,11 @@ func NewMappedFileReader(data []byte, opts ...Option) (*FileReader, error) { var ( cfg = newConfig(opts...) f = FileReader{ - r: &mappedReaderImpl{data: data}, + r: &mappedReaderImpl{ + data: data, + maxMetadataSize: cfg.maxMetadataSize, + maxBodySize: cfg.maxBodySize, + }, mem: cfg.alloc, } ) @@ -197,7 +248,11 @@ func NewFileReader(r ReadAtSeeker, opts ...Option) (*FileReader, error) { var ( cfg = newConfig(opts...) f = FileReader{ - r: &basicReaderImpl{r: r}, + r: &basicReaderImpl{ + r: r, + maxMetadataSize: cfg.maxMetadataSize, + maxBodySize: cfg.maxBodySize, + }, memo: dictutils.NewMemo(), mem: cfg.alloc, } @@ -385,7 +440,7 @@ func (f *FileReader) Record(i int) (arrow.Record, error) { // caller and must call Release() to free the memory. This method is safe to // call concurrently. func (f *FileReader) RecordBatchAt(i int) (arrow.RecordBatch, error) { - if i < 0 || i > f.NumRecords() { + if i < 0 || i >= f.NumRecords() { panic("arrow/ipc: record index out of bounds") } @@ -918,6 +973,9 @@ func (blk mappedFileBlock) NewMessage() (*Message, error) { // messages produced prior to version 0.15.0 prefix = 4 } + if int(blk.meta)-prefix < 4 { + return nil, fmt.Errorf("arrow/ipc: invalid file block metadata length %d for prefix length %d", blk.meta, prefix) + } meta = memory.NewBufferBytes(metaBytes[prefix:]) body = memory.NewBufferBytes(buf[blk.meta : int64(blk.meta)+blk.body]) diff --git a/arrow/ipc/ipc.go b/arrow/ipc/ipc.go index ca2e1ecdc..744a59e7b 100644 --- a/arrow/ipc/ipc.go +++ b/arrow/ipc/ipc.go @@ -72,14 +72,27 @@ type config struct { noAutoSchema bool emitDictDeltas bool minSpaceSavings float64 + maxMetadataSize int64 + maxBodySize int64 } +const ( + defaultMaxMetadataSize int64 = 64 * 1024 * 1024 + defaultMaxBodySize int64 = 256 * 1024 * 1024 +) + +// Option is a functional option to configure opening or creating Arrow files +// and streams. +type Option func(*config) + func newConfig(opts ...Option) *config { cfg := &config{ alloc: memory.NewGoAllocator(), codec: -1, // uncompressed ensureNativeEndian: true, compressNP: 1, + maxMetadataSize: defaultMaxMetadataSize, + maxBodySize: defaultMaxBodySize, } for _, opt := range opts { @@ -89,9 +102,23 @@ func newConfig(opts ...Option) *config { return cfg } -// Option is a functional option to configure opening or creating Arrow files -// and streams. -type Option func(*config) +// WithMetadataSizeLimit sets the largest IPC message metadata size readers +// will accept. The default is 64 MiB. Values less than or equal to zero disable +// the limit. +func WithMetadataSizeLimit(size int64) Option { + return func(cfg *config) { + cfg.maxMetadataSize = size + } +} + +// WithBodySizeLimit sets the largest IPC message body size readers will +// accept. The default is 256 MiB. Values less than or equal to zero disable +// the limit. +func WithBodySizeLimit(size int64) Option { + return func(cfg *config) { + cfg.maxBodySize = size + } +} // WithFooterOffset specifies the Arrow footer position in bytes. func WithFooterOffset(offset int64) Option { diff --git a/arrow/ipc/message.go b/arrow/ipc/message.go index a2c4e3703..df1ed8d9c 100644 --- a/arrow/ipc/message.go +++ b/arrow/ipc/message.go @@ -149,8 +149,10 @@ type messageReader struct { refCount atomic.Int64 msg *Message - mem memory.Allocator - header [4]byte + mem memory.Allocator + maxMetadataSize int64 + maxBodySize int64 + header [4]byte } // NewMessageReader returns a reader that reads messages from an input stream. @@ -160,7 +162,12 @@ func NewMessageReader(r io.Reader, opts ...Option) MessageReader { opt(cfg) } - mr := &messageReader{r: r, mem: cfg.alloc} + mr := &messageReader{ + r: r, + mem: cfg.alloc, + maxMetadataSize: cfg.maxMetadataSize, + maxBodySize: cfg.maxBodySize, + } mr.refCount.Add(1) return mr } @@ -188,9 +195,16 @@ func (r *messageReader) Release() { // Message returns the current message that has been extracted from the // underlying stream. // It is valid until the next call to Message. -func (r *messageReader) Message() (*Message, error) { +func (r *messageReader) Message() (msg *Message, err error) { + defer func() { + if recovered := recover(); recovered != nil { + msg = nil + err = fmt.Errorf("arrow/ipc: invalid message metadata: %v", recovered) + } + }() + buf := r.header[:] - _, err := io.ReadFull(r.r, buf) + _, err = io.ReadFull(r.r, buf) if err != nil { return nil, fmt.Errorf("arrow/ipc: could not read continuation indicator: %w", err) } @@ -218,6 +232,12 @@ func (r *messageReader) Message() (*Message, error) { // messages produced prior to version 0.15.0 msgLen = int32(cid) } + if msgLen < 4 { + return nil, fmt.Errorf("arrow/ipc: invalid message metadata length %d", msgLen) + } + if r.maxMetadataSize > 0 && int64(msgLen) > r.maxMetadataSize { + return nil, fmt.Errorf("arrow/ipc: message metadata length %d exceeds limit %d", msgLen, r.maxMetadataSize) + } buf = make([]byte, msgLen) _, err = io.ReadFull(r.r, buf) @@ -227,6 +247,13 @@ func (r *messageReader) Message() (*Message, error) { meta := flatbuf.GetRootAsMessage(buf, 0) bodyLen := meta.BodyLength() + maxInt := int64(^uint(0) >> 1) + if bodyLen < 0 || bodyLen > maxInt { + return nil, fmt.Errorf("arrow/ipc: invalid message body length %d", bodyLen) + } + if r.maxBodySize > 0 && bodyLen > r.maxBodySize { + return nil, fmt.Errorf("arrow/ipc: message body length %d exceeds limit %d", bodyLen, r.maxBodySize) + } body := memory.NewResizableBuffer(r.mem) defer body.Release() diff --git a/arrow/ipc/message_test.go b/arrow/ipc/message_test.go index a69738e8f..896a5bc4e 100644 --- a/arrow/ipc/message_test.go +++ b/arrow/ipc/message_test.go @@ -18,15 +18,117 @@ package ipc import ( "bytes" + "encoding/binary" "errors" + "fmt" "io" + "math" + "strconv" "testing" "github.com/apache/arrow-go/v18/arrow" "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/internal/flatbuf" "github.com/apache/arrow-go/v18/arrow/memory" + flatbuffers "github.com/google/flatbuffers/go" + "github.com/stretchr/testify/require" ) +type rejectingAllocator struct{} + +func (rejectingAllocator) Allocate(int) []byte { + panic("unexpected allocation") +} + +func (rejectingAllocator) Reallocate(int, []byte) []byte { + panic("unexpected allocation") +} + +func (rejectingAllocator) Free([]byte) {} + +func messageReaderInput(t *testing.T, bodyLen int64) *bytes.Buffer { + t.Helper() + + b := flatbuffers.NewBuilder(0) + metadata := writeMessageFB(b, memory.DefaultAllocator, flatbuf.MessageHeaderNONE, 0, bodyLen, arrow.Metadata{}) + defer metadata.Release() + + var input bytes.Buffer + require.NoError(t, binary.Write(&input, binary.LittleEndian, uint32(kIPCContToken))) + require.NoError(t, binary.Write(&input, binary.LittleEndian, uint32(metadata.Len()))) + _, err := input.Write(metadata.Bytes()) + require.NoError(t, err) + return &input +} + +func TestMessageReaderRejectsInvalidMetadataLengths(t *testing.T) { + for _, length := range []uint32{1, 3, 0x7fffffff, 0x80000000, 0xfffffffe} { + t.Run(fmt.Sprint(length), func(t *testing.T) { + var input bytes.Buffer + require.NoError(t, binary.Write(&input, binary.LittleEndian, length)) + + r := NewMessageReader(&input) + defer r.Release() + _, err := r.Message() + require.ErrorContains(t, err, "message metadata length") + }) + } +} + +func TestMessageReaderRejectsInvalidBodyLengths(t *testing.T) { + type testCase struct { + name string + length int64 + message string + } + tests := []testCase{ + {name: "negative", length: -1, message: "invalid message body length"}, + {name: "over default limit", length: 1 << 30, message: "message body length 1073741824 exceeds limit 268435456"}, + } + if strconv.IntSize == 32 { + tests = append(tests, testCase{ + name: "over max int", length: int64(math.MaxInt32) + 1, message: "invalid message body length", + }) + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r := NewMessageReader(messageReaderInput(t, tt.length), WithAllocator(rejectingAllocator{})) + defer r.Release() + + _, err := r.Message() + require.ErrorContains(t, err, tt.message) + }) + } +} + +func TestBodySizeLimitConfig(t *testing.T) { + require.Equal(t, defaultMaxBodySize, newConfig().maxBodySize) + require.Zero(t, newConfig(WithBodySizeLimit(0)).maxBodySize) + require.Equal(t, int64(1024), newConfig(WithBodySizeLimit(1024)).maxBodySize) +} + +func TestValidateFileBlock(t *testing.T) { + tests := []struct { + name string + offset, body, size int64 + meta int32 + }{ + {"negative offset", -1, 0, 16, 8}, + {"short metadata", 0, 0, 16, 3}, + {"negative body", 0, -1, 16, 8}, + {"past end", 8, 9, 24, 8}, + {"overflow", 1, math.MaxInt64, math.MaxInt64, 8}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Error(t, validateFileBlock(tt.offset, tt.meta, tt.body, tt.size, 0, 0)) + }) + } + require.NoError(t, validateFileBlock(8, 8, 8, 24, 0, 0)) +} + func TestMessageReaderBodyInAllocator(t *testing.T) { mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) defer mem.AssertSize(t, 0) diff --git a/arrow/ipc/metadata.go b/arrow/ipc/metadata.go index c139f95ee..7a6844397 100644 --- a/arrow/ipc/metadata.go +++ b/arrow/ipc/metadata.go @@ -114,6 +114,9 @@ func (blk fileBlock) NewMessage() (*Message, error) { // messages produced prior to version 0.15.0 prefix = 4 } + if int(blk.meta)-prefix < 4 { + return nil, fmt.Errorf("arrow/ipc: invalid file block metadata length %d for prefix length %d", blk.meta, prefix) + } // drop buf-size already known from blk.Meta meta = memory.SliceBuffer(meta, prefix, int(blk.meta)-prefix)