dnim/rust-postgres
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Overhaul connection APIs

Browse Source 
* `Connection` is now parameterized over the stream type, which can be
    any `AsyncRead + AsyncWrite`.
* The `TlsMode` enum is now a trait, and `NoTls`, `PreferTls`, and
    `RequireTls` are types implementing that trait.
* The `TlsConnect` trait no longer involves trait objects, and returns
    channel binding info alongside the stream type rather than requiring
    the stream to implement an additional trait.
* The `connect` free function and `ConnectParams` type is gone in favor
    of a `Builder` type. It takes a pre-connected stream rather than
    automatically opening a TCP or Unix socket connection.

Notably, we no longer have any dependency on the Tokio runtime. We do
use the `tokio-codec` and `tokio-io` crates, but those don't actually
depend on mio/tokio-reactor/etc. This means we can work with other
futures-based networking stacks.

We will almost certainly add back a convenience API that offers
something akin to the old logic to open a TCP/Unix connection
automatically but that will be worked out in a follow up PR.
This commit is contained in:
Steven Fackler
2018-11-26 22:45:14 -08:00
parent 0e60d80d4b
commit 08b4020534
20 changed files with 976 additions and 1083 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.circleci/config.yml
View File

@@ -22,7 +22,7 @@ version: 2
jobs:
build:
docker:
- image: rust:1.26.2
- image: rust:1.30.1
environment:
RUSTFLAGS: -D warnings
- image: sfackler/rust-postgres-test:4

2
docker/Dockerfile
View File

@@ -1,3 +1,3 @@
FROM postgres:11-beta1
FROM postgres:11
COPY sql_setup.sh /docker-entrypoint-initdb.d/

1
tokio-postgres-native-tls/Cargo.toml
View File

@@ -4,7 +4,6 @@ version = "0.1.0"
authors = ["Steven Fackler <sfackler@gmail.com>"]
[dependencies]
bytes = "0.4"
futures = "0.1"
native-tls = "0.2"
tokio-io = "0.1"

111
tokio-postgres-native-tls/src/lib.rs
View File

@@ -1,106 +1,71 @@
extern crate bytes;
extern crate futures;
extern crate native_tls;
extern crate tokio_io;
extern crate tokio_postgres;
extern crate tokio_tls;
#[macro_use]
extern crate futures;
#[cfg(test)]
extern crate tokio;
use bytes::{Buf, BufMut};
use futures::{Future, Poll};
use std::error::Error;
use std::io::{self, Read, Write};
use futures::{Async, Future, Poll};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_postgres::tls::{Socket, TlsConnect, TlsStream};
use tokio_postgres::{ChannelBinding, TlsConnect};
use tokio_tls::{Connect, TlsStream};
#[cfg(test)]
mod test;
pub struct TlsConnector {
connector: tokio_tls::TlsConnector,
domain: String,
}
impl TlsConnector {
pub fn new() -> Result<TlsConnector, native_tls::Error> {
pub fn new(domain: &str) -> Result<TlsConnector, native_tls::Error> {
let connector = native_tls::TlsConnector::new()?;
Ok(TlsConnector::with_connector(connector))
Ok(TlsConnector::with_connector(connector, domain))
}
pub fn with_connector(connector: native_tls::TlsConnector) -> TlsConnector {
pub fn with_connector(connector: native_tls::TlsConnector, domain: &str) -> TlsConnector {
TlsConnector {
connector: tokio_tls::TlsConnector::from(connector),
domain: domain.to_string(),
}
}
}
impl TlsConnect for TlsConnector {
fn connect(
&self,
domain: &str,
socket: Socket,
) -> Box<Future<Item = Box<TlsStream>, Error = Box<Error + Sync + Send>> + Sync + Send> {
let f = self
.connector
.connect(domain, socket)
.map(|s| {
let s: Box<TlsStream> = Box::new(SslStream(s));
s
}).map_err(|e| {
let e: Box<Error + Sync + Send> = Box::new(e);
e
});
Box::new(f)
impl<S> TlsConnect<S> for TlsConnector
where
S: AsyncRead + AsyncWrite,
{
type Stream = TlsStream<S>;
type Error = native_tls::Error;
type Future = TlsConnectFuture<S>;
fn connect(self, stream: S) -> TlsConnectFuture<S> {
TlsConnectFuture(self.connector.connect(&self.domain, stream))
}
}
struct SslStream(tokio_tls::TlsStream<Socket>);
pub struct TlsConnectFuture<S>(Connect<S>);
impl Read for SslStream {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.read(buf)
}
}
impl AsyncRead for SslStream {
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
self.0.prepare_uninitialized_buffer(buf)
}
fn read_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: BufMut,
{
self.0.read_buf(buf)
}
}
impl Write for SslStream {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.0.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.0.flush()
}
}
impl AsyncWrite for SslStream {
fn shutdown(&mut self) -> Poll<(), io::Error> {
self.0.shutdown()
}
fn write_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: Buf,
{
self.0.write_buf(buf)
}
}
impl TlsStream for SslStream {
fn tls_server_end_point(&self) -> Option<Vec<u8>> {
self.0.get_ref().tls_server_end_point().unwrap_or(None)
impl<S> Future for TlsConnectFuture<S>
where
S: AsyncRead + AsyncWrite,
{
type Item = (TlsStream<S>, ChannelBinding);
type Error = native_tls::Error;
fn poll(&mut self) -> Poll<(TlsStream<S>, ChannelBinding), native_tls::Error> {
let stream = try_ready!(self.0.poll());
let mut channel_binding = ChannelBinding::new();
if let Some(buf) = stream.get_ref().tls_server_end_point().unwrap_or(None) {
channel_binding = channel_binding.tls_server_end_point(buf);
}
Ok(Async::Ready((stream, channel_binding)))
}
}

37
tokio-postgres-native-tls/src/test.rs
View File

@@ -1,17 +1,24 @@
use futures::{Future, Stream};
use native_tls::{self, Certificate};
use tokio::net::TcpStream;
use tokio::runtime::current_thread::Runtime;
use tokio_postgres::{self, TlsMode};
use tokio_postgres::{self, PreferTls, RequireTls, TlsMode};
use TlsConnector;
fn smoke_test(url: &str, tls: TlsMode) {
fn smoke_test<T>(builder: &tokio_postgres::Builder, tls: T)
where
T: TlsMode<TcpStream>,
T::Stream: 'static,
{
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(url.parse().unwrap(), tls);
let handshake = TcpStream::connect(&"127.0.0.1:5433".parse().unwrap())
.map_err(|e| panic!("{}", e))
.and_then(|s| builder.connect(s, tls));
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
runtime.spawn(connection);
let prepare = client.prepare("SELECT 1::INT4");
let statement = runtime.block_on(prepare).unwrap();
@@ -33,10 +40,11 @@ fn require() {
Certificate::from_pem(include_bytes!("../../test/server.crt")).unwrap(),
).build()
.unwrap();
let connector = TlsConnector::with_connector(connector);
smoke_test(
"postgres://ssl_user@localhost:5433/postgres",
TlsMode::Require(Box::new(connector)),
tokio_postgres::Builder::new()
.user("ssl_user")
.database("postgres"),
RequireTls(TlsConnector::with_connector(connector, "localhost")),
);
}
@@ -47,10 +55,11 @@ fn prefer() {
Certificate::from_pem(include_bytes!("../../test/server.crt")).unwrap(),
).build()
.unwrap();
let connector = TlsConnector::with_connector(connector);
smoke_test(
"postgres://ssl_user@localhost:5433/postgres",
TlsMode::Prefer(Box::new(connector)),
tokio_postgres::Builder::new()
.user("ssl_user")
.database("postgres"),
PreferTls(TlsConnector::with_connector(connector, "localhost")),
);
}
@@ -61,9 +70,11 @@ fn scram_user() {
Certificate::from_pem(include_bytes!("../../test/server.crt")).unwrap(),
).build()
.unwrap();
let connector = TlsConnector::with_connector(connector);
smoke_test(
"postgres://scram_user:password@localhost:5433/postgres",
TlsMode::Require(Box::new(connector)),
tokio_postgres::Builder::new()
.user("scram_user")
.password("password")
.database("postgres"),
RequireTls(TlsConnector::with_connector(connector, "localhost")),
);
}

3
tokio-postgres-openssl/Cargo.toml
View File

@@ -4,11 +4,10 @@ version = "0.1.0"
authors = ["Steven Fackler <sfackler@gmail.com>"]
[dependencies]
bytes = "0.4"
futures = "0.1"
openssl = "0.10"
tokio-io = "0.1"
tokio-openssl = "0.2"
tokio-openssl = "0.3"
tokio-postgres = { version = "0.3", path = "../tokio-postgres" }
[dev-dependencies]

147
tokio-postgres-openssl/src/lib.rs
View File

@@ -1,141 +1,76 @@
extern crate bytes;
extern crate futures;
extern crate openssl;
extern crate tokio_io;
extern crate tokio_openssl;
extern crate tokio_postgres;
#[macro_use]
extern crate futures;
#[cfg(test)]
extern crate tokio;
use bytes::{Buf, BufMut};
use futures::{Future, IntoFuture, Poll};
use openssl::error::ErrorStack;
use openssl::ssl::{ConnectConfiguration, SslConnector, SslMethod, SslRef};
use std::error::Error;
use std::io::{self, Read, Write};
use futures::{Async, Future, Poll};
use openssl::ssl::{ConnectConfiguration, HandshakeError, SslRef};
use std::fmt::Debug;
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_openssl::ConnectConfigurationExt;
use tokio_postgres::tls::{Socket, TlsConnect, TlsStream};
use tokio_openssl::{ConnectAsync, ConnectConfigurationExt, SslStream};
use tokio_postgres::{ChannelBinding, TlsConnect};
#[cfg(test)]
mod test;
pub struct TlsConnector {
connector: SslConnector,
callback: Box<Fn(&mut ConnectConfiguration) -> Result<(), ErrorStack> + Sync + Send>,
ssl: ConnectConfiguration,
domain: String,
}
impl TlsConnector {
pub fn new() -> Result<TlsConnector, ErrorStack> {
let connector = SslConnector::builder(SslMethod::tls())?.build();
Ok(TlsConnector::with_connector(connector))
}
pub fn with_connector(connector: SslConnector) -> TlsConnector {
pub fn new(ssl: ConnectConfiguration, domain: &str) -> TlsConnector {
TlsConnector {
connector,
callback: Box::new(|_| Ok(())),
ssl,
domain: domain.to_string(),
}
}
}
pub fn set_callback<F>(&mut self, f: F)
where
F: Fn(&mut ConnectConfiguration) -> Result<(), ErrorStack> + 'static + Sync + Send,
{
self.callback = Box::new(f);
impl<S> TlsConnect<S> for TlsConnector
where
S: AsyncRead + AsyncWrite + Debug + 'static + Sync + Send,
{
type Stream = SslStream<S>;
type Error = HandshakeError<S>;
type Future = TlsConnectFuture<S>;
fn connect(self, stream: S) -> TlsConnectFuture<S> {
TlsConnectFuture(self.ssl.connect_async(&self.domain, stream))
}
}
impl TlsConnect for TlsConnector {
fn connect(
&self,
domain: &str,
socket: Socket,
) -> Box<Future<Item = Box<TlsStream>, Error = Box<Error + Sync + Send>> + Sync + Send> {
let f = self
.connector
.configure()
.and_then(|mut ssl| (self.callback)(&mut ssl).map(|_| ssl))
.map_err(|e| {
let e: Box<Error + Sync + Send> = Box::new(e);
e
})
.into_future()
.and_then({
let domain = domain.to_string();
move |ssl| {
ssl.connect_async(&domain, socket)
.map(|s| {
let s: Box<TlsStream> = Box::new(SslStream(s));
s
})
.map_err(|e| {
let e: Box<Error + Sync + Send> = Box::new(e);
e
})
}
});
Box::new(f)
}
}
pub struct TlsConnectFuture<S>(ConnectAsync<S>);
struct SslStream(tokio_openssl::SslStream<Socket>);
impl<S> Future for TlsConnectFuture<S>
where
S: AsyncRead + AsyncWrite + Debug + 'static + Sync + Send,
{
type Item = (SslStream<S>, ChannelBinding);
type Error = HandshakeError<S>;
impl Read for SslStream {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.read(buf)
}
}
fn poll(&mut self) -> Poll<(SslStream<S>, ChannelBinding), HandshakeError<S>> {
let stream = try_ready!(self.0.poll());
impl AsyncRead for SslStream {
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
self.0.prepare_uninitialized_buffer(buf)
}
fn read_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: BufMut,
{
self.0.read_buf(buf)
}
}
impl Write for SslStream {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.0.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.0.flush()
}
}
impl AsyncWrite for SslStream {
fn shutdown(&mut self) -> Poll<(), io::Error> {
self.0.shutdown()
}
fn write_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: Buf,
{
self.0.write_buf(buf)
}
}
impl TlsStream for SslStream {
fn tls_unique(&self) -> Option<Vec<u8>> {
let f = if self.0.get_ref().ssl().session_reused() {
let f = if stream.get_ref().ssl().session_reused() {
SslRef::peer_finished
} else {
SslRef::finished
};
let len = f(self.0.get_ref().ssl(), &mut []);
let mut buf = vec![0; len];
f(self.0.get_ref().ssl(), &mut buf);
let len = f(stream.get_ref().ssl(), &mut []);
let mut tls_unique = vec![0; len];
f(stream.get_ref().ssl(), &mut tls_unique);
Some(buf)
Ok(Async::Ready((
stream,
ChannelBinding::new().tls_unique(tls_unique),
)))
}
}

40
tokio-postgres-openssl/src/test.rs
View File

@@ -1,17 +1,24 @@
use futures::{Future, Stream};
use openssl::ssl::{SslConnector, SslMethod};
use tokio::net::TcpStream;
use tokio::runtime::current_thread::Runtime;
use tokio_postgres::{self, TlsMode};
use tokio_postgres::{self, PreferTls, RequireTls, TlsMode};
use TlsConnector;
fn smoke_test(url: &str, tls: TlsMode) {
fn smoke_test<T>(builder: &tokio_postgres::Builder, tls: T)
where
T: TlsMode<TcpStream>,
T::Stream: 'static,
{
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(url.parse().unwrap(), tls);
let handshake = TcpStream::connect(&"127.0.0.1:5433".parse().unwrap())
.map_err(|e| panic!("{}", e))
.and_then(|s| builder.connect(s, tls));
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
runtime.spawn(connection);
let prepare = client.prepare("SELECT 1::INT4");
let statement = runtime.block_on(prepare).unwrap();
@@ -30,10 +37,12 @@ fn smoke_test(url: &str, tls: TlsMode) {
fn require() {
let mut builder = SslConnector::builder(SslMethod::tls()).unwrap();
builder.set_ca_file("../test/server.crt").unwrap();
let connector = TlsConnector::with_connector(builder.build());
let ctx = builder.build();
smoke_test(
"postgres://ssl_user@localhost:5433/postgres",
TlsMode::Require(Box::new(connector)),
tokio_postgres::Builder::new()
.user("ssl_user")
.database("postgres"),
RequireTls(TlsConnector::new(ctx.configure().unwrap(), "localhost")),
);
}
@@ -41,10 +50,12 @@ fn require() {
fn prefer() {
let mut builder = SslConnector::builder(SslMethod::tls()).unwrap();
builder.set_ca_file("../test/server.crt").unwrap();
let connector = TlsConnector::with_connector(builder.build());
let ctx = builder.build();
smoke_test(
"postgres://ssl_user@localhost:5433/postgres",
TlsMode::Prefer(Box::new(connector)),
tokio_postgres::Builder::new()
.user("ssl_user")
.database("postgres"),
PreferTls(TlsConnector::new(ctx.configure().unwrap(), "localhost")),
);
}
@@ -52,9 +63,12 @@ fn prefer() {
fn scram_user() {
let mut builder = SslConnector::builder(SslMethod::tls()).unwrap();
builder.set_ca_file("../test/server.crt").unwrap();
let connector = TlsConnector::with_connector(builder.build());
let ctx = builder.build();
smoke_test(
"postgres://scram_user:password@localhost:5433/postgres",
TlsMode::Require(Box::new(connector)),
tokio_postgres::Builder::new()
.user("scram_user")
.password("password")
.database("postgres"),
RequireTls(TlsConnector::new(ctx.configure().unwrap(), "localhost")),
);
}

7
tokio-postgres/Cargo.toml
View File

@@ -37,7 +37,6 @@ bytes = "0.4"
fallible-iterator = "0.1.3"
futures = "0.1.7"
futures-cpupool = "0.1"
lazy_static = "1.0"
log = "0.4"
phf = "=0.7.22"
postgres-protocol = { version = "0.3.0", path = "../postgres-protocol" }
@@ -45,11 +44,7 @@ postgres-shared = { version = "0.4.0", path = "../postgres-shared" }
state_machine_future = "0.1.7"
tokio-codec = "0.1"
tokio-io = "0.1"
tokio-tcp = "0.1"
tokio-timer = "0.2"
[target.'cfg(unix)'.dependencies]
tokio-uds = "0.2.1"
void = "1.0"
[dev-dependencies]
tokio = "0.1.7"

55
tokio-postgres/src/builder.rs Normal file
View File

@@ -0,0 +1,55 @@
use std::collections::HashMap;
use tokio_io::{AsyncRead, AsyncWrite};
use proto::ConnectFuture;
use {Connect, TlsMode};
#[derive(Clone)]
pub struct Builder {
params: HashMap<String, String>,
password: Option<String>,
}
impl Builder {
pub fn new() -> Builder {
let mut params = HashMap::new();
params.insert("client_encoding".to_string(), "UTF8".to_string());
params.insert("timezone".to_string(), "GMT".to_string());
Builder {
params,
password: None,
}
}
pub fn user(&mut self, user: &str) -> &mut Builder {
self.param("user", user)
}
pub fn database(&mut self, database: &str) -> &mut Builder {
self.param("database", database)
}
pub fn param(&mut self, key: &str, value: &str) -> &mut Builder {
self.params.insert(key.to_string(), value.to_string());
self
}
pub fn password(&mut self, password: &str) -> &mut Builder {
self.password = Some(password.to_string());
self
}
pub fn connect<S, T>(&self, stream: S, tls_mode: T) -> Connect<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
Connect(ConnectFuture::new(
stream,
tls_mode,
self.password.clone(),
self.params.clone(),
))
}
}

13
tokio-postgres/src/error/mod.rs
View File

@@ -5,7 +5,6 @@ use postgres_protocol::message::backend::{ErrorFields, ErrorResponseBody};
use std::error;
use std::fmt;
use std::io;
use tokio_timer;
pub use self::sqlstate::*;
@@ -348,8 +347,6 @@ enum Kind {
MissingUser,
MissingPassword,
UnsupportedAuthentication,
Connect,
Timer,
Authentication,
}
@@ -396,8 +393,6 @@ impl error::Error for Error {
Kind::MissingUser => "username not provided",
Kind::MissingPassword => "password not provided",
Kind::UnsupportedAuthentication => "unsupported authentication method requested",
Kind::Connect => "error connecting to server",
Kind::Timer => "timer error",
Kind::Authentication => "authentication error",
}
}
@@ -489,14 +484,6 @@ impl Error {
Error::new(Kind::Tls, Some(e))
}
pub(crate) fn connect(e: io::Error) -> Error {
Error::new(Kind::Connect, Some(Box::new(e)))
}
pub(crate) fn timer(e: tokio_timer::Error) -> Error {
Error::new(Kind::Timer, Some(Box::new(e)))
}
pub(crate) fn io(e: io::Error) -> Error {
Error::new(Kind::Io, Some(Box::new(e)))
}

62
tokio-postgres/src/lib.rs
View File

@@ -7,21 +7,15 @@ extern crate postgres_protocol;
extern crate postgres_shared;
extern crate tokio_codec;
extern crate tokio_io;
extern crate tokio_tcp;
extern crate tokio_timer;
extern crate void;
#[macro_use]
extern crate futures;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
#[macro_use]
extern crate state_machine_future;
#[cfg(unix)]
extern crate tokio_uds;
use bytes::Bytes;
use futures::{Async, Future, Poll, Stream};
use postgres_shared::rows::RowIndex;
@@ -37,14 +31,16 @@ pub use postgres_shared::{params, types};
#[doc(inline)]
pub use postgres_shared::{CancelData, Notification};
use error::{DbError, Error};
use params::ConnectParams;
use tls::{TlsConnect, TlsStream};
pub use builder::*;
pub use error::*;
use proto::CancelFuture;
pub use tls::*;
use types::{FromSql, ToSql, Type};
mod builder;
pub mod error;
mod proto;
pub mod tls;
mod tls;
fn next_statement() -> String {
static ID: AtomicUsize = AtomicUsize::new(0);
@@ -56,18 +52,12 @@ fn next_portal() -> String {
format!("p{}", ID.fetch_add(1, Ordering::SeqCst))
}
pub enum TlsMode {
None,
Prefer(Box<TlsConnect>),
Require(Box<TlsConnect>),
}
pub fn cancel_query(params: ConnectParams, tls: TlsMode, cancel_data: CancelData) -> CancelQuery {
CancelQuery(proto::CancelFuture::new(params, tls, cancel_data))
}
pub fn connect(params: ConnectParams, tls: TlsMode) -> Handshake {
Handshake(proto::HandshakeFuture::new(params, tls))
pub fn cancel_query<S, T>(stream: S, tls_mode: T, cancel_data: CancelData) -> CancelQuery<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
CancelQuery(CancelFuture::new(stream, tls_mode, cancel_data))
}
pub struct Client(proto::Client);
@@ -165,9 +155,16 @@ pub enum AsyncMessage {
}
#[must_use = "futures do nothing unless polled"]
pub struct CancelQuery(proto::CancelFuture);
pub struct CancelQuery<S, T>(proto::CancelFuture<S, T>)
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>;
impl Future for CancelQuery {
impl<S, T> Future for CancelQuery<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
type Item = ();
type Error = Error;
@@ -177,13 +174,20 @@ impl Future for CancelQuery {
}
#[must_use = "futures do nothing unless polled"]
pub struct Handshake(proto::HandshakeFuture);
pub struct Connect<S, T>(proto::ConnectFuture<S, T>)
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>;
impl Future for Handshake {
type Item = (Client, Connection<Box<TlsStream>>);
impl<S, T> Future for Connect<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
type Item = (Client, Connection<T::Stream>);
type Error = Error;
fn poll(&mut self) -> Poll<(Client, Connection<Box<TlsStream>>), Error> {
fn poll(&mut self) -> Poll<(Client, Connection<T::Stream>), Error> {
let (client, connection) = try_ready!(self.0.poll());
Ok(Async::Ready((Client(client), Connection(connection))))

43
tokio-postgres/src/proto/cancel.rs
View File

@@ -2,35 +2,42 @@ use futures::{Future, Poll};
use postgres_protocol::message::frontend;
use state_machine_future::RentToOwn;
use tokio_io::io::{self, Flush, WriteAll};
use tokio_io::{AsyncRead, AsyncWrite};
use error::Error;
use params::ConnectParams;
use proto::connect::ConnectFuture;
use tls::TlsStream;
use proto::TlsFuture;
use {CancelData, TlsMode};
#[derive(StateMachineFuture)]
pub enum Cancel {
pub enum Cancel<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
#[state_machine_future(start, transitions(SendingCancel))]
Start {
future: ConnectFuture,
future: TlsFuture<S, T>,
cancel_data: CancelData,
},
#[state_machine_future(transitions(FlushingCancel))]
SendingCancel {
future: WriteAll<Box<TlsStream>, Vec<u8>>,
future: WriteAll<T::Stream, Vec<u8>>,
},
#[state_machine_future(transitions(Finished))]
FlushingCancel { future: Flush<Box<TlsStream>> },
FlushingCancel { future: Flush<T::Stream> },
#[state_machine_future(ready)]
Finished(()),
#[state_machine_future(error)]
Failed(Error),
}
impl PollCancel for Cancel {
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
let stream = try_ready!(state.future.poll());
impl<S, T> PollCancel<S, T> for Cancel<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
let (stream, _) = try_ready!(state.future.poll());
let mut buf = vec![];
frontend::cancel_request(
@@ -45,8 +52,8 @@ impl PollCancel for Cancel {
}
fn poll_sending_cancel<'a>(
state: &'a mut RentToOwn<'a, SendingCancel>,
) -> Poll<AfterSendingCancel, Error> {
state: &'a mut RentToOwn<'a, SendingCancel<S, T>>,
) -> Poll<AfterSendingCancel<S, T>, Error> {
let (stream, _) = try_ready_closed!(state.future.poll());
transition!(FlushingCancel {
@@ -55,15 +62,19 @@ impl PollCancel for Cancel {
}
fn poll_flushing_cancel<'a>(
state: &'a mut RentToOwn<'a, FlushingCancel>,
state: &'a mut RentToOwn<'a, FlushingCancel<S, T>>,
) -> Poll<AfterFlushingCancel, Error> {
try_ready_closed!(state.future.poll());
transition!(Finished(()))
}
}
impl CancelFuture {
pub fn new(params: ConnectParams, mode: TlsMode, cancel_data: CancelData) -> CancelFuture {
Cancel::start(ConnectFuture::new(params, mode), cancel_data)
impl<S, T> CancelFuture<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
pub fn new(stream: S, tls_mode: T, cancel_data: CancelData) -> CancelFuture<S, T> {
Cancel::start(TlsFuture::new(stream, tls_mode), cancel_data)
}
}

540
tokio-postgres/src/proto/connect.rs
View File

@@ -1,288 +1,330 @@
use futures::{Async, Future, Poll};
use futures_cpupool::{CpuFuture, CpuPool};
use fallible_iterator::FallibleIterator;
use futures::sink;
use futures::sync::mpsc;
use futures::{Future, Poll, Sink, Stream};
use postgres_protocol::authentication;
use postgres_protocol::authentication::sasl::{self, ScramSha256};
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use state_machine_future::RentToOwn;
use std::error::Error as StdError;
use std::collections::HashMap;
use std::io;
use std::net::{SocketAddr, ToSocketAddrs};
use std::time::{Duration, Instant};
use std::vec;
use tokio_io::io::{read_exact, write_all, ReadExact, WriteAll};
use tokio_tcp::{self, TcpStream};
use tokio_timer::Delay;
use tokio_codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};
#[cfg(unix)]
use tokio_uds::{self, UnixStream};
use params::{ConnectParams, Host};
use proto::socket::Socket;
use tls::{self, TlsConnect, TlsStream};
use {Error, TlsMode};
lazy_static! {
static ref DNS_POOL: CpuPool = CpuPool::new(2);
}
use proto::{Client, Connection, PostgresCodec, TlsFuture};
use {CancelData, ChannelBinding, Error, TlsMode};
#[derive(StateMachineFuture)]
pub enum Connect {
#[state_machine_future(start)]
#[cfg_attr(
unix,
state_machine_future(transitions(ResolvingDns, ConnectingUnix))
)]
#[cfg_attr(not(unix), state_machine_future(transitions(ResolvingDns)))]
Start { params: ConnectParams, tls: TlsMode },
#[state_machine_future(transitions(ConnectingTcp))]
ResolvingDns {
future: CpuFuture<vec::IntoIter<SocketAddr>, io::Error>,
timeout: Option<Duration>,
params: ConnectParams,
tls: TlsMode,
pub enum Connect<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
#[state_machine_future(start, transitions(SendingStartup))]
Start {
future: TlsFuture<S, T>,
password: Option<String>,
params: HashMap<String, String>,
},
#[state_machine_future(transitions(PreparingSsl))]
ConnectingTcp {
addrs: vec::IntoIter<SocketAddr>,
future: tokio_tcp::ConnectFuture,
timeout: Option<(Duration, Delay)>,
params: ConnectParams,
tls: TlsMode,
#[state_machine_future(transitions(ReadingAuth))]
SendingStartup {
future: sink::Send<Framed<T::Stream, PostgresCodec>>,
user: String,
password: Option<String>,
channel_binding: ChannelBinding,
},
#[cfg(unix)]
#[state_machine_future(transitions(PreparingSsl))]
ConnectingUnix {
future: tokio_uds::ConnectFuture,
timeout: Option<Delay>,
params: ConnectParams,
tls: TlsMode,
#[state_machine_future(transitions(ReadingInfo, SendingPassword, SendingSasl))]
ReadingAuth {
stream: Framed<T::Stream, PostgresCodec>,
user: String,
password: Option<String>,
channel_binding: ChannelBinding,
},
#[state_machine_future(transitions(Ready, SendingSsl))]
PreparingSsl {
socket: Socket,
params: ConnectParams,
tls: TlsMode,
#[state_machine_future(transitions(ReadingAuthCompletion))]
SendingPassword {
future: sink::Send<Framed<T::Stream, PostgresCodec>>,
},
#[state_machine_future(transitions(ReadingSsl))]
SendingSsl {
future: WriteAll<Socket, Vec<u8>>,
params: ConnectParams,
connector: Box<TlsConnect>,
required: bool,
#[state_machine_future(transitions(ReadingSasl))]
SendingSasl {
future: sink::Send<Framed<T::Stream, PostgresCodec>>,
scram: ScramSha256,
},
#[state_machine_future(transitions(ConnectingTls, Ready))]
ReadingSsl {
future: ReadExact<Socket, [u8; 1]>,
params: ConnectParams,
connector: Box<TlsConnect>,
required: bool,
#[state_machine_future(transitions(SendingSasl, ReadingAuthCompletion))]
ReadingSasl {
stream: Framed<T::Stream, PostgresCodec>,
scram: ScramSha256,
},
#[state_machine_future(transitions(Ready))]
ConnectingTls {
future:
Box<Future<Item = Box<TlsStream>, Error = Box<StdError + Sync + Send>> + Sync + Send>,
params: ConnectParams,
#[state_machine_future(transitions(ReadingInfo))]
ReadingAuthCompletion {
stream: Framed<T::Stream, PostgresCodec>,
},
#[state_machine_future(transitions(Finished))]
ReadingInfo {
stream: Framed<T::Stream, PostgresCodec>,
cancel_data: Option<CancelData>,
parameters: HashMap<String, String>,
},
#[state_machine_future(ready)]
Ready(Box<TlsStream>),
Finished((Client, Connection<T::Stream>)),
#[state_machine_future(error)]
Failed(Error),
}
impl PollConnect for Connect {
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
let state = state.take();
let timeout = state.params.connect_timeout();
let port = state.params.port();
match state.params.host().clone() {
Host::Tcp(host) => transition!(ResolvingDns {
future: DNS_POOL.spawn_fn(move || (&*host, port).to_socket_addrs()),
params: state.params,
tls: state.tls,
timeout,
}),
#[cfg(unix)]
Host::Unix(mut path) => {
path.push(format!(".s.PGSQL.{}", port));
transition!(ConnectingUnix {
future: UnixStream::connect(path),
timeout: timeout.map(|t| Delay::new(Instant::now() + t)),
params: state.params,
tls: state.tls,
})
},
#[cfg(not(unix))]
Host::Unix(_) => {
Err(Error::connect(io::Error::new(
io::ErrorKind::Other,
"unix sockets are not supported on this platform",
)))
},
}
}
fn poll_resolving_dns<'a>(
state: &'a mut RentToOwn<'a, ResolvingDns>,
) -> Poll<AfterResolvingDns, Error> {
let mut addrs = try_ready!(state.future.poll().map_err(Error::connect));
let state = state.take();
let addr = match addrs.next() {
Some(addr) => addr,
None => {
return Err(Error::connect(io::Error::new(
io::ErrorKind::Other,
"resolved to 0 addresses",
)))
}
};
transition!(ConnectingTcp {
addrs,
future: TcpStream::connect(&addr),
timeout: state.timeout.map(|t| (t, Delay::new(Instant::now() + t))),
params: state.params,
tls: state.tls,
})
}
fn poll_connecting_tcp<'a>(
state: &'a mut RentToOwn<'a, ConnectingTcp>,
) -> Poll<AfterConnectingTcp, Error> {
let socket = loop {
let error = match state.future.poll() {
Ok(Async::Ready(socket)) => break socket,
Ok(Async::NotReady) => match state.timeout {
Some((_, ref mut delay)) => {
try_ready!(delay.poll().map_err(Error::timer));
io::Error::new(io::ErrorKind::TimedOut, "connection timed out")
}
None => return Ok(Async::NotReady),
},
Err(e) => e,
};
let addr = match state.addrs.next() {
Some(addr) => addr,
None => return Err(Error::connect(error)),
};
state.future = TcpStream::connect(&addr);
if let Some((timeout, ref mut delay)) = state.timeout {
delay.reset(Instant::now() + timeout);
}
};
// Our read/write patterns may trigger Nagle's algorithm since we're pipelining which
// we don't want. Each individual write should be a full command we want the backend to
// see immediately.
socket.set_nodelay(true).map_err(Error::connect)?;
let state = state.take();
transition!(PreparingSsl {
socket: Socket::Tcp(socket),
params: state.params,
tls: state.tls,
})
}
#[cfg(unix)]
fn poll_connecting_unix<'a>(
state: &'a mut RentToOwn<'a, ConnectingUnix>,
) -> Poll<AfterConnectingUnix, Error> {
match state.future.poll().map_err(Error::connect)? {
Async::Ready(socket) => {
let state = state.take();
transition!(PreparingSsl {
socket: Socket::Unix(socket),
params: state.params,
tls: state.tls,
})
}
Async::NotReady => match state.timeout {
Some(ref mut delay) => {
try_ready!(delay.poll().map_err(Error::timer));
Err(Error::connect(io::Error::new(
io::ErrorKind::TimedOut,
"connection timed out",
)))
}
None => Ok(Async::NotReady),
},
}
}
fn poll_preparing_ssl<'a>(
state: &'a mut RentToOwn<'a, PreparingSsl>,
) -> Poll<AfterPreparingSsl, Error> {
let state = state.take();
let (connector, required) = match state.tls {
TlsMode::None => {
transition!(Ready(Box::new(state.socket)));
}
TlsMode::Prefer(connector) => (connector, false),
TlsMode::Require(connector) => (connector, true),
};
impl<S, T> PollConnect<S, T> for Connect<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
let (stream, channel_binding) = try_ready!(state.future.poll());
let mut state = state.take();
let mut buf = vec![];
frontend::ssl_request(&mut buf);
transition!(SendingSsl {
future: write_all(state.socket, buf),
params: state.params,
connector,
required,
frontend::startup_message(state.params.iter().map(|(k, v)| (&**k, &**v)), &mut buf)
.map_err(Error::encode)?;
let stream = Framed::new(stream, PostgresCodec);
let user = state
.params
.remove("user")
.ok_or_else(Error::missing_user)?;
transition!(SendingStartup {
future: stream.send(buf),
user,
password: state.password,
channel_binding,
})
}
fn poll_sending_ssl<'a>(
state: &'a mut RentToOwn<'a, SendingSsl>,
) -> Poll<AfterSendingSsl, Error> {
let (stream, _) = try_ready_closed!(state.future.poll());
fn poll_sending_startup<'a>(
state: &'a mut RentToOwn<'a, SendingStartup<S, T>>,
) -> Poll<AfterSendingStartup<S, T>, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::io));
let state = state.take();
transition!(ReadingSsl {
future: read_exact(stream, [0]),
params: state.params,
connector: state.connector,
required: state.required,
transition!(ReadingAuth {
stream,
user: state.user,
password: state.password,
channel_binding: state.channel_binding,
})
}
fn poll_reading_ssl<'a>(
state: &'a mut RentToOwn<'a, ReadingSsl>,
) -> Poll<AfterReadingSsl, Error> {
let (stream, buf) = try_ready_closed!(state.future.poll());
fn poll_reading_auth<'a>(
state: &'a mut RentToOwn<'a, ReadingAuth<S, T>>,
) -> Poll<AfterReadingAuth<S, T>, Error> {
let message = try_ready!(state.stream.poll().map_err(Error::io));
let state = state.take();
match buf[0] {
b'S' => {
let future = match state.params.host() {
Host::Tcp(domain) => state.connector.connect(domain, tls::Socket(stream)),
Host::Unix(_) => {
return Err(Error::tls("TLS over unix sockets not supported".into()))
}
};
transition!(ConnectingTls {
future,
params: state.params,
match message {
Some(Message::AuthenticationOk) => transition!(ReadingInfo {
stream: state.stream,
cancel_data: None,
parameters: HashMap::new(),
}),
Some(Message::AuthenticationCleartextPassword) => {
let pass = state.password.ok_or_else(Error::missing_password)?;
let mut buf = vec![];
frontend::password_message(&pass, &mut buf).map_err(Error::encode)?;
transition!(SendingPassword {
future: state.stream.send(buf)
})
}
b'N' if !state.required => transition!(Ready(Box::new(stream))),
b'N' => Err(Error::tls("TLS was required but not supported".into())),
_ => Err(Error::unexpected_message()),
Some(Message::AuthenticationMd5Password(body)) => {
let pass = state.password.ok_or_else(Error::missing_password)?;
let output =
authentication::md5_hash(state.user.as_bytes(), pass.as_bytes(), body.salt());
let mut buf = vec![];
frontend::password_message(&output, &mut buf).map_err(Error::encode)?;
transition!(SendingPassword {
future: state.stream.send(buf)
})
}
Some(Message::AuthenticationSasl(body)) => {
let pass = state.password.ok_or_else(Error::missing_password)?;
let mut has_scram = false;
let mut has_scram_plus = false;
let mut mechanisms = body.mechanisms();
while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
match mechanism {
sasl::SCRAM_SHA_256 => has_scram = true,
sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
_ => {}
}
}
let channel_binding = if let Some(tls_unique) = state.channel_binding.tls_unique {
Some(sasl::ChannelBinding::tls_unique(tls_unique))
} else if let Some(tls_server_end_point) =
state.channel_binding.tls_server_end_point
{
Some(sasl::ChannelBinding::tls_server_end_point(
tls_server_end_point,
))
} else {
None
};
let (channel_binding, mechanism) = if has_scram_plus {
match channel_binding {
Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
}
} else if has_scram {
match channel_binding {
Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
}
} else {
return Err(Error::unsupported_authentication());
};
let mut scram = ScramSha256::new(pass.as_bytes(), channel_binding);
let mut buf = vec![];
frontend::sasl_initial_response(mechanism, scram.message(), &mut buf)
.map_err(Error::encode)?;
transition!(SendingSasl {
future: state.stream.send(buf),
scram,
})
}
Some(Message::AuthenticationKerberosV5)
| Some(Message::AuthenticationScmCredential)
| Some(Message::AuthenticationGss)
| Some(Message::AuthenticationSspi) => Err(Error::unsupported_authentication()),
Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
Some(_) => Err(Error::unexpected_message()),
None => Err(Error::closed()),
}
}
fn poll_connecting_tls<'a>(
state: &'a mut RentToOwn<'a, ConnectingTls>,
) -> Poll<AfterConnectingTls, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::tls));
transition!(Ready(stream))
fn poll_sending_password<'a>(
state: &'a mut RentToOwn<'a, SendingPassword<S, T>>,
) -> Poll<AfterSendingPassword<S, T>, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::io));
transition!(ReadingAuthCompletion { stream })
}
fn poll_sending_sasl<'a>(
state: &'a mut RentToOwn<'a, SendingSasl<S, T>>,
) -> Poll<AfterSendingSasl<S, T>, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::io));
let state = state.take();
transition!(ReadingSasl {
stream,
scram: state.scram,
})
}
fn poll_reading_sasl<'a>(
state: &'a mut RentToOwn<'a, ReadingSasl<S, T>>,
) -> Poll<AfterReadingSasl<S, T>, Error> {
let message = try_ready!(state.stream.poll().map_err(Error::io));
let mut state = state.take();
match message {
Some(Message::AuthenticationSaslContinue(body)) => {
state
.scram
.update(body.data())
.map_err(Error::authentication)?;
let mut buf = vec![];
frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
transition!(SendingSasl {
future: state.stream.send(buf),
scram: state.scram,
})
}
Some(Message::AuthenticationSaslFinal(body)) => {
state
.scram
.finish(body.data())
.map_err(Error::authentication)?;
transition!(ReadingAuthCompletion {
stream: state.stream
})
}
Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
Some(_) => Err(Error::unexpected_message()),
None => Err(Error::closed()),
}
}
fn poll_reading_auth_completion<'a>(
state: &'a mut RentToOwn<'a, ReadingAuthCompletion<S, T>>,
) -> Poll<AfterReadingAuthCompletion<S, T>, Error> {
let message = try_ready!(state.stream.poll().map_err(Error::io));
let state = state.take();
match message {
Some(Message::AuthenticationOk) => transition!(ReadingInfo {
stream: state.stream,
cancel_data: None,
parameters: HashMap::new()
}),
Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
Some(_) => Err(Error::unexpected_message()),
None => Err(Error::closed()),
}
}
fn poll_reading_info<'a>(
state: &'a mut RentToOwn<'a, ReadingInfo<S, T>>,
) -> Poll<AfterReadingInfo<S, T>, Error> {
loop {
let message = try_ready!(state.stream.poll().map_err(Error::io));
match message {
Some(Message::BackendKeyData(body)) => {
state.cancel_data = Some(CancelData {
process_id: body.process_id(),
secret_key: body.secret_key(),
});
}
Some(Message::ParameterStatus(body)) => {
state.parameters.insert(
body.name().map_err(Error::parse)?.to_string(),
body.value().map_err(Error::parse)?.to_string(),
);
}
Some(Message::ReadyForQuery(_)) => {
let state = state.take();
let cancel_data = state.cancel_data.ok_or_else(|| {
Error::parse(io::Error::new(
io::ErrorKind::InvalidData,
"BackendKeyData message missing",
))
})?;
let (sender, receiver) = mpsc::unbounded();
let client = Client::new(sender);
let connection =
Connection::new(state.stream, cancel_data, state.parameters, receiver);
transition!(Finished((client, connection)))
}
Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
Some(Message::NoticeResponse(_)) => {}
Some(_) => return Err(Error::unexpected_message()),
None => return Err(Error::closed()),
}
}
}
}
impl ConnectFuture {
pub fn new(params: ConnectParams, tls: TlsMode) -> ConnectFuture {
Connect::start(params, tls)
impl<S, T> ConnectFuture<S, T>
where
S: AsyncRead + AsyncWrite,
T: TlsMode<S>,
{
pub fn new(
stream: S,
tls_mode: T,
password: Option<String>,
params: HashMap<String, String>,
) -> ConnectFuture<S, T> {
Connect::start(TlsFuture::new(stream, tls_mode), password, params)
}
}

328
tokio-postgres/src/proto/handshake.rs
View File

@@ -1,328 +0,0 @@
use fallible_iterator::FallibleIterator;
use futures::sink;
use futures::sync::mpsc;
use futures::{Future, Poll, Sink, Stream};
use postgres_protocol::authentication;
use postgres_protocol::authentication::sasl::{self, ChannelBinding, ScramSha256};
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use state_machine_future::RentToOwn;
use std::collections::HashMap;
use std::io;
use tokio_codec::Framed;
use params::{ConnectParams, User};
use proto::client::Client;
use proto::codec::PostgresCodec;
use proto::connect::ConnectFuture;
use proto::connection::Connection;
use tls::TlsStream;
use {CancelData, Error, TlsMode};
#[derive(StateMachineFuture)]
pub enum Handshake {
#[state_machine_future(start, transitions(SendingStartup))]
Start {
future: ConnectFuture,
params: ConnectParams,
},
#[state_machine_future(transitions(ReadingAuth))]
SendingStartup {
future: sink::Send<Framed<Box<TlsStream>, PostgresCodec>>,
user: User,
},
#[state_machine_future(transitions(ReadingInfo, SendingPassword, SendingSasl))]
ReadingAuth {
stream: Framed<Box<TlsStream>, PostgresCodec>,
user: User,
},
#[state_machine_future(transitions(ReadingAuthCompletion))]
SendingPassword {
future: sink::Send<Framed<Box<TlsStream>, PostgresCodec>>,
},
#[state_machine_future(transitions(ReadingSasl))]
SendingSasl {
future: sink::Send<Framed<Box<TlsStream>, PostgresCodec>>,
scram: ScramSha256,
},
#[state_machine_future(transitions(SendingSasl, ReadingAuthCompletion))]
ReadingSasl {
stream: Framed<Box<TlsStream>, PostgresCodec>,
scram: ScramSha256,
},
#[state_machine_future(transitions(ReadingInfo))]
ReadingAuthCompletion {
stream: Framed<Box<TlsStream>, PostgresCodec>,
},
#[state_machine_future(transitions(Finished))]
ReadingInfo {
stream: Framed<Box<TlsStream>, PostgresCodec>,
cancel_data: Option<CancelData>,
parameters: HashMap<String, String>,
},
#[state_machine_future(ready)]
Finished((Client, Connection<Box<TlsStream>>)),
#[state_machine_future(error)]
Failed(Error),
}
impl PollHandshake for Handshake {
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
let stream = try_ready!(state.future.poll());
let state = state.take();
let user = match state.params.user() {
Some(user) => user.clone(),
None => return Err(Error::missing_user()),
};
let mut buf = vec![];
{
let options = state
.params
.options()
.iter()
.map(|&(ref key, ref value)| (&**key, &**value));
let client_encoding = Some(("client_encoding", "UTF8"));
let timezone = Some(("timezone", "GMT"));
let user = Some(("user", user.name()));
let database = state.params.database().map(|s| ("database", s));
frontend::startup_message(
options
.chain(client_encoding)
.chain(timezone)
.chain(user)
.chain(database),
&mut buf,
).map_err(Error::encode)?;
}
let stream = Framed::new(stream, PostgresCodec);
transition!(SendingStartup {
future: stream.send(buf),
user,
})
}
fn poll_sending_startup<'a>(
state: &'a mut RentToOwn<'a, SendingStartup>,
) -> Poll<AfterSendingStartup, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::io));
let state = state.take();
transition!(ReadingAuth {
stream,
user: state.user,
})
}
fn poll_reading_auth<'a>(
state: &'a mut RentToOwn<'a, ReadingAuth>,
) -> Poll<AfterReadingAuth, Error> {
let message = try_ready!(state.stream.poll().map_err(Error::io));
let state = state.take();
match message {
Some(Message::AuthenticationOk) => transition!(ReadingInfo {
stream: state.stream,
cancel_data: None,
parameters: HashMap::new(),
}),
Some(Message::AuthenticationCleartextPassword) => {
let pass = state.user.password().ok_or_else(Error::missing_password)?;
let mut buf = vec![];
frontend::password_message(pass, &mut buf).map_err(Error::encode)?;
transition!(SendingPassword {
future: state.stream.send(buf)
})
}
Some(Message::AuthenticationMd5Password(body)) => {
let pass = state.user.password().ok_or_else(Error::missing_password)?;
let output = authentication::md5_hash(
state.user.name().as_bytes(),
pass.as_bytes(),
body.salt(),
);
let mut buf = vec![];
frontend::password_message(&output, &mut buf).map_err(Error::encode)?;
transition!(SendingPassword {
future: state.stream.send(buf)
})
}
Some(Message::AuthenticationSasl(body)) => {
let pass = state.user.password().ok_or_else(Error::missing_password)?;
let mut has_scram = false;
let mut has_scram_plus = false;
let mut mechanisms = body.mechanisms();
while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
match mechanism {
sasl::SCRAM_SHA_256 => has_scram = true,
sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
_ => {}
}
}
let channel_binding = state
.stream
.get_ref()
.tls_unique()
.map(ChannelBinding::tls_unique)
.or_else(|| {
state
.stream
.get_ref()
.tls_server_end_point()
.map(ChannelBinding::tls_server_end_point)
});
let (channel_binding, mechanism) = if has_scram_plus {
match channel_binding {
Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
None => (ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
}
} else if has_scram {
match channel_binding {
Some(_) => (ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
None => (ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
}
} else {
return Err(Error::unsupported_authentication());
};
let mut scram = ScramSha256::new(pass.as_bytes(), channel_binding);
let mut buf = vec![];
frontend::sasl_initial_response(mechanism, scram.message(), &mut buf)
.map_err(Error::encode)?;
transition!(SendingSasl {
future: state.stream.send(buf),
scram,
})
}
Some(Message::AuthenticationKerberosV5)
| Some(Message::AuthenticationScmCredential)
| Some(Message::AuthenticationGss)
| Some(Message::AuthenticationSspi) => Err(Error::unsupported_authentication()),
Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
Some(_) => Err(Error::unexpected_message()),
None => Err(Error::closed()),
}
}
fn poll_sending_password<'a>(
state: &'a mut RentToOwn<'a, SendingPassword>,
) -> Poll<AfterSendingPassword, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::io));
transition!(ReadingAuthCompletion { stream })
}
fn poll_sending_sasl<'a>(
state: &'a mut RentToOwn<'a, SendingSasl>,
) -> Poll<AfterSendingSasl, Error> {
let stream = try_ready!(state.future.poll().map_err(Error::io));
let state = state.take();
transition!(ReadingSasl {
stream,
scram: state.scram
})
}
fn poll_reading_sasl<'a>(
state: &'a mut RentToOwn<'a, ReadingSasl>,
) -> Poll<AfterReadingSasl, Error> {
let message = try_ready!(state.stream.poll().map_err(Error::io));
let mut state = state.take();
match message {
Some(Message::AuthenticationSaslContinue(body)) => {
state
.scram
.update(body.data())
.map_err(Error::authentication)?;
let mut buf = vec![];
frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
transition!(SendingSasl {
future: state.stream.send(buf),
scram: state.scram,
})
}
Some(Message::AuthenticationSaslFinal(body)) => {
state
.scram
.finish(body.data())
.map_err(Error::authentication)?;
transition!(ReadingAuthCompletion {
stream: state.stream,
})
}
Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
Some(_) => Err(Error::unexpected_message()),
None => Err(Error::closed()),
}
}
fn poll_reading_auth_completion<'a>(
state: &'a mut RentToOwn<'a, ReadingAuthCompletion>,
) -> Poll<AfterReadingAuthCompletion, Error> {
let message = try_ready!(state.stream.poll().map_err(Error::io));
let state = state.take();
match message {
Some(Message::AuthenticationOk) => transition!(ReadingInfo {
stream: state.stream,
cancel_data: None,
parameters: HashMap::new(),
}),
Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
Some(_) => Err(Error::unexpected_message()),
None => Err(Error::closed()),
}
}
fn poll_reading_info<'a>(
state: &'a mut RentToOwn<'a, ReadingInfo>,
) -> Poll<AfterReadingInfo, Error> {
loop {
let message = try_ready!(state.stream.poll().map_err(Error::io));
match message {
Some(Message::BackendKeyData(body)) => {
state.cancel_data = Some(CancelData {
process_id: body.process_id(),
secret_key: body.secret_key(),
});
}
Some(Message::ParameterStatus(body)) => {
state.parameters.insert(
body.name().map_err(Error::parse)?.to_string(),
body.value().map_err(Error::parse)?.to_string(),
);
}
Some(Message::ReadyForQuery(_)) => {
let state = state.take();
let cancel_data = state.cancel_data.ok_or_else(|| {
Error::parse(io::Error::new(
io::ErrorKind::InvalidData,
"BackendKeyData message missing",
))
})?;
let (sender, receiver) = mpsc::unbounded();
let client = Client::new(sender);
let connection =
Connection::new(state.stream, cancel_data, state.parameters, receiver);
transition!(Finished((client, connection)))
}
Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
Some(Message::NoticeResponse(_)) => {}
Some(_) => return Err(Error::unexpected_message()),
None => return Err(Error::closed()),
}
}
}
}
impl HandshakeFuture {
pub fn new(params: ConnectParams, tls: TlsMode) -> HandshakeFuture {
Handshake::start(ConnectFuture::new(params.clone(), tls), params)
}
}

7
tokio-postgres/src/proto/mod.rs
View File

@@ -27,14 +27,13 @@ mod connection;
mod copy_in;
mod copy_out;
mod execute;
mod handshake;
mod portal;
mod prepare;
mod query;
mod row;
mod simple_query;
mod socket;
mod statement;
mod tls;
mod transaction;
mod typeinfo;
mod typeinfo_composite;
@@ -44,16 +43,16 @@ pub use proto::bind::BindFuture;
pub use proto::cancel::CancelFuture;
pub use proto::client::Client;
pub use proto::codec::PostgresCodec;
pub use proto::connect::ConnectFuture;
pub use proto::connection::Connection;
pub use proto::copy_in::CopyInFuture;
pub use proto::copy_out::CopyOutStream;
pub use proto::execute::ExecuteFuture;
pub use proto::handshake::HandshakeFuture;
pub use proto::portal::Portal;
pub use proto::prepare::PrepareFuture;
pub use proto::query::QueryStream;
pub use proto::row::Row;
pub use proto::simple_query::SimpleQueryFuture;
pub use proto::socket::Socket;
pub use proto::statement::Statement;
pub use proto::tls::TlsFuture;
pub use proto::transaction::TransactionFuture;

84
tokio-postgres/src/proto/socket.rs
View File

@@ -1,84 +0,0 @@
use bytes::{Buf, BufMut};
use futures::Poll;
use std::io::{self, Read, Write};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_tcp::TcpStream;
#[cfg(unix)]
use tokio_uds::UnixStream;
pub enum Socket {
Tcp(TcpStream),
#[cfg(unix)]
Unix(UnixStream),
}
impl Read for Socket {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
Socket::Tcp(stream) => stream.read(buf),
#[cfg(unix)]
Socket::Unix(stream) => stream.read(buf),
}
}
}
impl AsyncRead for Socket {
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
match self {
Socket::Tcp(stream) => stream.prepare_uninitialized_buffer(buf),
#[cfg(unix)]
Socket::Unix(stream) => stream.prepare_uninitialized_buffer(buf),
}
}
fn read_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: BufMut,
{
match self {
Socket::Tcp(stream) => stream.read_buf(buf),
#[cfg(unix)]
Socket::Unix(stream) => stream.read_buf(buf),
}
}
}
impl Write for Socket {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
match self {
Socket::Tcp(stream) => stream.write(buf),
#[cfg(unix)]
Socket::Unix(stream) => stream.write(buf),
}
}
fn flush(&mut self) -> io::Result<()> {
match self {
Socket::Tcp(stream) => stream.flush(),
#[cfg(unix)]
Socket::Unix(stream) => stream.flush(),
}
}
}
impl AsyncWrite for Socket {
fn shutdown(&mut self) -> Poll<(), io::Error> {
match self {
Socket::Tcp(stream) => stream.shutdown(),
#[cfg(unix)]
Socket::Unix(stream) => stream.shutdown(),
}
}
fn write_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: Buf,
{
match self {
Socket::Tcp(stream) => stream.write_buf(buf),
#[cfg(unix)]
Socket::Unix(stream) => stream.write_buf(buf),
}
}
}

97
tokio-postgres/src/proto/tls.rs Normal file
View File

@@ -0,0 +1,97 @@
use futures::{Future, Poll};
use postgres_protocol::message::frontend;
use state_machine_future::RentToOwn;
use tokio_io::io::{self, ReadExact, WriteAll};
use tokio_io::{AsyncRead, AsyncWrite};
use {ChannelBinding, Error, TlsMode};
#[derive(StateMachineFuture)]
pub enum Tls<S, T>
where
T: TlsMode<S>,
S: AsyncRead + AsyncWrite,
{
#[state_machine_future(start, transitions(SendingTls, ConnectingTls))]
Start { stream: S, tls_mode: T },
#[state_machine_future(transitions(ReadingTls))]
SendingTls {
future: WriteAll<S, Vec<u8>>,
tls_mode: T,
},
#[state_machine_future(transitions(ConnectingTls))]
ReadingTls {
future: ReadExact<S, [u8; 1]>,
tls_mode: T,
},
#[state_machine_future(transitions(Ready))]
ConnectingTls { future: T::Future },
#[state_machine_future(ready)]
Ready((T::Stream, ChannelBinding)),
#[state_machine_future(error)]
Failed(Error),
}
impl<S, T> PollTls<S, T> for Tls<S, T>
where
T: TlsMode<S>,
S: AsyncRead + AsyncWrite,
{
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
let state = state.take();
if state.tls_mode.request_tls() {
let mut buf = vec![];
frontend::ssl_request(&mut buf);
transition!(SendingTls {
future: io::write_all(state.stream, buf),
tls_mode: state.tls_mode,
})
} else {
transition!(ConnectingTls {
future: state.tls_mode.handle_tls(false, state.stream),
})
}
}
fn poll_sending_tls<'a>(
state: &'a mut RentToOwn<'a, SendingTls<S, T>>,
) -> Poll<AfterSendingTls<S, T>, Error> {
let (stream, _) = try_ready!(state.future.poll().map_err(Error::io));
let state = state.take();
transition!(ReadingTls {
future: io::read_exact(stream, [0]),
tls_mode: state.tls_mode,
})
}
fn poll_reading_tls<'a>(
state: &'a mut RentToOwn<'a, ReadingTls<S, T>>,
) -> Poll<AfterReadingTls<S, T>, Error> {
let (stream, buf) = try_ready!(state.future.poll().map_err(Error::io));
let state = state.take();
let use_tls = buf[0] == b'S';
transition!(ConnectingTls {
future: state.tls_mode.handle_tls(use_tls, stream)
})
}
fn poll_connecting_tls<'a>(
state: &'a mut RentToOwn<'a, ConnectingTls<S, T>>,
) -> Poll<AfterConnectingTls<S, T>, Error> {
let t = try_ready!(state.future.poll().map_err(|e| Error::tls(e.into())));
transition!(Ready(t))
}
}
impl<S, T> TlsFuture<S, T>
where
T: TlsMode<S>,
S: AsyncRead + AsyncWrite,
{
pub fn new(stream: S, tls_mode: T) -> TlsFuture<S, T> {
Tls::start(stream, tls_mode)
}
}

271
tokio-postgres/src/tls.rs
View File

@@ -1,83 +1,274 @@
use bytes::{Buf, BufMut};
use futures::{Future, Poll};
use futures::future::{self, FutureResult};
use futures::{Async, Future, Poll};
use std::error::Error;
use std::fmt;
use std::io::{self, Read, Write};
use tokio_io::{AsyncRead, AsyncWrite};
use void::Void;
use proto;
pub struct ChannelBinding {
pub(crate) tls_server_end_point: Option<Vec<u8>>,
pub(crate) tls_unique: Option<Vec<u8>>,
}
pub struct Socket(pub(crate) proto::Socket);
impl ChannelBinding {
pub fn new() -> ChannelBinding {
ChannelBinding {
tls_server_end_point: None,
tls_unique: None,
}
}
impl Read for Socket {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.0.read(buf)
pub fn tls_server_end_point(mut self, tls_server_end_point: Vec<u8>) -> ChannelBinding {
self.tls_server_end_point = Some(tls_server_end_point);
self
}
pub fn tls_unique(mut self, tls_unique: Vec<u8>) -> ChannelBinding {
self.tls_unique = Some(tls_unique);
self
}
}
impl AsyncRead for Socket {
pub trait TlsMode<S> {
type Stream: AsyncRead + AsyncWrite;
type Error: Into<Box<Error + Sync + Send>>;
type Future: Future<Item = (Self::Stream, ChannelBinding), Error = Self::Error>;
fn request_tls(&self) -> bool;
fn handle_tls(self, use_tls: bool, stream: S) -> Self::Future;
}
pub trait TlsConnect<S> {
type Stream: AsyncRead + AsyncWrite;
type Error: Into<Box<Error + Sync + Send>>;
type Future: Future<Item = (Self::Stream, ChannelBinding), Error = Self::Error>;
fn connect(self, stream: S) -> Self::Future;
}
#[derive(Debug, Copy, Clone)]
pub struct NoTls;
impl<S> TlsMode<S> for NoTls
where
S: AsyncRead + AsyncWrite,
{
type Stream = S;
type Error = Void;
type Future = FutureResult<(S, ChannelBinding), Void>;
fn request_tls(&self) -> bool {
false
}
fn handle_tls(self, use_tls: bool, stream: S) -> FutureResult<(S, ChannelBinding), Void> {
debug_assert!(!use_tls);
future::ok((stream, ChannelBinding::new()))
}
}
#[derive(Debug, Copy, Clone)]
pub struct PreferTls<T>(pub T);
impl<T, S> TlsMode<S> for PreferTls<T>
where
T: TlsConnect<S>,
S: AsyncRead + AsyncWrite,
{
type Stream = MaybeTlsStream<T::Stream, S>;
type Error = T::Error;
type Future = PreferTlsFuture<T::Future, S>;
fn request_tls(&self) -> bool {
true
}
fn handle_tls(self, use_tls: bool, stream: S) -> PreferTlsFuture<T::Future, S> {
let f = if use_tls {
PreferTlsFutureInner::Tls(self.0.connect(stream))
} else {
PreferTlsFutureInner::Raw(Some(stream))
};
PreferTlsFuture(f)
}
}
enum PreferTlsFutureInner<F, S> {
Tls(F),
Raw(Option<S>),
}
pub struct PreferTlsFuture<F, S>(PreferTlsFutureInner<F, S>);
impl<F, S, T> Future for PreferTlsFuture<F, S>
where
F: Future<Item = (T, ChannelBinding)>,
{
type Item = (MaybeTlsStream<T, S>, ChannelBinding);
type Error = F::Error;
fn poll(&mut self) -> Poll<(MaybeTlsStream<T, S>, ChannelBinding), F::Error> {
match &mut self.0 {
PreferTlsFutureInner::Tls(f) => {
let (stream, channel_binding) = try_ready!(f.poll());
Ok(Async::Ready((MaybeTlsStream::Tls(stream), channel_binding)))
}
PreferTlsFutureInner::Raw(s) => Ok(Async::Ready((
MaybeTlsStream::Raw(s.take().expect("future polled after completion")),
ChannelBinding::new(),
))),
}
}
}
pub enum MaybeTlsStream<T, U> {
Tls(T),
Raw(U),
}
impl<T, U> Read for MaybeTlsStream<T, U>
where
T: Read,
U: Read,
{
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
MaybeTlsStream::Tls(s) => s.read(buf),
MaybeTlsStream::Raw(s) => s.read(buf),
}
}
}
impl<T, U> AsyncRead for MaybeTlsStream<T, U>
where
T: AsyncRead,
U: AsyncRead,
{
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
self.0.prepare_uninitialized_buffer(buf)
match self {
MaybeTlsStream::Tls(s) => s.prepare_uninitialized_buffer(buf),
MaybeTlsStream::Raw(s) => s.prepare_uninitialized_buffer(buf),
}
}
fn read_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: BufMut,
{
self.0.read_buf(buf)
match self {
MaybeTlsStream::Tls(s) => s.read_buf(buf),
MaybeTlsStream::Raw(s) => s.read_buf(buf),
}
}
}
impl Write for Socket {
impl<T, U> Write for MaybeTlsStream<T, U>
where
T: Write,
U: Write,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.0.write(buf)
match self {
MaybeTlsStream::Tls(s) => s.write(buf),
MaybeTlsStream::Raw(s) => s.write(buf),
}
}
fn flush(&mut self) -> io::Result<()> {
self.0.flush()
match self {
MaybeTlsStream::Tls(s) => s.flush(),
MaybeTlsStream::Raw(s) => s.flush(),
}
}
}
impl AsyncWrite for Socket {
impl<T, U> AsyncWrite for MaybeTlsStream<T, U>
where
T: AsyncWrite,
U: AsyncWrite,
{
fn shutdown(&mut self) -> Poll<(), io::Error> {
self.0.shutdown()
match self {
MaybeTlsStream::Tls(s) => s.shutdown(),
MaybeTlsStream::Raw(s) => s.shutdown(),
}
}
fn write_buf<B>(&mut self, buf: &mut B) -> Poll<usize, io::Error>
where
B: Buf,
{
self.0.write_buf(buf)
match self {
MaybeTlsStream::Tls(s) => s.write_buf(buf),
MaybeTlsStream::Raw(s) => s.write_buf(buf),
}
}
}
pub trait TlsConnect: Sync + Send {
fn connect(
&self,
domain: &str,
socket: Socket,
) -> Box<Future<Item = Box<TlsStream>, Error = Box<Error + Sync + Send>> + Sync + Send>;
}
#[derive(Debug, Copy, Clone)]
pub struct RequireTls<T>(pub T);
pub trait TlsStream: 'static + Sync + Send + AsyncRead + AsyncWrite {
/// Returns the data associated with the `tls-unique` channel binding type as described in
/// [RFC 5929], if supported.
///
/// An implementation only needs to support at most one of this or `tls_server_end_point`.
///
/// [RFC 5929]: https://tools.ietf.org/html/rfc5929
fn tls_unique(&self) -> Option<Vec<u8>> {
None
impl<T, S> TlsMode<S> for RequireTls<T>
where
T: TlsConnect<S>,
{
type Stream = T::Stream;
type Error = Box<Error + Sync + Send>;
type Future = RequireTlsFuture<T::Future>;
fn request_tls(&self) -> bool {
true
}
/// Returns the data associated with the `tls-server-end-point` channel binding type as
/// described in [RFC 5929], if supported.
///
/// An implementation only needs to support at most one of this or `tls_unique`.
///
/// [RFC 5929]: https://tools.ietf.org/html/rfc5929
fn tls_server_end_point(&self) -> Option<Vec<u8>> {
None
fn handle_tls(self, use_tls: bool, stream: S) -> RequireTlsFuture<T::Future> {
let f = if use_tls {
Ok(self.0.connect(stream))
} else {
Err(TlsUnsupportedError(()).into())
};
RequireTlsFuture { f: Some(f) }
}
}
impl TlsStream for proto::Socket {}
#[derive(Debug)]
pub struct TlsUnsupportedError(());
impl fmt::Display for TlsUnsupportedError {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.write_str("TLS was required but not supported by the server")
}
}
impl Error for TlsUnsupportedError {}
pub struct RequireTlsFuture<T> {
f: Option<Result<T, Box<Error + Sync + Send>>>,
}
impl<T> Future for RequireTlsFuture<T>
where
T: Future,
T::Error: Into<Box<Error + Sync + Send>>,
{
type Item = T::Item;
type Error = Box<Error + Sync + Send>;
fn poll(&mut self) -> Poll<T::Item, Box<Error + Sync + Send>> {
match self.f.take().expect("future polled after completion") {
Ok(mut f) => match f.poll().map_err(Into::into)? {
Async::Ready(r) => Ok(Async::Ready(r)),
Async::NotReady => {
self.f = Some(Ok(f));
Ok(Async::NotReady)
}
},
Err(e) => Err(e),
}
}
}

209
tokio-postgres/tests/test.rs
View File

@@ -12,18 +12,28 @@ use futures::stream;
use futures::sync::mpsc;
use std::error::Error;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio::prelude::*;
use tokio::runtime::current_thread::Runtime;
use tokio::timer::Delay;
use tokio_postgres::error::SqlState;
use tokio_postgres::types::{Kind, Type};
use tokio_postgres::{AsyncMessage, TlsMode};
use tokio_postgres::{AsyncMessage, Client, Connection, NoTls};
fn smoke_test(url: &str) {
fn connect(
builder: &tokio_postgres::Builder,
) -> impl Future<Item = (Client, Connection<TcpStream>), Error = tokio_postgres::Error> {
let builder = builder.clone();
TcpStream::connect(&"127.0.0.1:5433".parse().unwrap())
.map_err(|e| panic!("{}", e))
.and_then(move |s| builder.connect(s, NoTls))
}
fn smoke_test(builder: &tokio_postgres::Builder) {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(url.parse().unwrap(), TlsMode::None);
let handshake = connect(builder);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -46,9 +56,10 @@ fn plain_password_missing() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://pass_user@localhost:5433".parse().unwrap(),
TlsMode::None,
let handshake = connect(
tokio_postgres::Builder::new()
.user("pass_user")
.database("postgres"),
);
runtime.block_on(handshake).err().unwrap();
}
@@ -58,9 +69,11 @@ fn plain_password_wrong() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://pass_user:foo@localhost:5433".parse().unwrap(),
TlsMode::None,
let handshake = connect(
tokio_postgres::Builder::new()
.user("pass_user")
.password("foo")
.database("postgres"),
);
match runtime.block_on(handshake) {
Ok(_) => panic!("unexpected success"),
@@ -71,7 +84,12 @@ fn plain_password_wrong() {
#[test]
fn plain_password_ok() {
smoke_test("postgres://pass_user:password@localhost:5433/postgres");
smoke_test(
tokio_postgres::Builder::new()
.user("pass_user")
.password("password")
.database("postgres"),
);
}
#[test]
@@ -79,9 +97,10 @@ fn md5_password_missing() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://md5_user@localhost:5433".parse().unwrap(),
TlsMode::None,
let handshake = connect(
tokio_postgres::Builder::new()
.user("md5_user")
.database("postgres"),
);
runtime.block_on(handshake).err().unwrap();
}
@@ -91,9 +110,11 @@ fn md5_password_wrong() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://md5_user:foo@localhost:5433".parse().unwrap(),
TlsMode::None,
let handshake = connect(
tokio_postgres::Builder::new()
.user("md5_user")
.password("foo")
.database("postgres"),
);
match runtime.block_on(handshake) {
Ok(_) => panic!("unexpected success"),
@@ -104,7 +125,12 @@ fn md5_password_wrong() {
#[test]
fn md5_password_ok() {
smoke_test("postgres://md5_user:password@localhost:5433/postgres");
smoke_test(
tokio_postgres::Builder::new()
.user("md5_user")
.password("password")
.database("postgres"),
);
}
#[test]
@@ -112,9 +138,10 @@ fn scram_password_missing() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://scram_user@localhost:5433".parse().unwrap(),
TlsMode::None,
let handshake = connect(
tokio_postgres::Builder::new()
.user("scram_user")
.database("postgres"),
);
runtime.block_on(handshake).err().unwrap();
}
@@ -124,9 +151,11 @@ fn scram_password_wrong() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://scram_user:foo@localhost:5433".parse().unwrap(),
TlsMode::None,
let handshake = connect(
tokio_postgres::Builder::new()
.user("scram_user")
.password("foo")
.database("postgres"),
);
match runtime.block_on(handshake) {
Ok(_) => panic!("unexpected success"),
@@ -137,7 +166,12 @@ fn scram_password_wrong() {
#[test]
fn scram_password_ok() {
smoke_test("postgres://scram_user:password@localhost:5433/postgres");
smoke_test(
tokio_postgres::Builder::new()
.user("scram_user")
.password("password")
.database("postgres"),
);
}
#[test]
@@ -145,11 +179,9 @@ fn pipelined_prepare() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -167,11 +199,9 @@ fn insert_select() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -203,11 +233,9 @@ fn query_portal() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -246,11 +274,9 @@ fn cancel_query() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let cancel_data = connection.cancel_data();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -265,11 +291,10 @@ fn cancel_query() {
let cancel = Delay::new(Instant::now() + Duration::from_millis(100))
.then(|r| {
r.unwrap();
tokio_postgres::cancel_query(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
cancel_data,
)
TcpStream::connect(&"127.0.0.1:5433".parse().unwrap())
}).then(|r| {
let s = r.unwrap();
tokio_postgres::cancel_query(s, NoTls, cancel_data)
}).then(|r| {
r.unwrap();
Ok::<(), ()>(())
@@ -283,11 +308,9 @@ fn custom_enum() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -320,11 +343,9 @@ fn custom_domain() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -346,11 +367,9 @@ fn custom_array() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -373,11 +392,9 @@ fn custom_composite() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -413,11 +430,9 @@ fn custom_range() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -442,11 +457,9 @@ fn custom_simple() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, connection) = runtime.block_on(handshake).unwrap();
let (mut client, connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -463,11 +476,9 @@ fn notifications() {
let _ = env_logger::try_init();
let mut runtime = Runtime::new().unwrap();
let handshake = tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
);
let (mut client, mut connection) = runtime.block_on(handshake).unwrap();
let (mut client, mut connection) = runtime
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let (tx, rx) = mpsc::unbounded();
let connection = future::poll_fn(move || {
@@ -512,10 +523,8 @@ fn transaction_commit() {
let mut runtime = Runtime::new().unwrap();
let (mut client, connection) = runtime
.block_on(tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
)).unwrap();
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -547,10 +556,8 @@ fn transaction_abort() {
let mut runtime = Runtime::new().unwrap();
let (mut client, connection) = runtime
.block_on(tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
)).unwrap();
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -584,10 +591,8 @@ fn copy_in() {
let mut runtime = Runtime::new().unwrap();
let (mut client, connection) = runtime
.block_on(tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
)).unwrap();
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -628,10 +633,8 @@ fn copy_in_error() {
let mut runtime = Runtime::new().unwrap();
let (mut client, connection) = runtime
.block_on(tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
)).unwrap();
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();
@@ -668,10 +671,8 @@ fn copy_out() {
let mut runtime = Runtime::new().unwrap();
let (mut client, connection) = runtime
.block_on(tokio_postgres::connect(
"postgres://postgres@localhost:5433".parse().unwrap(),
TlsMode::None,
)).unwrap();
.block_on(connect(tokio_postgres::Builder::new().user("postgres")))
.unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.handle().spawn(connection).unwrap();