diff --git a/modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java b/modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java index c59ad3ebd83fc..e9b2a81f65815 100644 --- a/modules/benchmarks/src/main/java/org/apache/ignite/internal/benchmarks/jmh/direct/JmhDirectMessageReaderBenchmark.java +++ b/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(); } diff --git a/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java b/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java index 3473fc8a05bae..bde4e531b3407 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/direct/DirectMessageReader.java +++ b/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 state; @@ -514,13 +518,41 @@ private 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 { diff --git a/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java b/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java index b92f6000fdd7a..67bd768342552 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/direct/state/DirectMessageState.java +++ b/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 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)); diff --git a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java index 0ebeeeab274bc..ca39336175af1 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java +++ b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessage.java @@ -126,7 +126,18 @@ 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]; @@ -134,8 +145,9 @@ private byte[] uncompress() { 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); diff --git a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java index 81c97506d2e80..c1a42085bc95e 100644 --- a/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java +++ b/modules/core/src/main/java/org/apache/ignite/internal/managers/communication/CompressedMessageSerializer.java @@ -85,6 +85,9 @@ public class CompressedMessageSerializer implements MessageSerializer(msg.dataSize / CHUNK_SIZE + 1); + msg.chunks = new ArrayList<>(Math.min(msg.dataSize / CHUNK_SIZE + 1, 1024)); msg.chunks.add(msg.chunk); diff --git a/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java b/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java index 340fc02f38d35..98f1478c2b3a4 100644 --- a/modules/core/src/test/java/org/apache/ignite/internal/managers/communication/CompressedMessageTest.java +++ b/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> partHistSuppliers = new HashMap<>();