use std::fmt::{Debug, Formatter};
use std::net::SocketAddr;
use std::sync::Arc;
#[derive(Clone)]
pub struct ConnectionMetadata {
is_proxied: bool,
remote_addr: Option<SocketAddr>,
local_addr: Option<SocketAddr>,
poison_fn: Arc<dyn Fn() + Send + Sync>,
}
impl ConnectionMetadata {
pub fn poison(&self) {
tracing::info!("smithy connection was poisoned");
(self.poison_fn)()
}
#[deprecated(
since = "1.1.0",
note = "`ConnectionMetadata::new` is deprecated in favour of `ConnectionMetadata::builder`."
)]
pub fn new(
is_proxied: bool,
remote_addr: Option<SocketAddr>,
poison: impl Fn() + Send + Sync + 'static,
) -> Self {
Self {
is_proxied,
remote_addr,
local_addr: None,
poison_fn: Arc::new(poison),
}
}
pub fn builder() -> ConnectionMetadataBuilder {
ConnectionMetadataBuilder::new()
}
pub fn remote_addr(&self) -> Option<SocketAddr> {
self.remote_addr
}
pub fn local_addr(&self) -> Option<SocketAddr> {
self.local_addr
}
}
impl Debug for ConnectionMetadata {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SmithyConnection")
.field("is_proxied", &self.is_proxied)
.field("remote_addr", &self.remote_addr)
.field("local_addr", &self.local_addr)
.finish()
}
}
#[derive(Default)]
pub struct ConnectionMetadataBuilder {
is_proxied: Option<bool>,
remote_addr: Option<SocketAddr>,
local_addr: Option<SocketAddr>,
poison_fn: Option<Arc<dyn Fn() + Send + Sync>>,
}
impl Debug for ConnectionMetadataBuilder {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConnectionMetadataBuilder")
.field("is_proxied", &self.is_proxied)
.field("remote_addr", &self.remote_addr)
.field("local_addr", &self.local_addr)
.finish()
}
}
impl ConnectionMetadataBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn proxied(mut self, proxied: bool) -> Self {
self.set_proxied(Some(proxied));
self
}
pub fn set_proxied(&mut self, proxied: Option<bool>) -> &mut Self {
self.is_proxied = proxied;
self
}
pub fn remote_addr(mut self, remote_addr: SocketAddr) -> Self {
self.set_remote_addr(Some(remote_addr));
self
}
pub fn set_remote_addr(&mut self, remote_addr: Option<SocketAddr>) -> &mut Self {
self.remote_addr = remote_addr;
self
}
pub fn local_addr(mut self, local_addr: SocketAddr) -> Self {
self.set_local_addr(Some(local_addr));
self
}
pub fn set_local_addr(&mut self, local_addr: Option<SocketAddr>) -> &mut Self {
self.local_addr = local_addr;
self
}
pub fn poison_fn(mut self, poison_fn: impl Fn() + Send + Sync + 'static) -> Self {
self.set_poison_fn(Some(poison_fn));
self
}
pub fn set_poison_fn(
&mut self,
poison_fn: Option<impl Fn() + Send + Sync + 'static>,
) -> &mut Self {
self.poison_fn =
poison_fn.map(|poison_fn| Arc::new(poison_fn) as Arc<dyn Fn() + Send + Sync>);
self
}
pub fn build(self) -> ConnectionMetadata {
ConnectionMetadata {
is_proxied: self
.is_proxied
.expect("is_proxied should be set for ConnectionMetadata"),
remote_addr: self.remote_addr,
local_addr: self.local_addr,
poison_fn: self
.poison_fn
.expect("poison_fn should be set for ConnectionMetadata"),
}
}
}
#[cfg(test)]
mod tests {
use std::{
net::{IpAddr, Ipv6Addr},
sync::Mutex,
};
use super::*;
const TEST_SOCKET_ADDR: SocketAddr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 100);
#[test]
#[should_panic]
fn builder_panic_missing_proxied() {
ConnectionMetadataBuilder::new()
.poison_fn(|| {})
.local_addr(TEST_SOCKET_ADDR)
.remote_addr(TEST_SOCKET_ADDR)
.build();
}
#[test]
#[should_panic]
fn builder_panic_missing_poison_fn() {
ConnectionMetadataBuilder::new()
.proxied(true)
.local_addr(TEST_SOCKET_ADDR)
.remote_addr(TEST_SOCKET_ADDR)
.build();
}
#[test]
fn builder_all_fields_successful() {
let mutable_flag = Arc::new(Mutex::new(false));
let connection_metadata = ConnectionMetadataBuilder::new()
.proxied(true)
.local_addr(TEST_SOCKET_ADDR)
.remote_addr(TEST_SOCKET_ADDR)
.poison_fn({
let mutable_flag = Arc::clone(&mutable_flag);
move || {
let mut guard = mutable_flag.lock().unwrap();
*guard = !*guard;
}
})
.build();
assert!(connection_metadata.is_proxied);
assert_eq!(connection_metadata.remote_addr(), Some(TEST_SOCKET_ADDR));
assert_eq!(connection_metadata.local_addr(), Some(TEST_SOCKET_ADDR));
assert!(!(*mutable_flag.lock().unwrap()));
connection_metadata.poison();
assert!(*mutable_flag.lock().unwrap());
}
#[test]
fn builder_optional_fields_translate() {
let metadata1 = ConnectionMetadataBuilder::new()
.proxied(true)
.poison_fn(|| {})
.build();
assert_eq!(metadata1.local_addr(), None);
assert_eq!(metadata1.remote_addr(), None);
let metadata2 = ConnectionMetadataBuilder::new()
.proxied(true)
.poison_fn(|| {})
.local_addr(TEST_SOCKET_ADDR)
.build();
assert_eq!(metadata2.local_addr(), Some(TEST_SOCKET_ADDR));
assert_eq!(metadata2.remote_addr(), None);
let metadata3 = ConnectionMetadataBuilder::new()
.proxied(true)
.poison_fn(|| {})
.remote_addr(TEST_SOCKET_ADDR)
.build();
assert_eq!(metadata3.local_addr(), None);
assert_eq!(metadata3.remote_addr(), Some(TEST_SOCKET_ADDR));
}
}