SteNicholas commented on code in PR #3689: URL: https://github.com/apache/celeborn/pull/3689#discussion_r3400349580
########## client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java: ########## @@ -0,0 +1,71 @@ +/* + * 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(); + // The encrypted payload format is: [4-byte plaintext length][ciphertext...]. + // So the maximum valid decrypted length is length - 4 (the ciphertext portion). + // A value outside this range indicates corruption or a wrong key. + if (decryptedLength < 0 || decryptedLength > length - 4) { Review Comment: The bound `decryptedLength > length - 4` is loose: this length field is written in **cleartext** before the cipher stream, so a corrupted prefix value that is `>= 0` and `<= length - 4` passes validation, and `readFully(new byte[decryptedLength])` then silently returns a truncated plaintext instead of failing. Because the prefix is outside the cipher it also can't detect a wrong key. If the prefix is kept, consider asserting it equals the exact expected plaintext length rather than a loose upper bound. ########## client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java: ########## @@ -304,7 +320,32 @@ public static <K, C> CelebornShuffleReader<K, C> createColumnarShuffleReader( context, conf, metrics, - shuffleIdTracker); + shuffleIdTracker, + cryptoHandler); Review Comment: **Blocking: encryption is silently dropped for Spark 3.5 / 4 columnar shuffle.** `createColumnarShuffleReader` always passes `cryptoHandler` as the 10th `invoke()` arg, but the columnar reader in `spark-3.5-columnar-shuffle` and `spark-4-columnar-shuffle` was **not** updated with this parameter (only `spark-3-columnar-shuffle` was). Both still have the 9-arg constructor and both depend on this modified `celeborn-client-spark-3`. When `DynConstructors` binds the 9-arg fallback `.impl` (kept just above "for older columnar-shuffle modules"), `newInstanceChecked` **silently truncates** the extra arg: ```java if (args.length > ctor.getParameterCount()) { return ctor.newInstance(Arrays.copyOfRange(args, 0, ctor.getParameterCount())); } ``` So with `spark.io.encryption.enabled=true` + columnar shuffle on Spark 3.5/4, the writer encrypts but the reader is constructed with `Optional.empty()` and never decrypts → ciphertext is fed to the decompressor/deserializer → corruption / `SHUFFLE_DATA_LOST`, with no error. Suggest adding the `cryptoHandler` param to the spark-3.5 and spark-4 `CelebornColumnarShuffleReader` (and updating their test call sites), or making the reflection fail-fast instead of truncating. The new round-trip test won't catch this (it builds `CelebornInputStream` directly, not via the columnar reader). ########## client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java: ########## @@ -0,0 +1,71 @@ +/* + * 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 { Review Comment: AES/CTR (Spark's default IO-encryption transform) gives confidentiality but **no authentication**, and `celeborn.client.shuffle.integrityCheck.enabled` defaults to `false`. So a wrong key or a single bit-flip in ciphertext decrypts to garbage of the correct length with no exception, and by default nothing downstream detects it. `testDecryptWithWrongKeyFails` documents this ("CryptoStreamUtils may return garbage instead of throwing"). At minimum, a doc note that corruption/wrong-key detection requires enabling the shuffle integrity check would help. Two smaller points on this class: - `encrypt`/`decrypt` allocate a fresh `ByteArrayOutputStream`/`ByteArrayInputStream` + crypto stream + a new `byte[]` per batch on the push/read hot path; consider reusing Spark's `SerializerManager.wrapStream` and/or a reusable buffer. - The 4-byte cleartext length prefix is redundant under CTR (ciphertext length == plaintext length and the stream signals EOF); it adds 4 bytes/batch plus a cross-method format invariant. ########## client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java: ########## @@ -823,17 +836,38 @@ 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. + // encryptedSize tracks the on-wire (encrypted) byte count for metrics; size is + // reassigned to the decrypted length so downstream decompression and limit logic + // operate on the correct plaintext size. + int encryptedSize = size; + if (cryptoHandler.isPresent()) { Review Comment: Two efficiency notes on the encrypted read branch: - Decryption runs here, **before** the `attemptId == attempts[mapId]` and `batchSet.contains(batchId)` dedup checks below, so duplicate / stale-attempt batches are fully decrypted (cipher run + new `byte[]`) and then discarded. Decrypt is the expensive part — consider evaluating the dedup/attempt checks before decrypting. - `compressedBuf = decrypted` / `rawDataBuf = decrypted` reassign to the fresh array returned by `decrypt()` every batch, discarding the pre-allocated grow-only buffers `init()` set up. A decrypt-into-caller-buffer API would preserve the non-crypto path's buffer reuse. ########## client/src/main/java/org/apache/celeborn/client/ShuffleClient.java: ########## @@ -102,12 +115,14 @@ public static ShuffleClient get( _instance = new ShuffleClientImpl(appUniqueId, conf, userIdentifier); _instance.setupLifecycleManagerRef(driverHost, port); _instance.setExtension(extension); + _instance.setupCryptoHandler(cryptoHandler); Review Comment: The singleton binds the crypto handler only when `_instance` is first created. If the singleton is initialized first via a path that passes `Optional.empty()` — the 5/6-arg `get()` overloads, or a Spark call before `SparkEnv.get()`/the IO key is available — a later encryption-enabled `getReader`/`getWriter` reuses the handler-less instance and never installs the handler, so that executor reads/writes without decrypt/encrypt while the rest of the job assumes encryption. Same shape as `setExtension`, but here the failure is silent plaintext / undecryptable data rather than a missing tag. Worth confirming the first `get()` on an executor always comes from an encryption-aware path. ########## client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java: ########## @@ -188,6 +192,7 @@ private static final class CelebornInputStreamImpl extends CelebornInputStream { private final Map<String, LocationPushFailedBatches> failedBatches; + private byte[] encryptedBuf; Review Comment: `close()` nulls `compressedBuf`, `rawDataBuf`, `decompressor`, etc. for GC but not this new `encryptedBuf`. With encryption on and high partition fan-out, each closed-but-still-referenced stream retains its `encryptedBuf` (possibly grown for a large batch). Add `encryptedBuf = null;` to `close()` for consistency. ########## client/src/main/java/org/apache/celeborn/client/ShuffleClient.java: ########## @@ -150,6 +165,8 @@ public static void printReadStats(Logger logger) { String.format("%.2f", (localReadCount * 1.0d / totalReadCount) * 100)); } + public abstract void setupCryptoHandler(Optional<CryptoHandler> cryptoHandler); Review Comment: The `CryptoHandler` hook lives in shared `client/`, but only the Spark client ever supplies one — Flink/MR/Tez share `ShuffleClientImpl.pushOrMergeData` / `CelebornInputStream` and never call `setupCryptoHandler` (they use the 5/6-arg `get()`), so enabling encryption is a silent no-op for those engines. Reasonable scope for a "Spark impl" PR, but since the abstraction sits in the shared module, a doc note (or a warn/guard when the config is set on a non-Spark engine) would prevent operators from assuming cluster-wide coverage. ########## client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java: ########## @@ -101,6 +102,8 @@ public class ShuffleClientImpl extends ShuffleClient { protected byte[] extension; + private Optional<CryptoHandler> cryptoHandler = Optional.empty(); Review Comment: `cryptoHandler` is a non-volatile instance field, written by `setupCryptoHandler` (under `synchronized(ShuffleClient.class)` during singleton init) and read on the data-pusher push threads (`pushOrMergeData`) and read threads. Safe today via the init-once happens-before, but if `setupCryptoHandler` is ever called on an already-running client, push threads have no visibility guarantee and could observe a stale `Optional.empty()`. Marking it `volatile` would make the intent explicit and future-proof. ########## client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java: ########## @@ -791,6 +801,9 @@ private boolean moveToNextChunk() throws IOException { private void init() { int bufferSize = conf.clientFetchBufferSize(); + if (cryptoHandler.isPresent()) { + encryptedBuf = new byte[bufferSize]; Review Comment: `encryptedBuf` is sized to `clientFetchBufferSize()` here, **before** the `bufferSize += headerLen` adjustment that `compressedBuf`/`rawDataBuf` get just below. The on-wire encrypted size is `4 + IV(16) + compressedSize`, and `compressedSize` can reach `bufferSize + headerLen`, so `size > encryptedBuf.length` in `fillBuffer` fires on nearly every batch when compression is enabled → a reallocation per batch on the read hot path. Size it with the same `headerLen` (plus crypto) allowance the other buffers use. ########## client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java: ########## @@ -304,7 +320,32 @@ public static <K, C> CelebornShuffleReader<K, C> createColumnarShuffleReader( context, conf, metrics, - shuffleIdTracker); + shuffleIdTracker, + cryptoHandler); + } + + /** Overload for callers that do not use encryption at rest. */ + public static <K, C> CelebornShuffleReader<K, C> createColumnarShuffleReader( Review Comment: This no-crypto overload, the kept 9-arg fallback `.impl` above, and the extra `this(...)` aux constructors on the spark-3 reader all exist only to avoid touching the spark-3.5/spark-4 columnar modules. If those modules are updated for the `cryptoHandler` param (see the blocking comment), this overload and the fallback `.impl` can both be dropped, leaving a single reflective constructor. (`setupCryptoHandler`'s `!= null` guard in `ShuffleClientImpl` is also dead — all callers pass a non-null `Optional`.) -- 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]
