fix(arrow/ipc): validate file block framing and lengths

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")
+}

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
 }

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 {