Add `Forwarder::from_stream`

Make `Forwarder` more flexible by adding `Forwarder::from_stream`
which allows the consumer to bring their own socket connection.

1 files changed, 16 insertions, 5 deletions

diff --git a/mullvad-encrypted-dns-proxy/src/forwarder.rs b/mullvad-encrypted-dns-proxy/src/forwarder.rs
index f607d10978..0cac1f6072 100644
--- a/mullvad-encrypted-dns-proxy/src/forwarder.rs
+++ b/mullvad-encrypted-dns-proxy/src/forwarder.rs
@@ -19,10 +19,12 @@ pub struct Forwarder<S> {
     stream: S,
 }
 
-impl Forwarder<TcpStream> {
-    /// Create a forwarder that will connect to a given proxy endpoint.
-    pub async fn connect(proxy_config: &crate::config::ProxyConfig) -> io::Result<Self> {
-        let server_connection = TcpStream::connect(proxy_config.addr).await?;
+impl<S> Forwarder<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
+    /// Create a [`Forwarder`] with a connected `stream` to an encrypted DNS proxy server
+    pub fn from_stream(proxy_config: &crate::config::ProxyConfig, stream: S) -> io::Result<Self> {
         let (read_obfuscator, write_obfuscator) =
             if let Some(obfuscation_config) = &proxy_config.obfuscation {
                 (
@@ -36,9 +38,18 @@ impl Forwarder<TcpStream> {
         Ok(Self {
             read_obfuscator,
             write_obfuscator,
-            stream: server_connection,
+            stream,
         })
     }
+}
+
+/// Forward TCP traffic over various proxy configurations.
+impl Forwarder<TcpStream> {
+    /// Create a forwarder that will connect to a given proxy endpoint.
+    pub async fn connect(proxy_config: &crate::config::ProxyConfig) -> io::Result<Self> {
+        let server_connection = TcpStream::connect(proxy_config.addr).await?;
+        Self::from_stream(proxy_config, server_connection)
+    }
 
     /// Forwards traffic from the client stream to the remote proxy, obfuscating and deobfuscating
     /// it in the process.