mirror of
https://github.com/fafhrd91/actix-net
synced 2025-09-01 01:27:00 +02:00
add tests for custom resolver
This commit is contained in:
@@ -1,135 +1,147 @@
|
||||
use std::{
|
||||
collections::{vec_deque, VecDeque},
|
||||
fmt,
|
||||
iter::{FromIterator, FusedIterator},
|
||||
iter::{self, FromIterator as _},
|
||||
mem,
|
||||
net::SocketAddr,
|
||||
};
|
||||
|
||||
/// Connect request
|
||||
/// Parse a host into parts (hostname and port).
|
||||
pub trait Address: Unpin + 'static {
|
||||
/// Host name of the request
|
||||
fn host(&self) -> &str;
|
||||
/// Get hostname part.
|
||||
fn hostname(&self) -> &str;
|
||||
|
||||
/// Port of the request
|
||||
fn port(&self) -> Option<u16>;
|
||||
/// Get optional port part.
|
||||
fn port(&self) -> Option<u16> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Address for String {
|
||||
fn host(&self) -> &str {
|
||||
fn hostname(&self) -> &str {
|
||||
&self
|
||||
}
|
||||
|
||||
fn port(&self) -> Option<u16> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
impl Address for &'static str {
|
||||
fn host(&self) -> &str {
|
||||
fn hostname(&self) -> &str {
|
||||
self
|
||||
}
|
||||
|
||||
fn port(&self) -> Option<u16> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect request
|
||||
#[derive(Eq, PartialEq, Debug, Hash)]
|
||||
#[derive(Debug, Eq, PartialEq, Hash)]
|
||||
pub(crate) enum ConnectAddrs {
|
||||
None,
|
||||
One(SocketAddr),
|
||||
Multi(VecDeque<SocketAddr>),
|
||||
}
|
||||
|
||||
impl ConnectAddrs {
|
||||
pub(crate) fn is_none(&self) -> bool {
|
||||
matches!(self, Self::None)
|
||||
}
|
||||
|
||||
pub(crate) fn is_some(&self) -> bool {
|
||||
!self.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ConnectAddrs {
|
||||
fn default() -> Self {
|
||||
Self::None
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Option<SocketAddr>> for ConnectAddrs {
|
||||
fn from(addr: Option<SocketAddr>) -> Self {
|
||||
match addr {
|
||||
Some(addr) => ConnectAddrs::One(addr),
|
||||
None => ConnectAddrs::None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Connection info.
|
||||
#[derive(Debug, PartialEq, Eq, Hash)]
|
||||
pub struct Connect<T> {
|
||||
pub(crate) req: T,
|
||||
pub(crate) port: u16,
|
||||
pub(crate) addr: ConnectAddrs,
|
||||
}
|
||||
|
||||
#[derive(Eq, PartialEq, Debug, Hash)]
|
||||
pub(crate) enum ConnectAddrs {
|
||||
One(Option<SocketAddr>),
|
||||
Multi(VecDeque<SocketAddr>),
|
||||
}
|
||||
|
||||
impl ConnectAddrs {
|
||||
pub(crate) fn is_none(&self) -> bool {
|
||||
matches!(*self, Self::One(None))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ConnectAddrs {
|
||||
fn default() -> Self {
|
||||
Self::One(None)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Address> Connect<T> {
|
||||
/// Create `Connect` instance by splitting the string by ':' and convert the second part to u16
|
||||
pub fn new(req: T) -> Connect<T> {
|
||||
let (_, port) = parse(req.host());
|
||||
let (_, port) = parse_host(req.hostname());
|
||||
|
||||
Connect {
|
||||
req,
|
||||
port: port.unwrap_or(0),
|
||||
addr: ConnectAddrs::One(None),
|
||||
addr: ConnectAddrs::None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create new `Connect` instance from host and address. Connector skips name resolution stage
|
||||
/// for such connect messages.
|
||||
pub fn with(req: T, addr: SocketAddr) -> Connect<T> {
|
||||
pub fn with_addr(req: T, addr: SocketAddr) -> Connect<T> {
|
||||
Connect {
|
||||
req,
|
||||
port: 0,
|
||||
addr: ConnectAddrs::One(Some(addr)),
|
||||
addr: ConnectAddrs::One(addr),
|
||||
}
|
||||
}
|
||||
|
||||
/// Use port if address does not provide one.
|
||||
///
|
||||
/// By default it set to 0
|
||||
/// Default value is 0.
|
||||
pub fn set_port(mut self, port: u16) -> Self {
|
||||
self.port = port;
|
||||
self
|
||||
}
|
||||
|
||||
/// Use address.
|
||||
/// Set address.
|
||||
pub fn set_addr(mut self, addr: Option<SocketAddr>) -> Self {
|
||||
self.addr = ConnectAddrs::One(addr);
|
||||
self.addr = ConnectAddrs::from(addr);
|
||||
self
|
||||
}
|
||||
|
||||
/// Use addresses.
|
||||
/// Set list of addresses.
|
||||
pub fn set_addrs<I>(mut self, addrs: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = SocketAddr>,
|
||||
{
|
||||
let mut addrs = VecDeque::from_iter(addrs);
|
||||
self.addr = if addrs.len() < 2 {
|
||||
ConnectAddrs::One(addrs.pop_front())
|
||||
ConnectAddrs::from(addrs.pop_front())
|
||||
} else {
|
||||
ConnectAddrs::Multi(addrs)
|
||||
};
|
||||
self
|
||||
}
|
||||
|
||||
/// Host name
|
||||
pub fn host(&self) -> &str {
|
||||
self.req.host()
|
||||
/// Get hostname.
|
||||
pub fn hostname(&self) -> &str {
|
||||
self.req.hostname()
|
||||
}
|
||||
|
||||
/// Port of the request
|
||||
/// Get request port.
|
||||
pub fn port(&self) -> u16 {
|
||||
self.req.port().unwrap_or(self.port)
|
||||
}
|
||||
|
||||
/// Pre-resolved addresses of the request.
|
||||
/// Get resolved request addresses.
|
||||
pub fn addrs(&self) -> ConnectAddrsIter<'_> {
|
||||
match self.addr {
|
||||
ConnectAddrs::None => ConnectAddrsIter::None,
|
||||
ConnectAddrs::One(addr) => ConnectAddrsIter::One(addr),
|
||||
ConnectAddrs::Multi(ref addrs) => ConnectAddrsIter::Multi(addrs.iter()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Takes pre-resolved addresses of the request.
|
||||
/// Take resolved request addresses.
|
||||
pub fn take_addrs(&mut self) -> ConnectAddrsIter<'static> {
|
||||
match std::mem::take(&mut self.addr) {
|
||||
match mem::take(&mut self.addr) {
|
||||
ConnectAddrs::None => ConnectAddrsIter::None,
|
||||
ConnectAddrs::One(addr) => ConnectAddrsIter::One(addr),
|
||||
ConnectAddrs::Multi(addrs) => ConnectAddrsIter::MultiOwned(addrs.into_iter()),
|
||||
}
|
||||
@@ -144,14 +156,15 @@ impl<T: Address> From<T> for Connect<T> {
|
||||
|
||||
impl<T: Address> fmt::Display for Connect<T> {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}:{}", self.host(), self.port())
|
||||
write!(f, "{}:{}", self.hostname(), self.port())
|
||||
}
|
||||
}
|
||||
|
||||
/// Iterator over addresses in a [`Connect`] request.
|
||||
#[derive(Clone)]
|
||||
pub enum ConnectAddrsIter<'a> {
|
||||
One(Option<SocketAddr>),
|
||||
None,
|
||||
One(SocketAddr),
|
||||
Multi(vec_deque::Iter<'a, SocketAddr>),
|
||||
MultiOwned(vec_deque::IntoIter<SocketAddr>),
|
||||
}
|
||||
@@ -161,7 +174,8 @@ impl Iterator for ConnectAddrsIter<'_> {
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
match *self {
|
||||
Self::One(ref mut addr) => addr.take(),
|
||||
Self::None => None,
|
||||
Self::One(addr) => Some(addr),
|
||||
Self::Multi(ref mut iter) => iter.next().copied(),
|
||||
Self::MultiOwned(ref mut iter) => iter.next(),
|
||||
}
|
||||
@@ -169,8 +183,8 @@ impl Iterator for ConnectAddrsIter<'_> {
|
||||
|
||||
fn size_hint(&self) -> (usize, Option<usize>) {
|
||||
match *self {
|
||||
Self::One(None) => (0, Some(0)),
|
||||
Self::One(Some(_)) => (1, Some(1)),
|
||||
Self::None => (0, Some(0)),
|
||||
Self::One(_) => (1, Some(1)),
|
||||
Self::Multi(ref iter) => iter.size_hint(),
|
||||
Self::MultiOwned(ref iter) => iter.size_hint(),
|
||||
}
|
||||
@@ -183,23 +197,9 @@ impl fmt::Debug for ConnectAddrsIter<'_> {
|
||||
}
|
||||
}
|
||||
|
||||
impl ExactSizeIterator for ConnectAddrsIter<'_> {}
|
||||
impl iter::ExactSizeIterator for ConnectAddrsIter<'_> {}
|
||||
|
||||
impl FusedIterator for ConnectAddrsIter<'_> {}
|
||||
|
||||
fn parse(host: &str) -> (&str, Option<u16>) {
|
||||
let mut parts_iter = host.splitn(2, ':');
|
||||
if let Some(host) = parts_iter.next() {
|
||||
let port_str = parts_iter.next().unwrap_or("");
|
||||
if let Ok(port) = port_str.parse::<u16>() {
|
||||
(host, Some(port))
|
||||
} else {
|
||||
(host, None)
|
||||
}
|
||||
} else {
|
||||
(host, None)
|
||||
}
|
||||
}
|
||||
impl iter::FusedIterator for ConnectAddrsIter<'_> {}
|
||||
|
||||
pub struct Connection<T, U> {
|
||||
io: U,
|
||||
@@ -224,25 +224,25 @@ impl<T, U> Connection<T, U> {
|
||||
}
|
||||
|
||||
/// Replace inclosed object, return new Stream and old object
|
||||
pub fn replace<Y>(self, io: Y) -> (U, Connection<T, Y>) {
|
||||
pub fn replace_io<Y>(self, io: Y) -> (U, Connection<T, Y>) {
|
||||
(self.io, Connection { io, req: self.req })
|
||||
}
|
||||
|
||||
/// Returns a shared reference to the underlying stream.
|
||||
pub fn get_ref(&self) -> &U {
|
||||
pub fn io_ref(&self) -> &U {
|
||||
&self.io
|
||||
}
|
||||
|
||||
/// Returns a mutable reference to the underlying stream.
|
||||
pub fn get_mut(&mut self) -> &mut U {
|
||||
pub fn io_mut(&mut self) -> &mut U {
|
||||
&mut self.io
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Address, U> Connection<T, U> {
|
||||
/// Get request
|
||||
/// Get hostname.
|
||||
pub fn host(&self) -> &str {
|
||||
&self.req.host()
|
||||
self.req.hostname()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -265,3 +265,31 @@ impl<T, U: fmt::Debug> fmt::Debug for Connection<T, U> {
|
||||
write!(f, "Stream {{{:?}}}", self.io)
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_host(host: &str) -> (&str, Option<u16>) {
|
||||
let mut parts_iter = host.splitn(2, ':');
|
||||
|
||||
match parts_iter.next() {
|
||||
Some(hostname) => {
|
||||
let port_str = parts_iter.next().unwrap_or("");
|
||||
let port = port_str.parse::<u16>().ok();
|
||||
(hostname, port)
|
||||
}
|
||||
|
||||
None => (host, None),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_host_parser() {
|
||||
assert_eq!(parse_host("example.com"), ("example.com", None));
|
||||
assert_eq!(parse_host("example.com:8080"), ("example.com", Some(8080)));
|
||||
assert_eq!(parse_host("example:8080"), ("example", Some(8080)));
|
||||
assert_eq!(parse_host("example.com:false"), ("example.com", None));
|
||||
assert_eq!(parse_host("example.com:false:false"), ("example.com", None));
|
||||
}
|
||||
}
|
||||
|
Reference in New Issue
Block a user