Avoid extra copy when processing TLS handshake messages

- move handshake hash updates (outbound) to a more sensible place

6 files changed, 104 insertions, 58 deletions

diff --git a/crypto/src/crypto/tls/ByteQueue.cs b/crypto/src/crypto/tls/ByteQueue.cs
index f9398bbaf..c23ec2fbf 100644
--- a/crypto/src/crypto/tls/ByteQueue.cs
+++ b/crypto/src/crypto/tls/ByteQueue.cs
@@ -1,4 +1,5 @@
 using System;
+using System.IO;
 
 namespace Org.BouncyCastle.Crypto.Tls
 {
@@ -57,36 +58,14 @@ namespace Org.BouncyCastle.Crypto.Tls
             this.databuf = new byte[capacity];
         }
 
-        /// <summary>Read data from the buffer.</summary>
-        /// <param name="buf">The buffer where the read data will be copied to.</param>
-        /// <param name="offset">How many bytes to skip at the beginning of buf.</param>
-        /// <param name="len">How many bytes to read at all.</param>
-        /// <param name="skip">How many bytes from our data to skip.</param>
-        public void Read(
-            byte[]	buf,
-            int		offset,
-            int		len,
-            int		skip)
-        {
-            if ((buf.Length - offset) < len)
-            {
-                throw new ArgumentException("Buffer size of " + buf.Length + " is too small for a read of " + len + " bytes");
-            }
-            if ((available - skip) < len)
-            {
-                throw new InvalidOperationException("Not enough data to read");
-            }
-            Array.Copy(databuf, skipped + skip, buf, offset, len);
-        }
-
         /// <summary>Add some data to our buffer.</summary>
         /// <param name="data">A byte-array to read data from.</param>
         /// <param name="offset">How many bytes to skip at the beginning of the array.</param>
         /// <param name="len">How many bytes to read from the array.</param>
         public void AddData(
-            byte[]	data,
-            int		offset,
-            int		len)
+            byte[] data,
+            int offset,
+            int len)
         {
             if ((skipped + available + len) > databuf.Length)
             {
@@ -108,6 +87,64 @@ namespace Org.BouncyCastle.Crypto.Tls
             available += len;
         }
 
+        /// <summary>The number of bytes which are available in this buffer.</summary>
+        public int Available
+        {
+            get { return available; }
+        }
+
+        /// <summary>Copy some bytes from the beginning of the data to the provided <c cref="Stream">Stream</c>.</summary>
+        /// <param name="output">The <c cref="Stream">Stream</c> to copy the bytes to.</param>
+        /// <param name="length">How many bytes to copy.</param>
+		/// <exception cref="InvalidOperationException">If insufficient data is available.</exception>
+		/// <exception cref="IOException">If there is a problem copying the data.</exception>
+        public void CopyTo(Stream output, int length)
+        {
+            if (length > available)
+                throw new InvalidOperationException("Cannot copy " + length + " bytes, only got " + available);
+
+            output.Write(databuf, skipped, length);
+        }
+
+        /// <summary>Read data from the buffer.</summary>
+        /// <param name="buf">The buffer where the read data will be copied to.</param>
+        /// <param name="offset">How many bytes to skip at the beginning of buf.</param>
+        /// <param name="len">How many bytes to read at all.</param>
+        /// <param name="skip">How many bytes from our data to skip.</param>
+        public void Read(
+            byte[]	buf,
+            int		offset,
+            int		len,
+            int		skip)
+        {
+            if ((buf.Length - offset) < len)
+            {
+                throw new ArgumentException("Buffer size of " + buf.Length + " is too small for a read of " + len + " bytes");
+            }
+            if ((available - skip) < len)
+            {
+                throw new InvalidOperationException("Not enough data to read");
+            }
+            Array.Copy(databuf, skipped + skip, buf, offset, len);
+        }
+
+        /// <summary>Return a <c cref="MemoryStream">MemoryStream</c> over some bytes at the beginning of the data.</summary>
+        /// <param name="length">How many bytes will be readable.</param>
+        /// <returns>A <c cref="MemoryStream">MemoryStream</c> over the data.</returns>
+		/// <exception cref="InvalidOperationException">If insufficient data is available.</exception>
+        public MemoryStream ReadFrom(int length)
+        {
+            if (length > available)
+                throw new InvalidOperationException("Cannot read " + length + " bytes, only got " + available);
+
+            int position = skipped;
+
+            available -= length;
+            skipped += length;
+
+            return new MemoryStream(databuf, position, length, false);
+        }
+
         /// <summary>Remove some bytes from our data from the beginning.</summary>
         /// <param name="i">How many bytes to remove.</param>
         public void RemoveData(
@@ -137,11 +174,5 @@ namespace Org.BouncyCastle.Crypto.Tls
             RemoveData(buf, 0, len, skip);
             return buf;
         }
-
-        /// <summary>The number of bytes which are available in this buffer.</summary>
-        public int Available
-        {
-            get { return available; }
-        }
     }
 }
diff --git a/crypto/src/crypto/tls/RecordStream.cs b/crypto/src/crypto/tls/RecordStream.cs
index dc66d0df7..cddd5e6d7 100644
--- a/crypto/src/crypto/tls/RecordStream.cs
+++ b/crypto/src/crypto/tls/RecordStream.cs
@@ -2,6 +2,7 @@ using System;
 using System.IO;
 
 using Org.BouncyCastle.Utilities;
+using Org.BouncyCastle.Utilities.IO;
 
 namespace Org.BouncyCastle.Crypto.Tls
 {
@@ -24,6 +25,7 @@ namespace Org.BouncyCastle.Crypto.Tls
         private MemoryStream mBuffer = new MemoryStream();
 
         private TlsHandshakeHash mHandshakeHash = null;
+        private readonly BaseOutputStream mHandshakeHashUpdater;
 
         private ProtocolVersion mReadVersion = null, mWriteVersion = null;
         private bool mRestrictReadVersion = true;
@@ -37,6 +39,7 @@ namespace Org.BouncyCastle.Crypto.Tls
             this.mOutput = output;
             this.mReadCompression = new TlsNullCompression();
             this.mWriteCompression = this.mReadCompression;
+            this.mHandshakeHashUpdater = new HandshakeHashUpdateStream(this);
         }
 
         internal virtual void Init(TlsContext context)
@@ -257,11 +260,6 @@ namespace Org.BouncyCastle.Crypto.Tls
             if (plaintextLength < 1 && type != ContentType.application_data)
                 throw new TlsFatalAlert(AlertDescription.internal_error);
 
-            if (type == ContentType.handshake)
-            {
-                UpdateHandshakeData(plaintext, plaintextOffset, plaintextLength);
-            }
-
             Stream cOut = mWriteCompression.Compress(mBuffer);
 
             byte[] ciphertext;
@@ -308,6 +306,11 @@ namespace Org.BouncyCastle.Crypto.Tls
             get { return mHandshakeHash; }
         }
 
+        internal virtual Stream HandshakeHashUpdater
+        {
+            get { return mHandshakeHashUpdater; }
+        }
+
         internal virtual TlsHandshakeHash PrepareToFinish()
         {
             TlsHandshakeHash result = mHandshakeHash;
@@ -315,11 +318,6 @@ namespace Org.BouncyCastle.Crypto.Tls
             return result;
         }
 
-        internal virtual void UpdateHandshakeData(byte[] message, int offset, int len)
-        {
-            mHandshakeHash.BlockUpdate(message, offset, len);
-        }
-
         internal virtual void SafeClose()
         {
             try
@@ -371,5 +369,20 @@ namespace Org.BouncyCastle.Crypto.Tls
             if (length > limit)
                 throw new TlsFatalAlert(alertDescription);
         }
+
+        private class HandshakeHashUpdateStream
+            : BaseOutputStream
+        {
+            private readonly RecordStream mOuter;
+            public HandshakeHashUpdateStream(RecordStream mOuter)
+            {
+                this.mOuter = mOuter;
+            }
+
+            public override void Write(byte[] buf, int off, int len)
+            {
+                mOuter.mHandshakeHash.BlockUpdate(buf, off, len);
+            }
+        }
     }
 }
diff --git a/crypto/src/crypto/tls/TlsClientProtocol.cs b/crypto/src/crypto/tls/TlsClientProtocol.cs
index 6352a893f..d24d13bb5 100644
--- a/crypto/src/crypto/tls/TlsClientProtocol.cs
+++ b/crypto/src/crypto/tls/TlsClientProtocol.cs
@@ -135,10 +135,8 @@ namespace Org.BouncyCastle.Crypto.Tls
             get { return mTlsClient; }
         }
 
-        protected override void HandleHandshakeMessage(byte type, byte[] data)
+        protected override void HandleHandshakeMessage(byte type, MemoryStream buf)
         {
-            MemoryStream buf = new MemoryStream(data, false);
-
             if (this.mResumedSession)
             {
                 if (type != HandshakeType.finished || this.mConnectionState != CS_SERVER_HELLO)
diff --git a/crypto/src/crypto/tls/TlsProtocol.cs b/crypto/src/crypto/tls/TlsProtocol.cs
index c6c3375d9..eb0ad1101 100644
--- a/crypto/src/crypto/tls/TlsProtocol.cs
+++ b/crypto/src/crypto/tls/TlsProtocol.cs
@@ -114,7 +114,7 @@ namespace Org.BouncyCastle.Crypto.Tls
         {
         }
 
-        protected abstract void HandleHandshakeMessage(byte type, byte[] buf);
+        protected abstract void HandleHandshakeMessage(byte type, MemoryStream buf);
 
         protected virtual void HandleWarningMessage(byte description)
         {
@@ -291,18 +291,14 @@ namespace Org.BouncyCastle.Crypto.Tls
                     byte[] beginning = new byte[4];
                     mHandshakeQueue.Read(beginning, 0, 4, 0);
                     byte type = TlsUtilities.ReadUint8(beginning, 0);
-                    int len = TlsUtilities.ReadUint24(beginning, 1);
+                    int length = TlsUtilities.ReadUint24(beginning, 1);
+                    int totalLength = 4 + length;
 
                     /*
                      * Check if we have enough bytes in the buffer to read the full message.
                      */
-                    if (mHandshakeQueue.Available >= (len + 4))
+                    if (mHandshakeQueue.Available >= totalLength)
                     {
-                        /*
-                         * Read the message.
-                         */
-                        byte[] buf = mHandshakeQueue.RemoveData(len, 4);
-
                         CheckReceivedChangeCipherSpec(mConnectionState == CS_END || type == HandshakeType.finished);
 
                         /*
@@ -325,12 +321,15 @@ namespace Org.BouncyCastle.Crypto.Tls
                                 this.mExpectedVerifyData = CreateVerifyData(!ctx.IsServer);
                             }
 
-                            mRecordStream.UpdateHandshakeData(beginning, 0, 4);
-                            mRecordStream.UpdateHandshakeData(buf, 0, len);
+                            mHandshakeQueue.CopyTo(mRecordStream.HandshakeHashUpdater, totalLength);
                             break;
                         }
                         }
 
+                        mHandshakeQueue.RemoveData(4);
+
+                        MemoryStream buf = mHandshakeQueue.ReadFrom(length);
+
                         /*
                          * Now, parse the message.
                          */
@@ -629,6 +628,8 @@ namespace Org.BouncyCastle.Crypto.Tls
 
         protected virtual void WriteHandshakeMessage(byte[] buf, int off, int len)
         {
+            mRecordStream.HandshakeHashUpdater.Write(buf, off, len);
+
             while (len > 0)
             {
                 // Fragment data according to the current fragment limit.
diff --git a/crypto/src/crypto/tls/TlsServerProtocol.cs b/crypto/src/crypto/tls/TlsServerProtocol.cs
index fca685295..6642f43f4 100644
--- a/crypto/src/crypto/tls/TlsServerProtocol.cs
+++ b/crypto/src/crypto/tls/TlsServerProtocol.cs
@@ -124,10 +124,8 @@ namespace Org.BouncyCastle.Crypto.Tls
             get { return mTlsServer; }
         }
 
-        protected override void HandleHandshakeMessage(byte type, byte[] data)
+        protected override void HandleHandshakeMessage(byte type, MemoryStream buf)
         {
-            MemoryStream buf = new MemoryStream(data);
-
             switch (type)
             {
             case HandshakeType.client_hello:
diff --git a/crypto/src/util/io/BaseOutputStream.cs b/crypto/src/util/io/BaseOutputStream.cs
index a0608d111..0dbe821de 100644
--- a/crypto/src/util/io/BaseOutputStream.cs
+++ b/crypto/src/util/io/BaseOutputStream.cs
@@ -60,5 +60,10 @@ namespace Org.BouncyCastle.Utilities.IO
 		{
 			Write(buffer, 0, buffer.Length);
 		}
-	}
+
+        public override void WriteByte(byte b)
+        {
+            Write(new byte[]{ b }, 0, 1);
+        }
+    }
 }