From dde219a717cfbb0f85741a304fa83ccdb818831b Mon Sep 17 00:00:00 2001
From: Pierre Hubert <pierre.git@communiquons.org>
Date: Tue, 30 Aug 2022 15:09:39 +0200
Subject: [PATCH] Close Websocket when upstream TCP connection is closed

---
 tcp_relay_server/src/relay_ws.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tcp_relay_server/src/relay_ws.rs b/tcp_relay_server/src/relay_ws.rs
index 913c2d2..1fcfb33 100644
--- a/tcp_relay_server/src/relay_ws.rs
+++ b/tcp_relay_server/src/relay_ws.rs
@@ -4,6 +4,7 @@ use std::time::{Duration, Instant};
 use actix::{Actor, ActorContext, AsyncContext, Handler, Message, StreamHandler};
 use actix_web::{Error, HttpRequest, HttpResponse, web};
 use actix_web_actors::ws;
+use actix_web_actors::ws::{CloseCode, CloseReason};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
 use tokio::net::TcpStream;
@@ -20,6 +21,9 @@ const CLIENT_TIMEOUT: Duration = Duration::from_secs(60);
 #[rtype(result = "bool")]
 pub struct DataForWebSocket(Vec<u8>);
 
+#[derive(Message)]
+#[rtype(result = "()")]
+pub struct TCPReadEndClosed;
 
 /// Define HTTP actor
 struct RelayWS {
@@ -72,6 +76,7 @@ impl Actor for RelayWS {
                     Ok(l) => {
                         if l == 0 {
                             log::info!("Got empty read. Closing read end...");
+                            addr.do_send(TCPReadEndClosed);
                             return;
                         }
 
@@ -129,6 +134,16 @@ impl Handler<DataForWebSocket> for RelayWS {
     }
 }
 
+impl Handler<TCPReadEndClosed> for RelayWS {
+    type Result = ();
+
+    fn handle(&mut self, _msg: TCPReadEndClosed, ctx: &mut Self::Context) -> Self::Result {
+        ctx.close(Some(CloseReason {
+            code: CloseCode::Away,
+            description: Some("TCP read end closed.".to_string()),
+        }));
+    }
+}
 
 #[derive(serde::Deserialize)]
 pub struct WebSocketQuery {