Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 63 additions & 5 deletions arrow/ipc/file_reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"errors"
"fmt"
"io"
"math"

"github.com/apache/arrow-go/v18/arrow"
"github.com/apache/arrow-go/v18/arrow/array"
Expand Down Expand Up @@ -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) {
Expand All @@ -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(),
Expand All @@ -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(),
Expand All @@ -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) {
Expand All @@ -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(),
Expand All @@ -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(),
Expand Down Expand Up @@ -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,
}
)
Expand All @@ -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,
}
Expand Down Expand Up @@ -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")
}

Expand Down Expand Up @@ -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])
Expand Down
33 changes: 30 additions & 3 deletions arrow/ipc/ipc.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down
37 changes: 32 additions & 5 deletions arrow/ipc/message.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
}
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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)
Expand All @@ -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 {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By default maxBodySize is math.MaxInt (set in ipc.go), so on 64-bit a malformed stream advertising a huge-but-addressable BodyLength (e.g. 1<<40) passes this check and reaches body.Resize(int(bodyLen)) at L260 — a ~1 TiB allocation attempt — before io.ReadFull can discover the stream is truncated. That's an allocation-based DoS from a single malformed message on the default reader.

Suggest a sane default cap (cf. #953, which defaults to 64 MiB compressed / 256 MiB uncompressed) with an explicit opt-out for 0, and/or bounding the allocation to the bytes actually available. Also worth adding tests for negative BodyLength, BodyLength > MaxInt, and a huge-but-addressable body over the default limit — the current tests cover metadata-length prefixes and file blocks but not stream body rejection.

return nil, fmt.Errorf("arrow/ipc: message body length %d exceeds limit %d", bodyLen, r.maxBodySize)
}

body := memory.NewResizableBuffer(r.mem)
defer body.Release()
Expand Down
102 changes: 102 additions & 0 deletions arrow/ipc/message_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading
Loading