[SPARK-47318] Adding a flag to support old AuthEngine protocol

common/network-common/src/main/java/org/apache/spark/network/crypto/AuthEngine.java

@@ -47,11 +47,13 @@ class AuthEngine implements Closeable {
   private static final String MAC_ALGORITHM = "HMACSHA256";
   private static final int AES_GCM_KEY_SIZE_BYTES = 16;
   private static final byte[] EMPTY_TRANSCRIPT = new byte[0];
+  private static final int UNSAFE_SKIP_HKDF_VERSION = 1;
 
   private final String appId;
   private final byte[] preSharedSecret;
   private final TransportConf conf;
   private final Properties cryptoConf;
+  private final boolean unsafeSkipFinalHkdf;
 
   private byte[] clientPrivateKey;
   private TransportCipher sessionCipher;
@@ -63,6 +65,8 @@ class AuthEngine implements Closeable {
     this.preSharedSecret = preSharedSecret.getBytes(UTF_8);
     this.conf = conf;
     this.cryptoConf = conf.cryptoConf();
+    // This is for backward compatibility with older versions of this protocol which did not perform a final HKDF
+    this.unsafeSkipFinalHkdf = conf.authEngineVersion() == UNSAFE_SKIP_HKDF_VERSION;
   }
 
   @VisibleForTesting
@@ -202,7 +206,8 @@ private TransportCipher generateTransportCipher(
       byte[] sharedSecret,
       boolean isClient,
       byte[] transcript) throws GeneralSecurityException {
-    byte[] derivedKey = Hkdf.computeHkdf(
+    byte[] derivedKey = unsafeSkipFinalHkdf ? sharedSecret :  // This is for backwards compatibility
+      Hkdf.computeHkdf(
         MAC_ALGORITHM,
         sharedSecret,
         transcript,
@@ -231,7 +236,7 @@ private TransportCipher generateTransportCipher(
 
   private byte[] getTranscript(AuthMessage... encryptedPublicKeys) {
     ByteBuf transcript = Unpooled.buffer(
-        Arrays.stream(encryptedPublicKeys).mapToInt(k -> k.encodedLength()).sum());
+        Arrays.stream(encryptedPublicKeys).mapToInt(AuthMessage::encodedLength).sum());
     Arrays.stream(encryptedPublicKeys).forEachOrdered(k -> k.encode(transcript));
     return transcript.array();
   }

common/network-common/src/main/java/org/apache/spark/network/crypto/README.md

@@ -1,14 +1,6 @@
 Forward Secure Auth Protocol v2.0
 ==============================================
 
-Deprecation Notice
-------------------
-This is a bespoke key exchange protocol that was implemented before Spark supported TLS (aka SSL) for RPC
-calls. It is recommended that Spark users upgrade to using TLS for RPC calls between Spark processes.
-
-See the [Spark security documentation](https://github.com/apache/spark/blob/master/docs/security.md#ssl-encryption)
-for more information on how to configure TLS.
-
 Summary
 -------
 

common/network-common/src/main/java/org/apache/spark/network/util/TransportConf.java

@@ -213,6 +213,11 @@ public boolean encryptionEnabled() {
     return conf.getBoolean("spark.network.crypto.enabled", false);
   }
 
+  /**
+   * Version number to be used by the AuthEngine key agreement protocol. Value values are 1 or 2.
+   */
+  public int authEngineVersion() { return conf.getInt("spark.network.crypto.authEngineVersion", 2); }
+
   /**
    * The cipher transformation to use for encrypting session data.
    */

common/network-common/src/test/java/org/apache/spark/network/crypto/AuthEngineSuite.java

@@ -20,6 +20,7 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.WritableByteChannel;
 import java.security.GeneralSecurityException;
+import java.util.Collections;
 import java.util.Random;
 
 import com.google.crypto.tink.subtle.Hex;
@@ -49,6 +50,9 @@ public class AuthEngineSuite {
       "e25e9d5c5a380b8e6d1a184692fac065ed84f8592c18e9629f9c636809dca2ffc041f20346eb53db78738" +
       "08ecad08b46b5ee3ff";
   private static final String derivedKey = "2d6e7a9048c8265c33a8f3747bfcc84c";
+  // This key would have been derived from an older version of the protocol that did not run a final HKDF pass
+  private static final String unsafeDerivedKey = "31963f15a320d5c90333f7ecf5cf3a31c7eaf151de07fef8494663a9f47cfd31";
+
   private static final String inputIv = "fc6a5dc8b90a9dad8f54f08b51a59ed2";
   private static final String outputIv = "a72709baf00785cad6329ce09f631f71";
   private static TransportConf conf;
@@ -179,6 +183,27 @@ public void testFixedChallengeResponse() throws Exception {
     }
   }
 
+  public void testFixedChallengeResponseUnsafeVersion() throws Exception {
+    MapConfigProvider configProvider = new MapConfigProvider(Collections.singletonMap(
+      "spark.network.crypto.authEngineVersion", "1"));
+    TransportConf v1Conf = new TransportConf("rpc", configProvider);
+
+    try (AuthEngine client = new AuthEngine("appId", "secret", v1Conf)) {
+      byte[] clientPrivateKey = Hex.decode(clientPrivate);
+      client.setClientPrivateKey(clientPrivateKey);
+      AuthMessage clientChallenge =
+              AuthMessage.decodeMessage(ByteBuffer.wrap(Hex.decode(clientChallengeHex)));
+      AuthMessage serverResponse =
+              AuthMessage.decodeMessage(ByteBuffer.wrap(Hex.decode(serverResponseHex)));
+      // Verify that the client will accept an old transcript.
+      client.deriveSessionCipher(clientChallenge, serverResponse);
+      TransportCipher clientCipher = client.sessionCipher();
+      assertEquals(Hex.encode(clientCipher.getKey().getEncoded()), unsafeDerivedKey);
+      assertEquals(Hex.encode(clientCipher.getInputIv()), inputIv);
+      assertEquals(Hex.encode(clientCipher.getOutputIv()), outputIv);
+    }
+  }
+
   @Test
   public void testMismatchedSecret() throws Exception {
     try (AuthEngine client = new AuthEngine("appId", "secret", conf);