From 4714773370f273f45d543c2e31d1cab5c2afd7be Mon Sep 17 00:00:00 2001
From: Alexey Galakhov <agalakhov@snapview.de>
Date: Tue, 31 Jan 2017 21:05:51 +0100
Subject: [PATCH] close: implement complete close handshake

---
 src/error.rs        |  4 ++++
 src/protocol/mod.rs | 29 +++++++++++++++++++++++------
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 5a51d36..9126f09 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -16,6 +16,8 @@ pub type Result<T> = result::Result<T, Error>;
 /// Possible WebSocket errors
 #[derive(Debug)]
 pub enum Error {
+    /// WebSocket connection closed (normally)
+    ConnectionClosed,
     /// Input-output error
     Io(io::Error),
     /// Buffer capacity exhausted
@@ -33,6 +35,7 @@ pub enum Error {
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
+            Error::ConnectionClosed => write!(f, "Connection closed"),
             Error::Io(ref err) => write!(f, "IO error: {}", err),
             Error::Capacity(ref msg) => write!(f, "Space limit exceeded: {}", msg),
             Error::Protocol(ref msg) => write!(f, "WebSocket protocol error: {}", msg),
@@ -46,6 +49,7 @@ impl fmt::Display for Error {
 impl ErrorTrait for Error {
     fn description(&self) -> &str {
         match *self {
+            Error::ConnectionClosed => "",
             Error::Io(ref err) => err.description(),
             Error::Capacity(ref msg) => msg.borrow(),
             Error::Protocol(ref msg) => msg.borrow(),
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 1dc9b63..8d12f9f 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -82,7 +82,8 @@ impl<Stream> WebSocket<Stream>
         match self.state {
             WebSocketState::Active => {
                 self.state = WebSocketState::ClosedByUs;
-                // TODO
+                let frame = Frame::close(None);
+                self.send_queue.push_back(frame);
             }
             _ => {
                 // already closed, nothing to do
@@ -220,7 +221,6 @@ impl<Stream> WebSocket<Stream>
             } // match opcode
 
         } else {
-            //Ok(None) // TODO handle EOF?
             Err(Error::Protocol("Connection reset without closing handshake".into()))
         }
     }
@@ -240,25 +240,26 @@ impl<Stream> WebSocket<Stream>
                     Frame::close(None)
                 };
                 self.send_queue.push_back(reply);
+                Ok(())
             }
             WebSocketState::ClosedByPeer => {
                 // It is already closed, just ignore.
+                Ok(())
             }
             WebSocketState::ClosedByUs => {
                 // We received a reply.
                 match self.role {
                     Role::Client => {
                         // Client waits for the server to close the connection.
+                        Ok(())
                     }
                     Role::Server => {
                         // Server closes the connection.
-                        // TODO
+                        Err(Error::ConnectionClosed)
                     }
                 }
             }
         }
-        //unimplemented!()
-        Ok(())
     }
 
     /// Received a ping frame.
@@ -294,7 +295,23 @@ impl<Stream> WebSocket<Stream>
         while let Some(data) = self.send_queue.pop_front() {
             self.send_one_frame(data)?;
         }
-        Ok(())
+
+        // If we're closing and there is nothing to send anymore, we should close the connection.
+        match self.state {
+            WebSocketState::ClosedByPeer if self.send_queue.is_empty() => {
+                // The underlying TCP connection, in most normal cases, SHOULD be closed
+                // first by the server, so that it holds the TIME_WAIT state and not the
+                // client (as this would prevent it from re-opening the connection for 2
+                // maximum segment lifetimes (2MSL), while there is no corresponding
+                // server impact as a TIME_WAIT connection is immediately reopened upon
+                // a new SYN with a higher seq number). (RFC 6455)
+                match self.role {
+                    Role::Client => Ok(()),
+                    Role::Server => Err(Error::ConnectionClosed),
+                }
+            }
+            _ => Ok(()),
+        }
     }
 
     /// Send a single pending frame.