[CELEBORN-2329][CIP22] Encryption at Rest Spark Impl

client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java

@@ -19,15 +19,25 @@
 
 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 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()) {
+      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()));
+  }
 }

client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCryptoHandler.java

@@ -0,0 +1,86 @@
+/*
+ * 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;
+
+  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 = 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][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;
+    }
+  }
+}

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"));
+    }
+  }
+}

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)

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(