From 399fbee1e4db49d8c51fb70ed0916d0743b26c04 Mon Sep 17 00:00:00 2001
From: Bart Willems <bwillems@protonmail.com>
Date: Thu, 25 Jun 2020 06:20:27 +0200
Subject: [PATCH] send the correct exit code when closing the websocket (#335)

Signed-off-by: Bart Willems <bwillems@protonmail.com>
---
 websocket-autobahn/src/main.rs       | 4 ++++
 websocket-chat-broker/src/session.rs | 5 ++---
 websocket-chat/src/main.rs           | 3 ++-
 websocket-tcp-chat/src/main.rs       | 3 ++-
 websocket/src/main.rs                | 3 ++-
 5 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/websocket-autobahn/src/main.rs b/websocket-autobahn/src/main.rs
index 85522be8..d42dc243 100644
--- a/websocket-autobahn/src/main.rs
+++ b/websocket-autobahn/src/main.rs
@@ -25,6 +25,10 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WebSocket {
                 ws::Message::Text(text) => ctx.text(text),
                 ws::Message::Binary(bin) => ctx.binary(bin),
                 ws::Message::Ping(bytes) => ctx.pong(&bytes),
+                ws::Message::Close(reason) => {
+                    ctx.close(reason);
+                    ctx.stop();
+                }
                 _ => {}
             }
         } else {
diff --git a/websocket-chat-broker/src/session.rs b/websocket-chat-broker/src/session.rs
index 86b57838..a5029548 100644
--- a/websocket-chat-broker/src/session.rs
+++ b/websocket-chat-broker/src/session.rs
@@ -151,11 +151,10 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsChatSession {
                 }
                 self.send_msg(msg);
             }
-
-            ws::Message::Close(_) => {
+            ws::Message::Close(reason) => {
+                ctx.close(reason);
                 ctx.stop();
             }
-
             _ => {}
         }
     }
diff --git a/websocket-chat/src/main.rs b/websocket-chat/src/main.rs
index a9b7f6e8..3b28b7fe 100644
--- a/websocket-chat/src/main.rs
+++ b/websocket-chat/src/main.rs
@@ -182,7 +182,8 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsChatSession {
                 }
             }
             ws::Message::Binary(_) => println!("Unexpected binary"),
-            ws::Message::Close(_) => {
+            ws::Message::Close(reason) => {
+                ctx.close(reason);
                 ctx.stop();
             }
             ws::Message::Continuation(_) => {
diff --git a/websocket-tcp-chat/src/main.rs b/websocket-tcp-chat/src/main.rs
index 0d8f55b8..b85e4436 100644
--- a/websocket-tcp-chat/src/main.rs
+++ b/websocket-tcp-chat/src/main.rs
@@ -184,7 +184,8 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsChatSession {
                 }
             }
             ws::Message::Binary(_) => println!("Unexpected binary"),
-            ws::Message::Close(_) => {
+            ws::Message::Close(reason) => {
+                ctx.close(reason);
                 ctx.stop();
             }
             _ => (),
diff --git a/websocket/src/main.rs b/websocket/src/main.rs
index e1c77901..efa05fa4 100644
--- a/websocket/src/main.rs
+++ b/websocket/src/main.rs
@@ -59,7 +59,8 @@ impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for MyWebSocket {
             }
             Ok(ws::Message::Text(text)) => ctx.text(text),
             Ok(ws::Message::Binary(bin)) => ctx.binary(bin),
-            Ok(ws::Message::Close(_)) => {
+            Ok(ws::Message::Close(reason)) => {
+                ctx.close(reason);
                 ctx.stop();
             }
             _ => ctx.stop(),