Copilot commented on code in PR #3689: URL: https://github.com/apache/celeborn/pull/3689#discussion_r3340985743
########## 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: `decryptedLength` is taken directly from the encrypted payload and used to allocate a byte array. With corrupted/malicious input this can trigger very large allocations (executor OOM) before the crypto stream read fails. Add basic sanity checks (e.g., encrypted payload must be at least 4 bytes and decryptedLength must be within a reasonable bound such as `<= length - 4`). ########## 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; Review Comment: `size` is overwritten with the decrypted payload length, but later `callback.incBytesRead/incDuplicateBytesRead(BATCH_HEADER_SIZE + size)` still uses `size`. With encryption enabled this under-reports bytes read (the on-wire/encrypted size from the chunk header), which can skew shuffle read metrics and any logic based on them. Keep both encryptedSize (for metrics) and decryptedSize (for decompression/limit) instead of reusing `size` for both meanings. ########## 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 exercise a non-zero offset (it passes `offset = 0`), so it won't catch bugs in the offset handling. Populate the plaintext at a non-zero offset and pass that offset into `encrypt`. -- 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]
