TLS: Process CCM packet directly

3 files changed, 41 insertions, 11 deletions

diff --git a/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs b/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs
index 4965c92bd..82f1382bd 100644
--- a/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs
+++ b/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs
@@ -6,11 +6,11 @@ using Org.BouncyCastle.Crypto.Parameters;
 
 namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
 {
-    internal sealed class BcTlsAeadCipherImpl
+    internal class BcTlsAeadCipherImpl
         : TlsAeadCipherImpl
     {
         private readonly bool m_isEncrypting;
-        private readonly IAeadCipher m_cipher;
+        internal readonly IAeadCipher m_cipher;
 
         private KeyParameter key;
 
@@ -42,7 +42,7 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
             return m_cipher.GetOutputSize(inputLength);
         }
 
-        public int DoFinal(byte[] input, int inputOffset, int inputLength, byte[] output, int outputOffset)
+        public virtual int DoFinal(byte[] input, int inputOffset, int inputLength, byte[] output, int outputOffset)
         {
             int len = m_cipher.ProcessBytes(input, inputOffset, inputLength, output, outputOffset);
 
diff --git a/crypto/src/tls/crypto/impl/bc/BcTlsCcmImpl.cs b/crypto/src/tls/crypto/impl/bc/BcTlsCcmImpl.cs
new file mode 100644
index 000000000..641ad0ab0
--- /dev/null
+++ b/crypto/src/tls/crypto/impl/bc/BcTlsCcmImpl.cs
@@ -0,0 +1,31 @@
+using System;
+
+using Org.BouncyCastle.Crypto;
+using Org.BouncyCastle.Crypto.Modes;
+
+namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
+{
+    internal class BcTlsCcmImpl
+        : BcTlsAeadCipherImpl
+    {
+        internal BcTlsCcmImpl(CcmBlockCipher cipher, bool isEncrypting)
+            : base(cipher, isEncrypting)
+        {
+        }
+
+        public override int DoFinal(byte[] input, int inputOffset, int inputLength, byte[] output, int outputOffset)
+        {
+            if (!(m_cipher is CcmBlockCipher ccm))
+                throw new InvalidOperationException();
+
+            try
+            {
+                return ccm.ProcessPacket(input, inputOffset, inputLength, output, outputOffset);
+            }
+            catch (InvalidCipherTextException e)
+            {
+                throw new TlsFatalAlert(AlertDescription.bad_record_mac, e);
+            }
+        }
+    }
+}
diff --git a/crypto/src/tls/crypto/impl/bc/BcTlsCrypto.cs b/crypto/src/tls/crypto/impl/bc/BcTlsCrypto.cs
index 0cad3e10d..39df32ed8 100644
--- a/crypto/src/tls/crypto/impl/bc/BcTlsCrypto.cs
+++ b/crypto/src/tls/crypto/impl/bc/BcTlsCrypto.cs
@@ -540,8 +540,8 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
         protected virtual TlsAeadCipher CreateCipher_Aes_Ccm(TlsCryptoParameters cryptoParams, int cipherKeySize,
             int macSize)
         {
-            BcTlsAeadCipherImpl encrypt = new BcTlsAeadCipherImpl(CreateAeadCipher_Aes_Ccm(), true);
-            BcTlsAeadCipherImpl decrypt = new BcTlsAeadCipherImpl(CreateAeadCipher_Aes_Ccm(), false);
+            var encrypt = new BcTlsCcmImpl(CreateAeadCipher_Aes_Ccm(), true);
+            var decrypt = new BcTlsCcmImpl(CreateAeadCipher_Aes_Ccm(), false);
 
             return new TlsAeadCipher(cryptoParams, encrypt, decrypt, cipherKeySize, macSize, TlsAeadCipher.AEAD_CCM);
         }
@@ -587,8 +587,8 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
 
         protected virtual TlsAeadCipher CreateCipher_SM4_Ccm(TlsCryptoParameters cryptoParams)
         {
-            BcTlsAeadCipherImpl encrypt = new BcTlsAeadCipherImpl(CreateAeadCipher_SM4_Ccm(), true);
-            BcTlsAeadCipherImpl decrypt = new BcTlsAeadCipherImpl(CreateAeadCipher_SM4_Ccm(), false);
+            var encrypt = new BcTlsCcmImpl(CreateAeadCipher_SM4_Ccm(), true);
+            var decrypt = new BcTlsCcmImpl(CreateAeadCipher_SM4_Ccm(), false);
 
             return new TlsAeadCipher(cryptoParams, encrypt, decrypt, 16, 16, TlsAeadCipher.AEAD_CCM);
         }
@@ -637,18 +637,17 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
             return new SM4Engine();
         }
 
-        protected virtual IAeadCipher CreateCcmMode(IBlockCipher engine)
+        protected virtual CcmBlockCipher CreateCcmMode(IBlockCipher engine)
         {
             return new CcmBlockCipher(engine);
         }
 
         protected virtual IAeadCipher CreateGcmMode(IBlockCipher engine)
         {
-            // TODO Consider allowing custom configuration of multiplier
             return new GcmBlockCipher(engine);
         }
 
-        protected virtual IAeadCipher CreateAeadCipher_Aes_Ccm()
+        protected virtual CcmBlockCipher CreateAeadCipher_Aes_Ccm()
         {
             return CreateCcmMode(CreateAesEngine());
         }
@@ -668,7 +667,7 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
             return CreateGcmMode(CreateCamelliaEngine());
         }
 
-        protected virtual IAeadCipher CreateAeadCipher_SM4_Ccm()
+        protected virtual CcmBlockCipher CreateAeadCipher_SM4_Ccm()
         {
             return CreateCcmMode(CreateSM4Engine());
         }