diff --git a/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java b/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java index 36f6ea449b..1789dc4ffa 100644 --- a/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java +++ b/vector/src/main/java/org/apache/arrow/vector/ipc/message/MessageSerializer.java @@ -708,9 +708,16 @@ public static MessageMetadataResult readMessage(ReadChannel in) throws IOExcepti // ARROW-6313, if the first 4 bytes are continuation message, read the next 4 for the length if (in.readFully(buffer) == 4) { messageLength = MessageSerializer.bytesToInt(buffer.array()); + } else { + // Stream ended right after the continuation token, treat as end-of-stream. + return null; } } + if (messageLength < 0) { + throw new IOException("Invalid message length: " + messageLength); + } + // Length of 0 indicates end of stream if (messageLength != 0) { diff --git a/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java b/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java index b529ca645a..ae68f2a95f 100644 --- a/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java +++ b/vector/src/test/java/org/apache/arrow/vector/ipc/MessageSerializerTest.java @@ -20,6 +20,7 @@ import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -121,6 +122,23 @@ public void testWriteMessageBufferAligned() throws IOException { assertEquals(0, result.getInt()); } + @Test + public void testReadMessageWithContinuationTokenThenEndOfStream() throws IOException { + // A continuation token (0xFFFFFFFF) with no length prefix following it. + byte[] bytes = new byte[] {(byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF}; + ReadChannel in = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))); + assertNull(MessageSerializer.readMessage(in)); + } + + @Test + public void testReadMessageRejectsNegativeLength() { + // Little-endian prefix encoding a negative message length (-5), which is not the + // continuation token. + byte[] bytes = new byte[] {(byte) 0xFB, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF}; + ReadChannel in = new ReadChannel(Channels.newChannel(new ByteArrayInputStream(bytes))); + assertThrows(IOException.class, () -> MessageSerializer.readMessage(in)); + } + @Test public void testSchemaMessageSerialization() throws IOException { Schema schema = testSchema();