From d69439ff24649631b32d7d94a408a6cee3f59f94 Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Wed, 22 Dec 2021 11:33:55 +0100
Subject: [PATCH] Abstract EppConnection over stream type

---
 src/client.rs     |  5 ++++-
 src/connection.rs | 13 ++++++++-----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 47d4259..b37c9c7 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -35,6 +35,9 @@
 
 use std::net::SocketAddr;
 
+use tokio::net::TcpStream;
+use tokio_rustls::client::TlsStream;
+
 use crate::common::{Certificate, NoExtension, PrivateKey};
 use crate::connection::EppConnection;
 use crate::error::Error;
@@ -47,7 +50,7 @@ use crate::xml::EppXml;
 /// Once initialized, the EppClient instance can serialize EPP requests to XML and send them
 /// to the registry and deserialize the XML responses from the registry to local types
 pub struct EppClient {
-    connection: EppConnection,
+    connection: EppConnection<TlsStream<TcpStream>>,
 }
 
 impl EppClient {
diff --git a/src/connection.rs b/src/connection.rs
index 2876b9f..dc763b0 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -6,7 +6,8 @@ use std::sync::Arc;
 use std::{io, str, u32};
 
 use rustls::{OwnedTrustAnchor, RootCertStore};
-use tokio::{io::AsyncReadExt, io::AsyncWriteExt, net::TcpStream};
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+use tokio::net::TcpStream;
 use tokio_rustls::{client::TlsStream, rustls::ClientConfig, TlsConnector};
 use tracing::{debug, info};
 
@@ -14,20 +15,20 @@ use crate::common::{Certificate, PrivateKey};
 use crate::error::Error;
 
 /// EPP Connection struct with some metadata for the connection
-pub(crate) struct EppConnection {
+pub(crate) struct EppConnection<IO> {
     registry: String,
-    stream: TlsStream<TcpStream>,
+    stream: IO,
     pub greeting: String,
 }
 
-impl EppConnection {
+impl EppConnection<TlsStream<TcpStream>> {
     /// Create an EppConnection instance with the stream to the registry
     pub(crate) async fn connect(
         registry: String,
         addr: SocketAddr,
         hostname: &str,
         identity: Option<(Vec<Certificate>, PrivateKey)>,
-    ) -> Result<EppConnection, Error> {
+    ) -> Result<Self, Error> {
         let mut stream = epp_connect(addr, hostname, identity).await?;
 
         let mut buf = vec![0u8; 4096];
@@ -42,7 +43,9 @@ impl EppConnection {
             greeting,
         })
     }
+}
 
+impl<IO: AsyncRead + AsyncWrite + Unpin> EppConnection<IO> {
     /// Constructs an EPP XML request in the required form and sends it to the server
     async fn send_epp_request(&mut self, content: &str) -> Result<(), Error> {
         let len = content.len();