This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 1e715b2cf2 [CELEBORN-2329][CIP22] Encryption at Rest Spark Impl
1e715b2cf2 is described below
commit 1e715b2cf2936a3704c6a44da463628f7e737f35
Author: Aravind Patnam <[email protected]>
AuthorDate: Wed Jul 8 13:13:12 2026 +0800
[CELEBORN-2329][CIP22] Encryption at Rest Spark Impl
### What changes were proposed in this pull request?
Adds EAR support for Spark side.
See more details in
[doc](https://docs.google.com/document/d/1xBrLtpb8bk8CdJENiM3aLJCbFThRKLor8uAjSsfXoxE/edit?usp=sharing).
### Why are the changes needed?
See more details in
[doc](https://docs.google.com/document/d/1xBrLtpb8bk8CdJENiM3aLJCbFThRKLor8uAjSsfXoxE/edit?usp=sharing).
### Does this PR resolve a correctness bug?
- [ ] Yes.
### Does this PR introduce _any_ user-facing change?
- [ ] Yes.
### How was this patch tested?
Unit tests and tested in production internally.
Closes #3689 from akpatnam25/ear-spark-impl.
Lead-authored-by: Aravind Patnam <[email protected]>
Co-authored-by: Aravind Patnam <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../spark/shuffle/celeborn/SparkCommonUtils.java | 29 ++
.../spark/shuffle/celeborn/SparkCryptoHandler.java | 91 ++++++
.../shuffle/celeborn/SparkCryptoHandlerSuiteJ.java | 165 ++++++++++
.../celeborn/CelebornColumnarShuffleReader.scala | 9 +-
.../CelebornColumnarShuffleReaderSuite.scala | 8 +-
.../celeborn/CelebornColumnarShuffleReader.scala | 9 +-
.../CelebornColumnarShuffleReaderSuite.scala | 8 +-
.../shuffle/celeborn/SparkShuffleManager.java | 28 +-
.../apache/spark/shuffle/celeborn/SparkUtils.java | 11 +-
.../shuffle/celeborn/CelebornShuffleReader.scala | 35 ++-
.../celeborn/CelebornShuffleReaderSuite.scala | 2 +-
.../celeborn/CelebornColumnarShuffleReader.scala | 9 +-
.../CelebornColumnarShuffleReaderSuite.scala | 8 +-
.../apache/celeborn/client/DummyShuffleClient.java | 5 +
.../org/apache/celeborn/client/ShuffleClient.java | 24 ++
.../apache/celeborn/client/ShuffleClientImpl.java | 32 +-
.../celeborn/client/read/CelebornInputStream.java | 162 ++++++----
.../celeborn/client/security/CryptoHandler.java | 26 ++
.../CelebornInputStreamCryptoRoundTripSuiteJ.java | 334 +++++++++++++++++++++
.../read/CelebornInputStreamPeerFailoverTest.java | 8 +-
20 files changed, 924 insertions(+), 79 deletions(-)
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
index 84d74f8c14..08e217390c 100644
---
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
@@ -19,15 +19,25 @@ package org.apache.spark.shuffle.celeborn;
import java.util.Collections;
import java.util.Map;
+import java.util.Optional;
+
+import scala.Option;
import org.apache.spark.SparkConf;
+import org.apache.spark.SparkEnv;
import org.apache.spark.TaskContext;
+import org.apache.spark.internal.config.package$;
import org.apache.spark.memory.SparkOutOfMemoryError;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.reflect.DynConstructors;
import org.apache.celeborn.reflect.DynMethods;
public class SparkCommonUtils {
+ private static final Logger logger =
LoggerFactory.getLogger(SparkCommonUtils.class);
+
public static void validateAttemptConfig(SparkConf conf) throws
IllegalArgumentException {
int DEFAULT_MAX_CONSECUTIVE_STAGE_ATTEMPTS = 4;
int maxStageAttempts =
@@ -96,4 +106,23 @@ public class SparkCommonUtils {
}
}
}
+
+ 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()) {
+ logger.warn(
+ "IO encryption is enabled (spark.io.encryption.enabled=true) but the
IO encryption key "
+ + "is not available from the SecurityManager. Shuffle data will
be written as "
+ + "plaintext. Ensure the SecurityManager provides an IO
encryption key.");
+ return Optional.empty();
+ }
+ return Optional.of(new SparkCryptoHandler(conf, key.get()));
+ }
}
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java
new file mode 100644
index 0000000000..bf3922d484
--- /dev/null
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java
@@ -0,0 +1,91 @@
+/*
+ * 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 java.util.Properties;
+
+import org.apache.spark.SparkConf;
+import org.apache.spark.security.CryptoStreamUtils;
+
+import org.apache.celeborn.client.security.CryptoHandler;
+
+public class SparkCryptoHandler implements CryptoHandler {
+ // On-wire format: [4-byte plaintext length][16-byte IV][ciphertext].
+ // The minimum overhead (length prefix + IV) added by the crypto stream.
+ private static final int CRYPTO_OVERHEAD_BYTES =
+ Integer.BYTES + CryptoStreamUtils.IV_LENGTH_IN_BYTES();
+
+ private final SparkConf sparkConf;
+ private final byte[] key;
+ // Each push thread reuses its own ByteArrayOutputStream to avoid per-batch
allocation.
+ // The internal buffer grows to the high-water-mark once per thread and is
reset() each call.
+ private final ThreadLocal<ByteArrayOutputStream> encryptBaos =
+ ThreadLocal.withInitial(ByteArrayOutputStream::new);
+
+ public SparkCryptoHandler(SparkConf sparkConf, byte[] key) {
+ // Pre-filter sparkConf to only crypto-relevant keys so that
+ // CryptoStreamUtils.toCryptoConf() does not scan the full SparkConf on
every batch.
+ Properties cryptoProps = CryptoStreamUtils.toCryptoConf(sparkConf);
+ SparkConf minimalConf = new SparkConf(false);
+ String prefix =
CryptoStreamUtils.SPARK_IO_ENCRYPTION_COMMONS_CONFIG_PREFIX();
+ for (String propKey : cryptoProps.stringPropertyNames()) {
+ minimalConf.set(prefix + propKey, cryptoProps.getProperty(propKey));
+ }
+ this.sparkConf = minimalConf;
+ this.key = key;
+ }
+
+ @Override
+ public byte[] encrypt(byte[] input, int offset, int length) throws
IOException {
+ ByteArrayOutputStream baos = encryptBaos.get();
+ baos.reset();
+ 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][16-byte
IV][ciphertext].
+ // The minimum on-wire overhead is CRYPTO_OVERHEAD_BYTES (4 + 16 = 20), so
the maximum
+ // valid plaintext length is length - 20. A value outside this range
indicates corruption
+ // or a wrong key.
+ if (decryptedLength < 0 || decryptedLength > length -
CRYPTO_OVERHEAD_BYTES) {
+ 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;
+ }
+ }
+}
diff --git
a/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandlerSuiteJ.java
b/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandlerSuiteJ.java
new file mode 100644
index 0000000000..6baa62626f
--- /dev/null
+++
b/client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandlerSuiteJ.java
@@ -0,0 +1,165 @@
+/*
+ * 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.nio.ByteBuffer;
+import java.nio.ByteOrder;
+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();
+ int offset = 10;
+ byte[] padded = new byte[offset + actual.length + 20];
+ System.arraycopy(actual, 0, padded, offset, actual.length);
+
+ byte[] encrypted = handler.encrypt(padded, offset, actual.length);
+ byte[] decrypted = handler.decrypt(encrypted, 0, encrypted.length);
+
+ assertArrayEquals(actual, decrypted);
+ }
+
+ @Test
+ public void testDecryptWithWrongKeyFails() throws IOException {
+ byte[] plaintext = "secret data".getBytes();
+ byte[] encrypted = handler.encrypt(plaintext, 0, plaintext.length);
+
+ byte[] wrongKey = new byte[16];
+ new SecureRandom().nextBytes(wrongKey);
+ SparkConf sparkConf = new SparkConf(false);
+ sparkConf.set(package$.MODULE$.IO_ENCRYPTION_ENABLED(), true);
+ CryptoHandler wrongHandler = new SparkCryptoHandler(sparkConf, wrongKey);
+
+ byte[] decrypted = null;
+ try {
+ decrypted = wrongHandler.decrypt(encrypted, 0, encrypted.length);
+ } catch (IOException e) {
+ // acceptable — some implementations throw on wrong key
+ return;
+ }
+ // CryptoStreamUtils may return garbage instead of throwing
+ assertFalse(
+ "Decryption with wrong key should not produce original plaintext",
+ Arrays.equals(plaintext, decrypted));
+ }
+
+ @Test
+ public void testLargeData() throws IOException {
+ byte[] plaintext = new byte[64 * 1024]; // 64KB
+ new SecureRandom().nextBytes(plaintext);
+
+ byte[] encrypted = handler.encrypt(plaintext, 0, plaintext.length);
+ byte[] decrypted = handler.decrypt(encrypted, 0, encrypted.length);
+
+ assertArrayEquals(plaintext, decrypted);
+ }
+
+ @Test
+ public void testEmptyData() throws IOException {
+ byte[] encrypted = handler.encrypt(new byte[0], 0, 0);
+
+ byte[] decrypted = handler.decrypt(encrypted, 0, encrypted.length);
+ assertEquals(0, decrypted.length);
+ }
+
+ /**
+ * Verifies that the decrypt bounds check uses {@code length - 20} (4-byte
length prefix + 16-byte
+ * IV), not the previous {@code length - 4}. A crafted payload whose
embedded length value is
+ * between {@code length - 19} and {@code length - 5} (inclusive) must be
rejected.
+ */
+ @Test
+ public void
testDecryptRejectsCraftedLengthBetweenLengthMinus4AndLengthMinus20()
+ throws IOException {
+ // Construct a minimal on-wire buffer: [4-byte length][16-byte IV][0-byte
ciphertext].
+ // Total = 20 bytes. Embed a plaintext length of 1 — valid under the old
(length-4)
+ // guard (1 <= 20-4=16) but invalid under the corrected (length-20) guard
(1 > 20-20=0).
+ int totalLen = 20; // 4 (length prefix) + 16 (IV) + 0 (ciphertext)
+ byte[] crafted = new byte[totalLen];
+ ByteBuffer.wrap(crafted).order(ByteOrder.BIG_ENDIAN).putInt(1); // claim 1
byte of plaintext
+
+ try {
+ handler.decrypt(crafted, 0, crafted.length);
+ fail("Expected IOException for crafted length > length - 20");
+ } catch (IOException e) {
+ assertTrue(
+ "Exception message should mention decrypted length",
+ e.getMessage().contains("decrypted length") ||
e.getMessage().contains("Invalid"));
+ }
+ }
+}
diff --git
a/client-spark/spark-3-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
b/client-spark/spark-3-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
index fd888fb9dc..c32dfaf2eb 100644
---
a/client-spark/spark-3-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
+++
b/client-spark/spark-3-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
@@ -17,12 +17,15 @@
package org.apache.spark.shuffle.celeborn
+import java.util.Optional
+
import org.apache.spark.{ShuffleDependency, TaskContext}
import org.apache.spark.serializer.SerializerInstance
import org.apache.spark.shuffle.ShuffleReadMetricsReporter
import org.apache.spark.sql.execution.UnsafeRowSerializer
import org.apache.spark.sql.execution.columnar.{CelebornBatchBuilder,
CelebornColumnarBatchSerializer}
+import org.apache.celeborn.client.security.CryptoHandler
import org.apache.celeborn.common.CelebornConf
class CelebornColumnarShuffleReader[K, C](
@@ -34,7 +37,8 @@ class CelebornColumnarShuffleReader[K, C](
context: TaskContext,
conf: CelebornConf,
metrics: ShuffleReadMetricsReporter,
- shuffleIdTracker: ExecutorShuffleIdTracker)
+ shuffleIdTracker: ExecutorShuffleIdTracker,
+ cryptoHandler: Optional[CryptoHandler] = Optional.empty())
extends CelebornShuffleReader[K, C](
handle,
startPartition,
@@ -44,7 +48,8 @@ class CelebornColumnarShuffleReader[K, C](
context,
conf,
metrics,
- shuffleIdTracker) {
+ shuffleIdTracker,
+ cryptoHandler) {
override def newSerializerInstance(dep: ShuffleDependency[K, _, C]):
SerializerInstance = {
val schema = CustomShuffleDependencyUtils.getSchema(dep)
diff --git
a/client-spark/spark-3-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
b/client-spark/spark-3-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
index d0f4462be3..edc67d9833 100644
---
a/client-spark/spark-3-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
+++
b/client-spark/spark-3-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
@@ -17,6 +17,8 @@
package org.apache.spark.shuffle.celeborn
+import java.util.Optional
+
import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
import org.apache.spark.serializer.{KryoSerializer, KryoSerializerInstance}
import org.apache.spark.sql.execution.UnsafeRowSerializer
@@ -58,7 +60,8 @@ class CelebornColumnarShuffleReaderSuite {
taskContext,
new CelebornConf(),
null,
- new ExecutorShuffleIdTracker())
+ new ExecutorShuffleIdTracker(),
+ Optional.empty())
assert(shuffleReader.getClass ==
classOf[CelebornColumnarShuffleReader[Int, String]])
} finally {
if (shuffleClient != null) {
@@ -92,7 +95,8 @@ class CelebornColumnarShuffleReaderSuite {
taskContext,
new CelebornConf(),
null,
- new ExecutorShuffleIdTracker())
+ new ExecutorShuffleIdTracker(),
+ Optional.empty())
val shuffleDependency = Mockito.mock(classOf[ShuffleDependency[Int,
String, String]])
Mockito.when(shuffleDependency.shuffleId).thenReturn(0)
Mockito.when(shuffleDependency.serializer).thenReturn(new KryoSerializer(
diff --git
a/client-spark/spark-3.5-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
b/client-spark/spark-3.5-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
index fd888fb9dc..c32dfaf2eb 100644
---
a/client-spark/spark-3.5-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
+++
b/client-spark/spark-3.5-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
@@ -17,12 +17,15 @@
package org.apache.spark.shuffle.celeborn
+import java.util.Optional
+
import org.apache.spark.{ShuffleDependency, TaskContext}
import org.apache.spark.serializer.SerializerInstance
import org.apache.spark.shuffle.ShuffleReadMetricsReporter
import org.apache.spark.sql.execution.UnsafeRowSerializer
import org.apache.spark.sql.execution.columnar.{CelebornBatchBuilder,
CelebornColumnarBatchSerializer}
+import org.apache.celeborn.client.security.CryptoHandler
import org.apache.celeborn.common.CelebornConf
class CelebornColumnarShuffleReader[K, C](
@@ -34,7 +37,8 @@ class CelebornColumnarShuffleReader[K, C](
context: TaskContext,
conf: CelebornConf,
metrics: ShuffleReadMetricsReporter,
- shuffleIdTracker: ExecutorShuffleIdTracker)
+ shuffleIdTracker: ExecutorShuffleIdTracker,
+ cryptoHandler: Optional[CryptoHandler] = Optional.empty())
extends CelebornShuffleReader[K, C](
handle,
startPartition,
@@ -44,7 +48,8 @@ class CelebornColumnarShuffleReader[K, C](
context,
conf,
metrics,
- shuffleIdTracker) {
+ shuffleIdTracker,
+ cryptoHandler) {
override def newSerializerInstance(dep: ShuffleDependency[K, _, C]):
SerializerInstance = {
val schema = CustomShuffleDependencyUtils.getSchema(dep)
diff --git
a/client-spark/spark-3.5-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
b/client-spark/spark-3.5-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
index d0f4462be3..edc67d9833 100644
---
a/client-spark/spark-3.5-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
+++
b/client-spark/spark-3.5-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
@@ -17,6 +17,8 @@
package org.apache.spark.shuffle.celeborn
+import java.util.Optional
+
import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
import org.apache.spark.serializer.{KryoSerializer, KryoSerializerInstance}
import org.apache.spark.sql.execution.UnsafeRowSerializer
@@ -58,7 +60,8 @@ class CelebornColumnarShuffleReaderSuite {
taskContext,
new CelebornConf(),
null,
- new ExecutorShuffleIdTracker())
+ new ExecutorShuffleIdTracker(),
+ Optional.empty())
assert(shuffleReader.getClass ==
classOf[CelebornColumnarShuffleReader[Int, String]])
} finally {
if (shuffleClient != null) {
@@ -92,7 +95,8 @@ class CelebornColumnarShuffleReaderSuite {
taskContext,
new CelebornConf(),
null,
- new ExecutorShuffleIdTracker())
+ new ExecutorShuffleIdTracker(),
+ Optional.empty())
val shuffleDependency = Mockito.mock(classOf[ShuffleDependency[Int,
String, String]])
Mockito.when(shuffleDependency.shuffleId).thenReturn(0)
Mockito.when(shuffleDependency.serializer).thenReturn(new KryoSerializer(
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
index ed7865e19f..f528934b39 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -19,6 +19,7 @@ package org.apache.spark.shuffle.celeborn;
import java.io.IOException;
import java.util.Objects;
+import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.spark.*;
@@ -33,6 +34,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.ShuffleMode;
import org.apache.celeborn.reflect.DynMethods;
@@ -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) {
+ cryptoHandler = SparkCommonUtils.getCryptoHandler(conf);
+ } else {
+ return Optional.empty();
+ }
+ }
+ return cryptoHandler;
+ }
+
public SparkShuffleManager(SparkConf conf, boolean isDriver) {
if (conf.getBoolean(SQLConf.LOCAL_SHUFFLE_READER_ENABLED().key(), true)) {
logger.warn(
@@ -288,7 +307,8 @@ public class SparkShuffleManager implements ShuffleManager {
h.lifecycleManagerPort(),
celebornConf,
h.userIdentifier(),
- h.extension());
+ h.extension(),
+ getCryptoHandler());
if (h.stageRerunEnabled()) {
SparkUtils.addFailureListenerIfBarrierTask(shuffleClient, context,
h);
}
@@ -445,7 +465,8 @@ public class SparkShuffleManager implements ShuffleManager {
context,
celebornConf,
metrics,
- shuffleIdTracker);
+ shuffleIdTracker,
+ getCryptoHandler());
} else {
return new CelebornShuffleReader<>(
h,
@@ -456,7 +477,8 @@ public class SparkShuffleManager implements ShuffleManager {
context,
celebornConf,
metrics,
- shuffleIdTracker);
+ shuffleIdTracker,
+ getCryptoHandler());
}
}
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
index 613ac7f999..89cb7c3792 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
@@ -24,6 +24,7 @@ import java.io.ObjectOutputStream;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
@@ -66,6 +67,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.exception.CelebornRuntimeException;
@@ -282,7 +284,8 @@ public class SparkUtils {
TaskContext.class,
CelebornConf.class,
ShuffleReadMetricsReporter.class,
- ExecutorShuffleIdTracker.class)
+ ExecutorShuffleIdTracker.class,
+ Optional.class)
.build();
}
@@ -295,7 +298,8 @@ public class SparkUtils {
TaskContext context,
CelebornConf conf,
ShuffleReadMetricsReporter metrics,
- ExecutorShuffleIdTracker shuffleIdTracker) {
+ ExecutorShuffleIdTracker shuffleIdTracker,
+ Optional<CryptoHandler> cryptoHandler) {
return ColumnarShuffleReaderConstructorHolder.INSTANCE.invoke(
null,
handle,
@@ -306,7 +310,8 @@ public class SparkUtils {
context,
conf,
metrics,
- shuffleIdTracker);
+ shuffleIdTracker,
+ cryptoHandler);
}
// Added in SPARK-32920, for Spark 3.2 and above
diff --git
a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
index 155fc08861..32863c6a54 100644
---
a/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
+++
b/client-spark/spark-3/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReader.scala
@@ -18,7 +18,7 @@
package org.apache.spark.shuffle.celeborn
import java.io.IOException
-import java.util.{ArrayList => JArrayList, HashMap => JHashMap, Map => JMap,
Set => JSet}
+import java.util.{ArrayList => JArrayList, HashMap => JHashMap, Map => JMap,
Optional, Set => JSet}
import java.util.concurrent.{ConcurrentHashMap, ThreadPoolExecutor,
TimeoutException, TimeUnit}
import java.util.concurrent.atomic.AtomicReference
import java.util.function.BiFunction
@@ -40,6 +40,7 @@ import org.apache.spark.util.collection.ExternalSorter
import org.apache.celeborn.client.{ClientUtils, ShuffleClient}
import org.apache.celeborn.client.ShuffleClientImpl.ReduceFileGroups
import org.apache.celeborn.client.read.{CelebornInputStream, MetricsCallback}
+import org.apache.celeborn.client.security.CryptoHandler
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.exception.{CelebornBroadcastException,
CelebornIOException, CelebornRuntimeException, PartitionUnRetryAbleException}
import org.apache.celeborn.common.network.client.TransportClient
@@ -59,7 +60,8 @@ class CelebornShuffleReader[K, C](
conf: CelebornConf,
metrics: ShuffleReadMetricsReporter,
shuffleIdTracker: ExecutorShuffleIdTracker,
- needDecompress: Boolean)
+ needDecompress: Boolean,
+ cryptoHandler: Optional[CryptoHandler])
extends ShuffleReader[K, C] with Logging {
def this(
@@ -81,7 +83,31 @@ class CelebornShuffleReader[K, C](
conf,
metrics,
shuffleIdTracker,
- true)
+ true,
+ Optional.empty())
+
+ def this(
+ handle: CelebornShuffleHandle[K, _, C],
+ startPartition: Int,
+ endPartition: Int,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ context: TaskContext,
+ conf: CelebornConf,
+ metrics: ShuffleReadMetricsReporter,
+ shuffleIdTracker: ExecutorShuffleIdTracker,
+ cryptoHandler: Optional[CryptoHandler]) = this(
+ handle,
+ startPartition,
+ endPartition,
+ startMapIndex,
+ endMapIndex,
+ context,
+ conf,
+ metrics,
+ shuffleIdTracker,
+ true,
+ cryptoHandler)
private val dep = handle.dependency
@@ -92,7 +118,8 @@ class CelebornShuffleReader[K, C](
handle.lifecycleManagerPort,
conf,
handle.userIdentifier,
- handle.extension)
+ handle.extension,
+ cryptoHandler)
private val exceptionRef = new AtomicReference[IOException]
private val stageRerunEnabled = handle.stageRerunEnabled
diff --git
a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
index d2cec3abf8..901d7797b5 100644
---
a/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
+++
b/client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleReaderSuite.scala
@@ -60,7 +60,7 @@ class CelebornShuffleReaderSuite extends AnyFunSuite {
val tmpFile = Files.createTempFile("test", ".tmp").toFile
mockStatic(classOf[ShuffleClient]).when(() =>
- ShuffleClient.get(any(), any(), any(), any(), any(), any())).thenReturn(
+ ShuffleClient.get(any(), any(), any(), any(), any(), any(),
any())).thenReturn(
new DummyShuffleClient(conf, tmpFile))
val shuffleReader =
diff --git
a/client-spark/spark-4-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
b/client-spark/spark-4-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
index fd888fb9dc..c32dfaf2eb 100644
---
a/client-spark/spark-4-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
+++
b/client-spark/spark-4-columnar-shuffle/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReader.scala
@@ -17,12 +17,15 @@
package org.apache.spark.shuffle.celeborn
+import java.util.Optional
+
import org.apache.spark.{ShuffleDependency, TaskContext}
import org.apache.spark.serializer.SerializerInstance
import org.apache.spark.shuffle.ShuffleReadMetricsReporter
import org.apache.spark.sql.execution.UnsafeRowSerializer
import org.apache.spark.sql.execution.columnar.{CelebornBatchBuilder,
CelebornColumnarBatchSerializer}
+import org.apache.celeborn.client.security.CryptoHandler
import org.apache.celeborn.common.CelebornConf
class CelebornColumnarShuffleReader[K, C](
@@ -34,7 +37,8 @@ class CelebornColumnarShuffleReader[K, C](
context: TaskContext,
conf: CelebornConf,
metrics: ShuffleReadMetricsReporter,
- shuffleIdTracker: ExecutorShuffleIdTracker)
+ shuffleIdTracker: ExecutorShuffleIdTracker,
+ cryptoHandler: Optional[CryptoHandler] = Optional.empty())
extends CelebornShuffleReader[K, C](
handle,
startPartition,
@@ -44,7 +48,8 @@ class CelebornColumnarShuffleReader[K, C](
context,
conf,
metrics,
- shuffleIdTracker) {
+ shuffleIdTracker,
+ cryptoHandler) {
override def newSerializerInstance(dep: ShuffleDependency[K, _, C]):
SerializerInstance = {
val schema = CustomShuffleDependencyUtils.getSchema(dep)
diff --git
a/client-spark/spark-4-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
b/client-spark/spark-4-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
index d0f4462be3..edc67d9833 100644
---
a/client-spark/spark-4-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
+++
b/client-spark/spark-4-columnar-shuffle/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornColumnarShuffleReaderSuite.scala
@@ -17,6 +17,8 @@
package org.apache.spark.shuffle.celeborn
+import java.util.Optional
+
import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
import org.apache.spark.serializer.{KryoSerializer, KryoSerializerInstance}
import org.apache.spark.sql.execution.UnsafeRowSerializer
@@ -58,7 +60,8 @@ class CelebornColumnarShuffleReaderSuite {
taskContext,
new CelebornConf(),
null,
- new ExecutorShuffleIdTracker())
+ new ExecutorShuffleIdTracker(),
+ Optional.empty())
assert(shuffleReader.getClass ==
classOf[CelebornColumnarShuffleReader[Int, String]])
} finally {
if (shuffleClient != null) {
@@ -92,7 +95,8 @@ class CelebornColumnarShuffleReaderSuite {
taskContext,
new CelebornConf(),
null,
- new ExecutorShuffleIdTracker())
+ new ExecutorShuffleIdTracker(),
+ Optional.empty())
val shuffleDependency = Mockito.mock(classOf[ShuffleDependency[Int,
String, String]])
Mockito.when(shuffleDependency.shuffleId).thenReturn(0)
Mockito.when(shuffleDependency.serializer).thenReturn(new KryoSerializer(
diff --git
a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
index 8390a9c1b5..86e6a2d191 100644
--- a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -30,6 +30,7 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
@@ -41,6 +42,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.client.TransportClientFactory;
@@ -79,6 +81,9 @@ public class DummyShuffleClient extends ShuffleClient {
this.shuffleIntegrityCheckEnabled =
conf.clientShuffleIntegrityCheckEnabled();
}
+ @Override
+ public void setupCryptoHandler(Optional<CryptoHandler> cryptoHandler) {}
+
@Override
public void setupLifecycleManagerRef(String host, int port) {}
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
index 7035478ebf..8bc1a8911c 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -35,6 +35,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.identity.UserIdentifier;
@@ -90,6 +91,18 @@ public abstract class ShuffleClient {
CelebornConf conf,
UserIdentifier userIdentifier,
byte[] extension) {
+ return ShuffleClient.get(
+ appUniqueId, driverHost, port, conf, userIdentifier, extension,
Optional.empty());
+ }
+
+ public static ShuffleClient get(
+ String appUniqueId,
+ String driverHost,
+ int port,
+ CelebornConf conf,
+ UserIdentifier userIdentifier,
+ byte[] extension,
+ Optional<CryptoHandler> cryptoHandler) {
if (null == _instance || !initialized) {
synchronized (ShuffleClient.class) {
if (null == _instance) {
@@ -102,16 +115,25 @@ public abstract class ShuffleClient {
_instance = new ShuffleClientImpl(appUniqueId, conf, userIdentifier);
_instance.setupLifecycleManagerRef(driverHost, port);
_instance.setExtension(extension);
+ _instance.setupCryptoHandler(cryptoHandler);
initialized = true;
} else if (!initialized) {
_instance.shutdown();
_instance = new ShuffleClientImpl(appUniqueId, conf, userIdentifier);
_instance.setupLifecycleManagerRef(driverHost, port);
_instance.setExtension(extension);
+ _instance.setupCryptoHandler(cryptoHandler);
initialized = true;
}
}
}
+ // Apply the crypto handler even when the singleton is already
initialized. This handles
+ // the case where SparkEnv was transiently unavailable during the first
init call (causing
+ // an empty handler to be stored), so that encryption is correctly applied
on retry.
+ // setupCryptoHandler is a volatile write and safe to call without the
lock.
+ if (cryptoHandler != null && cryptoHandler.isPresent()) {
+ _instance.setupCryptoHandler(cryptoHandler);
+ }
return _instance;
}
@@ -150,6 +172,8 @@ public abstract class ShuffleClient {
String.format("%.2f", (localReadCount * 1.0d / totalReadCount) * 100));
}
+ public abstract void setupCryptoHandler(Optional<CryptoHandler>
cryptoHandler);
+
public abstract void setupLifecycleManagerRef(String host, int port);
public abstract void setupLifecycleManagerRef(RpcEndpointRef endpointRef);
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index 358bc227a1..46c1189ca4 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -43,6 +43,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.compress.Compressor;
import org.apache.celeborn.client.read.CelebornInputStream;
import org.apache.celeborn.client.read.MetricsCallback;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornBroadcastException;
import org.apache.celeborn.common.exception.CelebornIOException;
@@ -101,6 +102,8 @@ public class ShuffleClientImpl extends ShuffleClient {
protected byte[] extension;
+ private volatile Optional<CryptoHandler> cryptoHandler = Optional.empty();
+
// key: appShuffleIdentifier, value: shuffleId
protected Map<String, Tuple2<Integer, Boolean>> shuffleIdCache =
JavaUtils.newConcurrentHashMap();
@@ -1054,6 +1057,24 @@ public class ShuffleClientImpl extends ShuffleClient {
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);
+ if (logger.isDebugEnabled()) {
+ logger.debug(
+ "Encrypted shuffle data for shuffle {} map {} partition {}: {}
bytes -> {} bytes.",
+ shuffleId,
+ mapId,
+ partitionId,
+ length,
+ encrypted.length);
+ }
+ data = encrypted;
+ offset = 0;
+ length = encrypted.length;
+ }
+
final byte[] body = new byte[BATCH_HEADER_SIZE + length];
Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET, mapId);
Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 4, attemptId);
@@ -2039,7 +2060,8 @@ public class ShuffleClientImpl extends ShuffleClient {
partitionId,
exceptionMaker,
metricsCallback,
- needDecompress);
+ needDecompress,
+ cryptoHandler);
}
}
@@ -2104,6 +2126,14 @@ public class ShuffleClientImpl extends ShuffleClient {
this.extension = extension;
}
+ @Override
+ public void setupCryptoHandler(Optional<CryptoHandler> cryptoHandler) {
+ this.cryptoHandler = cryptoHandler != null ? cryptoHandler :
Optional.empty();
+ if (this.cryptoHandler.isPresent()) {
+ logger.info("IO encryption enabled for shuffle data (encryption at
rest).");
+ }
+ }
+
boolean mapperEnded(int shuffleId, int mapId) {
return (mapperEndMap.containsKey(shuffleId) &&
mapperEndMap.get(shuffleId).contains(mapId))
|| isStageEnded(shuffleId);
diff --git
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
index 37e0be3e37..d76b1eb9d3 100644
---
a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
+++
b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
@@ -39,6 +39,7 @@ import org.apache.celeborn.client.ClientUtils;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.client.compress.Decompressor;
import
org.apache.celeborn.client.read.checkpoint.PartitionReaderCheckpointMetadata;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.CommitMetadata;
import org.apache.celeborn.common.exception.CelebornIOException;
@@ -74,7 +75,8 @@ public abstract class CelebornInputStream extends InputStream
{
int partitionId,
ExceptionMaker exceptionMaker,
MetricsCallback metricsCallback,
- boolean needDecompress)
+ boolean needDecompress,
+ Optional<CryptoHandler> cryptoHandler)
throws IOException {
if (locations == null || locations.isEmpty()) {
return emptyInputStream;
@@ -106,7 +108,8 @@ public abstract class CelebornInputStream extends
InputStream {
metricsCallback,
needDecompress,
startMapIndex,
- endMapIndex);
+ endMapIndex,
+ cryptoHandler);
} else {
return new CelebornInputStreamImpl(
conf,
@@ -131,7 +134,8 @@ public abstract class CelebornInputStream extends
InputStream {
metricsCallback,
needDecompress,
-1,
- -1);
+ -1,
+ cryptoHandler);
}
}
}
@@ -188,6 +192,7 @@ public abstract class CelebornInputStream extends
InputStream {
private final Map<String, LocationPushFailedBatches> failedBatches;
+ private byte[] encryptedBuf;
private byte[] compressedBuf;
private byte[] rawDataBuf;
private Decompressor decompressor;
@@ -223,6 +228,7 @@ public abstract class CelebornInputStream extends
InputStream {
private int shuffleId;
private int partitionId;
private ExceptionMaker exceptionMaker;
+ private Optional<CryptoHandler> cryptoHandler;
private boolean closed = false;
private boolean integrityChecked = false;
private final CommitMetadata aggregatedActualCommitMetadata = new
CommitMetadata();
@@ -250,7 +256,8 @@ public abstract class CelebornInputStream extends
InputStream {
MetricsCallback metricsCallback,
boolean needDecompress,
int numberOfSubPartitions,
- int currentIndexOfSubPartition)
+ int currentIndexOfSubPartition,
+ Optional<CryptoHandler> cryptoHandler)
throws IOException {
this(
conf,
@@ -275,7 +282,8 @@ public abstract class CelebornInputStream extends
InputStream {
metricsCallback,
needDecompress,
numberOfSubPartitions,
- currentIndexOfSubPartition);
+ currentIndexOfSubPartition,
+ cryptoHandler);
}
CelebornInputStreamImpl(
@@ -301,7 +309,8 @@ public abstract class CelebornInputStream extends
InputStream {
MetricsCallback metricsCallback,
boolean needDecompress,
int numberOfSubPartitions,
- int currentIndexOfSubPartition)
+ int currentIndexOfSubPartition,
+ Optional<CryptoHandler> cryptoHandler)
throws IOException {
this.conf = conf;
this.clientFactory = clientFactory;
@@ -337,6 +346,7 @@ public abstract class CelebornInputStream extends
InputStream {
this.retryWaitMs =
conf.networkIoRetryWaitMs(TransportModuleConstants.DATA_MODULE);
this.callback = metricsCallback;
this.exceptionMaker = exceptionMaker;
+ this.cryptoHandler = cryptoHandler;
this.partitionId = partitionId;
this.appShuffleId = appShuffleId;
this.shuffleId = shuffleId;
@@ -730,6 +740,7 @@ public abstract class CelebornInputStream extends
InputStream {
compressedBuf = null;
rawDataBuf = null;
+ encryptedBuf = null;
batchesRead = null;
locations = null;
attempts = null;
@@ -790,9 +801,15 @@ public abstract class CelebornInputStream extends
InputStream {
private void init() {
int bufferSize = conf.clientFetchBufferSize();
+ int headerLen = shouldDecompress ?
Decompressor.getCompressionHeaderLength(conf) : 0;
+ if (cryptoHandler.isPresent()) {
+ // Pre-size to include the crypto overhead (e.g. SparkCryptoHandler
adds a 4-byte length
+ // prefix and a 16-byte IV = 20 bytes) so the buffer is large enough
for the first batch
+ // without an immediate reallocation.
+ encryptedBuf = new byte[bufferSize + headerLen + 64];
+ }
if (shouldDecompress) {
- int headerLen = Decompressor.getCompressionHeaderLength(conf);
bufferSize += headerLen;
compressedBuf = new byte[bufferSize];
decompressor = Decompressor.getDecompressor(conf);
@@ -823,65 +840,104 @@ public abstract class CelebornInputStream extends
InputStream {
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) {
+ if (decrypted.length <
Decompressor.getCompressionHeaderLength(conf)) {
+ throw new IOException(
+ "Decrypted batch too short to contain compression header: "
+ + decrypted.length
+ + " bytes (shuffleId="
+ + shuffleId
+ + ", partitionId="
+ + partitionId
+ + ")");
+ }
+ compressedBuf = decrypted;
+ } else {
+ rawDataBuf = decrypted;
+ }
+ } else if (shouldDecompress) {
if (size > compressedBuf.length) {
compressedBuf = new byte[size];
}
-
currentChunk.readBytes(compressedBuf, 0, size);
} else {
if (size > rawDataBuf.length) {
rawDataBuf = new byte[size];
}
-
currentChunk.readBytes(rawDataBuf, 0, size);
}
- // de-duplicate
- if (attemptId == attempts[mapId]) {
- 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);
- continue;
- }
- }
- }
- Set<Integer> batchSet = batchesRead.computeIfAbsent(mapId, k ->
new HashSet<>());
- if (!batchSet.contains(batchId)) {
- batchSet.add(batchId);
- callback.incBytesRead(BATCH_HEADER_SIZE + size);
- if (shouldDecompress) {
- // decompress data
- int originalLength =
decompressor.getOriginalLen(compressedBuf);
- if (rawDataBuf.length < originalLength) {
- rawDataBuf = new byte[originalLength];
- }
- limit = decompressor.decompress(compressedBuf, rawDataBuf, 0);
- } else {
- limit = size;
- }
- if (shuffleIntegrityCheckEnabled) {
-
aggregatedActualCommitMetadata.addDataWithOffsetAndLength(rawDataBuf, 0, limit);
- }
- position = 0;
- hasData = true;
- break;
- } else {
- callback.incDuplicateBytesRead(BATCH_HEADER_SIZE + size);
- logger.debug(
- "Skip duplicated batch: mapId {}, attemptId {}, batchId {}.",
- mapId,
- attemptId,
- batchId);
+ batchSet.add(batchId);
+ callback.incBytesRead(BATCH_HEADER_SIZE + encryptedSize);
+ if (shouldDecompress) {
+ // decompress data
+ int originalLength = decompressor.getOriginalLen(compressedBuf);
+ if (rawDataBuf.length < originalLength) {
+ rawDataBuf = new byte[originalLength];
}
+ limit = decompressor.decompress(compressedBuf, rawDataBuf, 0);
+ } else {
+ limit = size;
+ }
+ if (shuffleIntegrityCheckEnabled) {
+
aggregatedActualCommitMetadata.addDataWithOffsetAndLength(rawDataBuf, 0, limit);
}
+ position = 0;
+ hasData = true;
+ break;
}
if (!hasData) {
diff --git
a/client/src/main/java/org/apache/celeborn/client/security/CryptoHandler.java
b/client/src/main/java/org/apache/celeborn/client/security/CryptoHandler.java
new file mode 100644
index 0000000000..38da34fa79
--- /dev/null
+++
b/client/src/main/java/org/apache/celeborn/client/security/CryptoHandler.java
@@ -0,0 +1,26 @@
+/*
+ * 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.security;
+
+import java.io.IOException;
+
+public interface CryptoHandler {
+ byte[] encrypt(byte[] input, int offset, int length) throws IOException;
+
+ byte[] decrypt(byte[] input, int offset, int length) throws IOException;
+}
diff --git
a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamCryptoRoundTripSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamCryptoRoundTripSuiteJ.java
new file mode 100644
index 0000000000..06655b887b
--- /dev/null
+++
b/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 {
+ 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 {
+ return createStreamWithClient(
+ batchBuf, needDecompress, cryptoHandler, conf,
mock(ShuffleClient.class));
+ }
+
+ /**
+ * Like {@link #createStream} but with a caller-supplied ShuffleClient mock,
so tests can verify
+ * interactions such as {@code readReducerPartitionEnd}.
+ */
+ private CelebornInputStream createStreamWithClient(
+ ByteBuf batchBuf,
+ boolean needDecompress,
+ Optional<CryptoHandler> cryptoHandler,
+ CelebornConf conf,
+ ShuffleClient shuffleClient)
+ 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);
+
+ // 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 {
+ // Verify that EAR + shuffle integrity check work together: the checksum
must be computed
+ // over the *decrypted* plaintext, not the ciphertext. We capture the
crc32/bytes passed to
+ // readReducerPartitionEnd and assert they match an independently-computed
plaintext checksum.
+ byte[] plaintext = "integrity check should pass after
decryption".getBytes();
+ CelebornConf conf = new CelebornConf();
+ conf.set(CelebornConf.CLIENT_SHUFFLE_INTEGRITY_CHECK_ENABLED().key(),
"true");
+ XorCryptoHandler handler = new XorCryptoHandler((byte) 0x7F);
+
+ // Independently compute the expected checksum over the plaintext bytes.
+ CommitMetadata expected = new CommitMetadata();
+ expected.addDataWithOffsetAndLength(plaintext, 0, plaintext.length);
+
+ ByteBuf batchBuf = buildBatch(plaintext, false, handler, conf);
+
+ // createStream passes shuffleId=1, partitionId=0, startMapIndex=0,
endMapIndex=100
+ ShuffleClient shuffleClient = mock(ShuffleClient.class);
+ try (CelebornInputStream stream =
+ createStreamWithClient(batchBuf, false, Optional.of(handler), conf,
shuffleClient)) {
+ assertArrayEquals(plaintext, readAll(stream));
+ }
+
+ // Verify readReducerPartitionEnd was called with the checksum over
plaintext, not ciphertext.
+ ArgumentCaptor<Integer> crcCaptor = ArgumentCaptor.forClass(Integer.class);
+ ArgumentCaptor<Long> bytesCaptor = ArgumentCaptor.forClass(Long.class);
+ verify(shuffleClient)
+ .readReducerPartitionEnd(
+ anyInt(), anyInt(), anyInt(), anyInt(), crcCaptor.capture(),
bytesCaptor.capture());
+ assertEquals(
+ "checksum must be over plaintext", expected.getChecksum(), (int)
crcCaptor.getValue());
+ assertEquals(
+ "byte count must match plaintext length",
+ expected.getBytes(),
+ (long) bytesCaptor.getValue());
+ }
+
+ @Test
+ public void testLargePayloadEncryptDecrypt() throws IOException,
InterruptedException {
+ // 128 KB payload exercises buffer-boundary handling in fillBuffer()
+ byte[] plaintext = new byte[128 * 1024];
+ for (int i = 0; i < plaintext.length; i++) plaintext[i] = (byte) (i % 251);
+ CelebornConf conf = new CelebornConf();
+ XorCryptoHandler handler = new XorCryptoHandler((byte) 0xAB);
+
+ ByteBuf batchBuf = buildBatch(plaintext, false, handler, conf);
+ try (CelebornInputStream stream = createStream(batchBuf, false,
Optional.of(handler), conf)) {
+ assertArrayEquals(plaintext, readAll(stream));
+ }
+ }
+}
diff --git
a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
index a7dd7db317..456486b9e8 100644
---
a/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
+++
b/client/src/test/java/org/apache/celeborn/client/read/CelebornInputStreamPeerFailoverTest.java
@@ -32,6 +32,7 @@ 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 java.util.concurrent.atomic.AtomicInteger;
@@ -39,6 +40,7 @@ import org.junit.Before;
import org.junit.Test;
import org.apache.celeborn.client.ShuffleClient;
+import org.apache.celeborn.client.security.CryptoHandler;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.network.client.TransportClient;
@@ -174,7 +176,8 @@ public class CelebornInputStreamPeerFailoverTest {
0,
null,
new TestMetricsCallback(),
- false);
+ false,
+ Optional.<CryptoHandler>empty());
}
private void createInputStream(String primaryHost, String replicaHost)
throws IOException {
@@ -209,7 +212,8 @@ public class CelebornInputStreamPeerFailoverTest {
0,
null,
new TestMetricsCallback(),
- false);
+ false,
+ Optional.<CryptoHandler>empty());
}
/**
