1- using System . Globalization ;
2- using System . Linq ;
1+ using System ;
2+ using System . Globalization ;
3+ using System . Security . Cryptography ;
34
45using Org . BouncyCastle . Crypto . Agreement ;
56using Org . BouncyCastle . Crypto . Generators ;
1213
1314namespace Renci . SshNet . Security
1415{
16+ #pragma warning disable SYSLIB5006 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
1517 internal sealed class KeyExchangeMLKem768X25519Sha256 : KeyExchangeEC
1618 {
19+ private readonly MLKemAlgorithm _mlkem768 = MLKemAlgorithm . MLKem768 ;
20+ private MLKem _mlkem ;
1721 private MLKemDecapsulator _mlkemDecapsulator ;
1822 private X25519Agreement _x25519Agreement ;
1923
@@ -45,12 +49,22 @@ public override void Start(Session session, KeyExchangeInitMessage message, bool
4549
4650 Session . KeyExchangeHybridReplyMessageReceived += Session_KeyExchangeHybridReplyMessageReceived ;
4751
48- var mlkem768KeyPairGenerator = new MLKemKeyPairGenerator ( ) ;
49- mlkem768KeyPairGenerator . Init ( new MLKemKeyGenerationParameters ( CryptoAbstraction . SecureRandom , MLKemParameters . ml_kem_768 ) ) ;
50- var mlkem768KeyPair = mlkem768KeyPairGenerator . GenerateKeyPair ( ) ;
52+ if ( MLKem . IsSupported )
53+ {
54+ _mlkem = MLKem . GenerateKey ( _mlkem768 ) ;
55+ _clientExchangeValue = _mlkem . ExportEncapsulationKey ( ) ;
56+ }
57+ else
58+ {
59+ var mlkem768KeyPairGenerator = new MLKemKeyPairGenerator ( ) ;
60+ mlkem768KeyPairGenerator . Init ( new MLKemKeyGenerationParameters ( CryptoAbstraction . SecureRandom , MLKemParameters . ml_kem_768 ) ) ;
61+ var mlkem768KeyPair = mlkem768KeyPairGenerator . GenerateKeyPair ( ) ;
62+
63+ _mlkemDecapsulator = new MLKemDecapsulator ( MLKemParameters . ml_kem_768 ) ;
64+ _mlkemDecapsulator . Init ( mlkem768KeyPair . Private ) ;
5165
52- _mlkemDecapsulator = new MLKemDecapsulator ( MLKemParameters . ml_kem_768 ) ;
53- _mlkemDecapsulator . Init ( mlkem768KeyPair . Private ) ;
66+ _clientExchangeValue = ( ( MLKemPublicKeyParameters ) mlkem768KeyPair . Public ) . GetEncoded ( ) ;
67+ }
5468
5569 var x25519KeyPairGenerator = new X25519KeyPairGenerator ( ) ;
5670 x25519KeyPairGenerator . Init ( new X25519KeyGenerationParameters ( CryptoAbstraction . SecureRandom ) ) ;
@@ -59,10 +73,9 @@ public override void Start(Session session, KeyExchangeInitMessage message, bool
5973 _x25519Agreement = new X25519Agreement ( ) ;
6074 _x25519Agreement . Init ( x25519KeyPair . Private ) ;
6175
62- var mlkem768PublicKey = ( ( MLKemPublicKeyParameters ) mlkem768KeyPair . Public ) . GetEncoded ( ) ;
63- var x25519PublicKey = ( ( X25519PublicKeyParameters ) x25519KeyPair . Public ) . GetEncoded ( ) ;
76+ Array . Resize ( ref _clientExchangeValue , _mlkem768 . EncapsulationKeySizeInBytes + X25519PublicKeyParameters . KeySize ) ;
6477
65- _clientExchangeValue = mlkem768PublicKey . Concat ( x25519PublicKey ) ;
78+ ( ( X25519PublicKeyParameters ) x25519KeyPair . Public ) . Encode ( _clientExchangeValue , _mlkem768 . EncapsulationKeySizeInBytes ) ;
6679
6780 SendMessage ( new KeyExchangeHybridInitMessage ( _clientExchangeValue ) ) ;
6881 }
@@ -114,21 +127,39 @@ private void HandleServerHybridReply(byte[] hostKey, byte[] serverExchangeValue,
114127 _hostKey = hostKey ;
115128 _signature = signature ;
116129
117- if ( serverExchangeValue . Length != _mlkemDecapsulator . EncapsulationLength + _x25519Agreement . AgreementSize )
130+ if ( serverExchangeValue . Length != _mlkem768 . CiphertextSizeInBytes + _x25519Agreement . AgreementSize )
118131 {
119132 throw new SshConnectionException (
120133 string . Format ( CultureInfo . CurrentCulture , "Bad S_Reply length: {0}." , serverExchangeValue . Length ) ,
121134 DisconnectReason . KeyExchangeFailed ) ;
122135 }
123136
124- var secret = new byte [ _mlkemDecapsulator . SecretLength + _x25519Agreement . AgreementSize ] ;
137+ var secret = new byte [ _mlkem768 . SharedSecretSizeInBytes + _x25519Agreement . AgreementSize ] ;
125138
126- _mlkemDecapsulator . Decapsulate ( serverExchangeValue , 0 , _mlkemDecapsulator . EncapsulationLength , secret , 0 , _mlkemDecapsulator . SecretLength ) ;
139+ if ( MLKem . IsSupported )
140+ {
141+ _mlkem . Decapsulate ( serverExchangeValue . AsSpan ( 0 , _mlkem768 . CiphertextSizeInBytes ) , secret . AsSpan ( 0 , _mlkem768 . SharedSecretSizeInBytes ) ) ;
142+ }
143+ else
144+ {
145+ _mlkemDecapsulator . Decapsulate ( serverExchangeValue , 0 , _mlkemDecapsulator . EncapsulationLength , secret , 0 , _mlkemDecapsulator . SecretLength ) ;
146+ }
127147
128- var x25519PublicKey = new X25519PublicKeyParameters ( serverExchangeValue , _mlkemDecapsulator . EncapsulationLength ) ;
129- _x25519Agreement . CalculateAgreement ( x25519PublicKey , secret , _mlkemDecapsulator . SecretLength ) ;
148+ var x25519PublicKey = new X25519PublicKeyParameters ( serverExchangeValue , _mlkem768 . CiphertextSizeInBytes ) ;
149+ _x25519Agreement . CalculateAgreement ( x25519PublicKey , secret , _mlkem768 . SharedSecretSizeInBytes ) ;
130150
131151 SharedKey = CryptoAbstraction . HashSHA256 ( secret ) ;
132152 }
153+
154+ protected override void Dispose ( bool disposing )
155+ {
156+ if ( disposing )
157+ {
158+ _mlkem ? . Dispose ( ) ;
159+ }
160+
161+ base . Dispose ( disposing ) ;
162+ }
133163 }
164+ #pragma warning restore SYSLIB5006 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
134165}
0 commit comments