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

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>
+    </dependency>
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>

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

@@ -19,11 +19,17 @@
 
 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.apache.celeborn.client.security.CryptoHandler;
 import org.apache.celeborn.reflect.DynConstructors;
 import org.apache.celeborn.reflect.DynMethods;
 
@@ -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()) {
+      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,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) {
+      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,138 @@
+/*
+ * 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.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);
+  }
+}

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(

client-spark/spark-3-shaded/pom.xml

@@ -76,6 +76,7 @@
               <include>com.google.guava:failureaccess</include>
               <include>io.netty:*</include>
               <include>org.apache.commons:commons-lang3</include>
+              <include>org.apache.commons:commons-crypto</include>
               <include>org.roaringbitmap:RoaringBitmap</include>
               <include>commons-io:commons-io</include>
             </includes>

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)

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(