diff --git a/src/engine/terminal.rs b/src/engine/terminal.rs index 467afe4..76c480f 100644 --- a/src/engine/terminal.rs +++ b/src/engine/terminal.rs @@ -50,16 +50,21 @@ impl TerminalServer { async fn handle_client(mut reader: OwnedReadHalf) -> tokio::io::Result<()> { let input = tokio::spawn(async move { - let mut buffer = [0u8; 2048]; - loop { - let size = reader.read(&mut buffer).await?; + let mut len_buf = [0u8; size_of::()]; + let size = reader.read_exact(&mut len_buf).await?; - if size == 0 { + let len = usize::from_le_bytes(len_buf); + + if size == 0 || len == 0 { break; } - println!("Received {} bytes", size); + let mut buffer = vec![0u8; len]; + + reader.read_exact(&mut buffer).await?; + + println!("Received {} bytes", len); } Ok::<(), tokio::io::Error>(()) @@ -79,10 +84,7 @@ impl TerminalServer { let mut clients = self.clients.lock().await; for i in (0..clients.len()).rev() { - if let Err(e) = clients[i].write(&msg).await { - clients.remove(i); - println!("{e:?}"); - } else if let Err(e) = clients[i].flush().await { + if let Err(e) = Self::send_to_client(&mut clients, i, &msg).await { clients.remove(i); println!("{e:?}"); } @@ -90,4 +92,16 @@ impl TerminalServer { Ok(()) } + + pub async fn send_to_client( + clients: &mut tokio::sync::MutexGuard<'_, Vec>, + i: usize, + msg: &[u8], + ) -> tokio::io::Result<()> { + clients[i].write(&msg.len().to_le_bytes()).await?; + clients[i].write(msg).await?; + clients[i].flush().await?; + + Ok(()) + } } diff --git a/src/terminal/terminal.rs b/src/terminal/terminal.rs index 7c0f1f6..c58aea4 100644 --- a/src/terminal/terminal.rs +++ b/src/terminal/terminal.rs @@ -28,7 +28,10 @@ impl TerminalClient { &mut self, message: pulse_wire::terminal::TerminalClientMessage, ) -> tokio::io::Result<()> { - self.writer.write(&message.to_com()).await?; + let msg = message.to_com(); + self.writer.write(&msg.len().to_le_bytes()).await?; + self.writer.write(&msg).await?; + self.writer.flush().await?; Ok(()) } @@ -50,19 +53,21 @@ impl TerminalClient { tokio::spawn(async move { loop { - let mut buffer = vec![0u8; 4096]; + let mut len_buf = [0u8; size_of::()]; + reader + .read_exact(&mut len_buf) + .await + .expect("Failed to get header length"); - let len = reader - .read(&mut buffer) + let len = usize::from_le_bytes(len_buf); + + let mut buffer = vec![0u8; len]; + + reader + .read_exact(&mut buffer) .await .expect("Failed to read socket"); - if len == 0 { - break; - } - - buffer.truncate(len); - match TerminalServerMessage::from_com(&mut buffer) { TerminalServerMessage::WatchListUpdated(v) => { *watch_list.lock().await = v;