Copilot commented on code in PR #3689: URL: https://github.com/apache/celeborn/pull/3689#discussion_r3360083406
########## 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); Review Comment: `testEncryptWithOffset` currently calls `encrypt(padded, 0, actual.length)`, so it doesn't actually exercise the offset parameter. Adjust the test to place `actual` at a non-zero offset in a larger buffer and pass that offset to `encrypt`. ########## client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java: ########## @@ -293,7 +296,8 @@ public static <K, C> CelebornShuffleReader<K, C> createColumnarShuffleReader( TaskContext context, CelebornConf conf, ShuffleReadMetricsReporter metrics, - ExecutorShuffleIdTracker shuffleIdTracker) { + ExecutorShuffleIdTracker shuffleIdTracker, + Optional<CryptoHandler> cryptoHandler) { Review Comment: `SparkUtils.createColumnarShuffleReader` now requires an `Optional<CryptoHandler>` parameter, which breaks existing call sites that still use the old 9-arg signature (e.g. spark-3.5-columnar-shuffle / spark-4-columnar-shuffle tests). To keep source/binary compatibility, add an overload with the previous signature that delegates to this new method with `Optional.empty()`. ########## client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java: ########## @@ -280,7 +282,8 @@ private static class ColumnarShuffleReaderConstructorHolder { TaskContext.class, CelebornConf.class, ShuffleReadMetricsReporter.class, - ExecutorShuffleIdTracker.class) + ExecutorShuffleIdTracker.class, + Optional.class) Review Comment: `ColumnarShuffleReaderConstructorHolder` currently looks up only the constructor signature that includes `Optional` (crypto handler). Older columnar-shuffle modules (e.g. spark-3.5-columnar-shuffle / spark-4-columnar-shuffle) still define `CelebornColumnarShuffleReader` without that parameter, which will cause a runtime constructor lookup failure. Add a fallback `.impl(...)` without `Optional.class`; `DynConstructors.Ctor` will truncate extra args automatically when invoking. ########## client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java: ########## @@ -2090,6 +2108,14 @@ public void setExtension(byte[] extension) { this.extension = extension; } + @Override + public void setupCryptoHandler(Optional<CryptoHandler> cryptoHandler) { + this.cryptoHandler = cryptoHandler; + if (cryptoHandler.isPresent()) { + logger.info("IO encryption enabled for shuffle data (encryption at rest)."); + } Review Comment: `setupCryptoHandler` assigns the passed `Optional` directly. Since this is a public API, a null caller would leave `this.cryptoHandler` null and later `cryptoHandler.isPresent()` checks will throw NPE. Please defensively handle null by defaulting to `Optional.empty()` and log based on the stored field. ########## 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); + } Review Comment: `decrypt` trusts the unverified length prefix and allocates `new byte[decryptedLength]` with only a negative check. If the encrypted payload is corrupted or tampered with, this can trigger huge allocations / OOM. Add an upper-bound sanity check before allocating (e.g. `decryptedLength > length`). -- 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]
