Merge pull request #38 from mjgarton/fix_buffering_during_tls_neg

Fix buffering during switch to TLS

Author: jonhoo

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,112 @@ fn it_connects_tls_both() {
     .test(|_| {})
 }
 
+#[test]
+#[cfg(feature = "tls")]
+fn it_connects_tls_both_with_delayed_server_read() {
+    // This test is to ensure correctly handle the case when we read both the pre-TLS data as well
+    // as (at least part of) the TLS handshake into our the buffer.  When that happens, we need to
+    // ensure we correctly pass that TLS part of the data to rustls so that is can handle the TLS
+    // handshake properly.
+    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.