From 20242d19f73a366c4bda7d1e1153ca1e1ddabf9b Mon Sep 17 00:00:00 2001
From: Alexey Galakhov <agalakhov@snapview.de>
Date: Mon, 9 Jul 2018 17:00:42 +0200
Subject: [PATCH] config: protocol: add message size limiting

Signed-off-by: Alexey Galakhov <agalakhov@snapview.de>
---
 src/protocol/message.rs | 31 +++++++++++++++++++++++++--
 src/protocol/mod.rs     | 46 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/src/protocol/message.rs b/src/protocol/message.rs
index b2ee1a5..8ee3eac 100644
--- a/src/protocol/message.rs
+++ b/src/protocol/message.rs
@@ -3,7 +3,7 @@ use std::fmt;
 use std::result::Result as StdResult;
 use std::str;
 
-use error::Result;
+use error::{Result, Error};
 
 mod string_collect {
 
@@ -26,6 +26,11 @@ mod string_collect {
             }
         }
 
+        pub fn len(&self) -> usize {
+            self.data.len()
+                .saturating_add(self.incomplete.map(|i| i.buffer_len as usize).unwrap_or(0))
+        }
+
         pub fn extend<T: AsRef<[u8]>>(&mut self, tail: T) -> Result<()> {
             let mut input: &[u8] = tail.as_ref();
 
@@ -105,8 +110,29 @@ impl IncompleteMessage {
             }
         }
     }
+
+    /// Get the current filled size of the buffer.
+    pub fn len(&self) -> usize {
+        match self.collector {
+            IncompleteMessageCollector::Text(ref t) => t.len(),
+            IncompleteMessageCollector::Binary(ref b) => b.len(),
+        }
+    }
+
     /// Add more data to an existing message.
-    pub fn extend<T: AsRef<[u8]>>(&mut self, tail: T) -> Result<()> {
+    pub fn extend<T: AsRef<[u8]>>(&mut self, tail: T, size_limit: Option<usize>) -> Result<()> {
+        // Always have a max size. This ensures an error in case of concatenating two buffers
+        // of more than `usize::max_value()` bytes in total.
+        let max_size = size_limit.unwrap_or_else(usize::max_value);
+        let my_size = self.len();
+        let portion_size = tail.as_ref().len();
+        // Be careful about integer overflows here.
+        if my_size > max_size || portion_size > max_size - my_size {
+            return Err(Error::Capacity(
+                format!("Message too big: {} + {} > {}", my_size, portion_size, max_size).into()
+            ))
+        }
+
         match self.collector {
             IncompleteMessageCollector::Binary(ref mut v) => {
                 v.extend(tail.as_ref());
@@ -117,6 +143,7 @@ impl IncompleteMessage {
             }
         }
     }
+
     /// Convert an incomplete message into a complete one.
     pub fn complete(self) -> Result<Message> {
         match self.collector {
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 2e8e6c8..3436e20 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -33,12 +33,17 @@ pub struct WebSocketConfig {
     /// means here that the size of the queue is unlimited. The default value is the unlimited
     /// queue.
     pub max_send_queue: Option<usize>,
+    /// The maximum size of a message. `None` means no size limit. The default value is 64 megabytes
+    /// which should be reasonably big for all normal use-cases but small enough to prevent
+    /// memory eating by a malicious user.
+    pub max_message_size: Option<usize>,
 }
 
 impl Default for WebSocketConfig {
     fn default() -> Self {
         WebSocketConfig {
             max_send_queue: None,
+            max_message_size: Some(64 << 20),
         }
     }
 }
@@ -308,8 +313,7 @@ impl<Stream: Read + Write> WebSocket<Stream> {
                     match data {
                         OpData::Continue => {
                             if let Some(ref mut msg) = self.incomplete {
-                                // TODO if msg too big
-                                msg.extend(frame.into_data())?;
+                                msg.extend(frame.into_data(), self.config.max_message_size)?;
                             } else {
                                 return Err(Error::Protocol("Continue frame but nothing to continue".into()))
                             }
@@ -332,7 +336,7 @@ impl<Stream: Read + Write> WebSocket<Stream> {
                                     _ => panic!("Bug: message is not text nor binary"),
                                 };
                                 let mut m = IncompleteMessage::new(message_type);
-                                m.extend(frame.into_data())?;
+                                m.extend(frame.into_data(), self.config.max_message_size)?;
                                 m
                             };
                             if fin {
@@ -475,7 +479,7 @@ impl WebSocketState {
 
 #[cfg(test)]
 mod tests {
-    use super::{WebSocket, Role, Message};
+    use super::{WebSocket, Role, Message, WebSocketConfig};
 
     use std::io;
     use std::io::Cursor;
@@ -517,4 +521,38 @@ mod tests {
         assert_eq!(socket.read_message().unwrap(), Message::Binary(vec![0x01, 0x02, 0x03]));
     }
 
+
+    #[test]
+    fn size_limiting_text_fragmented() {
+        let incoming = Cursor::new(vec![
+            0x01, 0x07,
+            0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x2c, 0x20,
+            0x80, 0x06,
+            0x57, 0x6f, 0x72, 0x6c, 0x64, 0x21,
+        ]);
+        let limit = WebSocketConfig {
+            max_message_size: Some(10),
+            .. WebSocketConfig::default()
+        };
+        let mut socket = WebSocket::from_raw_socket(WriteMoc(incoming), Role::Client, Some(limit));
+        assert_eq!(socket.read_message().unwrap_err().to_string(),
+            "Space limit exceeded: Message too big: 7 + 6 > 10"
+        );
+    }
+
+    #[test]
+    fn size_limiting_binary() {
+        let incoming = Cursor::new(vec![
+            0x82, 0x03,
+            0x01, 0x02, 0x03,
+        ]);
+        let limit = WebSocketConfig {
+            max_message_size: Some(2),
+            .. WebSocketConfig::default()
+        };
+        let mut socket = WebSocket::from_raw_socket(WriteMoc(incoming), Role::Client, Some(limit));
+        assert_eq!(socket.read_message().unwrap_err().to_string(),
+            "Space limit exceeded: Message too big: 0 + 3 > 2"
+        );
+    }
 }