SteNicholas commented on code in PR #3689:
URL: https://github.com/apache/celeborn/pull/3689#discussion_r3464927095
##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java:
##########
@@ -91,6 +93,17 @@ public class SparkShuffleManager implements ShuffleManager {
private ExecutorShuffleIdTracker shuffleIdTracker = new
ExecutorShuffleIdTracker();
+ // The IO encryption key is fixed for the app lifetime. Lazily initialized
on first
+ // writer/reader call (not in the constructor) to ensure SparkEnv is
available.
+ private volatile Optional<CryptoHandler> cryptoHandler = null;
+
+ private Optional<CryptoHandler> getCryptoHandler() {
Review Comment:
`getCryptoHandler()` is an unsynchronized check-then-act on a `volatile
Optional` with `null` as the sentinel. Two notes:
1. **Race (minor):** concurrent `getReader`/`getWriter` task threads can
each observe `cryptoHandler == null` and independently call
`SparkCommonUtils.getCryptoHandler(conf)`, building redundant
`SparkCryptoHandler` instances and emitting duplicate `IO encryption enabled`
logs. The key is identical, so data stays correct — wasted allocation, not
corruption.
2. **Permanent empty memoization (latent):**
`SparkCommonUtils.getCryptoHandler` returns `Optional.empty()` when
`SparkEnv.get() == null` or the IO key is undefined. If the first reader/writer
call on an executor ever observes a transiently-unready `SparkEnv`, the empty
handler is cached for the JVM lifetime and that executor then reads/writes
shuffle **without** encryption while peers encrypt — a silent cross-executor
`Invalid decrypted length`/garbage-deserialize failure. The comment notes init
is deferred so SparkEnv is up by task time, so this is unlikely in practice;
but since the failure is silent and irreversible, consider not memoizing the
empty result (recompute until a present handler resolves).
This block is duplicated verbatim in spark-2 `SparkShuffleManager` — a
shared helper that caches internally would remove both copies.
##########
client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java:
##########
@@ -1054,6 +1057,22 @@ public int pushOrMergeData(
length = compressor.getCompressedTotalSize();
}
+ // Snapshot volatile field once to avoid a TOCTOU race between isPresent()
and get().
+ Optional<CryptoHandler> handler = cryptoHandler;
+ if (handler.isPresent()) {
+ byte[] encrypted = handler.get().encrypt(data, offset, length);
+ logger.debug(
Review Comment:
Minor: this encrypt-side `logger.debug(...)` is unguarded, while the
symmetric decrypt-side log in `CelebornInputStream` uses `if
(logger.isDebugEnabled())`. With debug off (the prod default) this still
evaluates/boxes 5 args on every push; pick one idiom for both sides (either
drop the read-side guard or add one here).
Non-blocking design note: encryption is applied per-batch in
`pushOrMergeData` after per-batch compression, so every batch carries its own
`CryptoOutputStream` header/IV + 4-byte length (~20B framing) and a fresh
cipher init, inflating on-wire bytes and CPU proportional to batch count rather
than data volume. Encrypting the partition byte stream once would be a deeper
seam if this overhead ever shows up in profiling.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -823,65 +839,94 @@ 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) {
+ // encryptedSize is the on-wire byte count (used for metrics); size
will be
+ // reassigned to the decrypted length after decryption.
+ int encryptedSize = size;
+
+ // Perform dedup/stale-attempt checks before decrypting to avoid
paying the
+ // crypto cost for batches that will be discarded anyway.
+ if (attemptId != attempts[mapId]) {
+ currentChunk.skipBytes(size);
+ continue;
+ }
+ if (readSkewPartitionWithoutMapRange) {
+ LocationPushFailedBatches locationPushFailedBatches =
+
this.failedBatches.get(currentReader.getLocation().getUniqueId());
+ if (null != locationPushFailedBatches) {
+ if (locationPushFailedBatches.contains(mapId, attemptId,
batchId)) {
+ logger.warn(
+ "Skip duplicated batch: mapId={}, attemptId={},
batchId={}",
+ mapId,
+ attemptId,
+ batchId);
+ currentChunk.skipBytes(size);
+ continue;
+ }
+ }
+ }
+ Set<Integer> batchSet = batchesRead.computeIfAbsent(mapId, k -> new
HashSet<>());
+ if (batchSet.contains(batchId)) {
+ callback.incDuplicateBytesRead(BATCH_HEADER_SIZE + encryptedSize);
+ logger.debug(
+ "Skip duplicated batch: mapId {}, attemptId {}, batchId {}.",
+ mapId,
+ attemptId,
+ batchId);
+ currentChunk.skipBytes(size);
+ continue;
+ }
+
+ // Batch is unique and from the correct attempt — now read and
optionally decrypt.
+ if (cryptoHandler.isPresent()) {
+ if (size > encryptedBuf.length) {
+ encryptedBuf = new byte[size];
+ }
+ currentChunk.readBytes(encryptedBuf, 0, size);
+ byte[] decrypted = cryptoHandler.get().decrypt(encryptedBuf, 0,
size);
+ if (logger.isDebugEnabled()) {
+ logger.debug(
+ "Decrypted shuffle data for shuffle {} partition {}: {}
bytes -> {} bytes.",
+ shuffleId,
+ partitionId,
+ size,
+ decrypted.length);
+ }
+ size = decrypted.length;
+ if (shouldDecompress) {
+ compressedBuf = decrypted;
Review Comment:
`decrypt()` returns a fresh array sized exactly to the plaintext, bound
straight into `compressedBuf` here with no check that `decrypted.length` is at
least the compression header size. On the non-crypto branches the payload is
`readBytes` into a `>= bufferSize` buffer; here a corrupt/truncated batch whose
4-byte length prefix decoded to a small-but-non-negative value (cf. the loose
`decryptedLength > length - 4` bound in `SparkCryptoHandler`) yields a short
`decrypted`, and the subsequent `decompressor.getOriginalLen(compressedBuf)`
reads past it — surfacing as an opaque `ArrayIndexOutOfBounds`/fetch failure
rather than a clear corruption error. A length guard before decompress would
make this diagnosable.
##########
client-spark/spark-3-shaded/pom.xml:
##########
@@ -76,6 +76,7 @@
<include>com.google.guava:failureaccess</include>
<include>io.netty:*</include>
<include>org.apache.commons:commons-lang3</include>
+ <include>org.apache.commons:commons-crypto</include>
Review Comment:
spark-2-shaded and spark-3-shaded both add this `commons-crypto` include,
but `client-spark/spark-4-shaded/pom.xml` was not updated — its `<includes>`
block (currently `commons-lang3`, `RoaringBitmap`, `commons-io`, …) has no
`commons-crypto`. Since `client-spark/common` now hard-depends on
`org.apache.commons:commons-crypto` at compile scope (via `SparkCryptoHandler`
→ `CryptoStreamUtils`), an encrypted shuffle on Spark 4 assembled from the
shaded jar would hit `NoClassDefFoundError` at the first encrypt/decrypt unless
commons-crypto is on the classpath transitively from spark-core. Please confirm
Spark 4 provides it transitively, or add the same include to spark-4-shaded for
parity.
##########
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 {
Review Comment:
`decrypt()` (and `encrypt()`) call
`CryptoStreamUtils.createCrypto*Stream(stream, sparkConf, …)` on every batch,
which rebuilds `CryptoParams` via `toCryptoConf(sparkConf)` — iterating the
entire `SparkConf` into a `Properties` — plus a fresh `SecretKeySpec` and
Commons-Crypto cipher init each time. The transformation, key spec, and crypto
`Properties` are app-lifetime constant; with thousands of batches per reducer
partition this per-batch conf scan + cipher init is pure overhead versus the
unencrypted path. Consider resolving the key/properties once in the constructor
and only reading the per-batch IV in the hot loop.
##########
client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamCryptoRoundTripSuiteJ.java:
##########
@@ -0,0 +1,298 @@
+/*
+ * 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 {
+ // Validate the buffer is large enough to hold the 4-byte length prefix
+ if (length < 4) {
+ throw new IOException("Encrypted buffer too short: " + length);
+ }
+ // 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;
duplicate() shares
+ // the underlying data without incrementing the ref count, so
the stream's
+ // single release correctly frees the buffer.
+ cb.onSuccess(0, new NettyManagedBuffer(batchBuf.duplicate()));
+ return null;
+ })
+ .when(client)
+ .fetchChunk(anyLong(), anyInt(), anyLong(),
any(ChunkReceivedCallback.class));
+
+ TransportClientFactory clientFactory = mock(TransportClientFactory.class);
+ when(clientFactory.createClient(anyString(), anyInt())).thenReturn(client);
+
+ ShuffleClient shuffleClient = mock(ShuffleClient.class);
+
+ // PRIMARY location pointing to a single HDD partition
+ PartitionLocation location =
+ new PartitionLocation(
+ 0, 0, "host1", 9001, 9002, 9003, 9004,
PartitionLocation.Mode.PRIMARY);
+ location.setStorageInfo(new StorageInfo(StorageInfo.Type.HDD, true,
"/mnt/disk1"));
+
+ ArrayList<PartitionLocation> locations = new ArrayList<>();
+ locations.add(location);
+ ArrayList<PbStreamHandler> handlers = new ArrayList<>();
+
handlers.add(PbStreamHandler.newBuilder().setStreamId(1L).setNumChunks(1).build());
+
+ return CelebornInputStream.create(
+ conf,
+ clientFactory,
+ SHUFFLE_KEY,
+ locations,
+ handlers,
+ new int[] {0},
+ new HashMap<>(),
+ new HashMap<>(),
+ 0,
+ 1L,
+ 0,
+ 100,
+ new ConcurrentHashMap<>(),
+ shuffleClient,
+ 1,
+ 1,
+ 0,
+ null,
+ new MetricsCallback() {
+ @Override
+ public void incBytesRead(long bytes) {}
+
+ @Override
+ public void incReadTime(long time) {}
+ },
+ needDecompress,
+ cryptoHandler);
+ }
+
+ private byte[] readAll(CelebornInputStream stream) throws IOException {
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ byte[] buf = new byte[4096];
+ int n;
+ while ((n = stream.read(buf)) != -1) {
+ baos.write(buf, 0, n);
+ }
+ return baos.toByteArray();
+ }
+
+ @Test
+ public void testEncryptDecryptRoundTrip() throws IOException,
InterruptedException {
+ byte[] plaintext = "hello, EAR round-trip without compression".getBytes();
+ CelebornConf conf = new CelebornConf();
+ XorCryptoHandler handler = new XorCryptoHandler((byte) 0x5A);
+
+ // Build an encrypted batch and read it back through CelebornInputStream
+ ByteBuf batchBuf = buildBatch(plaintext, false, handler, conf);
+ try (CelebornInputStream stream = createStream(batchBuf, false,
Optional.of(handler), conf)) {
+ assertArrayEquals(plaintext, readAll(stream));
+ }
+ }
+
+ @Test
+ public void testNoEncryptionRoundTrip() throws IOException,
InterruptedException {
+ byte[] plaintext = "unencrypted shuffle data sanity check".getBytes();
+ CelebornConf conf = new CelebornConf();
+
+ // Baseline: with no CryptoHandler the data flows through unchanged
+ ByteBuf batchBuf = buildBatch(plaintext, false, null, conf);
+ try (CelebornInputStream stream = createStream(batchBuf, false,
Optional.empty(), conf)) {
+ assertArrayEquals(plaintext, readAll(stream));
+ }
+ }
+
+ @Test
+ public void testCompressThenEncryptRoundTrip() throws IOException,
InterruptedException {
+ // Reproduce the compress-then-encrypt ordering used in ShuffleClientImpl.
+ byte[] plaintext = "shuffle data with compression and encryption enabled
for EAR".getBytes();
+ CelebornConf conf = new CelebornConf();
+ // Use LZ4 (default)
+ conf.set(CelebornConf.SHUFFLE_COMPRESSION_CODEC().key(), "lz4");
+ XorCryptoHandler handler = new XorCryptoHandler((byte) 0x3C);
+
+ // Writer: LZ4-compress then XOR-encrypt; Reader: decrypt then decompress
+ ByteBuf batchBuf = buildBatch(plaintext, true, handler, conf);
+ try (CelebornInputStream stream = createStream(batchBuf, true,
Optional.of(handler), conf)) {
+ assertArrayEquals(plaintext, readAll(stream));
+ }
+ }
+
+ @Test
+ public void testEncryptWithIntegrityCheckEnabled() throws IOException,
InterruptedException {
Review Comment:
This test sets `integrityCheck.enabled=true` and reads back through
`CelebornInputStream`, but it never actually exercises the integrity
comparison. `validateIntegrity()` only aggregates bytes locally and delegates
the real checksum/byte comparison to
`shuffleClient.readReducerPartitionEnd(...)`; here `shuffleClient` is a Mockito
`mock(ShuffleClient.class)`, so that call is a no-op stub that compares nothing
and never throws, and the test only asserts `decrypted == plaintext` with no
`verify(shuffleClient).readReducerPartitionEnd(...)`. A regression that
computed the integrity checksum over the wrong bytes (e.g. still-encrypted)
would still pass here — false confidence on the most fragile interaction this
PR introduces.
Consider stubbing `readReducerPartitionEnd` to assert the expected
checksum/bytes, or at least `verify(...)` it was invoked with them. Relatedly,
there's no test covering compression + encryption + integrity-check *together*,
which is the combination most likely to diverge from the write-side plaintext
CRC.
--
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]
