Copilot commented on code in PR #3689:
URL: https://github.com/apache/celeborn/pull/3689#discussion_r3418771752
##########
client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java:
##########
@@ -1054,6 +1057,20 @@ public int pushOrMergeData(
length = compressor.getCompressedTotalSize();
}
+ if (cryptoHandler.isPresent()) {
+ byte[] encrypted = cryptoHandler.get().encrypt(data, offset, length);
+ logger.debug(
+ "Encrypted shuffle data for shuffle {} map {} partition {}: {} bytes
-> {} bytes.",
+ shuffleId,
+ mapId,
+ partitionId,
+ length,
+ encrypted.length);
+ data = encrypted;
+ offset = 0;
+ length = encrypted.length;
+ }
Review Comment:
cryptoHandler is a volatile Optional; calling isPresent() and then get()
rereads the volatile and can theoretically race with setupCryptoHandler updates
(leading to NoSuchElementException). Reading it once into a local variable
avoids this and is clearer.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -790,9 +801,14 @@ private boolean moveToNextChunk() throws IOException {
private void init() {
int bufferSize = conf.clientFetchBufferSize();
+ int headerLen = shouldDecompress ?
Decompressor.getCompressionHeaderLength(conf) : 0;
+ if (cryptoHandler.isPresent()) {
+ // The encrypted payload is: IV(16) + ciphertext(compressedSize), where
+ // compressedSize can reach bufferSize + headerLen, so match the same
headroom.
+ encryptedBuf = new byte[bufferSize + headerLen];
Review Comment:
This comment hard-codes an encrypted payload layout ("IV(16) + ciphertext")
that doesn't match the actual CryptoHandler contract used here (e.g.,
SparkCryptoHandler prefixes a 4-byte plaintext length, and cipher streams may
add additional overhead). This can mislead future changes around buffer sizing.
##########
client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamCryptoRoundTripSuiteJ.java:
##########
@@ -0,0 +1,292 @@
+/*
+ * 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 org.apache.celeborn.client.read;
+
+import static org.junit.Assert.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import org.junit.Test;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.compress.Compressor;
+import org.apache.celeborn.client.security.CryptoHandler;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.buffer.NettyManagedBuffer;
+import org.apache.celeborn.common.network.client.ChunkReceivedCallback;
+import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientFactory;
+import org.apache.celeborn.common.network.protocol.TransportMessage;
+import org.apache.celeborn.common.protocol.MessageType;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PbStreamHandler;
+import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.common.unsafe.Platform;
+
+/**
+ * Integration-style round-trip tests for EAR (Encryption At Rest) wiring in
{@link
+ * CelebornInputStream}. These tests verify that the encrypt-on-write /
decrypt-on-read path works
+ * end-to-end, including interactions with compression and the shuffle
integrity check.
+ */
+public class CelebornInputStreamCryptoRoundTripSuiteJ {
+
+ private static final int BATCH_HEADER_SIZE = 16;
+ private static final String SHUFFLE_KEY = "app-1-1";
+
+ /**
+ * A minimal CryptoHandler for testing: the encrypted format is [4-byte
plaintext length
+ * (int)][XOR-encrypted payload]. This matches the structural contract of
SparkCryptoHandler so
+ * the bounds check (decryptedLength > length - 4) is also exercised.
+ */
+ static class XorCryptoHandler implements CryptoHandler {
+ private final byte key;
+
+ XorCryptoHandler(byte key) {
+ this.key = key;
+ }
+
+ @Override
+ public byte[] encrypt(byte[] input, int offset, int length) throws
IOException {
+ // Prefix with 4-byte plaintext length, then XOR-encrypt the payload
+ byte[] out = new byte[4 + length];
+ Platform.putInt(out, Platform.BYTE_ARRAY_OFFSET, length);
+ for (int i = 0; i < length; i++) {
+ out[4 + i] = (byte) (input[offset + i] ^ key);
+ }
+ return out;
+ }
+
+ @Override
+ public byte[] decrypt(byte[] input, int offset, int length) throws
IOException {
+ // Read the plaintext length from the 4-byte prefix
+ int decryptedLength = Platform.getInt(input, Platform.BYTE_ARRAY_OFFSET
+ offset);
+ // Validate bounds: the 4-byte prefix must fit inside the encrypted
buffer
+ if (decryptedLength < 0 || decryptedLength > length - 4) {
+ throw new IOException(
+ "Invalid decrypted length: " + decryptedLength + ", encrypted
length: " + length);
+ }
+ byte[] out = new byte[decryptedLength];
+ for (int i = 0; i < decryptedLength; i++) {
+ out[i] = (byte) (input[offset + 4 + i] ^ key);
+ }
+ return out;
+ }
+ }
+
+ /**
+ * Build a single batch ByteBuf as ShuffleClientImpl.pushOrMergeData does:
optionally compress,
+ * optionally encrypt, then prepend the 16-byte batch header.
+ */
+ private ByteBuf buildBatch(
+ byte[] plaintext, boolean compress, CryptoHandler cryptoHandler,
CelebornConf conf)
+ throws IOException {
+ byte[] data = plaintext;
+ int offset = 0;
+ int length = plaintext.length;
+
+ // Step 1: optionally compress (compress-then-encrypt ordering matches
ShuffleClientImpl)
+ if (compress) {
+ Compressor compressor = Compressor.getCompressor(conf);
+ compressor.compress(data, offset, length);
+ data = compressor.getCompressedBuffer();
+ offset = 0;
+ length = compressor.getCompressedTotalSize();
+ }
+
+ // Step 2: optionally encrypt the (possibly compressed) payload
+ if (cryptoHandler != null) {
+ data = cryptoHandler.encrypt(data, offset, length);
+ offset = 0;
+ length = data.length;
+ }
+
+ // Step 3: prepend the 16-byte batch header
[mapId|attemptId|batchId|payloadLen]
+ byte[] body = new byte[BATCH_HEADER_SIZE + length];
+ Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET, 0); // mapId
+ Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 4, 0); // attemptId
+ Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 8, 0); // batchId
+ Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 12, length); // payload
length
+ System.arraycopy(data, offset, body, BATCH_HEADER_SIZE, length);
+ return Unpooled.wrappedBuffer(body);
+ }
+
+ /**
+ * Create a CelebornInputStream backed by a mock TransportClient that serves
the given batchBuf as
+ * a single chunk.
+ */
+ private CelebornInputStream createStream(
+ ByteBuf batchBuf,
+ boolean needDecompress,
+ Optional<CryptoHandler> cryptoHandler,
+ CelebornConf conf)
+ throws IOException, InterruptedException {
+ TransportClient client = mock(TransportClient.class);
+ PbStreamHandler pbHandler =
+ PbStreamHandler.newBuilder().setStreamId(1L).setNumChunks(1).build();
+ // Encode the stream handler into an RPC response that CelebornInputStream
expects
+ ByteBuffer rpcResponse =
+ new TransportMessage(MessageType.STREAM_HANDLER,
pbHandler.toByteArray()).toByteBuffer();
+ when(client.sendRpcSync(any(ByteBuffer.class),
anyLong())).thenReturn(rpcResponse);
+ doNothing().when(client).sendRpc(any(ByteBuffer.class));
+ doAnswer(
+ invocation -> {
+ ChunkReceivedCallback cb = invocation.getArgument(3);
+ // Serve the pre-built batch buffer immediately as chunk 0
+ cb.onSuccess(0, new
NettyManagedBuffer(batchBuf.duplicate().retain()));
+ return null;
Review Comment:
The test retains the ByteBuf before wrapping it in NettyManagedBuffer, but
the stream will release the buffer only once; this leaves the original
ref-count unreleased and can trigger Netty leak detection. Since the test
doesn't use batchBuf after handing it to the stream, don't retain here (or
explicitly release batchBuf in a finally block).
##########
client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamCryptoRoundTripSuiteJ.java:
##########
@@ -0,0 +1,292 @@
+/*
+ * 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 org.apache.celeborn.client.read;
+
+import static org.junit.Assert.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
+import org.junit.Test;
+
+import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.compress.Compressor;
+import org.apache.celeborn.client.security.CryptoHandler;
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.network.buffer.NettyManagedBuffer;
+import org.apache.celeborn.common.network.client.ChunkReceivedCallback;
+import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientFactory;
+import org.apache.celeborn.common.network.protocol.TransportMessage;
+import org.apache.celeborn.common.protocol.MessageType;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PbStreamHandler;
+import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.common.unsafe.Platform;
+
+/**
+ * Integration-style round-trip tests for EAR (Encryption At Rest) wiring in
{@link
+ * CelebornInputStream}. These tests verify that the encrypt-on-write /
decrypt-on-read path works
+ * end-to-end, including interactions with compression and the shuffle
integrity check.
+ */
+public class CelebornInputStreamCryptoRoundTripSuiteJ {
+
+ private static final int BATCH_HEADER_SIZE = 16;
+ private static final String SHUFFLE_KEY = "app-1-1";
+
+ /**
+ * A minimal CryptoHandler for testing: the encrypted format is [4-byte
plaintext length
+ * (int)][XOR-encrypted payload]. This matches the structural contract of
SparkCryptoHandler so
+ * the bounds check (decryptedLength > length - 4) is also exercised.
+ */
+ static class XorCryptoHandler implements CryptoHandler {
+ private final byte key;
+
+ XorCryptoHandler(byte key) {
+ this.key = key;
+ }
+
+ @Override
+ public byte[] encrypt(byte[] input, int offset, int length) throws
IOException {
+ // Prefix with 4-byte plaintext length, then XOR-encrypt the payload
+ byte[] out = new byte[4 + length];
+ Platform.putInt(out, Platform.BYTE_ARRAY_OFFSET, length);
+ for (int i = 0; i < length; i++) {
+ out[4 + i] = (byte) (input[offset + i] ^ key);
+ }
+ return out;
+ }
+
+ @Override
+ public byte[] decrypt(byte[] input, int offset, int length) throws
IOException {
+ // Read the plaintext length from the 4-byte prefix
+ int decryptedLength = Platform.getInt(input, Platform.BYTE_ARRAY_OFFSET
+ offset);
+ // Validate bounds: the 4-byte prefix must fit inside the encrypted
buffer
+ if (decryptedLength < 0 || decryptedLength > length - 4) {
+ throw new IOException(
+ "Invalid decrypted length: " + decryptedLength + ", encrypted
length: " + length);
+ }
Review Comment:
XorCryptoHandler.decrypt reads the 4-byte length prefix via Platform.getInt
before validating that the encrypted buffer is at least 4 bytes long. Because
Platform uses unsafe access, a short/corrupted buffer could cause an
out-of-bounds read (potentially crashing the JVM) rather than a clean test
failure.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
