Copilot commented on code in PR #3689:
URL: https://github.com/apache/celeborn/pull/3689#discussion_r3256230206
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -823,17 +836,34 @@ private boolean fillBuffer() throws IOException {
int batchId = Platform.getInt(sizeBuf, Platform.BYTE_ARRAY_OFFSET +
8);
int size = Platform.getInt(sizeBuf, Platform.BYTE_ARRAY_OFFSET + 12);
- if (shouldDecompress) {
+ // Read and optionally decrypt data into the appropriate buffer
+ if (cryptoHandler.isPresent()) {
+ if (size > encryptedBuf.length) {
+ encryptedBuf = new byte[size];
+ }
+ currentChunk.readBytes(encryptedBuf, 0, size);
+ byte[] decrypted = cryptoHandler.get().decrypt(encryptedBuf, 0,
size);
+ logger.debug(
+ "Decrypted shuffle data for shuffle {} partition {}: {} bytes
-> {} bytes.",
+ shuffleId,
+ partitionId,
+ size,
+ decrypted.length);
+ size = decrypted.length;
+ if (shouldDecompress) {
+ compressedBuf = decrypted;
+ } else {
+ rawDataBuf = decrypted;
+ }
Review Comment:
In the cryptoHandler branch, `size` is overwritten with the decrypted
payload length. Later logic (e.g., bytes-read / duplicate-bytes metrics) uses
`size` to account for how many bytes were fetched from the chunk, which will
undercount network bytes once encryption is enabled (encrypted size !=
decrypted size). Keep the original encrypted size (from the batch header) in a
separate variable and use that for `incBytesRead`/`incDuplicateBytesRead`,
while continuing to use the decrypted size for `limit`/decompression.
##########
client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandlerSuiteJ.java:
##########
@@ -0,0 +1,136 @@
+/*
+ * 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.spark.shuffle.celeborn;
+
+import static org.junit.Assert.*;
+
+import java.io.IOException;
+import java.security.SecureRandom;
+import java.util.Arrays;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.internal.config.package$;
+import org.junit.Before;
+import org.junit.Test;
+
+import org.apache.celeborn.client.security.CryptoHandler;
+
+public class SparkCryptoHandlerSuiteJ {
+
+ private byte[] key;
+ private CryptoHandler handler;
+
+ @Before
+ public void setUp() {
+ key = new byte[16];
+ new SecureRandom().nextBytes(key);
+ SparkConf sparkConf = new SparkConf(false);
+ sparkConf.set(package$.MODULE$.IO_ENCRYPTION_ENABLED(), true);
+ handler = new SparkCryptoHandler(sparkConf, key);
+ }
+
+ @Test
+ public void testRoundTrip() throws IOException {
+ byte[] plaintext = "hello world, this is a test of encryption".getBytes();
+
+ byte[] encrypted = handler.encrypt(plaintext, 0, plaintext.length);
+ assertFalse(
+ "Encrypted output should differ from plaintext",
Arrays.equals(plaintext, encrypted));
+
+ byte[] decrypted = handler.decrypt(encrypted, 0, encrypted.length);
+ assertArrayEquals(plaintext, decrypted);
+ }
+
+ @Test
+ public void testEncryptedDiffersFromPlaintext() throws IOException {
+ byte[] plaintext = "deterministic test data for comparison".getBytes();
+
+ byte[] encrypted = handler.encrypt(plaintext, 0, plaintext.length);
+ assertFalse(
+ "Encrypted output should differ from plaintext",
Arrays.equals(plaintext, encrypted));
+ }
+
+ @Test
+ public void testSameDataEncryptsThenDecrypts() throws IOException {
+ byte[] plaintext = "same data encrypted twice".getBytes();
+
+ byte[] encrypted1 = handler.encrypt(plaintext, 0, plaintext.length);
+ byte[] encrypted2 = handler.encrypt(plaintext, 0, plaintext.length);
+
+ // Both should decrypt to the same plaintext
+ byte[] decrypted1 = handler.decrypt(encrypted1, 0, encrypted1.length);
+ byte[] decrypted2 = handler.decrypt(encrypted2, 0, encrypted2.length);
+
+ assertArrayEquals(plaintext, decrypted1);
+ assertArrayEquals(plaintext, decrypted2);
+ }
+
+ @Test
+ public void testEncryptWithOffset() throws IOException {
+ byte[] actual = "offset test data".getBytes();
+ byte[] padded = Arrays.copyOf(actual, actual.length + 20);
+
+ byte[] encrypted = handler.encrypt(padded, 0, actual.length);
+ byte[] decrypted = handler.decrypt(encrypted, 0, encrypted.length);
+
+ assertArrayEquals(actual, decrypted);
+ }
Review Comment:
`testEncryptWithOffset` doesn't actually test a non-zero offset: it calls
`encrypt(padded, 0, actual.length)`. Use a non-zero offset (and ensure trailing
padding isn’t included) so the test validates the offset/length handling of the
CryptoHandler implementation.
##########
client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.spark.shuffle.celeborn;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.security.CryptoStreamUtils;
+
+import org.apache.celeborn.client.security.CryptoHandler;
+
+public class SparkCryptoHandler implements CryptoHandler {
+ private final SparkConf sparkConf;
+ private final byte[] key;
+
+ public SparkCryptoHandler(SparkConf sparkConf, byte[] key) {
+ this.sparkConf = sparkConf;
+ this.key = key;
+ }
+
+ @Override
+ public byte[] encrypt(byte[] input, int offset, int length) throws
IOException {
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ DataOutputStream dos = new DataOutputStream(baos);
+ dos.writeInt(length);
+ try (OutputStream cos = CryptoStreamUtils.createCryptoOutputStream(dos,
sparkConf, key)) {
+ cos.write(input, offset, length);
+ }
+ return baos.toByteArray();
+ }
+
+ @Override
+ public byte[] decrypt(byte[] input, int offset, int length) throws
IOException {
+ ByteArrayInputStream bais = new ByteArrayInputStream(input, offset,
length);
+ DataInputStream dis = new DataInputStream(bais);
+ int decryptedLength = dis.readInt();
+ if (decryptedLength < 0) {
+ throw new IOException(
+ "Invalid decrypted length: " + decryptedLength + ", encrypted
length: " + length);
+ }
+ try (DataInputStream cis =
+ new DataInputStream(CryptoStreamUtils.createCryptoInputStream(dis,
sparkConf, key))) {
+ byte[] decrypted = new byte[decryptedLength];
+ cis.readFully(decrypted);
+ return decrypted;
Review Comment:
`decrypt()` trusts the first 4 bytes of the ciphertext as `decryptedLength`
and allocates an array of that size. With corrupted input or a wrong key, this
value can be a large positive int, leading to excessive allocation/OOM. Add an
upper-bound validation (e.g., `decryptedLength <= length - 4` and possibly also
a reasonable max from config) and fail fast with an IOException when the bound
is exceeded.
--
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]
