use std::{
fmt::{self, Formatter},
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
};
use ::http::Extensions;
#[cfg(feature = "tokio")]
pub use self::http::{HttpConnector, HttpInfo};
#[cfg(feature = "tokio")]
pub mod dns;
#[cfg(feature = "tokio")]
mod http;
pub mod proxy;
pub(crate) mod capture;
pub use capture::{capture_connection, CaptureConnection};
pub use self::sealed::Connect;
pub trait Connection {
fn connected(&self) -> Connected;
}
#[derive(Debug)]
pub struct Connected {
pub(super) alpn: Alpn,
pub(super) is_proxied: bool,
pub(super) extra: Option<Extra>,
pub(super) poisoned: PoisonPill,
}
#[derive(Clone)]
pub(crate) struct PoisonPill {
poisoned: Arc<AtomicBool>,
}
impl fmt::Debug for PoisonPill {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(
f,
"PoisonPill@{:p} {{ poisoned: {} }}",
self.poisoned,
self.poisoned.load(Ordering::Relaxed)
)
}
}
impl PoisonPill {
pub(crate) fn healthy() -> Self {
Self {
poisoned: Arc::new(AtomicBool::new(false)),
}
}
pub(crate) fn poison(&self) {
self.poisoned.store(true, Ordering::Relaxed)
}
pub(crate) fn poisoned(&self) -> bool {
self.poisoned.load(Ordering::Relaxed)
}
}
pub(super) struct Extra(Box<dyn ExtraInner>);
#[derive(Clone, Copy, Debug, PartialEq)]
pub(super) enum Alpn {
H2,
None,
}
impl Connected {
pub fn new() -> Connected {
Connected {
alpn: Alpn::None,
is_proxied: false,
extra: None,
poisoned: PoisonPill::healthy(),
}
}
pub fn proxy(mut self, is_proxied: bool) -> Connected {
self.is_proxied = is_proxied;
self
}
pub fn is_proxied(&self) -> bool {
self.is_proxied
}
pub fn extra<T: Clone + Send + Sync + 'static>(mut self, extra: T) -> Connected {
if let Some(prev) = self.extra {
self.extra = Some(Extra(Box::new(ExtraChain(prev.0, extra))));
} else {
self.extra = Some(Extra(Box::new(ExtraEnvelope(extra))));
}
self
}
pub fn get_extras(&self, extensions: &mut Extensions) {
if let Some(extra) = &self.extra {
extra.set(extensions);
}
}
pub fn negotiated_h2(mut self) -> Connected {
self.alpn = Alpn::H2;
self
}
pub fn is_negotiated_h2(&self) -> bool {
self.alpn == Alpn::H2
}
pub fn poison(&self) {
self.poisoned.poison();
tracing::debug!(
poison_pill = ?self.poisoned, "connection was poisoned. this connection will not be reused for subsequent requests"
);
}
pub(super) fn clone(&self) -> Connected {
Connected {
alpn: self.alpn,
is_proxied: self.is_proxied,
extra: self.extra.clone(),
poisoned: self.poisoned.clone(),
}
}
}
impl Extra {
pub(super) fn set(&self, res: &mut Extensions) {
self.0.set(res);
}
}
impl Clone for Extra {
fn clone(&self) -> Extra {
Extra(self.0.clone_box())
}
}
impl fmt::Debug for Extra {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Extra").finish()
}
}
trait ExtraInner: Send + Sync {
fn clone_box(&self) -> Box<dyn ExtraInner>;
fn set(&self, res: &mut Extensions);
}
#[derive(Clone)]
struct ExtraEnvelope<T>(T);
impl<T> ExtraInner for ExtraEnvelope<T>
where
T: Clone + Send + Sync + 'static,
{
fn clone_box(&self) -> Box<dyn ExtraInner> {
Box::new(self.clone())
}
fn set(&self, res: &mut Extensions) {
res.insert(self.0.clone());
}
}
struct ExtraChain<T>(Box<dyn ExtraInner>, T);
impl<T: Clone> Clone for ExtraChain<T> {
fn clone(&self) -> Self {
ExtraChain(self.0.clone_box(), self.1.clone())
}
}
impl<T> ExtraInner for ExtraChain<T>
where
T: Clone + Send + Sync + 'static,
{
fn clone_box(&self) -> Box<dyn ExtraInner> {
Box::new(self.clone())
}
fn set(&self, res: &mut Extensions) {
self.0.set(res);
res.insert(self.1.clone());
}
}
pub(super) mod sealed {
use std::error::Error as StdError;
use std::future::Future;
use ::http::Uri;
use hyper::rt::{Read, Write};
use super::Connection;
pub trait Connect: Sealed + Sized {
#[doc(hidden)]
type _Svc: ConnectSvc;
#[doc(hidden)]
fn connect(self, internal_only: Internal, dst: Uri) -> <Self::_Svc as ConnectSvc>::Future;
}
pub trait ConnectSvc {
type Connection: Read + Write + Connection + Unpin + Send + 'static;
type Error: Into<Box<dyn StdError + Send + Sync>>;
type Future: Future<Output = Result<Self::Connection, Self::Error>> + Unpin + Send + 'static;
fn connect(self, internal_only: Internal, dst: Uri) -> Self::Future;
}
impl<S, T> Connect for S
where
S: tower_service::Service<Uri, Response = T> + Send + 'static,
S::Error: Into<Box<dyn StdError + Send + Sync>>,
S::Future: Unpin + Send,
T: Read + Write + Connection + Unpin + Send + 'static,
{
type _Svc = S;
fn connect(self, _: Internal, dst: Uri) -> crate::service::Oneshot<S, Uri> {
crate::service::Oneshot::new(self, dst)
}
}
impl<S, T> ConnectSvc for S
where
S: tower_service::Service<Uri, Response = T> + Send + 'static,
S::Error: Into<Box<dyn StdError + Send + Sync>>,
S::Future: Unpin + Send,
T: Read + Write + Connection + Unpin + Send + 'static,
{
type Connection = T;
type Error = S::Error;
type Future = crate::service::Oneshot<S, Uri>;
fn connect(self, _: Internal, dst: Uri) -> Self::Future {
crate::service::Oneshot::new(self, dst)
}
}
impl<S, T> Sealed for S
where
S: tower_service::Service<Uri, Response = T> + Send,
S::Error: Into<Box<dyn StdError + Send + Sync>>,
S::Future: Unpin + Send,
T: Read + Write + Connection + Unpin + Send + 'static,
{
}
pub trait Sealed {}
#[allow(missing_debug_implementations)]
pub struct Internal;
}
#[cfg(test)]
mod tests {
use super::Connected;
#[derive(Clone, Debug, PartialEq)]
struct Ex1(usize);
#[derive(Clone, Debug, PartialEq)]
struct Ex2(&'static str);
#[derive(Clone, Debug, PartialEq)]
struct Ex3(&'static str);
#[test]
fn test_connected_extra() {
let c1 = Connected::new().extra(Ex1(41));
let mut ex = ::http::Extensions::new();
assert_eq!(ex.get::<Ex1>(), None);
c1.extra.as_ref().expect("c1 extra").set(&mut ex);
assert_eq!(ex.get::<Ex1>(), Some(&Ex1(41)));
}
#[test]
fn test_connected_extra_chain() {
let c1 = Connected::new()
.extra(Ex1(45))
.extra(Ex2("zoom"))
.extra(Ex3("pew pew"));
let mut ex1 = ::http::Extensions::new();
assert_eq!(ex1.get::<Ex1>(), None);
assert_eq!(ex1.get::<Ex2>(), None);
assert_eq!(ex1.get::<Ex3>(), None);
c1.extra.as_ref().expect("c1 extra").set(&mut ex1);
assert_eq!(ex1.get::<Ex1>(), Some(&Ex1(45)));
assert_eq!(ex1.get::<Ex2>(), Some(&Ex2("zoom")));
assert_eq!(ex1.get::<Ex3>(), Some(&Ex3("pew pew")));
let c2 = Connected::new()
.extra(Ex1(33))
.extra(Ex2("hiccup"))
.extra(Ex1(99));
let mut ex2 = ::http::Extensions::new();
c2.extra.as_ref().expect("c2 extra").set(&mut ex2);
assert_eq!(ex2.get::<Ex1>(), Some(&Ex1(99)));
assert_eq!(ex2.get::<Ex2>(), Some(&Ex2("hiccup")));
}
}