SteNicholas commented on code in PR #3689:
URL: https://github.com/apache/celeborn/pull/3689#discussion_r3501002827
##########
client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java:
##########
@@ -96,4 +102,19 @@ public static void throwSparkOutOfMemoryError() {
}
}
}
+
+ public static Optional<CryptoHandler> getCryptoHandler(SparkConf conf) {
+ if (!(Boolean) conf.get(package$.MODULE$.IO_ENCRYPTION_ENABLED())) {
+ return Optional.empty();
+ }
+ SparkEnv env = SparkEnv.get();
+ if (env == null) {
+ return Optional.empty();
+ }
+ Option<byte[]> key = env.securityManager().getIOEncryptionKey();
+ if (!key.isDefined()) {
Review Comment:
**Fail-open:** when `spark.io.encryption.enabled=true` but the IO key is
undefined (115) or `SparkEnv` is null (111), this returns `Optional.empty()`
with no warn/error/exception — so shuffle data is written and read as
**plaintext** while the user believes encryption-at-rest is on, with zero
signal.
A security control should fail-closed: when encryption is enabled and a
handler cannot be produced, throw (or at minimum `logWarn`/`logError` loudly)
rather than silently degrading. As written, `enabled == true` does not imply
`encrypted-or-error`.
##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java:
##########
@@ -91,6 +93,23 @@ 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() {
+ if (cryptoHandler == null) {
+ // Only cache when SparkEnv is ready. If it is transiently null (e.g.
called before
+ // the executor env is initialized), return empty without caching so the
next call retries.
+ if (SparkEnv.get() != null) {
Review Comment:
The `SparkEnv`-null branch returns `Optional.empty()` *without* caching
specifically so it can retry later — but that empty is passed straight into
`ShuffleClient.get()`, which is one-shot (see the `setupCryptoHandler`
comment). So the retry only refreshes this field; the already-initialized
singleton keeps the empty handler forever.
Also note: when `SparkEnv` is non-null but `getIOEncryptionKey()` is
transiently undefined, `SparkCommonUtils.getCryptoHandler` returns empty and
this code **caches** it (the no-cache guard only covers the `SparkEnv == null`
case), latching plaintext permanently. And `if (cryptoHandler == null)` is a
non-atomic check-then-set, so two threads can build two handlers (benign-same
key, but redundant).
##########
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:
**First-caller-wins:** `setupCryptoHandler(cryptoHandler)` is applied only
inside the one-time `_instance` init block (118 and 125). `_instance` is a
process-wide singleton; once `initialized == true`, every later `get(...,
Optional.of(handler))` returns the cached instance and **never re-applies the
handler**.
So whichever `get()` wins initialization decides crypto for the whole
executor JVM. If that first call carries `Optional.empty()` (see the
`SparkShuffleManager.getCryptoHandler` transient-empty path, and the 6-arg
`get()` overload that passes `Optional.empty()`), the singleton is latched with
no handler: as a writer it silently pushes plaintext; as a reader of
peer-encrypted data it feeds ciphertext into the decompressor → corruption /
FetchFailed.
Fix: apply the handler on every `get()` (it's idempotent), or reject an
empty handler when encryption is enabled, so an empty first-init can't
permanently disable crypto.
##########
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 4-byte plaintext length is written in **cleartext** (45, before the
crypto stream), and this bound only checks `decryptedLength > length - 4`. The
real on-wire layout is `[4-byte len][16-byte IV][ciphertext]`, so the true max
plaintext is `length - 20`, not `length - 4` — the guard (and the comment above
it) are off by the IV.
Because the length is unauthenticated and CTR has no AEAD, an **under-size**
corruption of those 4 bytes passes the check and `cis.readFully(new
byte[decryptedLength])` returns a **silently truncated** plaintext with no
exception (the over-size case throws `EOFException`). With the CommitMetadata
CRC disabled (EAR doesn't require it) this surfaces as garbled records
downstream rather than a clean failure. Tighten the bound to `length - 4 -
IV_LEN`, and consider authenticating the length (or relying on the crypto
layer's own framing) so the prefix can't be tampered undetected.
##########
client-spark/common/pom.xml:
##########
@@ -75,6 +75,10 @@
<artifactId>spark-sql_${scala.binary.version}</artifactId>
<scope>provided</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.commons</groupId>
+ <artifactId>commons-crypto</artifactId>
Review Comment:
`commons-crypto` is added here as a (default) **compile**-scope dependency
and bundled+relocated into `spark-3-shaded` (line 79) and `spark-4-shaded`
(line 81) — but no Celeborn code references `org.apache.commons.crypto`
(`SparkCryptoHandler` only calls Spark's
`org.apache.spark.security.CryptoStreamUtils`, which uses Spark's *own*
unshaded commons-crypto on the executor classpath at runtime).
So the bundled+relocated copy is dead weight (never loaded; the relocation
even moves the native-lib resource path), and the pinned version (`1.0.0` in
the parent, stale) is irrelevant. It's needed only for compile-time type
resolution of `CryptoStreamUtils`' return type → make it `provided` and drop
the two shade `<include>`s.
##########
client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamCryptoRoundTripSuiteJ.java:
##########
@@ -0,0 +1,334 @@
+/*
+ * 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.mockito.ArgumentCaptor;
+
+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.CommitMetadata;
+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 {
Review Comment:
Test-fidelity gap: the only end-to-end read-path crypto coverage uses this
`XorCryptoHandler`, whose format is `[4-byte len][XOR payload]` with **no IV**
— so `encrypted.length == 4 + plaintext.length` and the real
`SparkCryptoHandler`'s decrypt bound (`length-4` vs the true `length-20`) and
IV-stripping are never exercised through `CelebornInputStream`. Combined with
`DummyShuffleClient.setupCryptoHandler` being a no-op (so reader suites that
stub the client drop the handler entirely), the real handler has no integration
coverage of the read path. Consider an end-to-end test that drives the actual
`SparkCryptoHandler` through write→read.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -823,65 +839,104 @@ 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);
Review Comment:
A few hot-path efficiency points with encryption on (not correctness — the
round-trip is verified):
- `decrypt(...)` returns a fresh `byte[]` per batch that's then assigned
into `compressedBuf`/`rawDataBuf` (further down), so the long-lived reusable
read buffers are discarded every batch — the non-crypto path's near-zero
steady-state allocation becomes one+ allocation per batch.
- `init()` pre-sizes `encryptedBuf` to the *plaintext* `bufferSize +
headerLen` (809), but an encrypted batch is always ≥20 bytes larger (len + IV),
so `size > encryptedBuf.length` is hit on essentially every full-size batch and
it reallocates at 882 — the pre-size never fits.
- `Decompressor.getCompressionHeaderLength(conf)` is re-parsed per batch at
896 although `headerLen` already holds it from `init()` (804).
- Write side (`SparkCryptoHandler.encrypt`) copies three times: BAOS →
`toByteArray()` → arraycopy into `body`; the middle copy is avoidable with a
length-prefixed write into a sized buffer.
Per-batch `Cipher.init` + `SecureRandom` IV draw in `CryptoStreamUtils` is
also inherent to creating a fresh stream each call; worth measuring under high
batch throughput.
--
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]
