Fix buffering during switch to TLS

This fixes a race bug where the intereaction between our buffering and
the TLS handshake can cause problems.

There was an assert in the code to ensure we don't read "too far" into
the stream before we enter the TLS handshake.  However, it turns out
that sometimes (when the client sends the tls handshake information
quickly enough after the pre-tls data) we can end up hitting this
assertion.

This adds a test for this scenario, by using a new `DelayedReadRW` wrapper
in the server which reliably triggers this.

The fix introduces a new `PrependedReader` type that allows us to push
any extra data we may have buffered into it (together with the
underlying connection) before passing it in to rustls for the handshake.

Author: mjgarton

src/packet.rs

@@ -77,9 +77,11 @@ impl<W: Read + Write> PacketConn<W> {
 
     #[cfg(feature = "tls")]
     pub fn switch_to_tls(&mut self, config: std::sync::Arc<ServerConfig>) -> io::Result<()> {
-        assert_eq!(self.remaining, 0); // otherwise we've read ahead into the TLS handshake and will be in trouble.
-
-        self.rw.switch_to_tls(config)
+        let res = self
+            .rw
+            .switch_to_tls(config, &self.bytes[self.bytes.len() - self.remaining..]);
+        self.remaining = 0;
+        res
     }
 }
 

src/tls.rs

@@ -1,4 +1,4 @@
-use std::io;
+use std::io::{self, Chain, Cursor};
 use std::io::{Read, Write};
 use std::sync::Arc;
 
@@ -17,7 +17,7 @@ pub(crate) struct SwitchableConn<T: Read + Write>(Option<EitherConn<T>>);
 
 pub(crate) enum EitherConn<T: Read + Write> {
     Plain(T),
-    Tls(rustls::StreamOwned<ServerConnection, T>),
+    Tls(rustls::StreamOwned<ServerConnection, PrependedReader<T>>),
 }
 
 impl<T: Read + Write> Read for SwitchableConn<T> {
@@ -50,9 +50,16 @@ impl<T: Read + Write> SwitchableConn<T> {
         SwitchableConn(Some(EitherConn::Plain(rw)))
     }
 
-    pub fn switch_to_tls(&mut self, config: Arc<ServerConfig>) -> io::Result<()> {
+    pub fn switch_to_tls(
+        &mut self,
+        config: Arc<ServerConfig>,
+        to_prepend: &[u8],
+    ) -> io::Result<()> {
         let replacement = match self.0.take() {
-            Some(EitherConn::Plain(plain)) => Ok(EitherConn::Tls(create_stream(plain, config)?)),
+            Some(EitherConn::Plain(plain)) => Ok(EitherConn::Tls(create_stream(
+                PrependedReader::new(to_prepend, plain),
+                config,
+            )?)),
             Some(EitherConn::Tls(_)) => Err(io::Error::new(
                 io::ErrorKind::Other,
                 "tls variant found when plain was expected",
@@ -64,3 +71,48 @@ impl<T: Read + Write> SwitchableConn<T> {
         Ok(())
     }
 }
+
+pub(crate) struct PrependedReader<RW: Read + Write> {
+    inner: Chain<Cursor<Vec<u8>>, RW>,
+}
+
+impl<RW: Read + Write> PrependedReader<RW> {
+    fn new(prepended: &[u8], rw: RW) -> PrependedReader<RW> {
+        PrependedReader {
+            inner: Cursor::new(prepended.to_vec()).chain(rw),
+        }
+    }
+}
+
+impl<RW: Read + Write> Read for PrependedReader<RW> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.read(buf)
+    }
+}
+
+impl<RW: Read + Write> Write for PrependedReader<RW> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.inner.get_mut().1.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.inner.get_mut().1.flush()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::{Cursor, Read};
+
+    use super::PrependedReader;
+
+    #[test]
+    fn test_bufreader_replace() {
+        let mut rw = Cursor::new(vec![1, 2, 3]);
+        let mut br = PrependedReader::new(&[0, 1, 2], &mut rw);
+        let mut out = Vec::new();
+        br.read_to_end(&mut out).unwrap();
+
+        assert_eq!(&out, &[0, 1, 2, 1, 2, 3]);
+    }
+}

tests/main.rs

@@ -12,8 +12,11 @@ use mysql::SslOpts;
 use rustls::{Certificate, PrivateKey, ServerConfig};
 use std::error::Error;
 use std::io;
+use std::io::Read;
+use std::io::Write;
 use std::net;
 use std::thread;
+use std::time::Duration;
 
 use msql_srv::{
     Column, ErrorKind, InitWriter, MysqlIntermediary, MysqlShim, ParamParser, QueryResultWriter,
@@ -210,6 +213,108 @@ fn it_connects_tls_both() {
     .test(|_| {})
 }
 
+#[test]
+#[cfg(feature = "tls")]
+fn it_connects_tls_both_with_delayed_server_read() {
+    use std::{marker::PhantomData, sync::Arc};
+
+    struct MyShim<RW> {
+        ph: PhantomData<RW>,
+    }
+
+    impl<RW: Read + Write> MysqlShim<RW> for MyShim<RW> {
+        type Error = io::Error;
+
+        fn on_prepare(
+            &mut self,
+            _: &str,
+            _: StatementMetaWriter<'_, RW>,
+        ) -> Result<(), Self::Error> {
+            unreachable!()
+        }
+
+        fn on_execute(
+            &mut self,
+            _: u32,
+            _: ParamParser<'_>,
+            _: QueryResultWriter<'_, RW>,
+        ) -> Result<(), Self::Error> {
+            unreachable!()
+        }
+
+        fn on_close(&mut self, _: u32) {
+            unreachable!()
+        }
+
+        fn on_query(&mut self, _: &str, _: QueryResultWriter<'_, RW>) -> Result<(), Self::Error> {
+            unreachable!()
+        }
+
+        fn tls_config(&self) -> Option<Arc<ServerConfig>> {
+            let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]).unwrap();
+
+            Some(std::sync::Arc::new(
+                ServerConfig::builder()
+                    .with_safe_defaults()
+                    .with_no_client_auth()
+                    .with_single_cert(
+                        vec![Certificate(cert.serialize_der().unwrap())],
+                        PrivateKey(cert.get_key_pair().serialize_der()),
+                    )
+                    .unwrap(),
+            ))
+        }
+    }
+
+    let shim = MyShim {
+        ph: PhantomData::default(),
+    };
+
+    let listener = net::TcpListener::bind("127.0.0.1:0").unwrap();
+    let port = listener.local_addr().unwrap().port();
+    let jh = thread::spawn(move || {
+        let (s, _) = listener.accept().unwrap();
+        let s = DelayedReadRW {
+            s,
+            read_delay: Duration::from_millis(200),
+        };
+        MysqlIntermediary::run_on(shim, s)
+    });
+
+    let db = mysql::Conn::new(
+        OptsBuilder::default()
+            .ip_or_hostname(Some("localhost"))
+            .tcp_port(port)
+            .ssl_opts(Some(
+                SslOpts::default().with_danger_accept_invalid_certs(true),
+            )),
+    )
+    .unwrap();
+    drop(db);
+    jh.join().unwrap().unwrap();
+}
+
+struct DelayedReadRW<RW: Read + Write> {
+    s: RW,
+    read_delay: Duration,
+}
+
+impl<RW: Read + Write> Read for DelayedReadRW<RW> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        thread::sleep(self.read_delay);
+        self.s.read(buf)
+    }
+}
+
+impl<RW: Read + Write> Write for DelayedReadRW<RW> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.s.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.s.flush()
+    }
+}
 #[test]
 fn it_does_not_connect_tls_client_only() {
     // Client requesting tls fails as expected when server does not support it.