Allow the Connection to connect itself

This commit is contained in:
Dirkjan Ochtman 2022-02-04 22:38:45 +01:00 committed by masalachai
parent 95b62891c9
commit 36558c429c
4 changed files with 119 additions and 67 deletions

View File

@ -11,6 +11,7 @@ repository = "https://github.com/masalachai/epp-client"
default = ["tokio-rustls"]
[dependencies]
async-trait = "0.1.52"
celes = "2.1"
chrono = "0.4"
quick-xml = { version = "0.22", features = [ "serialize" ] }

View File

@ -4,17 +4,18 @@ use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncWrite};
use async_trait::async_trait;
use tokio::net::TcpStream;
#[cfg(feature = "tokio-rustls")]
use tokio_rustls::client::TlsStream;
#[cfg(feature = "tokio-rustls")]
use tokio_rustls::rustls::{ClientConfig, OwnedTrustAnchor, RootCertStore};
use tokio_rustls::rustls::{ClientConfig, OwnedTrustAnchor, RootCertStore, ServerName};
#[cfg(feature = "tokio-rustls")]
use tokio_rustls::TlsConnector;
use tracing::info;
use crate::common::{Certificate, NoExtension, PrivateKey};
pub use crate::connection::Connector;
use crate::connection::{self, EppConnection};
use crate::error::Error;
use crate::hello::{Greeting, GreetingDocument, HelloDocument};
@ -68,12 +69,12 @@ use crate::xml::EppXml;
/// Domain: eppdev.com, Available: 1
/// Domain: eppdev.net, Available: 1
/// ```
pub struct EppClient<IO> {
connection: EppConnection<IO>,
pub struct EppClient<C: Connector> {
connection: EppConnection<C>,
}
#[cfg(feature = "tokio-rustls")]
impl EppClient<TlsStream<TcpStream>> {
impl EppClient<RustlsConnector> {
/// Connect to the specified `addr` and `hostname` over TLS
///
/// The `registry` is used as a name in internal logging; `addr` provides the address to
@ -91,52 +92,16 @@ impl EppClient<TlsStream<TcpStream>> {
timeout: Duration,
) -> Result<Self, Error> {
info!("Connecting to server: {:?}", addr);
let mut roots = RootCertStore::empty();
roots.add_server_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.0.iter().map(|ta| {
OwnedTrustAnchor::from_subject_spki_name_constraints(
ta.subject,
ta.spki,
ta.name_constraints,
)
}));
let builder = ClientConfig::builder()
.with_safe_defaults()
.with_root_certificates(roots);
let config = match identity {
Some((certs, key)) => {
let certs = certs
.into_iter()
.map(|cert| tokio_rustls::rustls::Certificate(cert.0))
.collect();
builder
.with_single_cert(certs, tokio_rustls::rustls::PrivateKey(key.0))
.map_err(|e| Error::Other(e.into()))?
}
None => builder.with_no_client_auth(),
};
let domain = hostname.try_into().map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
format!("Invalid domain: {}", hostname),
)
})?;
let connector = TlsConnector::from(Arc::new(config));
let future = connector.connect(domain, TcpStream::connect(&addr).await?);
let stream = connection::timeout(timeout, future).await?;
Self::new(registry, stream, timeout).await
let connector = RustlsConnector::new(addr, hostname, identity)?;
Self::new(connector, registry, timeout).await
}
}
impl<IO: AsyncRead + AsyncWrite + Unpin> EppClient<IO> {
impl<C: Connector> EppClient<C> {
/// Create an `EppClient` from an already established connection
pub async fn new(registry: String, stream: IO, timeout: Duration) -> Result<Self, Error> {
pub async fn new(connector: C, registry: String, timeout: Duration) -> Result<Self, Error> {
Ok(Self {
connection: EppConnection::new(registry, stream, timeout).await?,
connection: EppConnection::new(connector, registry, timeout).await?,
})
}
@ -149,21 +114,22 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> EppClient<IO> {
Ok(GreetingDocument::deserialize(&response)?.data)
}
pub async fn transact<'c, 'e, C, E>(
pub async fn transact<'c, 'e, Cmd, Ext>(
&mut self,
data: impl Into<RequestData<'c, 'e, C, E>>,
data: impl Into<RequestData<'c, 'e, Cmd, Ext>>,
id: &str,
) -> Result<Response<C::Response, E::Response>, Error>
) -> Result<Response<Cmd::Response, Ext::Response>, Error>
where
C: Transaction<E> + Command + 'c,
E: Extension + 'e,
Cmd: Transaction<Ext> + Command + 'c,
Ext: Extension + 'e,
{
let data = data.into();
let epp_xml = <C as Transaction<E>>::serialize_request(data.command, data.extension, id)?;
let epp_xml =
<Cmd as Transaction<Ext>>::serialize_request(data.command, data.extension, id)?;
let response = self.connection.transact(&epp_xml).await?;
C::deserialize_response(&response)
Cmd::deserialize_response(&response)
}
/// Accepts raw EPP XML and returns the raw EPP XML response to it.
@ -209,3 +175,69 @@ impl<'c, 'e, C: Command, E: Extension> From<(&'c C, &'e E)> for RequestData<'c,
}
}
}
#[cfg(feature = "tokio-rustls")]
pub struct RustlsConnector {
inner: TlsConnector,
domain: ServerName,
addr: SocketAddr,
}
impl RustlsConnector {
pub fn new(
addr: SocketAddr,
hostname: &str,
identity: Option<(Vec<Certificate>, PrivateKey)>,
) -> Result<Self, Error> {
let mut roots = RootCertStore::empty();
roots.add_server_trust_anchors(webpki_roots::TLS_SERVER_ROOTS.0.iter().map(|ta| {
OwnedTrustAnchor::from_subject_spki_name_constraints(
ta.subject,
ta.spki,
ta.name_constraints,
)
}));
let builder = ClientConfig::builder()
.with_safe_defaults()
.with_root_certificates(roots);
let config = match identity {
Some((certs, key)) => {
let certs = certs
.into_iter()
.map(|cert| tokio_rustls::rustls::Certificate(cert.0))
.collect();
builder
.with_single_cert(certs, tokio_rustls::rustls::PrivateKey(key.0))
.map_err(|e| Error::Other(e.into()))?
}
None => builder.with_no_client_auth(),
};
let domain = hostname.try_into().map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
format!("Invalid domain: {}", hostname),
)
})?;
Ok(Self {
inner: TlsConnector::from(Arc::new(config)),
domain,
addr,
})
}
}
#[cfg(feature = "tokio-rustls")]
#[async_trait]
impl Connector for RustlsConnector {
type Connection = TlsStream<TcpStream>;
async fn connect(&self, timeout: Duration) -> Result<Self::Connection, Error> {
let stream = TcpStream::connect(&self.addr).await?;
let future = self.inner.connect(self.domain.clone(), stream);
connection::timeout(timeout, future).await
}
}

View File

@ -5,28 +5,29 @@ use std::future::Future;
use std::time::Duration;
use std::{io, str, u32};
use async_trait::async_trait;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tracing::{debug, info};
use crate::error::Error;
/// EPP Connection struct with some metadata for the connection
pub(crate) struct EppConnection<IO> {
pub(crate) struct EppConnection<C: Connector> {
registry: String,
stream: IO,
stream: C::Connection,
pub greeting: String,
timeout: Duration,
}
impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
impl<C: Connector> EppConnection<C> {
pub(crate) async fn new(
connector: C,
registry: String,
stream: IO,
timeout: Duration,
) -> Result<Self, Error> {
let mut this = Self {
registry,
stream,
stream: connector.connect(timeout).await?,
greeting: String::new(),
timeout,
};
@ -118,3 +119,10 @@ pub(crate) async fn timeout<T, E: Into<Error>>(
Err(_) => Err(Error::Timeout),
}
}
#[async_trait]
pub trait Connector {
type Connection: AsyncRead + AsyncWrite + Unpin;
async fn connect(&self, timeout: Duration) -> Result<Self::Connection, Error>;
}

View File

@ -3,6 +3,7 @@ use std::io::{self, Read, Write};
use std::str;
use std::time::Duration;
use async_trait::async_trait;
use regex::Regex;
use tokio_test::io::Builder;
@ -78,16 +79,26 @@ fn build_stream(units: &[&str]) -> Builder {
#[tokio::test]
async fn client() {
let _guard = log_to_stdout();
let stream = build_stream(&[
"response/greeting.xml",
"request/login.xml",
"response/login.xml",
"request/domain/check.xml",
"response/domain/check.xml",
])
.build();
let mut client = EppClient::new("test".into(), stream, Duration::from_secs(5))
struct FakeConnector;
#[async_trait]
impl epp_client::client::Connector for FakeConnector {
type Connection = tokio_test::io::Mock;
async fn connect(&self, _: Duration) -> Result<Self::Connection, epp_client::Error> {
Ok(build_stream(&[
"response/greeting.xml",
"request/login.xml",
"response/login.xml",
"request/domain/check.xml",
"response/domain/check.xml",
])
.build())
}
}
let mut client = EppClient::new(FakeConnector, "test".into(), Duration::from_secs(5))
.await
.unwrap();