From c5dfead7f999aa4dfe5ecd0a6282a5407ec4f4c4 Mon Sep 17 00:00:00 2001
From: Heinz Gies <hgies@wayfair.com>
Date: Tue, 10 Sep 2019 16:47:37 +0200
Subject: [PATCH] Improve close condition handling

---
 actix-http/src/ws/codec.rs |  2 +-
 actix-http/src/ws/frame.rs | 13 +++++++++----
 actix-http/src/ws/proto.rs | 26 ++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/actix-http/src/ws/codec.rs b/actix-http/src/ws/codec.rs
index 032db31ea..72055f7ba 100644
--- a/actix-http/src/ws/codec.rs
+++ b/actix-http/src/ws/codec.rs
@@ -184,7 +184,7 @@ impl Decoder for Codec {
                     OpCode::Bad => Err(ProtocolError::BadOpCode),
                     OpCode::Close => {
                         if let Some(ref pl) = payload {
-                            let close_reason = Parser::parse_close_payload(pl);
+                            let close_reason = Parser::parse_close_payload(pl)?;
                             Ok(Some(Frame::Close(close_reason)))
                         } else {
                             Ok(Some(Frame::Close(None)))
diff --git a/actix-http/src/ws/frame.rs b/actix-http/src/ws/frame.rs
index 9f93897e7..90b7f2153 100644
--- a/actix-http/src/ws/frame.rs
+++ b/actix-http/src/ws/frame.rs
@@ -140,18 +140,23 @@ impl Parser {
     }
 
     /// Parse the payload of a close frame.
-    pub fn parse_close_payload(payload: &[u8]) -> Option<CloseReason> {
+    pub fn parse_close_payload(payload: &[u8]) -> Result<Option<CloseReason>, ProtocolError> {
         if payload.len() >= 2 {
             let raw_code = u16::from_be_bytes(TryFrom::try_from(&payload[..2]).unwrap());
             let code = CloseCode::from(raw_code);
             let description = if payload.len() > 2 {
-                Some(String::from_utf8_lossy(&payload[2..]).into())
+                if let Ok(desc) = String::from_utf8(payload[2..].to_vec()) {
+                    Some(desc)
+                } else {
+                    return Err(ProtocolError::BadEncoding)
+                }
+                
             } else {
                 None
             };
-            Some(CloseReason { code, description })
+            Ok(Some(CloseReason { code, description }))
         } else {
-            None
+            Ok(None)
         }
     }
 
diff --git a/actix-http/src/ws/proto.rs b/actix-http/src/ws/proto.rs
index e14651a56..ae0489ffb 100644
--- a/actix-http/src/ws/proto.rs
+++ b/actix-http/src/ws/proto.rs
@@ -128,6 +128,20 @@ pub enum CloseCode {
     /// connect to a different IP (when multiple targets exist), or
     /// reconnect to the same IP when a user has performed an action.
     Again,
+    ///   3000-3999
+    ///
+    ///   Status codes in the range 3000-3999 are reserved for use by
+    ///   libraries, frameworks, and applications.  These status codes are
+    ///   registered directly with IANA.  The interpretation of these codes
+    ///   is undefined by this protocol.
+    Library(u16),
+    ///   4000-4999
+    ///
+    ///  Status codes in the range 4000-4999 are reserved for private use
+    ///  and thus can't be registered.  Such codes can be used by prior
+    ///  agreements between WebSocket applications.  The interpretation of
+    ///  these codes is undefined by this protocol.
+    Private(u16),
     #[doc(hidden)]
     Tls,
     #[doc(hidden)]
@@ -150,6 +164,8 @@ impl Into<u16> for CloseCode {
             Restart => 1012,
             Again => 1013,
             Tls => 1015,
+            Library(code) => code,
+            Private(code) => code,
             Other(code) => code,
         }
     }
@@ -171,6 +187,8 @@ impl From<u16> for CloseCode {
             1012 => Restart,
             1013 => Again,
             1015 => Tls,
+            3000..=3999 => Library(code),
+            4000..=4999 => Private(code),
             _ => Other(code),
         }
     }
@@ -293,6 +311,10 @@ mod test {
         assert_eq!(CloseCode::from(1013u16), CloseCode::Again);
         assert_eq!(CloseCode::from(1015u16), CloseCode::Tls);
         assert_eq!(CloseCode::from(2000u16), CloseCode::Other(2000));
+        assert_eq!(CloseCode::from(3000u16), CloseCode::Library(3000));
+        assert_eq!(CloseCode::from(3999u16), CloseCode::Library(3999));
+        assert_eq!(CloseCode::from(4000u16), CloseCode::Private(4000));
+        assert_eq!(CloseCode::from(4999u16), CloseCode::Private(4999));
     }
 
     #[test]
@@ -311,5 +333,9 @@ mod test {
         assert_eq!(1013u16, Into::<u16>::into(CloseCode::Again));
         assert_eq!(1015u16, Into::<u16>::into(CloseCode::Tls));
         assert_eq!(2000u16, Into::<u16>::into(CloseCode::Other(2000)));
+        assert_eq!(3000u16, Into::<u16>::into(CloseCode::Library(3000)));
+        assert_eq!(3999u16, Into::<u16>::into(CloseCode::Library(3999)));
+        assert_eq!(4000u16, Into::<u16>::into(CloseCode::Private(4000)));
+        assert_eq!(4999u16, Into::<u16>::into(CloseCode::Private(4999)));
     }
 }