wip

Author: isabelatkinson

Cargo.lock

@@ -84,6 +84,10 @@ opentelemetry = ["dep:opentelemetry"]
 # Capture backtraces in errors.  This can be slow, memory intensive, and very verbose.
 error-backtrace = []
 
+# Enable SOCKS5 proxy support.
+socks5-proxy = ["dep:fast-socks5"]
+
+
 [dependencies]
 base64 = "0.22"
 bitflags = "2"
@@ -93,6 +97,7 @@ chrono = { version = "0.4.32", default-features = false, optional = true, featur
 ] }
 derive_more = { version = "2", features = ["display", "from"] }
 derive-where = "1.2.7"
+fast-socks5 = { version = "0.10", optional = true }
 flate2 = { version = "1.0", optional = true }
 futures-io = "0.3.21"
 futures-core = "0.3.14"

Cargo.toml

@@ -75,6 +75,10 @@ const URI_OPTIONS: &[&str] = &[
     "maxpoolsize",
     "minpoolsize",
     "maxconnecting",
+    "proxyHost",
+    "proxyPort",
+    "proxyUsername",
+    "proxyPassword",
     "readconcernlevel",
     "readpreference",
     "readpreferencetags",
@@ -616,6 +620,18 @@ pub struct ClientOptions {
     #[cfg(feature = "opentelemetry")]
     pub tracing: Option<crate::otel::OpentelemetryOptions>,
 
+    #[cfg(feature = "socks5-proxy")]
+    pub proxy_host: Option<String>,
+
+    #[cfg(feature = "socks5-proxy")]
+    pub proxy_port: Option<u16>,
+
+    #[cfg(feature = "socks5-proxy")]
+    pub proxy_username: Option<String>,
+
+    #[cfg(feature = "socks5-proxy")]
+    pub proxy_password: Option<String>,
+
     /// Information from the SRV URI that generated these client options, if applicable.
     #[builder(setter(skip))]
     #[serde(skip)]

src/client/options.rs

@@ -164,6 +164,15 @@ impl ClientOptions {
             srv_service_name: conn_str.srv_service_name,
             #[cfg(feature = "opentelemetry")]
             tracing: None,
+            // isabeltodo propagate from connstr
+            #[cfg(feature = "socks5-proxy")]
+            proxy_host: None,
+            #[cfg(feature = "socks5-proxy")]
+            proxy_port: None,
+            #[cfg(feature = "socks5-proxy")]
+            proxy_username: None,
+            #[cfg(feature = "socks5-proxy")]
+            proxy_password: None,
         }
     }
 }

src/client/options/parse.rs

@@ -40,6 +40,9 @@ pub(crate) struct ConnectionEstablisher {
 
     #[cfg(test)]
     test_patch_reply: Option<fn(&mut Result<HelloReply>)>,
+
+    #[cfg(feature = "socks5-proxy")]
+    proxy: Option<Proxy>,
 }
 
 pub(crate) struct EstablisherOptions {
@@ -48,6 +51,64 @@ pub(crate) struct EstablisherOptions {
     connect_timeout: Option<Duration>,
     #[cfg(test)]
     pub(crate) test_patch_reply: Option<fn(&mut Result<HelloReply>)>,
+    #[cfg(feature = "socks5-proxy")]
+    proxy: Option<Proxy>,
+}
+
+#[derive(Clone)]
+#[cfg(feature = "socks5-proxy")]
+struct Proxy {
+    host: String,
+    port: u16,
+    authentication: Option<(String, String)>,
+}
+
+#[cfg(feature = "socks5-proxy")]
+impl Proxy {
+    fn from_client_options(options: &crate::options::ClientOptions) -> Option<Self> {
+        let host = options.proxy_host.as_ref()?.clone();
+        let port = options.proxy_port.unwrap_or(1080);
+        let authentication = match (&options.proxy_username, &options.proxy_password) {
+            (Some(username), Some(password)) => Some((username.clone(), password.clone())),
+            // ClientOptions::validate will return an error if the username and password are not
+            // provided together or both omitted.
+            _ => None,
+        };
+        Some(Self {
+            host,
+            port,
+            authentication,
+        })
+    }
+
+    async fn connect(
+        &self,
+        address: ServerAddress,
+        connect_timeout: Duration,
+    ) -> Result<AsyncStream> {
+        use fast_socks5::client::{Config, Socks5Stream};
+
+        let mut config = Config::default();
+        config.set_connect_timeout(connect_timeout.as_secs());
+
+        let stream = if let Some((username, password)) = self.authentication.as_ref() {
+            Socks5Stream::connect_with_password(
+                address.to_string(),
+                self.host.clone(),
+                self.port,
+                username.clone(),
+                password.clone(),
+                config,
+            )
+            .await
+        } else {
+            Socks5Stream::connect(address.to_string(), self.host.clone(), self.port, config).await
+        }
+        .map_err(|error| ErrorKind::ProxyConnect {
+            message: error.to_string(),
+        })?;
+        Ok(AsyncStream::Socks5(stream))
+    }
 }
 
 impl From<&TopologySpec> for EstablisherOptions {
@@ -58,6 +119,8 @@ impl From<&TopologySpec> for EstablisherOptions {
             connect_timeout: spec.options.connect_timeout,
             #[cfg(test)]
             test_patch_reply: None,
+            #[cfg(feature = "socks5-proxy")]
+            proxy: Proxy::from_client_options(&spec.options),
         }
     }
 }
@@ -113,6 +176,8 @@ impl ConnectionEstablisher {
         let address = pending_connection.address.clone();
         let cancellation_receiver = pending_connection.cancellation_receiver.take();
 
+        if let Some(ref proxy) = self.proxy {}
+
         let stream = self
             .make_stream(address)
             .await

src/cmap/establish.rs

@@ -577,7 +577,8 @@ impl Error {
             | ErrorKind::Authentication { .. }
             | ErrorKind::Custom(_)
             | ErrorKind::Shutdown
-            | ErrorKind::GridFs(_) => {}
+            | ErrorKind::GridFs(_)
+            | ErrorKind::ProxyConnect { .. } => {}
             #[cfg(feature = "in-use-encryption")]
             ErrorKind::Encryption(_) => {}
             #[cfg(feature = "bson-3")]
@@ -766,6 +767,12 @@ pub enum ErrorKind {
     /// A method was called on a client that was shut down.
     #[error("Client has been shut down")]
     Shutdown,
+
+    /// An error occurred when connecting to a proxy host.
+    #[error("An error occurred when connecting to a proxy host: {message}")]
+    #[non_exhaustive]
+    #[cfg(feature = "socks5-proxy")]
+    ProxyConnect { message: String },
 }
 
 impl ErrorKind {
@@ -812,6 +819,7 @@ impl ErrorKind {
             ErrorKind::Encryption(..) => "Encryption",
             ErrorKind::Custom(..) => "Custom",
             ErrorKind::Shutdown => "Shutdown",
+            ErrorKind::ProxyConnect { .. } => "ProxyConnect",
         }
     }
 }

src/error.rs

@@ -38,6 +38,10 @@ pub(crate) enum AsyncStream {
     /// A Unix domain socket connection.
     #[cfg(unix)]
     Unix(tokio::net::UnixStream),
+
+    /// A connection to a SOCKS5 proxy.
+    #[cfg(feature = "socks5-proxy")]
+    Socks5(fast_socks5::client::Socks5Stream<TcpStream>),
 }
 
 impl AsyncStream {
@@ -175,6 +179,10 @@ impl tokio::io::AsyncRead for AsyncStream {
             Self::Tls(ref mut inner) => tokio::io::AsyncRead::poll_read(Pin::new(inner), cx, buf),
             #[cfg(unix)]
             Self::Unix(ref mut inner) => tokio::io::AsyncRead::poll_read(Pin::new(inner), cx, buf),
+            #[cfg(feature = "socks5-proxy")]
+            Self::Socks5(ref mut inner) => {
+                tokio::io::AsyncRead::poll_read(Pin::new(inner), cx, buf)
+            }
         }
     }
 }
@@ -191,6 +199,8 @@ impl AsyncWrite for AsyncStream {
             Self::Tls(ref mut inner) => Pin::new(inner).poll_write(cx, buf),
             #[cfg(unix)]
             Self::Unix(ref mut inner) => AsyncWrite::poll_write(Pin::new(inner), cx, buf),
+            #[cfg(feature = "socks5-proxy")]
+            Self::Socks5(ref mut inner) => AsyncWrite::poll_write(Pin::new(inner), cx, buf),
         }
     }
 
@@ -201,6 +211,8 @@ impl AsyncWrite for AsyncStream {
             Self::Tls(ref mut inner) => Pin::new(inner).poll_flush(cx),
             #[cfg(unix)]
             Self::Unix(ref mut inner) => AsyncWrite::poll_flush(Pin::new(inner), cx),
+            #[cfg(feature = "socks5-proxy")]
+            Self::Socks5(ref mut inner) => AsyncWrite::poll_flush(Pin::new(inner), cx),
         }
     }
 
@@ -211,6 +223,8 @@ impl AsyncWrite for AsyncStream {
             Self::Tls(ref mut inner) => Pin::new(inner).poll_shutdown(cx),
             #[cfg(unix)]
             Self::Unix(ref mut inner) => Pin::new(inner).poll_shutdown(cx),
+            #[cfg(feature = "socks5-proxy")]
+            Self::Socks5(ref mut inner) => Pin::new(inner).poll_shutdown(cx),
         }
     }
 
@@ -225,6 +239,8 @@ impl AsyncWrite for AsyncStream {
             Self::Tls(ref mut inner) => Pin::new(inner).poll_write_vectored(cx, bufs),
             #[cfg(unix)]
             Self::Unix(ref mut inner) => Pin::new(inner).poll_write_vectored(cx, bufs),
+            #[cfg(feature = "socks5-proxy")]
+            Self::Socks5(ref mut inner) => Pin::new(inner).poll_write_vectored(cx, bufs),
         }
     }
 
@@ -235,6 +251,8 @@ impl AsyncWrite for AsyncStream {
             Self::Tls(ref inner) => inner.is_write_vectored(),
             #[cfg(unix)]
             Self::Unix(ref inner) => inner.is_write_vectored(),
+            #[cfg(feature = "socks5-proxy")]
+            Self::Socks5(ref inner) => inner.is_write_vectored(),
         }
     }
 }