From 9d797a513fadae2b05eebc6b4a311d38b4e7b05d Mon Sep 17 00:00:00 2001
From: Kasper Ziemianek <kasper.ziemianek@gmail.com>
Date: Mon, 4 Nov 2024 12:28:18 +0100
Subject: [PATCH] Process all pending WS messages (#3154)

* add ws server message handling tests

* fix

* ws message processing fix
---
 .../tls-websocket-server/src/connection.rs    |  5 +--
 .../tls-websocket-server/src/ws_server.rs     | 39 +++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/tee-worker/common/core/tls-websocket-server/src/connection.rs b/tee-worker/common/core/tls-websocket-server/src/connection.rs
index 6959cf0ba3..cae41f1199 100644
--- a/tee-worker/common/core/tls-websocket-server/src/connection.rs
+++ b/tee-worker/common/core/tls-websocket-server/src/connection.rs
@@ -71,10 +71,7 @@ where
 			Some(MaybeServerTlsStream::Rustls(s)) => {
 				let tls_session = &mut s.sess;
 				match tls_session.read_tls(&mut s.sock) {
-					Ok(r) =>
-						if r == 0 {
-							return ConnectionState::Closing
-						},
+					Ok(_) => {},
 					Err(err) => {
 						if let std::io::ErrorKind::WouldBlock = err.kind() {
 							debug!(
diff --git a/tee-worker/common/core/tls-websocket-server/src/ws_server.rs b/tee-worker/common/core/tls-websocket-server/src/ws_server.rs
index ba9d3a6030..cd7d8384cf 100644
--- a/tee-worker/common/core/tls-websocket-server/src/ws_server.rs
+++ b/tee-worker/common/core/tls-websocket-server/src/ws_server.rs
@@ -327,6 +327,45 @@ mod tests {
 	};
 	use url::Url;
 
+	#[test]
+	fn server_handles_multiple_messages() {
+		let _ = env_logger::builder().is_test(true).try_init();
+		let port: u16 = 21780;
+		let (server, handler) = create_server(vec![], port);
+
+		let server_clone = server.clone();
+		let server_join_handle = thread::spawn(move || server_clone.run());
+
+		let messages_number = 1000;
+
+		// Wait until server is up.
+		while !server.is_running().unwrap() {
+			thread::sleep(std::time::Duration::from_millis(50));
+		}
+
+		let client_join_handle = thread::spawn(move || {
+			let mut socket = connect_tls_client(get_server_addr(port).as_str());
+
+			for i in 0..messages_number {
+				socket
+					.write_message(Message::Text(format!("Request: {:}", i)))
+					.expect("client write message to be successful");
+			}
+			// We never read, just send a message and close the connection, despite the server
+			// trying to send a reply (which will fail).
+			socket.close(None).unwrap();
+			socket.write_pending().unwrap();
+		});
+
+		client_join_handle.join().unwrap();
+		// let's wait a while so it process all messages
+		thread::sleep(Duration::from_secs(2));
+		server.shut_down().unwrap();
+		server_join_handle.join().unwrap().unwrap();
+
+		assert_eq!(messages_number, handler.get_handled_messages().len());
+	}
+
 	#[test]
 	fn server_handles_multiple_connections() {
 		let _ = env_logger::builder().is_test(true).try_init();