warp-tech · Eugeny · Dec 27, 2024 · Dec 26, 2024 · Dec 26, 2024 · Dec 26, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -27,6 +27,8 @@ serde_json = "1.0"
 russh = { version = "0.49.0" }
 russh-keys = { version = "0.49.0" }
 tracing = "0.1"
+futures = "0.3"
+tokio-stream = { version = "0.1.17", features = ["net"] }
 
 [profile.release]
 lto = true

diff --git a/tests/test_ssh_proto.py b/tests/test_ssh_proto.py
@@ -167,7 +167,7 @@ def test_direct_tcpip(
             "-v",
             *common_args,
             "-L",
-            f"{local_port}:neverssl.com:80",
+            f"{local_port}:github.com:443",
             "-N",
             password="123",
         )
@@ -178,8 +178,8 @@ def test_direct_tcpip(
 
         s = requests.Session()
         retries = requests.adapters.Retry(total=5, backoff_factor=1)
-        s.mount("http://", requests.adapters.HTTPAdapter(max_retries=retries))
-        response = s.get(f"http://localhost:{local_port}", timeout=timeout)
+        s.mount("https://", requests.adapters.HTTPAdapter(max_retries=retries))
+        response = s.get(f"https://localhost:{local_port}", timeout=timeout, verify=False)
         assert response.status_code == 200
         ssh_client.kill()
 
@@ -193,32 +193,35 @@ def test_tcpip_forward(
         user, ssh_target = setup_user_and_target(
             processes, shared_wg, wg_c_ed25519_pubkey
         )
+        fw_port = alloc_port()
         pf_client = processes.start_ssh_client(
             f"{user.username}:{ssh_target.name}@localhost",
             "-p",
             str(shared_wg.ssh_port),
             "-v",
             *common_args,
             "-R",
-            "1234:neverssl.com:80",
+            f"{fw_port}:www.google.com:443",
             "-N",
             password="123",
         )
-        time.sleep(5)
+        # time.sleep(5)
         ssh_client = processes.start_ssh_client(
             f"{user.username}:{ssh_target.name}@localhost",
             "-p",
             str(shared_wg.ssh_port),
             "-v",
             *common_args,
             "curl",
-            "-v",
-            "http://localhost:1234",
+            "-vk",
+            "--http1.1",
+            "-H", "Host: www.google.com",
+            f"https://localhost:{fw_port}",
             password="123",
         )
         output = ssh_client.communicate(timeout=timeout)[0]
         assert ssh_client.returncode == 0
-        assert b"<html>" in output
+        assert b"</html>" in output
         pf_client.kill()
 
     def test_shell(

diff --git a/warpgate-admin/Cargo.toml b/warpgate-admin/Cargo.toml
@@ -9,7 +9,7 @@ anyhow = { version = "1.0", features = ["std"] }
 async-trait = "0.1"
 bytes.workspace = true
 chrono = { version = "0.4", default-features = false }
-futures = "0.3"
+futures.workspace = true
 hex = "0.4"
 mime_guess = { version = "2.0", default-features = false }
 poem = { version = "3.1", features = [

diff --git a/warpgate-common/Cargo.toml b/warpgate-common/Cargo.toml
@@ -13,7 +13,7 @@ chrono = { version = "0.4", default-features = false, features = ["serde"] }
 data-encoding.workspace = true
 delegate = "0.6"
 humantime-serde = "1.1"
-futures = "0.3"
+futures.workspace = true
 once_cell = "1.17"
 password-hash = "0.4"
 poem = { version = "3.1", features = ["rustls"] }
@@ -47,3 +47,4 @@ rustls = "0.23"
 rustls-pemfile = "1.0"
 webpki = "0.22"
 aho-corasick = "1.1.3"
+tokio-stream.workspace = true
diff --git a/warpgate-common/src/config/defaults.rs b/warpgate-common/src/config/defaults.rs
@@ -41,17 +41,17 @@ pub(crate) fn _default_database_url() -> Secret<String> {
 
 #[inline]
 pub(crate) fn _default_http_listen() -> ListenEndpoint {
-    ListenEndpoint(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 8888))
+    ListenEndpoint::from(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 8888))
 }
 
 #[inline]
 pub(crate) fn _default_mysql_listen() -> ListenEndpoint {
-    ListenEndpoint(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 33306))
+    ListenEndpoint::from(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 33306))
 }
 
 #[inline]
 pub(crate) fn _default_postgres_listen() -> ListenEndpoint {
-    ListenEndpoint(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 55432))
+    ListenEndpoint::from(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 55432))
 }
 
 #[inline]
@@ -75,7 +75,7 @@ pub(crate) fn _default_empty_vec<T>() -> Vec<T> {
 }
 
 pub(crate) fn _default_ssh_listen() -> ListenEndpoint {
-    ListenEndpoint(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 2222))
+    ListenEndpoint::from(SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), 2222))
 }
 
 pub(crate) fn _default_ssh_keys_path() -> String {

diff --git a/warpgate-common/src/error.rs b/warpgate-common/src/error.rs
@@ -37,6 +37,8 @@ pub enum WarpgateError {
     Sso(#[from] SsoError),
     #[error(transparent)]
     RusshKeys(#[from] russh_keys::Error),
+    #[error("I/O: {0}")]
+    Io(#[from] std::io::Error),
 
     #[error("Session end")]
     SessionEnd,

diff --git a/warpgate-common/src/types/listen_endpoint.rs b/warpgate-common/src/types/listen_endpoint.rs
@@ -1,16 +1,85 @@
 use std::fmt::Debug;
-use std::net::{SocketAddr, ToSocketAddrs};
-use std::ops::Deref;
+use std::io::ErrorKind;
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
 
+use futures::stream::{iter, FuturesUnordered};
+use futures::{Stream, StreamExt, TryStreamExt};
+use poem::listener::Listener;
 use serde::{Deserialize, Serialize};
+use tokio::net::{TcpListener, TcpStream};
+use tokio_stream::wrappers::TcpListenerStream;
+
+use crate::WarpgateError;
 
 #[derive(Clone)]
-pub struct ListenEndpoint(pub SocketAddr);
+pub struct ListenEndpoint(SocketAddr);
+
+impl ListenEndpoint {
+    pub fn addresses_to_listen_on(&self) -> Result<Vec<SocketAddr>, WarpgateError> {
+        // For [::], explicitly return both addresses so that we are not affected
+        // by the state of the ipv6only sysctl.
+        if self.0.ip() == Ipv6Addr::UNSPECIFIED {
+            let addr6 = SocketAddr::new(Ipv6Addr::UNSPECIFIED.into(), self.0.port());
+            let addr4 = SocketAddr::new(Ipv4Addr::UNSPECIFIED.into(), self.0.port());
+            let listener6 = std::net::TcpListener::bind(addr6)?;
+            let listener4 = std::net::TcpListener::bind(addr4);
+            let result = match listener4 {
+                Ok(_) => vec![addr4, addr6],
+                Err(e) if e.kind() == ErrorKind::AddrInUse => vec![addr6],
+                Err(e) => return Err(WarpgateError::Io(e)),
+            };
+            drop(listener6);
+            Ok(result)
+        } else {
+            Ok(vec![self.0])
+        }
+    }
+
+    pub async fn tcp_listeners(&self) -> Result<Vec<TcpListener>, WarpgateError> {
+        Ok(self
+            .addresses_to_listen_on()?
+            .into_iter()
+            .map(TcpListener::bind)
+            .collect::<FuturesUnordered<_>>()
+            .try_collect()
+            .await?)
+    }
+
+    pub async fn poem_listener(&self) -> Result<poem::listener::BoxListener, WarpgateError> {
+        let addrs = self.addresses_to_listen_on()?;
+        #[allow(clippy::unwrap_used)] // length known >=1
+        let (first, rest) = addrs.split_first().unwrap();
+        let mut listener: poem::listener::BoxListener =
+            poem::listener::TcpListener::bind(first.to_string()).boxed();
+        for addr in rest {
+            listener = listener
+                .combine(poem::listener::TcpListener::bind(addr.to_string()))
+                .boxed();
+        }
+
+        Ok(listener)
+    }
+
+    pub async fn tcp_accept_stream(
+        &self,
+    ) -> Result<impl Stream<Item = std::io::Result<TcpStream>>, WarpgateError> {
+        Ok(iter(
+            self.tcp_listeners()
+                .await?
+                .into_iter()
+                .map(TcpListenerStream::new),
+        )
+        .flatten_unordered(None))
+    }
+
+    pub fn port(&self) -> u16 {
+        self.0.port()
+    }
+}
 
-impl Deref for ListenEndpoint {
-    type Target = SocketAddr;
-    fn deref(&self) -> &Self::Target {
-        &self.0
+impl From<SocketAddr> for ListenEndpoint {
+    fn from(addr: SocketAddr) -> Self {
+        Self(addr)
     }
 }
 

diff --git a/warpgate-core/Cargo.toml b/warpgate-core/Cargo.toml
@@ -16,7 +16,7 @@ bytes.workspace = true
 chrono = { version = "0.4", default-features = false, features = ["serde"] }
 data-encoding.workspace = true
 humantime-serde = "1.1"
-futures = "0.3"
+futures.workspace = true
 once_cell = "1.17"
 packet = "0.1"
 password-hash = "0.4"

diff --git a/warpgate-core/src/protocols/mod.rs b/warpgate-core/src/protocols/mod.rs
@@ -1,10 +1,9 @@
 mod handle;
-use std::net::SocketAddr;
 
 use anyhow::Result;
 use async_trait::async_trait;
 pub use handle::{SessionHandle, WarpgateServerHandle};
-use warpgate_common::Target;
+use warpgate_common::{ListenEndpoint, Target};
 
 #[derive(Debug, thiserror::Error)]
 pub enum TargetTestError {
@@ -22,6 +21,6 @@ pub enum TargetTestError {
 
 #[async_trait]
 pub trait ProtocolServer {
-    async fn run(self, address: SocketAddr) -> Result<()>;
+    async fn run(self, address: ListenEndpoint) -> Result<()>;
     async fn test_target(&self, target: Target) -> Result<(), TargetTestError>;
 }
diff --git a/warpgate-protocol-http/Cargo.toml b/warpgate-protocol-http/Cargo.toml
@@ -11,7 +11,7 @@ chrono = { version = "0.4", default-features = false, features = ["serde"] }
 cookie = "0.17"
 data-encoding.workspace = true
 delegate = "0.6"
-futures = "0.3"
+futures.workspace = true
 http = "1.0"
 once_cell = "1.17"
 poem = { version = "3.1", features = [

diff --git a/warpgate-protocol-http/src/lib.rs b/warpgate-protocol-http/src/lib.rs
@@ -10,7 +10,6 @@ mod session;
 mod session_handle;
 
 use std::fmt::Debug;
-use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -21,7 +20,7 @@ pub use common::{SsoLoginState, PROTOCOL_NAME};
 use http::HeaderValue;
 use logging::{get_client_ip, log_request_error, log_request_result, span_for_request};
 use poem::endpoint::{EmbeddedFileEndpoint, EmbeddedFilesEndpoint};
-use poem::listener::{Listener, RustlsConfig, TcpListener};
+use poem::listener::{Listener, RustlsConfig};
 use poem::middleware::SetHeader;
 use poem::session::{CookieConfig, MemoryStorage, ServerSession, Session};
 use poem::web::Data;
@@ -31,7 +30,8 @@ use tokio::sync::Mutex;
 use tracing::*;
 use warpgate_admin::admin_api_app;
 use warpgate_common::{
-    Target, TargetOptions, TlsCertificateAndPrivateKey, TlsCertificateBundle, TlsPrivateKey,
+    ListenEndpoint, Target, TargetOptions, TlsCertificateAndPrivateKey, TlsCertificateBundle,
+    TlsPrivateKey,
 };
 use warpgate_core::{ProtocolServer, Services, TargetTestError};
 use warpgate_web::Assets;
@@ -59,7 +59,7 @@ fn make_session_storage() -> SharedSessionStorage {
 
 #[async_trait]
 impl ProtocolServer for HTTPProtocolServer {
-    async fn run(self, address: SocketAddr) -> Result<()> {
+    async fn run(self, address: ListenEndpoint) -> Result<()> {
         let admin_api_app = admin_api_app(&self.services).into_endpoint();
         let api_service = OpenApiService::new(
             crate::api::get(),
@@ -205,7 +205,9 @@ impl ProtocolServer for HTTPProtocolServer {
 
         info!(?address, "Listening");
         Server::new(
-            TcpListener::bind(address)
+            address
+                .poem_listener()
+                .await?
                 .rustls(RustlsConfig::new().fallback(certificate_and_key.into())),
         )
         .run(app)

diff --git a/warpgate-protocol-mysql/Cargo.toml b/warpgate-protocol-mysql/Cargo.toml
@@ -11,6 +11,7 @@ warpgate-db-entities = { version = "*", path = "../warpgate-db-entities" }
 warpgate-database-protocols = { version = "*", path = "../warpgate-database-protocols" }
 anyhow = { version = "1.0", features = ["std"] }
 async-trait = "0.1"
+futures.workspace = true
 tokio = { version = "1.20", features = ["tracing", "signal"] }
 tracing.workspace = true
 uuid = { version = "1.3", features = ["v4"] }