Split into BclImpl and BouncyCastleImpl

Author: scott-xu

src/Renci.SshNet/Security/KeyExchangeMLKem768X25519Sha256.BclImpl.cs

@@ -0,0 +1,36 @@
+using System;
+using System.Security.Cryptography;
+
+namespace Renci.SshNet.Security
+{
+    internal sealed partial class KeyExchangeMLKem768X25519Sha256
+    {
+        private sealed class MLKemBclImpl : Impl
+        {
+            private MLKem _mlkem;
+
+            public override byte[] GenerateClientPublicKey()
+            {
+                _mlkem = MLKem.GenerateKey(MLKemAlgorithm.MLKem768);
+                return _mlkem.ExportEncapsulationKey();
+            }
+
+            public override byte[] CalculateAgreement(byte[] serverPublicKey)
+            {
+                var mlkemSecret = new byte[MLKemAlgorithm.MLKem768.SharedSecretSizeInBytes];
+                _mlkem.Decapsulate(serverPublicKey.AsSpan(0, MLKemAlgorithm.MLKem768.CiphertextSizeInBytes), mlkemSecret);
+                return mlkemSecret;
+            }
+
+            protected override void Dispose(bool disposing)
+            {
+                if (disposing)
+                {
+                    _mlkem?.Dispose();
+                }
+
+                base.Dispose(disposing);
+            }
+        }
+    }
+}

src/Renci.SshNet/Security/KeyExchangeMLKem768X25519Sha256.BouncyCastleImpl.cs

@@ -0,0 +1,36 @@
+using Org.BouncyCastle.Crypto.Generators;
+using Org.BouncyCastle.Crypto.Kems;
+using Org.BouncyCastle.Crypto.Parameters;
+
+using Renci.SshNet.Abstractions;
+
+namespace Renci.SshNet.Security
+{
+    internal sealed partial class KeyExchangeMLKem768X25519Sha256
+    {
+        private sealed class MLKemBouncyCastleImpl : Impl
+        {
+            private MLKemDecapsulator _mlkemDecapsulator;
+
+            public override byte[] GenerateClientPublicKey()
+            {
+                var mlkem768KeyPairGenerator = new MLKemKeyPairGenerator();
+                mlkem768KeyPairGenerator.Init(new MLKemKeyGenerationParameters(CryptoAbstraction.SecureRandom, MLKemParameters.ml_kem_768));
+                var mlkem768KeyPair = mlkem768KeyPairGenerator.GenerateKeyPair();
+
+                _mlkemDecapsulator = new MLKemDecapsulator(MLKemParameters.ml_kem_768);
+                _mlkemDecapsulator.Init(mlkem768KeyPair.Private);
+
+                return ((MLKemPublicKeyParameters)mlkem768KeyPair.Public).GetEncoded();
+            }
+
+            public override byte[] CalculateAgreement(byte[] serverPublicKey)
+            {
+                var mlkemSecret = new byte[_mlkemDecapsulator.SecretLength];
+                _mlkemDecapsulator.Decapsulate(serverPublicKey, 0, _mlkemDecapsulator.EncapsulationLength, mlkemSecret, 0, _mlkemDecapsulator.SecretLength);
+
+                return mlkemSecret;
+            }
+        }
+    }
+}

src/Renci.SshNet/Security/KeyExchangeMLKem768X25519Sha256.cs

@@ -1,9 +1,6 @@
-using System;
-using System.Globalization;
+using System.Globalization;
 using System.Security.Cryptography;
 
-using Org.BouncyCastle.Crypto.Generators;
-using Org.BouncyCastle.Crypto.Kems;
 using Org.BouncyCastle.Crypto.Parameters;
 
 using Renci.SshNet.Abstractions;
@@ -12,11 +9,9 @@
 
 namespace Renci.SshNet.Security
 {
-    internal sealed class KeyExchangeMLKem768X25519Sha256 : KeyExchangeECCurve25519
+    internal sealed partial class KeyExchangeMLKem768X25519Sha256 : KeyExchangeECCurve25519
     {
-        private readonly MLKemAlgorithm _mlkem768 = MLKemAlgorithm.MLKem768;
-        private MLKem _mlkem;
-        private MLKemDecapsulator _mlkemDecapsulator;
+        private Impl _mlkemImpl;
 
         /// <summary>
         /// Gets algorithm name.
@@ -44,24 +39,17 @@ protected override void StartImpl()
 
             Session.KeyExchangeHybridReplyMessageReceived += Session_KeyExchangeHybridReplyMessageReceived;
 
-            byte[] mlkem768PublicKey;
             if (MLKem.IsSupported)
             {
-                _mlkem = MLKem.GenerateKey(_mlkem768);
-                mlkem768PublicKey = _mlkem.ExportEncapsulationKey();
+                _mlkemImpl = new MLKemBclImpl();
             }
             else
             {
-                var mlkem768KeyPairGenerator = new MLKemKeyPairGenerator();
-                mlkem768KeyPairGenerator.Init(new MLKemKeyGenerationParameters(CryptoAbstraction.SecureRandom, MLKemParameters.ml_kem_768));
-                var mlkem768KeyPair = mlkem768KeyPairGenerator.GenerateKeyPair();
-
-                _mlkemDecapsulator = new MLKemDecapsulator(MLKemParameters.ml_kem_768);
-                _mlkemDecapsulator.Init(mlkem768KeyPair.Private);
-
-                mlkem768PublicKey = ((MLKemPublicKeyParameters)mlkem768KeyPair.Public).GetEncoded();
+                _mlkemImpl = new MLKemBouncyCastleImpl();
             }
 
+            var mlkem768PublicKey = _mlkemImpl.GenerateClientPublicKey();
+
             var x25519PublicKey = _impl.GenerateClientPublicKey();
 
             _clientExchangeValue = mlkem768PublicKey.Concat(x25519PublicKey);
@@ -112,25 +100,16 @@ private void HandleServerHybridReply(byte[] hostKey, byte[] serverExchangeValue,
             _hostKey = hostKey;
             _signature = signature;
 
-            if (serverExchangeValue.Length != _mlkemDecapsulator.EncapsulationLength + X25519PublicKeyParameters.KeySize)
+            if (serverExchangeValue.Length != MLKemAlgorithm.MLKem768.CiphertextSizeInBytes + X25519PublicKeyParameters.KeySize)
             {
                 throw new SshConnectionException(
                     string.Format(CultureInfo.CurrentCulture, "Bad S_Reply length: {0}.", serverExchangeValue.Length),
                     DisconnectReason.KeyExchangeFailed);
             }
 
-            var mlkemSecret = new byte[_mlkemDecapsulator.SecretLength];
-
-            if (MLKem.IsSupported)
-            {
-                _mlkem.Decapsulate(serverExchangeValue.AsSpan(0, _mlkem768.CiphertextSizeInBytes), mlkemSecret);
-            }
-            else
-            {
-                _mlkemDecapsulator.Decapsulate(serverExchangeValue, 0, _mlkemDecapsulator.EncapsulationLength, mlkemSecret, 0, _mlkemDecapsulator.SecretLength);
-            }
+            var mlkemSecret = _mlkemImpl.CalculateAgreement(serverExchangeValue);
 
-            var x25519Agreement = _impl.CalculateAgreement(serverExchangeValue.Take(_mlkemDecapsulator.EncapsulationLength, X25519PublicKeyParameters.KeySize));
+            var x25519Agreement = _impl.CalculateAgreement(serverExchangeValue.Take(MLKemAlgorithm.MLKem768.CiphertextSizeInBytes, X25519PublicKeyParameters.KeySize));
 
             SharedKey = CryptoAbstraction.HashSHA256(mlkemSecret.Concat(x25519Agreement));
         }
@@ -139,7 +118,7 @@ protected override void Dispose(bool disposing)
         {
             if (disposing)
             {
-                _mlkem?.Dispose();
+                _mlkemImpl?.Dispose();
             }
 
             base.Dispose(disposing);