v0.2.6: preserve buffered lines and add request read timeouts

4 files changed, 53 insertions, 22 deletions

diff --git a/src/bin/client/connection.rs b/src/bin/client/connection.rs
index 93a7829..49695dd 100644
--- a/src/bin/client/connection.rs
+++ b/src/bin/client/connection.rs
@@ -66,6 +66,7 @@ pub fn connect(
     }
 
     let mut tls = tls;
+    tls.set_read_timeout(Some(Duration::from_secs(10)))?;
 
     let start = Instant::now();
     tls.write_line("PING")?;
@@ -124,6 +125,13 @@ impl TlsStream {
         self.inner.write_line(msg)
     }
 
+    pub fn set_read_timeout(
+        &self,
+        duration: Option<std::time::Duration>,
+    ) -> std::io::Result<()> {
+        self.inner.set_read_timeout(duration)
+    }
+
     pub fn try_clone(&self) -> std::io::Result<TlsStream> {
         Ok(TlsStream {
             inner: self.inner.try_clone()?,
diff --git a/src/bin/client/repl.rs b/src/bin/client/repl.rs
index 76b149b..033fb1a 100644
--- a/src/bin/client/repl.rs
+++ b/src/bin/client/repl.rs
@@ -18,6 +18,7 @@ use zagan_bet::client::database;
 use zagan_bet::client::ui::{
     display_width, format_pong, host_lines, input_cursor_column, lines_to_move_up,
 };
+use zagan_bet::protocol::io_failure_message;
 
 use crate::connection::TlsStream;
 
@@ -299,9 +300,9 @@ impl ClientState {
                 self.stream = Some(stream);
                 response
             }
-            Err(_) => {
+            Err(error) => {
                 self.stream = None;
-                "ERROR Connection lost".to_string()
+                io_failure_message(&error).to_string()
             }
         }
     }
diff --git a/src/bin/client/tls.rs b/src/bin/client/tls.rs
index ca99838..b118f63 100644
--- a/src/bin/client/tls.rs
+++ b/src/bin/client/tls.rs
@@ -19,6 +19,7 @@ use std::sync::Arc;
 use parking_lot::Mutex;
 use rustls::client::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
 use rustls::DigitallySignedStruct;
+use zagan_bet::protocol::LineBuffer;
 use zagan_bet::tofu::{decide, spki_fingerprint, TrustDecision};
 
 pub struct TlsConnector {
@@ -59,7 +60,11 @@ impl TlsConnector {
         conn.complete_io(&mut sock)?;
 
         Ok(TlsStream {
-            session: Arc::new(Mutex::new(Session { sock, tls: conn })),
+            session: Arc::new(Mutex::new(Session {
+                sock,
+                tls: conn,
+                lines: LineBuffer::new(),
+            })),
         })
     }
 
@@ -144,6 +149,7 @@ impl ServerCertVerifier for TofuVerifier {
 struct Session {
     sock: TcpStream,
     tls: rustls::ClientConnection,
+    lines: LineBuffer,
 }
 
 impl Session {
@@ -159,17 +165,25 @@ pub struct TlsStream {
 
 impl TlsStream {
     pub fn read_line(&mut self, buf: &mut String) -> std::io::Result<usize> {
-        let mut byte_buf = [0u8; 1024];
         let mut s = self.session.lock();
         loop {
+            if let Some(line) = s.lines.next_line() {
+                buf.push_str(&line);
+                buf.push('\n');
+                return Ok(line.len() + 1);
+            }
+
+            let mut byte_buf = [0u8; 4096];
             match s.tls.reader().read(&mut byte_buf) {
-                Ok(0) => return Ok(0),
-                Ok(n) => {
-                    buf.push_str(std::str::from_utf8(&byte_buf[..n]).unwrap_or(""));
-                    if buf.contains('\n') {
-                        return Ok(n);
+                Ok(0) => {
+                    if let Some(rest) = s.lines.take_rest() {
+                        let read = rest.len();
+                        buf.push_str(&rest);
+                        return Ok(read);
                     }
+                    return Ok(0);
                 }
+                Ok(n) => s.lines.push(&byte_buf[..n]),
                 Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => s.drive()?,
                 Err(e) => return Err(e),
             }
@@ -183,6 +197,14 @@ impl TlsStream {
         s.drive()
     }
 
+    /// Apply a read timeout to the underlying TCP socket.
+    pub fn set_read_timeout(
+        &self,
+        duration: Option<std::time::Duration>,
+    ) -> std::io::Result<()> {
+        self.session.lock().sock.set_read_timeout(duration)
+    }
+
     /// Clones share the same underlying TLS session.
     pub fn try_clone(&self) -> std::io::Result<TlsStream> {
         Ok(self.clone())
diff --git a/src/bin/server/tls.rs b/src/bin/server/tls.rs
index acf3315..62a25db 100644
--- a/src/bin/server/tls.rs
+++ b/src/bin/server/tls.rs
@@ -1,6 +1,8 @@
 use std::io::{Read, Write};
 use std::net::TcpStream;
 
+use zagan_bet::protocol::LineBuffer;
+
 #[derive(Clone)]
 pub struct TlsServer {
     acceptor: native_tls::TlsAcceptor,
@@ -42,27 +44,25 @@ impl TlsServer {
 
 pub struct TlsStream {
     stream: native_tls::TlsStream<TcpStream>,
+    lines: LineBuffer,
 }
 
 impl TlsStream {
     pub fn new(stream: native_tls::TlsStream<TcpStream>) -> Self {
-        Self { stream }
+        Self {
+            stream,
+            lines: LineBuffer::new(),
+        }
     }
 
     pub fn read_line(&mut self, s: &mut String) -> std::io::Result<usize> {
-        let mut buf = [0u8; 1024];
-        loop {
-            match self.stream.read(&mut buf) {
-                Ok(0) => return Ok(0),
-                Ok(n) => {
-                    let chunk = std::str::from_utf8(&buf[..n]).unwrap_or("");
-                    s.push_str(chunk);
-                    if s.contains('\n') {
-                        return Ok(n);
-                    }
-                }
-                Err(e) => return Err(e),
+        match self.lines.read_line(&mut self.stream)? {
+            Some(line) => {
+                s.push_str(&line);
+                s.push('\n');
+                Ok(line.len() + 1)
             }
+            None => Ok(0),
         }
     }