Process all pending WS messages (#3154)

* add ws server message handling tests

* fix

* ws message processing fix

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!(

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();