IGNITE-28853 Harden compressed payload deserialization against corrupted streams

modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java

@@ -100,7 +100,8 @@ public void setup() {
 
         boolean finished = writer.writeMessage(msg, true);
 
-        assert finished;
+        if (!finished)
+            throw new IllegalStateException("Message does not fit into the buffer.");
 
         buf.flip();
     }

modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java

@@ -23,6 +23,7 @@
 import java.util.Map;
 import java.util.UUID;
 import java.util.function.Function;
+import org.apache.ignite.IgniteException;
 import org.apache.ignite.internal.direct.state.DirectMessageState;
 import org.apache.ignite.internal.direct.state.DirectMessageStateItem;
 import org.apache.ignite.internal.direct.stream.DirectByteBufferStream;
@@ -49,6 +50,9 @@
  * Message reader implementation.
  */
 public class DirectMessageReader implements MessageReader {
+    /** Empty buffer to release the payload reference from {@link #tmpReader} streams between compressed reads. */
+    private static final ByteBuffer EMPTY_BUF = ByteBuffer.wrap(new byte[0]);
+
     /** State. */
     @GridToStringInclude
     private final DirectMessageState<StateItem> state;
@@ -514,13 +518,41 @@ private <T> T readCompressedMessageAndDeserialize(DirectByteBufferStream stream,
 
         tmpReader.setBuffer(ByteBuffer.wrap(msg0.uncompressed()));
 
-        T res = fun.apply(tmpReader);
+        T res;
+
+        boolean ok = false;
+
+        try {
+            res = fun.apply(tmpReader);
 
-        lastRead = tmpReader.state.item().stream.lastFinished();
+            // The payload buffer is always complete, so a partial read means a corrupted stream, not a lack of data.
+            if (!tmpReader.state.item().stream.lastFinished())
+                throw new IgniteException("Failed to deserialize compressed payload: uncompressed data ended " +
+                    "unexpectedly [dataSize=" + msg0.dataSize() + ']');
+
+            ok = true;
+        }
+        finally {
+            // Don't reuse a reader with half-read state.
+            if (!ok)
+                tmpReader = null;
+        }
+
+        // Don't pin the payload: the reader lives as long as the connection.
+        tmpReader.releasePayload();
+
+        lastRead = true;
 
         return res;
     }
 
+    /** Releases the payload buffer reference from all streams of this reader. */
+    private void releasePayload() {
+        buf = EMPTY_BUF;
+
+        state.forEachItem(item -> item.stream.setBuffer(EMPTY_BUF));
+    }
+
     /**
      */
     private static class StateItem implements DirectMessageStateItem {

modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java

@@ -19,6 +19,7 @@
 
 import java.lang.reflect.Array;
 import java.util.Arrays;
+import java.util.function.Consumer;
 import org.apache.ignite.internal.util.tostring.GridToStringExclude;
 import org.apache.ignite.internal.util.typedef.internal.S;
 import org.apache.ignite.lang.IgniteOutClosure;
@@ -98,6 +99,19 @@ public void reset() {
         stack[0].reset();
     }
 
+    /**
+     * @param c Closure to apply to every created item.
+     */
+    public void forEachItem(Consumer<T> c) {
+        // Items are created contiguously by forward().
+        for (T item : stack) {
+            if (item == null)
+                break;
+
+            c.accept(item);
+        }
+    }
+
     /** {@inheritDoc} */
     @Override public String toString() {
         return S.toString(DirectMessageState.class, this, "stack", Arrays.toString(stack));

modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java

@@ -126,16 +126,28 @@ private void compress(ByteBuffer buf) {
     /** @return Uncompressed data. */
     private byte[] uncompress() {
         if (chunks == null)
-            return null;
+            throw new IgniteException("Compressed stream is truncated [expected=" + dataSize + ", inflated=0]");
+
+        long compressedTotal = 0;
+
+        for (int i = 0; i < chunks.size(); i++)
+            compressedTotal += chunks.get(i).length;
+
+        // Raw deflate cannot expand beyond ~1032:1, a larger claim means a corrupted size header; also bounds the
+        // upfront allocation by ~1032x of the actually received bytes.
+        if (dataSize > compressedTotal * 1032L + 64)
+            throw new IgniteException("Invalid compressed message data size [dataSize=" + dataSize +
+                ", compressedBytes=" + compressedTotal + ']');
 
         byte[] data = new byte[dataSize];
 
         Inflater inflater = new Inflater(true);
 
         try {
             int off = 0;
+            int i = 0;
 
-            for (int i = 0; i < chunks.size() && off < dataSize; i++) {
+            for (; i < chunks.size() && off < dataSize; i++) {
                 inflater.setInput(chunks.get(i));
 
                 int n;
@@ -146,6 +158,19 @@ private byte[] uncompress() {
 
             if (off != dataSize)
                 throw new IgniteException("Compressed stream is truncated [expected=" + dataSize + ", inflated=" + off + ']');
+
+            // Any extra inflatable byte means the size header is understated.
+            byte[] probe = new byte[1];
+
+            while (true) {
+                if (inflater.inflate(probe, 0, 1) > 0)
+                    throw new IgniteException("Compressed stream is longer than expected [expected=" + dataSize + ']');
+
+                if (inflater.finished() || i == chunks.size())
+                    break;
+
+                inflater.setInput(chunks.get(i++));
+            }
         }
         catch (DataFormatException e) {
             throw new IgniteException(e);

modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java

@@ -85,6 +85,9 @@ public class CompressedMessageSerializer implements MessageSerializer<Compressed
                     if (!reader.isLastRead())
                         return false;
 
+                    if (msg.dataSize < 0)
+                        throw new IgniteException("Invalid compressed message data size: " + msg.dataSize);
+
                     if (msg.dataSize == 0)
                         return true;
 
@@ -112,7 +115,7 @@ public class CompressedMessageSerializer implements MessageSerializer<Compressed
                             "(stream is corrupted or the sender is incompatible).");
 
                     if (msg.chunks == null)
-                        msg.chunks = new ArrayList<>(msg.dataSize / CHUNK_SIZE + 1);
+                        msg.chunks = new ArrayList<>(Math.min(msg.dataSize / CHUNK_SIZE + 1, 1024));
 
                     msg.chunks.add(msg.chunk);
 

modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java

@@ -21,8 +21,10 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.Set;
 import java.util.UUID;
+import java.util.zip.Deflater;
 import org.apache.ignite.IgniteException;
 import org.apache.ignite.internal.CoreMessagesProvider;
 import org.apache.ignite.internal.direct.DirectMessageReader;
@@ -136,6 +138,110 @@ public void testReadFailsOnNullChunk() {
             "unexpected null chunk");
     }
 
+    /** Read must fail fast on a negative data size from the wire. */
+    @Test
+    public void testReadFailsOnNegativeDataSize() {
+        DirectMessageWriter writer = new DirectMessageWriter(MSG_FACTORY);
+
+        ByteBuffer buf = ByteBuffer.allocate(16);
+
+        writer.setBuffer(buf);
+
+        writer.writeInt(-5);
+
+        buf.flip();
+
+        DirectMessageReader reader = new DirectMessageReader(MSG_FACTORY, null);
+
+        reader.setBuffer(buf);
+
+        GridTestUtils.assertThrows(null,
+            () -> new CompressedMessageSerializer().readFrom(new CompressedMessage(), reader),
+            IgniteException.class,
+            "Invalid compressed message data size");
+    }
+
+    /** Uncompress must fail when dataSize > 0 but no chunks were received. */
+    @Test
+    public void testUncompressFailsWithoutChunks() {
+        CompressedMessage rcvd = new CompressedMessage();
+
+        rcvd.dataSize = 100;
+        rcvd.finalChunk = true;
+
+        GridTestUtils.assertThrows(null, rcvd::uncompressed, IgniteException.class, "truncated");
+    }
+
+    /** Uncompress must fail when the stream inflates to more bytes than the size header claims. */
+    @Test
+    public void testUncompressFailsOnUnderstatedDataSize() {
+        byte[] data = new byte[1000];
+
+        CompressedMessage sent = new CompressedMessage(ByteBuffer.wrap(data), Deflater.BEST_SPEED);
+
+        CompressedMessage rcvd = new CompressedMessage();
+
+        rcvd.dataSize = data.length - 1;
+        rcvd.chunks = sent.chunks;
+        rcvd.finalChunk = true;
+
+        GridTestUtils.assertThrows(null, rcvd::uncompressed, IgniteException.class, "longer than expected");
+    }
+
+    /** Same as {@link #testUncompressFailsOnUnderstatedDataSize()}, but with a multi-chunk compressed stream. */
+    @Test
+    public void testUncompressFailsOnUnderstatedDataSizeMultiChunk() {
+        byte[] data = new byte[CompressedMessage.CHUNK_SIZE * 3];
+
+        new Random(42).nextBytes(data);
+
+        CompressedMessage sent = new CompressedMessage(ByteBuffer.wrap(data), Deflater.BEST_SPEED);
+
+        assertTrue(sent.chunks.size() > 1);
+
+        CompressedMessage rcvd = new CompressedMessage();
+
+        rcvd.dataSize = CompressedMessage.CHUNK_SIZE / 2;
+        rcvd.chunks = sent.chunks;
+        rcvd.finalChunk = true;
+
+        GridTestUtils.assertThrows(null, rcvd::uncompressed, IgniteException.class, "longer than expected");
+    }
+
+    /** A complete envelope whose payload doesn't deserialize fully must fail instead of hanging as a partial read. */
+    @Test
+    public void testReadFailsOnTruncatedPayload() {
+        DirectMessageWriter writer = new DirectMessageWriter(MSG_FACTORY);
+
+        ByteBuffer tmpBuf = ByteBuffer.allocate(1 << 20);
+
+        writer.setBuffer(tmpBuf);
+
+        assertTrue(writer.writeMessage(fullMessage(), false));
+
+        tmpBuf.flip();
+
+        tmpBuf.limit(tmpBuf.limit() - 5); // Truncate the serialized message.
+
+        CompressedMessage compressedMsg = new CompressedMessage(tmpBuf, Deflater.BEST_SPEED);
+
+        DirectMessageWriter wireWriter = new DirectMessageWriter(MSG_FACTORY);
+
+        ByteBuffer wire = ByteBuffer.allocate(1 << 20);
+
+        wireWriter.setBuffer(wire);
+
+        assertTrue(wireWriter.writeMessage(compressedMsg, false));
+
+        wire.flip();
+
+        DirectMessageReader reader = new DirectMessageReader(MSG_FACTORY, null);
+
+        reader.setBuffer(wire);
+
+        GridTestUtils.assertThrows(null, () -> reader.readMessage(true), IgniteException.class, "ended unexpectedly");
+    }
+
     /** */
     private GridDhtPartitionsFullMessage fullMessage() {
         Map<UUID, Map<GroupPartitionIdPair, Long>> partHistSuppliers = new HashMap<>();