diff --git a/arrow/ipc/file_block_test.go b/arrow/ipc/file_block_test.go new file mode 100644 index 00000000..3e7e7d0c --- /dev/null +++ b/arrow/ipc/file_block_test.go @@ -0,0 +1,106 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ipc + +import ( + "bytes" + "encoding/binary" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "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" +) + +func testFramedMessage(t *testing.T, bodyLen int, legacy bool) ([]byte, int32, int64) { + t.Helper() + meta := writeMessageFB(flatbuffers.NewBuilder(0), memory.DefaultAllocator, + flatbuf.MessageHeaderNONE, 0, int64(bodyLen), arrow.Metadata{}) + defer meta.Release() + + padding := (8 - meta.Len()%8) % 8 + prefix := 8 + if legacy { + prefix = 4 + } + framed := make([]byte, prefix+meta.Len()+padding) + if legacy { + binary.LittleEndian.PutUint32(framed, uint32(meta.Len()+padding)) + } else { + binary.LittleEndian.PutUint32(framed, kIPCContToken) + binary.LittleEndian.PutUint32(framed[4:], uint32(meta.Len()+padding)) + } + copy(framed[prefix:], meta.Bytes()) + return framed, int32(len(framed)), int64(bodyLen) +} + +func TestFileBlockNewMessageValidatesFraming(t *testing.T) { + for _, legacy := range []bool{false, true} { + for _, mapped := range []bool{false, true} { + name := "continuation" + if legacy { + name = "legacy" + } + if mapped { + name += " mapped" + } else { + name += " reader" + } + + t.Run(name, func(t *testing.T) { + framed, metaLen, bodyLen := testFramedMessage(t, 4, legacy) + body := []byte{1, 2, 3, 4} + + newBlock := func(meta int32, blockBody int64) dataBlock { + data := append(append([]byte{}, framed...), body...) + if meta > metaLen { + data = append(data[:int(metaLen)], make([]byte, int(meta-metaLen))...) + data = append(data, body...) + } + if blockBody > bodyLen { + data = append(data, make([]byte, int(blockBody-bodyLen))...) + } + if mapped { + return mappedFileBlock{meta: meta, body: blockBody, data: data} + } + return fileBlock{meta: meta, body: blockBody, r: bytes.NewReader(data), mem: memory.DefaultAllocator} + } + + msg, err := newBlock(metaLen, bodyLen).NewMessage() + require.NoError(t, err) + require.EqualValues(t, bodyLen, msg.BodyLen()) + msg.Release() + + _, err = newBlock(metaLen+8, bodyLen).NewMessage() + require.ErrorContains(t, err, "metadata length prefix") + _, err = newBlock(metaLen-4, bodyLen).NewMessage() + require.ErrorContains(t, err, "metadata length prefix") + _, err = newBlock(metaLen, bodyLen+1).NewMessage() + require.ErrorContains(t, err, "body length") + _, err = newBlock(metaLen, bodyLen-1).NewMessage() + require.ErrorContains(t, err, "body length") + }) + } + } +} + +func TestValidateFileBlockRejectsUnalignedBody(t *testing.T) { + err := validateFileBlock(8, 8, 4, 24, 0, 0) + require.ErrorContains(t, err, "not a multiple of 8") +} diff --git a/arrow/ipc/file_reader.go b/arrow/ipc/file_reader.go index 3128a150..abf2a70f 100644 --- a/arrow/ipc/file_reader.go +++ b/arrow/ipc/file_reader.go @@ -75,6 +75,9 @@ func validateFileBlock(offset int64, meta int32, body, fileSize, maxMetadataSize if body < 0 { return fmt.Errorf("arrow/ipc: invalid file block body length %d", body) } + if body%8 != 0 { + return fmt.Errorf("arrow/ipc: file block body length %d is not a multiple of 8", body) + } if maxMetadataSize > 0 && int64(meta) > maxMetadataSize { return fmt.Errorf("arrow/ipc: file block metadata length %d exceeds limit %d", meta, maxMetadataSize) } @@ -974,21 +977,18 @@ func (blk mappedFileBlock) NewMessage() (*Message, error) { metaBytes := buf[:blk.meta] - prefix := 0 - switch binary.LittleEndian.Uint32(metaBytes) { - case 0: - case kIPCContToken: - prefix = 8 - default: - // ARROW-6314: backwards compatibility for reading old IPC - // 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) + prefix, err := validateFileBlockMetadata(metaBytes, blk.meta) + if err != nil { + return nil, err } meta = memory.NewBufferBytes(metaBytes[prefix:]) body = memory.NewBufferBytes(buf[blk.meta : int64(blk.meta)+blk.body]) - return NewMessage(meta, body), nil + msg := NewMessage(meta, body) + messageBodyLen := msg.BodyLen() + if messageBodyLen != blk.body { + msg.Release() + return nil, fmt.Errorf("arrow/ipc: file block body length %d does not match message body length %d", blk.body, messageBodyLen) + } + return msg, nil } diff --git a/arrow/ipc/metadata.go b/arrow/ipc/metadata.go index 7a684439..54cde676 100644 --- a/arrow/ipc/metadata.go +++ b/arrow/ipc/metadata.go @@ -75,6 +75,39 @@ func (blk fileBlock) Offset() int64 { return blk.offset } func (blk fileBlock) Meta() int32 { return blk.meta } func (blk fileBlock) Body() int64 { return blk.body } +func validateFileBlockMetadata(buf []byte, meta int32) (int, error) { + if len(buf) < 4 { + return 0, fmt.Errorf("arrow/ipc: file block metadata is too short: %d", len(buf)) + } + + var ( + prefix int + length uint32 + ) + switch binary.LittleEndian.Uint32(buf) { + case 0: + return 0, errors.New("arrow/ipc: unexpected end-of-stream marker in file block") + case kIPCContToken: + prefix = 8 + if len(buf) < prefix { + return 0, fmt.Errorf("arrow/ipc: file block metadata is too short for prefix length %d", prefix) + } + default: + // ARROW-6314: backwards compatibility for reading old IPC + // messages produced prior to version 0.15.0 + prefix = 4 + } + length = binary.LittleEndian.Uint32(buf[prefix-4:]) + + if int(meta)-prefix < 4 { + return 0, fmt.Errorf("arrow/ipc: invalid file block metadata length %d for prefix length %d", meta, prefix) + } + if int64(length) != int64(meta)-int64(prefix) { + return 0, fmt.Errorf("arrow/ipc: file block metadata length prefix %d does not match footer length %d", length, int64(meta)-int64(prefix)) + } + return prefix, nil +} + func fileBlocksToFB(b *flatbuffers.Builder, blocks []dataBlock, start startVecFunc) flatbuffers.UOffsetT { start(b, len(blocks)) for i := len(blocks) - 1; i >= 0; i-- { @@ -104,18 +137,9 @@ func (blk fileBlock) NewMessage() (*Message, error) { return nil, fmt.Errorf("arrow/ipc: could not read message metadata: %w", err) } - prefix := 0 - switch binary.LittleEndian.Uint32(buf) { - case 0: - case kIPCContToken: - prefix = 8 - default: - // ARROW-6314: backwards compatibility for reading old IPC - // 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) + prefix, err := validateFileBlockMetadata(buf, blk.meta) + if err != nil { + return nil, err } // drop buf-size already known from blk.Meta @@ -131,7 +155,13 @@ func (blk fileBlock) NewMessage() (*Message, error) { return nil, fmt.Errorf("arrow/ipc: could not read message body: %w", err) } - return NewMessage(meta, body), nil + msg := NewMessage(meta, body) + messageBodyLen := msg.BodyLen() + if messageBodyLen != blk.body { + msg.Release() + return nil, fmt.Errorf("arrow/ipc: file block body length %d does not match message body length %d", blk.body, messageBodyLen) + } + return msg, nil } func (blk fileBlock) section() io.Reader {