DTLS: Remove the need to reset AEAD ciphers

5 files changed, 52 insertions, 19 deletions

diff --git a/crypto/src/tls/TlsPeer.cs b/crypto/src/tls/TlsPeer.cs
index 53ab7ff40..ed7fa109a 100644
--- a/crypto/src/tls/TlsPeer.cs
+++ b/crypto/src/tls/TlsPeer.cs
@@ -136,6 +136,7 @@ namespace Org.BouncyCastle.Tls
         /// <returns>the <see cref="HeartbeatMode"/> value.</returns>
         short GetHeartbeatPolicy();
 
+        // TODO[api] Remove this and treat it as default 'true'
         /// <summary>Indicates whether a DTLS connection should ignore corrupt records (bad_record_mac) instead of
         /// failing the connection.</summary>
         /// <remarks>Called only once at the start of a connection and applies throughout.</remarks>
diff --git a/crypto/src/tls/crypto/impl/TlsAeadCipher.cs b/crypto/src/tls/crypto/impl/TlsAeadCipher.cs
index e34ca589d..972e93167 100644
--- a/crypto/src/tls/crypto/impl/TlsAeadCipher.cs
+++ b/crypto/src/tls/crypto/impl/TlsAeadCipher.cs
@@ -240,8 +240,9 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
                     output[outputPos + plaintextLength] = (byte)contentType;
                 }
 
-                m_encryptCipher.Init(nonce, m_macSize, additionalData);
-                outputPos += m_encryptCipher.DoFinal(output, outputPos, innerPlaintextLength, output, outputPos);
+                m_encryptCipher.Init(nonce, m_macSize, null);
+                outputPos += m_encryptCipher.DoFinal(additionalData, output, outputPos, innerPlaintextLength, output,
+                    outputPos);
             }
             catch (IOException)
             {
@@ -318,8 +319,9 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
                     output[outputPos + plaintext.Length] = (byte)contentType;
                 }
 
-                m_encryptCipher.Init(nonce, m_macSize, additionalData);
-                outputPos += m_encryptCipher.DoFinal(output, outputPos, innerPlaintextLength, output, outputPos);
+                m_encryptCipher.Init(nonce, m_macSize, null);
+                outputPos += m_encryptCipher.DoFinal(additionalData, output, outputPos, innerPlaintextLength, output,
+                    outputPos);
             }
             catch (IOException)
             {
@@ -366,6 +368,8 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
                 throw new TlsFatalAlert(AlertDescription.internal_error);
             }
 
+            m_decryptCipher.Init(nonce, m_macSize, null);
+
             int encryptionOffset = ciphertextOffset + m_record_iv_length;
             int encryptionLength = ciphertextLength - m_record_iv_length;
             int innerPlaintextLength = m_decryptCipher.GetOutputSize(encryptionLength);
@@ -376,17 +380,8 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
             int outputPos;
             try
             {
-                m_decryptCipher.Init(nonce, m_macSize, additionalData);
-                outputPos = m_decryptCipher.DoFinal(ciphertext, encryptionOffset, encryptionLength, ciphertext,
-                    encryptionOffset);
-            }
-            catch (TlsFatalAlert fatalAlert)
-            {
-                if (AlertDescription.bad_record_mac == fatalAlert.AlertDescription)
-                {
-                    m_decryptCipher.Reset();
-                }
-                throw;
+                outputPos = m_decryptCipher.DoFinal(additionalData, ciphertext, encryptionOffset, encryptionLength,
+                    ciphertext, encryptionOffset);
             }
             catch (IOException)
             {
@@ -394,7 +389,6 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
             }
             catch (Exception e)
             {
-                m_decryptCipher.Reset();
                 throw new TlsFatalAlert(AlertDescription.bad_record_mac, e);
             }
 
diff --git a/crypto/src/tls/crypto/impl/TlsAeadCipherImpl.cs b/crypto/src/tls/crypto/impl/TlsAeadCipherImpl.cs
index 0cd2923c2..c521041ff 100644
--- a/crypto/src/tls/crypto/impl/TlsAeadCipherImpl.cs
+++ b/crypto/src/tls/crypto/impl/TlsAeadCipherImpl.cs
@@ -17,6 +17,7 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
         void SetKey(ReadOnlySpan<byte> key);
 #endif
 
+        // TODO[api] Remove additionalData parameter
         /// <summary>Initialise the parameters for the AEAD operator.</summary>
         /// <param name="nonce">the nonce.</param>
         /// <param name="macSize">MAC size in bytes.</param>
@@ -29,6 +30,7 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
         /// <returns>the maximum size of the output.</returns>
         int GetOutputSize(int inputLength);
 
+        // TODO[api] Remove unused variant
         /// <summary>Perform the cipher encryption/decryption returning the output in output.</summary>
         /// <remarks>
         /// Note: we have to use DoFinal() here as it is the only way to guarantee output from the underlying cipher.
@@ -42,6 +44,22 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl
         /// <exception cref="IOException">in case of failure.</exception>
         int DoFinal(byte[] input, int inputOffset, int inputLength, byte[] output, int outputOffset);
 
+        /// <summary>Perform the cipher encryption/decryption returning the output in output.</summary>
+        /// <remarks>
+        /// Note: we have to use DoFinal() here as it is the only way to guarantee output from the underlying cipher.
+        /// </remarks>
+        /// <param name="additionalData">any additional data to be included in the MAC calculation.</param>
+        /// <param name="input">array holding input data to the cipher.</param>
+        /// <param name="inputOffset">offset into input array data starts at.</param>
+        /// <param name="inputLength">length of the input data in the array.</param>
+        /// <param name="output">array to hold the cipher output.</param>
+        /// <param name="outputOffset">offset into output array to start saving output.</param>
+        /// <returns>the amount of data written to output.</returns>
+        /// <exception cref="IOException">in case of failure.</exception>
+        int DoFinal(byte[] additionalData, byte[] input, int inputOffset, int inputLength, byte[] output,
+            int outputOffset);
+
+        // TODO[api] Remove
         void Reset();
     }
 }
diff --git a/crypto/src/tls/crypto/impl/bc/BcChaCha20Poly1305.cs b/crypto/src/tls/crypto/impl/bc/BcChaCha20Poly1305.cs
index 06a09bbb1..dc5672814 100644
--- a/crypto/src/tls/crypto/impl/bc/BcChaCha20Poly1305.cs
+++ b/crypto/src/tls/crypto/impl/bc/BcChaCha20Poly1305.cs
@@ -73,6 +73,18 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
             }
         }
 
+        public int DoFinal(byte[] additionalData, byte[] input, int inputOffset, int inputLength, byte[] output,
+            int outputOffset)
+        {
+            if (additionalData != null)
+            {
+                this.m_additionalDataLength += additionalData.Length;
+                UpdateMac(additionalData, 0, additionalData.Length);
+            }
+
+            return DoFinal(input, inputOffset, inputLength, output, outputOffset);
+        }
+
         public int GetOutputSize(int inputLength)
         {
             return m_isEncrypting ? inputLength + 16 : inputLength - 16;
@@ -98,8 +110,6 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
 
         public void Reset()
         {
-            m_cipher.Reset();
-            m_mac.Reset();
         }
 
         public void SetKey(byte[] key, int keyOff, int keyLen)
diff --git a/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs b/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs
index 82f1382bd..7dfa737d9 100644
--- a/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs
+++ b/crypto/src/tls/crypto/impl/bc/BcTlsAeadCipherImpl.cs
@@ -58,9 +58,19 @@ namespace Org.BouncyCastle.Tls.Crypto.Impl.BC
             return len;
         }
 
+        public virtual int DoFinal(byte[] additionalData, byte[] input, int inputOffset, int inputLength, byte[] output,
+            int outputOffset)
+        {
+            if (additionalData != null)
+            {
+                m_cipher.ProcessAadBytes(additionalData, 0, additionalData.Length);
+            }
+
+            return DoFinal(input, inputOffset, inputLength, output, outputOffset);
+        }
+
         public void Reset()
         {
-            m_cipher.Reset();
         }
     }
 }