This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch branch-3.5
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-3.5 by this push:
new b75a07db85e6 [SPARK-56227][3.5][CORE] Fix GcmTransportCipher to
correctly handle multiple messages per channel
b75a07db85e6 is described below
commit b75a07db85e6630725e9bbaa070e08c29906aa80
Author: Akira Ajisaka <[email protected]>
AuthorDate: Fri Jul 10 22:04:32 2026 -0700
[SPARK-56227][3.5][CORE] Fix GcmTransportCipher to correctly handle
multiple messages per channel
### What changes were proposed in this pull request?
Backport #55028 to branch-3.5.
### Why are the changes needed?
To successfully run Spark jobs on YARN with
`spark.network.crypto.cipher="AES/GCM/NoPadding"`
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Added new regression tests.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.6)
Closes #55621 from aajisaka/fix-rpc-encryption-branch-3.5.
Authored-by: Akira Ajisaka <[email protected]>
Signed-off-by: Holden Karau <[email protected]>
---
.../spark/network/crypto/GcmTransportCipher.java | 276 ++++++++---
.../spark/network/crypto/GcmAuthEngineSuite.java | 522 ++++++++++++++++++++-
2 files changed, 710 insertions(+), 88 deletions(-)
diff --git
a/common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java
b/common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java
index d3f1bf490d3a..5f6d6ef50207 100644
---
a/common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java
+++
b/common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java
@@ -22,9 +22,11 @@ import com.google.common.base.Preconditions;
import com.google.common.primitives.Longs;
import com.google.crypto.tink.subtle.*;
import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.util.ReferenceCounted;
+
import org.apache.spark.network.util.AbstractFileRegion;
import org.apache.spark.network.util.ByteBufferWriteableChannel;
@@ -41,6 +43,12 @@ public class GcmTransportCipher implements TransportCipher {
private static final int LENGTH_HEADER_BYTES = 8;
@VisibleForTesting
static final int CIPHERTEXT_BUFFER_SIZE = 32 * 1024; // 32KB
+ // Maximum plaintext bytes to accumulate before flushing to downstream
handlers, even
+ // if the GCM message is not yet complete. Bounds on-heap retention for
large shuffle
+ // blocks that route to disk via spark.maxRemoteBlockSizeFetchToMem, while
still
+ // reducing EventLoop callbacks to 1 for the common small-message case (<
1 MB).
+ @VisibleForTesting
+ static final int MAX_PLAINTEXT_BATCH_BYTES = 1024 * 1024; // 1 MB
private final SecretKeySpec aesKey;
public GcmTransportCipher(SecretKeySpec aesKey) {
@@ -82,25 +90,16 @@ public class GcmTransportCipher implements TransportCipher {
@VisibleForTesting
class EncryptionHandler extends ChannelOutboundHandlerAdapter {
- private final ByteBuffer plaintextBuffer;
- private final ByteBuffer ciphertextBuffer;
private final AesGcmHkdfStreaming aesGcmHkdfStreaming;
EncryptionHandler() throws InvalidAlgorithmParameterException {
aesGcmHkdfStreaming = getAesGcmHkdfStreaming();
- plaintextBuffer =
ByteBuffer.allocate(aesGcmHkdfStreaming.getPlaintextSegmentSize());
- ciphertextBuffer =
ByteBuffer.allocate(aesGcmHkdfStreaming.getCiphertextSegmentSize());
}
@Override
public void write(ChannelHandlerContext ctx, Object msg,
ChannelPromise promise)
throws Exception {
- GcmEncryptedMessage encryptedMessage = new GcmEncryptedMessage(
- aesGcmHkdfStreaming,
- msg,
- plaintextBuffer,
- ciphertextBuffer);
- ctx.write(encryptedMessage, promise);
+ ctx.write(new GcmEncryptedMessage(aesGcmHkdfStreaming, msg),
promise);
}
}
@@ -116,22 +115,39 @@ public class GcmTransportCipher implements
TransportCipher {
private final long encryptedCount;
GcmEncryptedMessage(AesGcmHkdfStreaming aesGcmHkdfStreaming,
- Object plaintextMessage,
- ByteBuffer plaintextBuffer,
- ByteBuffer ciphertextBuffer) throws
GeneralSecurityException {
+ Object plaintextMessage) throws
GeneralSecurityException {
Preconditions.checkArgument(
plaintextMessage instanceof ByteBuf || plaintextMessage
instanceof FileRegion,
"Unrecognized message type: %s",
plaintextMessage.getClass().getName());
this.plaintextMessage = plaintextMessage;
- this.plaintextBuffer = plaintextBuffer;
- this.ciphertextBuffer = ciphertextBuffer;
+ this.plaintextBuffer =
+
ByteBuffer.allocate(aesGcmHkdfStreaming.getPlaintextSegmentSize());
+ this.ciphertextBuffer =
+
ByteBuffer.allocate(aesGcmHkdfStreaming.getCiphertextSegmentSize());
// If the ciphertext buffer cannot be fully written the target,
transferTo may
// return with it containing some unwritten data. The initial call
we'll explicitly
// set its limit to 0 to indicate the first call to transferTo.
((Buffer) this.ciphertextBuffer).limit(0);
+
this.bytesToRead = getReadableBytes();
- this.encryptedCount =
- LENGTH_HEADER_BYTES +
aesGcmHkdfStreaming.expectedCiphertextSize(bytesToRead);
+ int plaintextSegmentSize =
aesGcmHkdfStreaming.getPlaintextSegmentSize();
+ int ciphertextSegmentSize =
aesGcmHkdfStreaming.getCiphertextSegmentSize();
+ int tagSize = ciphertextSegmentSize - plaintextSegmentSize;
+ long fullSegments = bytesToRead / plaintextSegmentSize;
+ long partialBytes = bytesToRead % plaintextSegmentSize;
+ // encryptedCount covers all bytes written to the wire:
+ // LENGTH_HEADER_BYTES - Spark's 8-byte big-endian framing
prefix, read by the
+ // receiver to know where this message
ends in the TCP stream.
+ // getHeaderLength() - Tink's streaming header (salt + nonce
prefix), written by
+ // the encrypter before any ciphertext
segments.
+ // segment bytes - full segments each padded with a
16-byte GCM auth tag,
+ // plus any partial final segment
similarly padded.
+ // NOTE: This formula is coupled to Tink's internal layout.
testEncryptedCountBoundary
+ // guards current behavior.
+ this.encryptedCount = LENGTH_HEADER_BYTES
+ + aesGcmHkdfStreaming.getHeaderLength()
+ + fullSegments * ciphertextSegmentSize
+ + (partialBytes > 0 ? partialBytes + tagSize : 0);
byte[] lengthAad = Longs.toByteArray(encryptedCount);
this.encrypter =
aesGcmHkdfStreaming.newStreamSegmentEncrypter(lengthAad);
this.headerByteBuffer = createHeaderByteBuffer();
@@ -297,27 +313,66 @@ public class GcmTransportCipher implements
TransportCipher {
private final ByteBuffer headerBuffer;
private final ByteBuffer ciphertextBuffer;
private final AesGcmHkdfStreaming aesGcmHkdfStreaming;
- private final StreamSegmentDecrypter decrypter;
+ private StreamSegmentDecrypter decrypter;
+ private final int headerLength;
private final int plaintextSegmentSize;
private boolean decrypterInit = false;
private boolean completed = false;
private int segmentNumber = 0;
private long expectedLength = -1;
private long ciphertextRead = 0;
+ // Accumulates decrypted segments and flushes them to downstream in
bounded batches.
+ // For small messages (total plaintext < MAX_PLAINTEXT_BATCH_BYTES) a
single
+ // ctx.fireChannelRead() fires at message completion, reducing N
per-segment
+ // EventLoop callbacks to 1 and preventing EventLoop starvation under
shuffle load.
+ // For large messages the accumulator flushes every
MAX_PLAINTEXT_BATCH_BYTES,
+ // bounding on-heap retention for blocks that route to disk. Null
between flushes;
+ // ownership transfers to downstream on fireChannelRead().
+ private CompositeByteBuf plaintextAccumulator = null;
DecryptionHandler() throws GeneralSecurityException {
aesGcmHkdfStreaming = getAesGcmHkdfStreaming();
+ headerLength = aesGcmHkdfStreaming.getHeaderLength();
expectedLengthBuffer = ByteBuffer.allocate(LENGTH_HEADER_BYTES);
- headerBuffer =
ByteBuffer.allocate(aesGcmHkdfStreaming.getHeaderLength());
+ headerBuffer = ByteBuffer.allocate(headerLength);
ciphertextBuffer =
ByteBuffer.allocate(aesGcmHkdfStreaming.getCiphertextSegmentSize());
decrypter = aesGcmHkdfStreaming.newStreamSegmentDecrypter();
plaintextSegmentSize =
aesGcmHkdfStreaming.getPlaintextSegmentSize();
}
- private boolean initalizeExpectedLength(ByteBuf ciphertextNettyBuf) {
+ /**
+ * Resets all per-message state so that the next incoming GCM message
can be
+ * decoded through the same channel handler instance. This must be
called after
+ * every successfully completed message because AesGcmHkdfStreaming is
a one-shot
+ * streaming primitive: each encrypted message carries its own random
IV and must
+ * be decrypted with a fresh StreamSegmentDecrypter.
+ */
+ private void resetForNextMessage() throws GeneralSecurityException {
+ expectedLength = -1;
+ ((Buffer) expectedLengthBuffer).clear();
+ ((Buffer) headerBuffer).clear();
+ ((Buffer) ciphertextBuffer).clear();
+ decrypterInit = false;
+ completed = false;
+ segmentNumber = 0;
+ ciphertextRead = 0;
+ decrypter = aesGcmHkdfStreaming.newStreamSegmentDecrypter();
+ plaintextAccumulator = null; // defensive; should already be null
after fireChannelRead
+ }
+
+ private boolean initializeExpectedLength(ByteBuf ciphertextNettyBuf) {
if (expectedLength < 0) {
- ciphertextNettyBuf.readBytes(expectedLengthBuffer);
+ // ByteBuf.readBytes(ByteBuffer) throws if fewer than
dst.remaining() bytes
+ // are available, so temporarily narrow the limit to what is
actually present.
+ int toRead = Math.min(ciphertextNettyBuf.readableBytes(),
+ expectedLengthBuffer.remaining());
+ if (toRead > 0) {
+ int savedLimit = expectedLengthBuffer.limit();
+ ((Buffer)
expectedLengthBuffer).limit(expectedLengthBuffer.position() + toRead);
+ ciphertextNettyBuf.readBytes(expectedLengthBuffer);
+ ((Buffer) expectedLengthBuffer).limit(savedLimit);
+ }
if (expectedLengthBuffer.hasRemaining()) {
// We did not read enough bytes to initialize the expected
length.
return false;
@@ -332,12 +387,22 @@ public class GcmTransportCipher implements
TransportCipher {
return true;
}
- private boolean initalizeDecrypter(ByteBuf ciphertextNettyBuf)
+ private boolean initializeDecrypter(ByteBuf ciphertextNettyBuf)
throws GeneralSecurityException {
// Check if the ciphertext header has been read. This contains
// the IV and other internal metadata.
if (!decrypterInit) {
- ciphertextNettyBuf.readBytes(headerBuffer);
+ // ByteBuf.readBytes(ByteBuffer) throws if fewer than
dst.remaining() bytes
+ // are available. Under TCP fragmentation the header can
arrive in multiple
+ // chunks, so temporarily narrow the limit to what is actually
present.
+ int toRead = Math.min(ciphertextNettyBuf.readableBytes(),
+ headerBuffer.remaining());
+ if (toRead > 0) {
+ int savedLimit = headerBuffer.limit();
+ ((Buffer) headerBuffer).limit(headerBuffer.position() +
toRead);
+ ciphertextNettyBuf.readBytes(headerBuffer);
+ ((Buffer) headerBuffer).limit(savedLimit);
+ }
if (headerBuffer.hasRemaining()) {
// We did not read enough bytes to initialize the header.
return false;
@@ -346,7 +411,7 @@ public class GcmTransportCipher implements TransportCipher {
byte[] lengthAad = Longs.toByteArray(expectedLength);
decrypter.init(headerBuffer, lengthAad);
decrypterInit = true;
- ciphertextRead += aesGcmHkdfStreaming.getHeaderLength();
+ ciphertextRead += headerLength;
if (expectedLength == ciphertextRead) {
// If the expected length is just the header, the
ciphertext is 0 length.
completed = true;
@@ -362,61 +427,130 @@ public class GcmTransportCipher implements
TransportCipher {
"Unrecognized message type: %s",
ciphertextMessage.getClass().getName());
ByteBuf ciphertextNettyBuf = (ByteBuf) ciphertextMessage;
- // The format of the output is:
+ // The format of each message is:
// [8 byte length][Internal IV and header][Ciphertext][Auth Tag]
+ //
+ // A single channelRead() call may deliver bytes from multiple
back-to-back
+ // GCM messages (common under shuffle load when TCP coalesces
writes). The
+ // outer loop processes as many complete messages as possible from
the buffer
+ // before releasing it, so that bytes belonging to the next
message are never
+ // discarded mid-stream.
try {
- if (!initalizeExpectedLength(ciphertextNettyBuf)) {
- // We have not read enough bytes to initialize the
expected length.
- return;
- }
- if (!initalizeDecrypter(ciphertextNettyBuf)) {
- // We have not read enough bytes to initialize a header,
needed to
- // initialize a decrypter.
- return;
- }
- int nettyBufReadableBytes = ciphertextNettyBuf.readableBytes();
- while (nettyBufReadableBytes > 0 && !completed) {
- // Read the ciphertext into the local buffer
- int readableBytes = Integer.min(
- nettyBufReadableBytes,
- ciphertextBuffer.remaining());
- int expectedRemaining = (int) (expectedLength -
ciphertextRead);
- int bytesToRead = Integer.min(readableBytes,
expectedRemaining);
- // The smallest ciphertext size is 16 bytes for the auth
tag
- ((Buffer) ciphertextBuffer).limit(
- ((Buffer) ciphertextBuffer).position() +
bytesToRead);
- ciphertextNettyBuf.readBytes(ciphertextBuffer);
- ciphertextRead += bytesToRead;
- // Check if this is the last segment
- if (ciphertextRead == expectedLength) {
- completed = true;
- } else if (ciphertextRead > expectedLength) {
- throw new IllegalStateException("Read more ciphertext
than expected.");
+ while (true) {
+ if (!initializeExpectedLength(ciphertextNettyBuf)) {
+ // We have not read enough bytes to initialize the
expected length.
+ break;
+ }
+ if (!initializeDecrypter(ciphertextNettyBuf)) {
+ // We have not read enough bytes to initialize a
header, needed to
+ // initialize a decrypter.
+ break;
+ }
+ int nettyBufReadableBytes =
ciphertextNettyBuf.readableBytes();
+ while (nettyBufReadableBytes > 0 && !completed) {
+ // Read the ciphertext into the local buffer
+ int readableBytes = Math.min(
+ nettyBufReadableBytes,
+ ciphertextBuffer.remaining());
+ int expectedRemaining = (int) (expectedLength -
ciphertextRead);
+ int bytesToRead = Math.min(readableBytes,
expectedRemaining);
+ // The smallest ciphertext size is 16 bytes for the
auth tag
+ ((Buffer) ciphertextBuffer).limit(
+ ciphertextBuffer.position() + bytesToRead);
+ ciphertextNettyBuf.readBytes(ciphertextBuffer);
+ ciphertextRead += bytesToRead;
+ // Check if this is the last segment
+ if (ciphertextRead == expectedLength) {
+ completed = true;
+ } else if (ciphertextRead > expectedLength) {
+ throw new IllegalStateException("Read more
ciphertext than expected.");
+ }
+ // If the ciphertext buffer is full, or this is the
last segment,
+ // then decrypt it and fire a read.
+ if (ciphertextBuffer.limit() ==
ciphertextBuffer.capacity() || completed) {
+ ByteBuffer plaintextBuffer =
ByteBuffer.allocate(plaintextSegmentSize);
+ ((Buffer) ciphertextBuffer).flip();
+ decrypter.decryptSegment(
+ ciphertextBuffer,
+ segmentNumber,
+ completed,
+ plaintextBuffer);
+ segmentNumber++;
+ // Clear the ciphertext buffer because it's been
read
+ ((Buffer) ciphertextBuffer).clear();
+ ((Buffer) plaintextBuffer).flip();
+ if (plaintextAccumulator == null) {
+ // Integer.MAX_VALUE disables Netty's internal
consolidation,
+ // which would copy all components into a
single buffer on
+ // access. The initial component array is
small regardless of
+ // this cap (min(16, maxNumComponents)
elements).
+ plaintextAccumulator =
+
Unpooled.compositeBuffer(Integer.MAX_VALUE);
+ }
+ // Zero-copy append: addComponent(true, ...)
increases writerIndex
+ // so the component is immediately readable from
the composite.
+ plaintextAccumulator.addComponent(
+ true,
Unpooled.wrappedBuffer(plaintextBuffer));
+ // Flush when the batch threshold is reached or
the message ends.
+ // Small messages flush once at completion; large
blocks flush
+ // every MAX_PLAINTEXT_BATCH_BYTES to cap on-heap
retention.
+ if (completed ||
+ plaintextAccumulator.readableBytes() >=
+ MAX_PLAINTEXT_BATCH_BYTES) {
+ flushAccumulator(ctx);
+ }
+ } else {
+ // Set the ciphertext buffer up to read the next
chunk
+ ((Buffer)
ciphertextBuffer).limit(ciphertextBuffer.capacity());
+ }
+ nettyBufReadableBytes =
ciphertextNettyBuf.readableBytes();
}
- // If the ciphertext buffer is full, or this is the last
segment,
- // then decrypt it and fire a read.
- if (ciphertextBuffer.limit() ==
ciphertextBuffer.capacity() || completed) {
- ByteBuffer plaintextBuffer =
ByteBuffer.allocate(plaintextSegmentSize);
- ((Buffer) ciphertextBuffer).flip();
- decrypter.decryptSegment(
- ciphertextBuffer,
- segmentNumber,
- completed,
- plaintextBuffer);
- segmentNumber++;
- // Clear the ciphertext buffer because it's been read
- ((Buffer) ciphertextBuffer).clear();
- ((Buffer) plaintextBuffer).flip();
-
ctx.fireChannelRead(Unpooled.wrappedBuffer(plaintextBuffer));
- } else {
- // Set the ciphertext buffer up to read the next chunk
- ((Buffer)
ciphertextBuffer).limit(ciphertextBuffer.capacity());
+ if (!completed) {
+ // Partial message: more bytes needed from the next
channelRead() call.
+ break;
}
- nettyBufReadableBytes = ciphertextNettyBuf.readableBytes();
+ // Flush any remaining accumulator not yet fired inside
the segment
+ // loop (e.g. a zero-length ciphertext where the loop
never ran).
+ if (plaintextAccumulator != null) {
+ flushAccumulator(ctx);
+ }
+ // Current message is fully decoded. Reset state so the
handler can
+ // decode the next independent GCM message on the same
channel.
+ resetForNextMessage();
+ if (ciphertextNettyBuf.readableBytes() == 0) {
+ break;
+ }
+ // Remaining bytes may belong to another message; loop to
process them.
}
} finally {
ciphertextNettyBuf.release();
}
}
+
+ private void flushAccumulator(ChannelHandlerContext ctx) {
+ CompositeByteBuf out = plaintextAccumulator;
+ // null the accumulator before firing
+ plaintextAccumulator = null;
+ ctx.fireChannelRead(out);
+ }
+
+ @Override
+ public void channelInactive(ChannelHandlerContext ctx) throws
Exception {
+ releaseAccumulator();
+ ctx.fireChannelInactive();
+ }
+
+ @Override
+ public void exceptionCaught(ChannelHandlerContext ctx, Throwable
cause) throws Exception {
+ releaseAccumulator();
+ super.exceptionCaught(ctx, cause);
+ }
+
+ private void releaseAccumulator() {
+ if (plaintextAccumulator != null) {
+ plaintextAccumulator.release();
+ plaintextAccumulator = null;
+ }
+ }
}
}
diff --git
a/common/network-common/src/test/java/org/apache/spark/network/crypto/GcmAuthEngineSuite.java
b/common/network-common/src/test/java/org/apache/spark/network/crypto/GcmAuthEngineSuite.java
index f25277aa1a99..37fd32f17a36 100644
---
a/common/network-common/src/test/java/org/apache/spark/network/crypto/GcmAuthEngineSuite.java
+++
b/common/network-common/src/test/java/org/apache/spark/network/crypto/GcmAuthEngineSuite.java
@@ -35,6 +35,7 @@ import java.nio.ByteBuffer;
import java.nio.channels.WritableByteChannel;
import java.util.Arrays;
+import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
import static org.mockito.Mockito.*;
@@ -94,13 +95,13 @@ public class GcmAuthEngineSuite extends AuthEngineSuite {
// Capture the decrypted values and verify them
ArgumentCaptor<ByteBuf> captorPlaintext =
ArgumentCaptor.forClass(ByteBuf.class);
decryptionHandler.channelRead(ctx, ciphertext);
- verify(ctx, times(2))
+ verify(ctx, times(1))
.fireChannelRead(captorPlaintext.capture());
- ByteBuf lastPlaintextSegment = captorPlaintext.getValue();
- assertEquals(plaintextSegmentSize/2,
- lastPlaintextSegment.readableBytes());
+ ByteBuf plaintext = captorPlaintext.getValue();
+ assertEquals(plaintextSegmentSize + (plaintextSegmentSize / 2),
+ plaintext.readableBytes());
assertEquals('c',
- lastPlaintextSegment.getByte((plaintextSegmentSize/2) - 10));
+ plaintext.getByte(plaintext.readableBytes() - 10));
}
}
@@ -223,13 +224,11 @@ public class GcmAuthEngineSuite extends AuthEngineSuite {
// Capture the decrypted values and verify them
ArgumentCaptor<ByteBuf> captorPlaintext =
ArgumentCaptor.forClass(ByteBuf.class);
decryptionHandler.channelRead(ctx, ciphertext);
- verify(ctx, times(2)).fireChannelRead(captorPlaintext.capture());
+ verify(ctx, times(1)).fireChannelRead(captorPlaintext.capture());
ByteBuf plaintext = captorPlaintext.getValue();
- // We expect this to be the last partial plaintext segment
- int expectedLength = totalSize % plaintextSegmentSize;
- assertEquals(expectedLength, plaintext.readableBytes());
- // This will be the "remainder" segment that is filled to 'c'
- assertEquals('c', plaintext.getByte(0));
+ assertEquals(totalSize, plaintext.readableBytes());
+ // 'c' starts at the second plaintext segment (offset
plaintextSegmentSize)
+ assertEquals('c', plaintext.getByte(plaintextSegmentSize));
}
}
@@ -285,12 +284,501 @@ public class GcmAuthEngineSuite extends AuthEngineSuite {
// Capture the decrypted values and verify them
ArgumentCaptor<ByteBuf> captorPlaintext =
ArgumentCaptor.forClass(ByteBuf.class);
decryptionHandler.channelRead(ctx, mockCiphertext);
- verify(ctx, times(2)).fireChannelRead(captorPlaintext.capture());
- ByteBuf lastPlaintextSegment = captorPlaintext.getValue();
- assertEquals(plaintextSegmentSize/2,
- lastPlaintextSegment.readableBytes());
- assertEquals('x',
- lastPlaintextSegment.getByte((plaintextSegmentSize/2) - 10));
+ verify(ctx, times(1)).fireChannelRead(captorPlaintext.capture());
+ ByteBuf plaintext = captorPlaintext.getValue();
+ assertEquals(plaintextSize, plaintext.readableBytes());
+ assertEquals('x', plaintext.getByte(plaintextSize - 10));
+ }
+ }
+
+ /**
+ * Verifies that the same DecryptionHandler instance correctly decodes
multiple independent
+ * GCM-encrypted messages sent over the same channel. This is the regression
test for the
+ * bug where DecryptionHandler.completed was never reset, causing every
message after the
+ * first to be silently dropped - which manifested as YARN container launch
failures.
+ */
+ @Test
+ public void testMultipleMessages() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ // --- First message ---
+ byte[] data1 = new byte[1024];
+ Arrays.fill(data1, (byte) 'A');
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> captor1 =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data1), promise);
+ verify(ctx).write(captor1.capture(), eq(promise));
+ ByteBuffer ct1 = ByteBuffer.allocate((int) captor1.getValue().count());
+ captor1.getValue().transferTo(new ByteBufferWriteableChannel(ct1), 0);
+ ((Buffer) ct1).flip();
+
+ ArgumentCaptor<ByteBuf> plaintextCaptor1 =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx, Unpooled.wrappedBuffer(ct1));
+ verify(ctx, atLeastOnce()).fireChannelRead(plaintextCaptor1.capture());
+ byte[] decrypted1 = new byte[data1.length];
+ int offset = 0;
+ for (ByteBuf segment : plaintextCaptor1.getAllValues()) {
+ int len = segment.readableBytes();
+ segment.readBytes(decrypted1, offset, len);
+ offset += len;
+ }
+ assertEquals(data1.length, offset);
+ assertArrayEquals(data1, decrypted1);
+
+ // --- Second message (same handler, different content) ---
+ reset(ctx);
+ byte[] data2 = new byte[2048];
+ Arrays.fill(data2, (byte) 'B');
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> captor2 =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data2), promise);
+ verify(ctx).write(captor2.capture(), eq(promise));
+ ByteBuffer ct2 = ByteBuffer.allocate((int) captor2.getValue().count());
+ captor2.getValue().transferTo(new ByteBufferWriteableChannel(ct2), 0);
+ ((Buffer) ct2).flip();
+
+ ArgumentCaptor<ByteBuf> plaintextCaptor2 =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx, Unpooled.wrappedBuffer(ct2));
+ verify(ctx, atLeastOnce()).fireChannelRead(plaintextCaptor2.capture());
+ byte[] decrypted2 = new byte[data2.length];
+ offset = 0;
+ for (ByteBuf segment : plaintextCaptor2.getAllValues()) {
+ int len = segment.readableBytes();
+ segment.readBytes(decrypted2, offset, len);
+ offset += len;
+ }
+ assertEquals(data2.length, offset);
+ assertArrayEquals(data2, decrypted2);
+ }
+ }
+
+ /**
+ * Verifies that multiple GCM-encrypted messages delivered inside a single
channelRead()
+ * call (TCP coalescing) are all decoded correctly. This is the regression
test for the
+ * IllegalStateException("Invalid expected ciphertext length") observed
under SparkSQL
+ * shuffle load: when Netty batches two messages into one ByteBuf, the old
code released
+ * the buffer after the first message, discarding remaining bytes. The next
channelRead()
+ * then read bytes from the middle of the second message as a length header.
+ */
+ @Test
+ public void testBatchedMessages() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ byte[] data1 = new byte[1024];
+ Arrays.fill(data1, (byte) 'A');
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> captor1 =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data1), promise);
+ verify(ctx).write(captor1.capture(), eq(promise));
+ ByteBuffer ct1 = ByteBuffer.allocate((int) captor1.getValue().count());
+ captor1.getValue().transferTo(new ByteBufferWriteableChannel(ct1), 0);
+ ((Buffer) ct1).flip();
+
+ reset(ctx);
+ byte[] data2 = new byte[2048];
+ Arrays.fill(data2, (byte) 'B');
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> captor2 =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data2), promise);
+ verify(ctx).write(captor2.capture(), eq(promise));
+ ByteBuffer ct2 = ByteBuffer.allocate((int) captor2.getValue().count());
+ captor2.getValue().transferTo(new ByteBufferWriteableChannel(ct2), 0);
+ ((Buffer) ct2).flip();
+
+ // Simulate TCP coalescing: deliver both ciphertexts in one
channelRead() call.
+ reset(ctx);
+ ByteBuf batched = Unpooled.wrappedBuffer(ct1, ct2);
+ ArgumentCaptor<ByteBuf> plaintextCaptor =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx, batched);
+ verify(ctx, atLeastOnce()).fireChannelRead(plaintextCaptor.capture());
+
+ byte[] decrypted = new byte[data1.length + data2.length];
+ int offset = 0;
+ for (ByteBuf segment : plaintextCaptor.getAllValues()) {
+ int len = segment.readableBytes();
+ segment.readBytes(decrypted, offset, len);
+ offset += len;
+ }
+ assertEquals(data1.length + data2.length, offset);
+ assertArrayEquals(data1, Arrays.copyOfRange(decrypted, 0, data1.length));
+ assertArrayEquals(data2, Arrays.copyOfRange(decrypted, data1.length,
decrypted.length));
+ }
+ }
+
+ /**
+ * Verifies that DecryptionHandler correctly handles a GCM message whose
framing header
+ * is split across two channelRead() calls. This is the regression test for
the
+ * IndexOutOfBoundsException in initializeDecrypter observed in
benchmarking: when only
+ * 4 bytes of the 24-byte GCM internal header arrived in one Netty buffer,
+ * ByteBuf.readBytes(ByteBuffer) threw because it requires all
dst.remaining() bytes to
+ * be available rather than performing a partial fill.
+ */
+ @Test
+ public void testSplitGcmHeader() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ byte[] data = new byte[1024];
+ Arrays.fill(data, (byte) 'X');
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> captor =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data), promise);
+ verify(ctx).write(captor.capture(), eq(promise));
+
+ ByteBuffer ciphertextBuffer = ByteBuffer.allocate((int)
captor.getValue().count());
+ captor.getValue().transferTo(new
ByteBufferWriteableChannel(ciphertextBuffer), 0);
+ ((Buffer) ciphertextBuffer).flip();
+ byte[] ciphertext = new byte[ciphertextBuffer.remaining()];
+ ciphertextBuffer.get(ciphertext);
+
+ // Split in the middle of the 24-byte GCM internal header:
+ // chunk1 = [8-byte length field][4 bytes of GCM header]
+ // chunk2 = [remaining 20 bytes of GCM header][full ciphertext]
+ int splitPoint = 8 + 4;
+ ByteBuf chunk1 = Unpooled.wrappedBuffer(ciphertext, 0, splitPoint);
+ ByteBuf chunk2 = Unpooled.wrappedBuffer(
+ ciphertext, splitPoint, ciphertext.length - splitPoint);
+
+ decryptionHandler.channelRead(ctx, chunk1);
+ // Only a partial header was delivered; no plaintext should be emitted
yet.
+ verify(ctx, never()).fireChannelRead(any());
+
+ ArgumentCaptor<ByteBuf> plaintextCaptor =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx, chunk2);
+ verify(ctx, atLeastOnce()).fireChannelRead(plaintextCaptor.capture());
+
+ byte[] decrypted = new byte[data.length];
+ int offset = 0;
+ for (ByteBuf segment : plaintextCaptor.getAllValues()) {
+ int len = segment.readableBytes();
+ segment.readBytes(decrypted, offset, len);
+ offset += len;
+ }
+ assertEquals(data.length, offset);
+ assertArrayEquals(data, decrypted);
+ }
+ }
+
+ /**
+ * Verifies that DecryptionHandler correctly handles a GCM message whose
8-byte Spark length
+ * prefix is split across two channelRead() calls. This exercises the
partial-fill path in
+ * {@code initializeExpectedLength}: when fewer than 8 bytes of the length
prefix arrive,
+ * the method narrows {@code expectedLengthBuffer}'s limit so that {@code
readBytes()} does
+ * not throw, accumulates the partial result, and returns {@code false} to
signal that the
+ * handler must wait for the rest. The second call completes the prefix and
proceeds normally.
+ */
+ @Test
+ public void testSplitLengthPrefix() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ byte[] data = new byte[1024];
+ Arrays.fill(data, (byte) 'Y');
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> captor =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data), promise);
+ verify(ctx).write(captor.capture(), eq(promise));
+
+ ByteBuffer ciphertextBuffer = ByteBuffer.allocate((int)
captor.getValue().count());
+ captor.getValue().transferTo(new
ByteBufferWriteableChannel(ciphertextBuffer), 0);
+ ((Buffer) ciphertextBuffer).flip();
+ byte[] ciphertext = new byte[ciphertextBuffer.remaining()];
+ ciphertextBuffer.get(ciphertext);
+
+ // Split in the middle of the 8-byte Spark length prefix:
+ // chunk1 = [4 bytes of length prefix]
+ // chunk2 = [remaining 4 bytes of length prefix][GCM header][ciphertext]
+ int splitPoint = 4;
+ ByteBuf chunk1 = Unpooled.wrappedBuffer(ciphertext, 0, splitPoint);
+ ByteBuf chunk2 = Unpooled.wrappedBuffer(
+ ciphertext, splitPoint, ciphertext.length - splitPoint);
+
+ decryptionHandler.channelRead(ctx, chunk1);
+ // Only a partial length prefix was delivered; no plaintext should be
emitted yet.
+ verify(ctx, never()).fireChannelRead(any());
+
+ ArgumentCaptor<ByteBuf> plaintextCaptor =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx, chunk2);
+ verify(ctx, atLeastOnce()).fireChannelRead(plaintextCaptor.capture());
+
+ byte[] decrypted = new byte[data.length];
+ int offset = 0;
+ for (ByteBuf segment : plaintextCaptor.getAllValues()) {
+ int len = segment.readableBytes();
+ segment.readBytes(decrypted, offset, len);
+ offset += len;
+ }
+ assertEquals(data.length, offset);
+ assertArrayEquals(data, decrypted);
+ }
+ }
+
+ /**
+ * Regression test for the encryptedCount miscalculation that caused shuffle
fetch stalls
+ * for plaintext sizes in (plaintextSegmentSize - getCiphertextOffset(),
plaintextSegmentSize]
+ * = (32728, 32752].
+ *
+ * Root cause: encryptedCount was computed using {@code
expectedCiphertextSize(P)} directly.
+ * Tink's formula internally adds {@code getCiphertextOffset()} = 24 to P
before dividing by
+ * plaintextSegmentSize to count segments. For P in (32728, 32752] this
predicted two
+ * ciphertext segments, but {@code transferTo()} writes the Tink header
separately and passes
+ * all P bytes to a single {@code encryptSegment()} call, producing only one
segment. The
+ * resulting encryptedCount was inflated by TAG_SIZE_IN_BYTES = 16, so
+ * {@code count() > transferred()} after all ciphertext was written.
Subsequent
+ * {@code transferTo()} calls returned 0 and the receiver stalled waiting
indefinitely for
+ * 16 bytes that were never sent.
+ */
+ @Test
+ public void testEncryptedCountBoundary() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ // plaintextSegmentSize = CIPHERTEXT_BUFFER_SIZE - TAG_SIZE = 32768 - 16
= 32752
+ // getCiphertextOffset() = 24 (Tink streaming header, written separately
by transferTo)
+ // Buggy range: P in (32728, 32752] - test lower boundary, midpoint, and
upper boundary.
+ int plaintextSegmentSize = GcmTransportCipher.CIPHERTEXT_BUFFER_SIZE -
16;
+ int[] plaintextSizes = {
+ plaintextSegmentSize - 23, // 32729: one above the lower boundary
+ plaintextSegmentSize - 12, // 32740: midpoint of the affected range
+ plaintextSegmentSize // 32752: exactly one full segment (upper
boundary)
+ };
+
+ for (int plaintextSize : plaintextSizes) {
+ reset(ctx);
+ byte[] data = new byte[plaintextSize];
+ Arrays.fill(data, (byte) 'Z');
+
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> encryptedCaptor
=
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data), promise);
+ verify(ctx).write(encryptedCaptor.capture(), eq(promise));
+
+ GcmTransportCipher.GcmEncryptedMessage encrypted =
encryptedCaptor.getValue();
+ ByteBuffer ciphertextBuf = ByteBuffer.allocate((int)
encrypted.count());
+ encrypted.transferTo(new ByteBufferWriteableChannel(ciphertextBuf), 0);
+
+ // Before the fix: count() was inflated by 16 bytes for these sizes, so
+ // transferred() < count() after all plaintext was consumed. The
channel stalled
+ // because subsequent transferTo() calls returned 0 instead of
completing.
+ assertEquals("count() != transferred() for plaintextSize=" +
plaintextSize,
+ encrypted.count(), encrypted.transferred());
+
+ // Verify the full round-trip also decrypts correctly.
+ reset(ctx);
+ ((Buffer) ciphertextBuf).flip();
+ ArgumentCaptor<ByteBuf> plaintextCaptor =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx,
Unpooled.wrappedBuffer(ciphertextBuf));
+ verify(ctx, times(1)).fireChannelRead(plaintextCaptor.capture());
+ ByteBuf plaintext = plaintextCaptor.getValue();
+ assertEquals(plaintextSize, plaintext.readableBytes());
+ assertEquals('Z', plaintext.getByte(0));
+ assertEquals('Z', plaintext.getByte(plaintextSize - 1));
+ }
+ }
+ }
+
+ /**
+ * Regression test for the executor heartbeat timeout caused by per-segment
+ * {@code ctx.fireChannelRead()} calls in {@code DecryptionHandler}. The old
code fired one
+ * EventLoop callback per 32 KB ciphertext segment; decoding a large shuffle
block produced
+ * many synchronous callbacks inside a single {@code processSelectedKeys()}
call,
+ * monopolising the Netty EventLoop and starving the executor-driver
heartbeat task.
+ *
+ * The fix accumulates all decrypted segments into a {@code
CompositeByteBuf} and reduce
+ * the number of {@code ctx.fireChannelRead()} calls. This test verifies
that a
+ * multi-segment plaintext below the 1 MB accumulator threshold
+ * ({@code MAX_PLAINTEXT_BATCH_BYTES}) produces exactly one {@code
fireChannelRead} call.
+ * Above that threshold the accumulator flushes mid-message; see
+ * {@code testLargeMessageBatchedFlush}.
+ */
+ @Test
+ public void testSingleFirePerMessage() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ // Use a 5-segment plaintext so the old per-segment fire (5 calls) is
clearly
+ // distinguishable from the expected single-fire behaviour (1 call).
+ int plaintextSegmentSize = GcmTransportCipher.CIPHERTEXT_BUFFER_SIZE -
16;
+ int plaintextSize = plaintextSegmentSize * 5;
+ byte[] data = new byte[plaintextSize];
+ Arrays.fill(data, (byte) 'M');
+
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> encryptedCaptor =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data), promise);
+ verify(ctx).write(encryptedCaptor.capture(), eq(promise));
+
+ GcmTransportCipher.GcmEncryptedMessage encrypted =
encryptedCaptor.getValue();
+ ByteBuffer ciphertextBuf = ByteBuffer.allocate((int) encrypted.count());
+ encrypted.transferTo(new ByteBufferWriteableChannel(ciphertextBuf), 0);
+ ((Buffer) ciphertextBuf).flip();
+
+ reset(ctx);
+ ArgumentCaptor<ByteBuf> plaintextCaptor =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx,
Unpooled.wrappedBuffer(ciphertextBuf));
+ // The old code fired once per 32 KB segment (5 times for this plaintext
size).
+ // The fix must fire exactly once for the whole message.
+ verify(ctx, times(1)).fireChannelRead(plaintextCaptor.capture());
+ ByteBuf plaintext = plaintextCaptor.getValue();
+ assertEquals(plaintextSize, plaintext.readableBytes());
+ assertEquals('M', plaintext.getByte(0));
+ assertEquals('M', plaintext.getByte(plaintextSize - 1));
+ }
+ }
+
+ /**
+ * Verifies that DecryptionHandler flushes plaintext in multiple batches
when the message
+ * plaintext exceeds {@code MAX_PLAINTEXT_BATCH_BYTES} (1 MB). The bounded
accumulation
+ * prevents full on-heap materialisation of large shuffle blocks that route
to disk via
+ * {@code spark.maxRemoteBlockSizeFetchToMem}. This test uses a 2 MB
plaintext, which
+ * crosses the 1 MB threshold mid-message (triggering the first flush) and
flushes once
+ * more at message completion, producing at least two {@code
fireChannelRead} calls whose
+ * concatenated output must reproduce the original plaintext exactly.
+ */
+ @Test
+ public void testLargeMessageBatchedFlush() throws Exception {
+ TransportConf gcmConf = getConf(2, false);
+ try (AuthEngine client = new AuthEngine("appId", "secret", gcmConf);
+ AuthEngine server = new AuthEngine("appId", "secret", gcmConf)) {
+ AuthMessage clientChallenge = client.challenge();
+ AuthMessage serverResponse = server.response(clientChallenge);
+ client.deriveSessionCipher(clientChallenge, serverResponse);
+ TransportCipher cipher = server.sessionCipher();
+ assert (cipher instanceof GcmTransportCipher);
+ GcmTransportCipher gcmTransportCipher = (GcmTransportCipher) cipher;
+
+ GcmTransportCipher.EncryptionHandler encryptionHandler =
+ gcmTransportCipher.getEncryptionHandler();
+ GcmTransportCipher.DecryptionHandler decryptionHandler =
+ gcmTransportCipher.getDecryptionHandler();
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelPromise promise = mock(ChannelPromise.class);
+
+ // 2 MB plaintext exceeds MAX_PLAINTEXT_BATCH_BYTES (1 MB): the
accumulator flushes
+ // when the threshold is crossed (after ~33 segments) and again at
completion,
+ // so at least two fireChannelRead() calls are expected rather than one.
+ int plaintextSize = 2 * GcmTransportCipher.MAX_PLAINTEXT_BATCH_BYTES;
+ byte[] data = new byte[plaintextSize];
+ Arrays.fill(data, (byte) 'L');
+
+ ArgumentCaptor<GcmTransportCipher.GcmEncryptedMessage> encryptedCaptor =
+
ArgumentCaptor.forClass(GcmTransportCipher.GcmEncryptedMessage.class);
+ encryptionHandler.write(ctx, Unpooled.wrappedBuffer(data), promise);
+ verify(ctx).write(encryptedCaptor.capture(), eq(promise));
+
+ GcmTransportCipher.GcmEncryptedMessage encrypted =
encryptedCaptor.getValue();
+ ByteBuffer ciphertextBuf = ByteBuffer.allocate((int) encrypted.count());
+ encrypted.transferTo(new ByteBufferWriteableChannel(ciphertextBuf), 0);
+ ((Buffer) ciphertextBuf).flip();
+
+ reset(ctx);
+ ArgumentCaptor<ByteBuf> plaintextCaptor =
ArgumentCaptor.forClass(ByteBuf.class);
+ decryptionHandler.channelRead(ctx,
Unpooled.wrappedBuffer(ciphertextBuf));
+ // At least two batches: one when the accumulator crosses 1 MB, one at
completion.
+ verify(ctx, atLeast(2)).fireChannelRead(plaintextCaptor.capture());
+
+ byte[] decrypted = new byte[plaintextSize];
+ int offset = 0;
+ for (ByteBuf segment : plaintextCaptor.getAllValues()) {
+ int len = segment.readableBytes();
+ segment.readBytes(decrypted, offset, len);
+ offset += len;
+ }
+ assertEquals(plaintextSize, offset);
+ assertArrayEquals(data, decrypted);
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
