Merge branch 'master' into further_tls_functionality

Author: mjgarton

Cargo.toml

@@ -31,16 +31,15 @@ nom = "7"
 mysql_common = { version = "0.28.0", features = ["chrono"] }
 byteorder = "1"
 chrono = "0.4"
-time = "0.2.25"
 rustls = {version = "0.20.0", optional=true}
 
 [dev-dependencies]
 postgres = "0.19.1"
-mysql = "18"
-mysql_async = "0.20.0"
+mysql = "22"
+mysql_async = "0.29.0"
 slab = "0.4.2"
-tokio = "0.1.19"
-futures = "0.1.26"
+tokio = { version = "1.15.0", features = ["full"] }
+futures = "0.3.0"
 rcgen = "0.8.14"
 tempfile = "3.3.0"
 native-tls = "0.2.8"

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
     }
 
     #[cfg(feature = "tls")]
@@ -103,15 +105,7 @@ impl<R: Read + Write> PacketConn<R> {
 
         loop {
             if self.remaining != 0 {
-                let bytes = {
-                    // NOTE: this is all sorts of unfortunate. what we really want to do is to give
-                    // &self.bytes[self.start..] to `packet()`, and the lifetimes should all work
-                    // out. however, without NLL, borrowck doesn't realize that self.bytes is no
-                    // longer borrowed after the match, and so can be mutated.
-                    let bytes = &self.bytes[self.start..];
-                    unsafe { ::std::slice::from_raw_parts(bytes.as_ptr(), bytes.len()) }
-                };
-                match packet(bytes) {
+                match packet(&self.bytes[self.start..]) {
                     Ok((rest, p)) => {
                         self.remaining = rest.len();
                         return Ok(Some(p));

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>(pub(crate) 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]);
+    }
+}