Add support for non-TLS API connections

1 files changed, 66 insertions, 33 deletions

diff --git a/mullvad-api/src/proxy.rs b/mullvad-api/src/proxy.rs
index 2f3764e7e6..fa1da913ba 100644
--- a/mullvad-api/src/proxy.rs
+++ b/mullvad-api/src/proxy.rs
@@ -1,8 +1,6 @@
-use crate::tls_stream::TlsStream;
 use futures::Stream;
-use hyper::client::connect::{Connected, Connection};
+use hyper::client::connect::Connected;
 use serde::{Deserialize, Serialize};
-use shadowsocks::relay::tcprelay::ProxyClientStream;
 use std::{
     fmt, io,
     net::SocketAddr,
@@ -14,7 +12,6 @@ use talpid_types::{net::openvpn::ShadowsocksProxySettings, ErrorExt};
 use tokio::{
     fs,
     io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf},
-    net::TcpStream,
 };
 
 const CURRENT_CONFIG_FILENAME: &str = "api-endpoint.json";
@@ -130,57 +127,93 @@ impl ApiConnectionMode {
     }
 }
 
-/// Stream that is either a regular TLS stream or TLS via shadowsocks
-pub enum ApiConnection {
-    Direct(Box<TlsStream<TcpStream>>),
-    Proxied(Box<TlsStream<ProxyClientStream<TcpStream>>>),
+/// Implements `hyper::client::connect::Connection` by wrapping a type.
+pub struct ConnectionDecorator<T: AsyncRead + AsyncWrite>(pub T);
+
+impl<T: AsyncRead + AsyncWrite + Unpin> AsyncRead for ConnectionDecorator<T> {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_read(cx, buf)
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite + Unpin> AsyncWrite for ConnectionDecorator<T> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut self.0).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_flush(cx)
+    }
+
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_shutdown(cx)
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite> hyper::client::connect::Connection for ConnectionDecorator<T> {
+    fn connected(&self) -> Connected {
+        Connected::new()
+    }
+}
+
+trait Connection: AsyncRead + AsyncWrite + Unpin + hyper::client::connect::Connection + Send {}
+
+impl<T: AsyncRead + AsyncWrite + Unpin + hyper::client::connect::Connection + Send> Connection
+    for T
+{
+}
+
+/// Stream that represents a Mullvad API connection
+pub struct ApiConnection(Box<dyn Connection>);
+
+impl ApiConnection {
+    pub fn new<
+        T: AsyncRead + AsyncWrite + Unpin + hyper::client::connect::Connection + Send + 'static,
+    >(
+        conn: Box<T>,
+    ) -> Self {
+        Self(conn)
+    }
 }
 
 impl AsyncRead for ApiConnection {
     fn poll_read(
-        self: Pin<&mut Self>,
+        mut self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<()>> {
-        match Pin::get_mut(self) {
-            ApiConnection::Direct(s) => Pin::new(s).poll_read(cx, buf),
-            ApiConnection::Proxied(s) => Pin::new(s).poll_read(cx, buf),
-        }
+        Pin::new(&mut self.0).poll_read(cx, buf)
     }
 }
 
 impl AsyncWrite for ApiConnection {
     fn poll_write(
-        self: Pin<&mut Self>,
+        mut self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
         buf: &[u8],
     ) -> Poll<io::Result<usize>> {
-        match Pin::get_mut(self) {
-            ApiConnection::Direct(s) => Pin::new(s).poll_write(cx, buf),
-            ApiConnection::Proxied(s) => Pin::new(s).poll_write(cx, buf),
-        }
+        Pin::new(&mut self.0).poll_write(cx, buf)
     }
 
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
-        match Pin::get_mut(self) {
-            ApiConnection::Direct(s) => Pin::new(s).poll_flush(cx),
-            ApiConnection::Proxied(s) => Pin::new(s).poll_flush(cx),
-        }
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_flush(cx)
     }
 
-    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
-        match Pin::get_mut(self) {
-            ApiConnection::Direct(s) => Pin::new(s).poll_shutdown(cx),
-            ApiConnection::Proxied(s) => Pin::new(s).poll_shutdown(cx),
-        }
+    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.0).poll_shutdown(cx)
     }
 }
 
-impl Connection for ApiConnection {
+impl hyper::client::connect::Connection for ApiConnection {
     fn connected(&self) -> Connected {
-        match self {
-            ApiConnection::Direct(s) => s.connected(),
-            ApiConnection::Proxied(s) => s.connected(),
-        }
+        self.0.connected()
     }
 }