Detect abnormal close messages for websocket streams

For #13

Author: Wesmania

Cargo.lock

@@ -13,7 +13,7 @@ mod test {
     use rand::Rng;
     use tokio::{io::{AsyncReadExt, AsyncWriteExt}, join, net::TcpStream, select};
     use tokio_util::sync::CancellationToken;
-    use tokio_websockets::{ClientBuilder, Message};
+    use tokio_websockets::{ClientBuilder, CloseCode, Message};
 
     use crate::{config::test::default_config, server::server::server_with_real_deps, util::test::{compare_bufs, get_file, setup_logging, sleep_ms}};
 
@@ -89,6 +89,32 @@ mod test {
         }
     }
 
+    async fn websocket_writer_end_with_normal_closure(port: u16, header: &[u8], file: Vec<u8>) {
+        let uri = format!("ws://127.0.0.1:{}", port).parse::<Uri>().unwrap();
+        let (mut client, _) = ClientBuilder::from_uri(uri).connect().await.unwrap();
+
+        random_sleep().await;
+        client.send(Message::binary(header.to_owned())).await.unwrap();
+        random_sleep().await;
+        for data in file.chunks(1000) {
+                client.send(Message::binary(data.to_owned())).await.unwrap();
+                random_sleep().await;
+        }
+        // Add a ping and pong we should ignore for good measure.
+        client.send(Message::ping("foo")).await.unwrap();
+        client.send(Message::pong("bar")).await.unwrap();
+        client.send(Message::close(Some(CloseCode::NORMAL_CLOSURE), "")).await.unwrap();
+    }
+
+    async fn websocket_writer_abnormal_closure(port: u16, header: &[u8]) {
+        let uri = format!("ws://127.0.0.1:{}", port).parse::<Uri>().unwrap();
+        let (mut client, _) = ClientBuilder::from_uri(uri).connect().await.unwrap();
+        random_sleep().await;
+        client.send(Message::binary(header.to_owned())).await.unwrap();
+        random_sleep().await;
+        client.send(Message::close(Some(CloseCode::GOING_AWAY), "")).await.unwrap();
+    }
+
     async fn websocket_reader(port: u16, header: &[u8]) -> Vec<u8> {
         let uri = format!("ws://127.0.0.1:{}", port).parse::<Uri>().unwrap();
         let (mut client, _) = ClientBuilder::from_uri(uri).connect().await.unwrap();
@@ -216,4 +242,100 @@ mod test {
 
         // TODO: verify stats saved to the database.
     }
+
+    // See issue #13.
+    #[cfg_attr(not(feature = "local_db_tests"), ignore)]
+    #[tokio::test(flavor = "multi_thread")]
+    async fn integration_test_server_websocket_normal_closure_message() {
+
+        setup_logging();
+
+        let replay_dir = tempfile::tempdir().unwrap();
+        let mut config = default_config();
+        config.server.port = None;
+        config.server.websocket_port = Some(0);
+        config.database.host = std::env::var("DB_HOST").unwrap();
+        config.database.port = str::parse::<u16>(&std::env::var("DB_PORT").unwrap()).unwrap();
+        config.storage.vault_path = replay_dir.path().to_str().unwrap().into();
+        config.server.connection_accept_timeout_s = Duration::from_millis(50);
+
+        let token = CancellationToken::new();
+        let (server, port_info) = server_with_real_deps(Arc::new(config), token.child_token()).await;
+
+        let sent_replay = get_file("example");
+        let replay_writer = websocket_writer_end_with_normal_closure(port_info.websocket.unwrap(), b"P/2001/foo\0", sent_replay.clone());
+
+        let exit_server = async move || {
+            tokio::time::sleep(Duration::from_secs(1)).await;
+            token.cancel();
+        };
+
+        let run = async move || {
+            let (_, _, _) = join! {
+               server.run(),
+               exit_server(),
+               replay_writer,
+            };
+        };
+        tokio::time::timeout(Duration::from_secs(2), run()).await.unwrap();
+
+        let mut replay_path: PathBuf = replay_dir.path().to_owned();
+        replay_path.push("0");
+        replay_path.push("0");
+        replay_path.push("0");
+        replay_path.push("20");
+        replay_path.push("2001.fafreplay");
+        let saved_replay = std::fs::read(&replay_path).unwrap();
+
+        // This file was copied from this test and manually verified to be correct.
+        // Zstd is deterministic for the same data + params + version, let's rely on that.
+        let reference_replay = get_file("integration/game_2001.fafreplay");
+        compare_bufs(saved_replay, reference_replay);
+    }
+
+    #[cfg_attr(not(feature = "local_db_tests"), ignore)]
+    #[tokio::test(flavor = "multi_thread")]
+    async fn integration_test_server_websocket_abnormal_closure() {
+
+        setup_logging();
+
+        let replay_dir = tempfile::tempdir().unwrap();
+        let mut config = default_config();
+        config.server.port = None;
+        config.server.websocket_port = Some(0);
+        config.database.host = std::env::var("DB_HOST").unwrap();
+        config.database.port = str::parse::<u16>(&std::env::var("DB_PORT").unwrap()).unwrap();
+        config.storage.vault_path = replay_dir.path().to_str().unwrap().into();
+        config.server.connection_accept_timeout_s = Duration::from_millis(50);
+
+        let token = CancellationToken::new();
+        let (server, port_info) = server_with_real_deps(Arc::new(config), token.child_token()).await;
+
+        let replay_writer = websocket_writer_abnormal_closure(port_info.websocket.unwrap(), b"P/2001/foo\0");
+
+        let exit_server = async move || {
+            tokio::time::sleep(Duration::from_secs(1)).await;
+            token.cancel();
+        };
+
+        let run = async move || {
+            let (_, _, _) = join! {
+               server.run(),
+               exit_server(),
+               replay_writer,
+            };
+        };
+        tokio::time::timeout(Duration::from_secs(2), run()).await.unwrap();
+
+        // The abnormal closure should've been interpreted as an error and not data.
+        // FIXME this doesn't check that the connection ended with an error...
+        // FIXME a test in replay.rs might be better.
+        let mut replay_path: PathBuf = replay_dir.path().to_owned();
+        replay_path.push("0");
+        replay_path.push("0");
+        replay_path.push("0");
+        replay_path.push("20");
+        replay_path.push("2001.fafreplay");
+        let _saved_replay = std::fs::read(&replay_path).expect_err("Replay should have been empty and not saved.");
+    }
 }

src/integration_tests.rs

@@ -7,7 +7,7 @@ use tokio_util::{
     bytes::Bytes,
     io::{CopyToBytes, SinkWriter, StreamReader},
 };
-use tokio_websockets::{Error, Message, ServerBuilder, WebSocketStream};
+use tokio_websockets::{CloseCode, Error, Message, ServerBuilder, WebSocketStream};
 
 use super::connection::{ReaderType, WriterType};
 
@@ -51,7 +51,19 @@ fn map_websocket_sink_to_bytes(sink: impl WebsocketSink) -> impl for<'a> Sink<&'
 fn map_websocket_stream_to_bytes(stream: impl WebsocketStream) -> impl Stream<Item = Result<Bytes, std::io::Error>> {
     stream.map(|i| {
         i.map_err(map_message_error_to_io_error)
-            .map(|mes| mes.into_payload().into())
+            .and_then(|mes| {
+                if let Some((code, reason)) = mes.as_close() {
+                    if code != CloseCode::NORMAL_CLOSURE {
+                        Err(std::io::Error::other(reason))
+                    } else {
+                        Ok(Bytes::new())
+                    }
+                } else if mes.is_ping() || mes.is_pong() {
+                    Ok(Bytes::new())
+                } else {
+                    Ok(mes.into_payload().into())
+                }
+            })
     })
 }