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

1 files changed, 29 insertions, 7 deletions

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