TLS support

Author: sfackler

postgres-tokio/Cargo.toml

@@ -3,6 +3,9 @@ name = "postgres-tokio"
 version = "0.1.0"
 authors = ["Steven Fackler <[email protected]>"]
 
+[features]
+with-openssl = ["tokio-openssl", "openssl"]
+
 [dependencies]
 fallible-iterator = "0.1.3"
 futures = "0.1.7"
@@ -12,3 +15,6 @@ postgres-protocol = "0.2"
 tokio-core = "0.1"
 tokio-dns-unofficial = "0.1"
 tokio-uds = "0.1"
+
+tokio-openssl = { version = "0.1", optional = true }
+openssl = { version = "0.9", optional = true }

postgres-tokio/src/lib.rs

@@ -7,6 +7,11 @@ extern crate tokio_core;
 extern crate tokio_dns;
 extern crate tokio_uds;
 
+#[cfg(feature = "tokio-openssl")]
+extern crate tokio_openssl;
+#[cfg(feature = "openssl")]
+extern crate openssl;
+
 use fallible_iterator::FallibleIterator;
 use futures::{Future, IntoFuture, BoxFuture, Stream, Sink, Poll, StartSend};
 use futures::future::Either;
@@ -31,13 +36,21 @@ use error::{ConnectError, Error, DbError};
 use params::{ConnectParams, IntoConnectParams};
 use stream::PostgresStream;
 use types::{Oid, Type, ToSql, SessionInfo, IsNull, FromSql, WrongType};
+use tls::Handshake;
 
 pub mod error;
 mod stream;
+pub mod tls;
 
 #[cfg(test)]
 mod test;
 
+pub enum TlsMode {
+    Require(Box<Handshake>),
+    Prefer(Box<Handshake>),
+    None,
+}
+
 #[derive(Debug, Copy, Clone)]
 pub struct CancelData {
     pub process_id: i32,
@@ -119,16 +132,18 @@ impl fmt::Debug for Connection {
 }
 
 impl Connection {
-    pub fn connect<T>(params: T, handle: &Handle) -> BoxFuture<Connection, ConnectError>
+    pub fn connect<T>(params: T,
+                      tls_mode: TlsMode,
+                      handle: &Handle)
+                      -> BoxFuture<Connection, ConnectError>
         where T: IntoConnectParams
     {
         let params = match params.into_connect_params() {
             Ok(params) => params,
             Err(e) => return futures::failed(ConnectError::ConnectParams(e)).boxed(),
         };
 
-        stream::connect(params.host(), params.port(), handle)
-            .map_err(ConnectError::Io)
+        stream::connect(params.host().clone(), params.port(), tls_mode, handle)
             .map(|s| {
                 let (sender, receiver) = mpsc::channel();
                 Connection(InnerConnection {

postgres-tokio/src/stream.rs

@@ -1,75 +1,126 @@
-use futures::{BoxFuture, Future, IntoFuture, Async};
+use futures::{BoxFuture, Future, IntoFuture, Async, Sink, Stream as FuturesStream};
+use futures::future::Either;
 use postgres_shared::params::Host;
 use postgres_protocol::message::backend::{self, ParseResult};
+use postgres_protocol::message::frontend;
 use std::io::{self, Read, Write};
 use tokio_core::io::{Io, Codec, EasyBuf, Framed};
 use tokio_core::net::TcpStream;
 use tokio_core::reactor::Handle;
 use tokio_dns;
 use tokio_uds::UnixStream;
 
-pub type PostgresStream = Framed<InnerStream, PostgresCodec>;
+use TlsMode;
+use error::ConnectError;
+use tls::TlsStream;
 
-pub fn connect(host: &Host,
-                port: u16,
-                handle: &Handle)
-                -> BoxFuture<PostgresStream, io::Error> {
-    match *host {
+pub type PostgresStream = Framed<Box<TlsStream>, PostgresCodec>;
+
+pub fn connect(host: Host,
+               port: u16,
+               tls_mode: TlsMode,
+               handle: &Handle)
+               -> BoxFuture<PostgresStream, ConnectError> {
+    let inner = match host {
         Host::Tcp(ref host) => {
-            tokio_dns::tcp_connect((&**host, port), handle.remote().clone())
-                .map(|s| InnerStream::Tcp(s).framed(PostgresCodec))
-                .boxed()
+            Either::A(tokio_dns::tcp_connect((&**host, port), handle.remote().clone())
+                .map(|s| Stream(InnerStream::Tcp(s))))
         }
         Host::Unix(ref host) => {
             let addr = host.join(format!(".s.PGSQL.{}", port));
-            UnixStream::connect(addr, handle)
-                .map(|s| InnerStream::Unix(s).framed(PostgresCodec))
-                .into_future()
-                .boxed()
+            Either::B(UnixStream::connect(addr, handle)
+                .map(|s| Stream(InnerStream::Unix(s)))
+                .into_future())
         }
-    }
+    };
+
+    let (required, mut handshaker) = match tls_mode {
+        TlsMode::Require(h) => (true, h),
+        TlsMode::Prefer(h) => (false, h),
+        TlsMode::None => {
+            return inner.map(|s| {
+                    let s: Box<TlsStream> = Box::new(s);
+                    s.framed(PostgresCodec)
+                })
+                .map_err(ConnectError::Io)
+                .boxed()
+        },
+    };
+
+    inner.map(|s| s.framed(SslCodec))
+        .and_then(|s| {
+            let mut buf = vec![];
+            frontend::ssl_request(&mut buf);
+            s.send(buf)
+        })
+        .and_then(|s| s.into_future().map_err(|e| e.0))
+        .map_err(ConnectError::Io)
+        .and_then(move |(m, s)| {
+            let s = s.into_inner();
+            match (m, required) {
+                (Some(b'N'), true) => {
+                    Either::A(Err(ConnectError::Tls("the server does not support TLS".into())).into_future())
+                }
+                (Some(b'N'), false) => {
+                    let s: Box<TlsStream> = Box::new(s);
+                    Either::A(Ok(s).into_future())
+                },
+                (None, _) => Either::A(Err(ConnectError::Io(io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected EOF"))).into_future()),
+                _ => {
+                    let host = match host {
+                        Host::Tcp(ref host) => host,
+                        Host::Unix(_) => unreachable!(),
+                    };
+                    Either::B(handshaker.handshake(host, s).map_err(ConnectError::Tls))
+                }
+            }
+        })
+        .map(|s| s.framed(PostgresCodec))
+        .boxed()
 }
 
-pub enum InnerStream {
+pub struct Stream(InnerStream);
+
+enum InnerStream {
     Tcp(TcpStream),
     Unix(UnixStream),
 }
 
-impl Read for InnerStream {
+impl Read for Stream {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        match *self {
+        match self.0 {
             InnerStream::Tcp(ref mut s) => s.read(buf),
             InnerStream::Unix(ref mut s) => s.read(buf),
         }
     }
 }
 
-impl Write for InnerStream {
+impl Write for Stream {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        match *self {
+        match self.0 {
             InnerStream::Tcp(ref mut s) => s.write(buf),
             InnerStream::Unix(ref mut s) => s.write(buf),
         }
     }
 
     fn flush(&mut self) -> io::Result<()> {
-        match *self {
+        match self.0 {
             InnerStream::Tcp(ref mut s) => s.flush(),
             InnerStream::Unix(ref mut s) => s.flush(),
         }
     }
 }
 
-impl Io for InnerStream {
+impl Io for Stream {
     fn poll_read(&mut self) -> Async<()> {
-        match *self {
+        match self.0 {
             InnerStream::Tcp(ref mut s) => s.poll_read(),
             InnerStream::Unix(ref mut s) => s.poll_read(),
         }
     }
 
     fn poll_write(&mut self) -> Async<()> {
-        match *self {
+        match self.0 {
             InnerStream::Tcp(ref mut s) => s.poll_write(),
             InnerStream::Unix(ref mut s) => s.poll_write(),
         }
@@ -98,3 +149,25 @@ impl Codec for PostgresCodec {
         Ok(())
     }
 }
+
+struct SslCodec;
+
+impl Codec for SslCodec {
+    type In = u8;
+    type Out = Vec<u8>;
+
+    fn decode(&mut self, buf: &mut EasyBuf) -> io::Result<Option<u8>> {
+        if buf.as_slice().is_empty() {
+            Ok(None)
+        } else {
+            let byte = buf.as_slice()[0];
+            buf.drain_to(1);
+            Ok(Some(byte))
+        }
+    }
+
+    fn encode(&mut self, msg: Vec<u8>, buf: &mut Vec<u8>) -> io::Result<()> {
+        buf.extend_from_slice(&msg);
+        Ok(())
+    }
+}