1
0
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:
Rob Ede
2021-01-26 08:05:19 +00:00
parent 636cef8868
commit eaefe21b98
13 changed files with 294 additions and 217 deletions

View File

@@ -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));
}
}