110 lines
3.0 KiB
Rust
110 lines
3.0 KiB
Rust
use futures::{try_ready, Future, Poll};
|
|
use state_machine_future::{transition, RentToOwn, StateMachineFuture};
|
|
use std::io;
|
|
|
|
use crate::config::{Host, SslMode};
|
|
use crate::proto::{CancelQueryRawFuture, ConnectSocketFuture};
|
|
use crate::{Config, Error, MakeTlsConnect, Socket};
|
|
|
|
#[derive(StateMachineFuture)]
|
|
pub enum CancelQuery<T>
|
|
where
|
|
T: MakeTlsConnect<Socket>,
|
|
{
|
|
#[state_machine_future(start, transitions(ConnectingSocket))]
|
|
Start {
|
|
tls: T,
|
|
idx: Option<usize>,
|
|
config: Config,
|
|
process_id: i32,
|
|
secret_key: i32,
|
|
},
|
|
#[state_machine_future(transitions(Canceling))]
|
|
ConnectingSocket {
|
|
future: ConnectSocketFuture,
|
|
mode: SslMode,
|
|
tls: T::TlsConnect,
|
|
process_id: i32,
|
|
secret_key: i32,
|
|
},
|
|
#[state_machine_future(transitions(Finished))]
|
|
Canceling {
|
|
future: CancelQueryRawFuture<Socket, T::TlsConnect>,
|
|
},
|
|
#[state_machine_future(ready)]
|
|
Finished(()),
|
|
#[state_machine_future(error)]
|
|
Failed(Error),
|
|
}
|
|
|
|
impl<T> PollCancelQuery<T> for CancelQuery<T>
|
|
where
|
|
T: MakeTlsConnect<Socket>,
|
|
{
|
|
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<T>>) -> Poll<AfterStart<T>, Error> {
|
|
let mut state = state.take();
|
|
|
|
let idx = state.idx.ok_or_else(|| {
|
|
Error::connect(io::Error::new(io::ErrorKind::InvalidInput, "unknown host"))
|
|
})?;
|
|
|
|
let hostname = match &state.config.0.host[idx] {
|
|
Host::Tcp(host) => &**host,
|
|
// postgres doesn't support TLS over unix sockets, so the choice here doesn't matter
|
|
#[cfg(unix)]
|
|
Host::Unix(_) => "",
|
|
};
|
|
let tls = state
|
|
.tls
|
|
.make_tls_connect(hostname)
|
|
.map_err(|e| Error::tls(e.into()))?;
|
|
|
|
transition!(ConnectingSocket {
|
|
mode: state.config.0.ssl_mode,
|
|
future: ConnectSocketFuture::new(state.config, idx),
|
|
tls,
|
|
process_id: state.process_id,
|
|
secret_key: state.secret_key,
|
|
})
|
|
}
|
|
|
|
fn poll_connecting_socket<'a>(
|
|
state: &'a mut RentToOwn<'a, ConnectingSocket<T>>,
|
|
) -> Poll<AfterConnectingSocket<T>, Error> {
|
|
let socket = try_ready!(state.future.poll());
|
|
let state = state.take();
|
|
|
|
transition!(Canceling {
|
|
future: CancelQueryRawFuture::new(
|
|
socket,
|
|
state.mode,
|
|
state.tls,
|
|
state.process_id,
|
|
state.secret_key
|
|
),
|
|
})
|
|
}
|
|
|
|
fn poll_canceling<'a>(
|
|
state: &'a mut RentToOwn<'a, Canceling<T>>,
|
|
) -> Poll<AfterCanceling, Error> {
|
|
try_ready!(state.future.poll());
|
|
transition!(Finished(()))
|
|
}
|
|
}
|
|
|
|
impl<T> CancelQueryFuture<T>
|
|
where
|
|
T: MakeTlsConnect<Socket>,
|
|
{
|
|
pub fn new(
|
|
tls: T,
|
|
idx: Option<usize>,
|
|
config: Config,
|
|
process_id: i32,
|
|
secret_key: i32,
|
|
) -> CancelQueryFuture<T> {
|
|
CancelQuery::start(tls, idx, config, process_id, secret_key)
|
|
}
|
|
}
|