Remove lazy_connect option

src/tcp_options.rs

@@ -25,10 +25,6 @@ pub struct TcpOptions {
     #[cfg_attr(feature = "clap", arg(long = "tcp-recv-timeout", value_parser = duration_secs_from_str))]
     pub recv_timeout: Option<Duration>,
 
-    /// If true, wait for the first incoming UDP datagram before connecting to the TCP endpoint.
-    #[cfg_attr(feature = "clap", arg(skip))]
-    pub lazy_connect: bool,
-
     /// If given, sets the SO_MARK option on the TCP socket.
     /// This exists only on Linux.
     #[cfg(target_os = "linux")]

src/udp2tcp.rs

@@ -5,41 +5,37 @@ use crate::logging::Redact;
 use std::fmt;
 use std::io;
 use std::net::SocketAddr;
-use tokio::net::{TcpSocket, TcpStream, UdpSocket};
+use tokio::net::{TcpSocket, UdpSocket};
 
 #[cfg(unix)]
 use std::os::unix::io::{AsRawFd, RawFd};
 
 #[derive(Debug)]
-pub enum ConnectError {
+pub enum Error {
     /// Failed to create the TCP socket.
     CreateTcpSocket(io::Error),
-    /// Failed to connect to TCP forward address.
-    ConnectTcp(io::Error),
     /// Failed to apply the given TCP socket options.
     ApplyTcpOptions(crate::tcp_options::ApplyTcpOptionsError),
     /// Failed to bind UDP socket locally.
     BindUdp(io::Error),
 }
 
-impl fmt::Display for ConnectError {
+impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        use ConnectError::*;
+        use Error::*;
         match self {
             CreateTcpSocket(_) => "Failed to create the TCP socket".fmt(f),
-            ConnectTcp(_) => "Failed to connect to TCP forward address".fmt(f),
             ApplyTcpOptions(e) => e.fmt(f),
             BindUdp(_) => "Failed to bind UDP socket locally".fmt(f),
         }
     }
 }
 
-impl std::error::Error for ConnectError {
+impl std::error::Error for Error {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
-        use ConnectError::*;
+        use Error::*;
         match self {
             CreateTcpSocket(e) => Some(e),
-            ConnectTcp(e) => Some(e),
             ApplyTcpOptions(e) => e.source(),
             BindUdp(e) => Some(e),
         }
@@ -50,7 +46,7 @@ impl std::error::Error for ConnectError {
 pub enum ForwardError {
     ReadUdp(io::Error),
     ConnectUdp(io::Error),
-    ConnectTcp(ConnectError),
+    ConnectTcp(io::Error),
 }
 
 impl fmt::Display for ForwardError {
@@ -59,7 +55,7 @@ impl fmt::Display for ForwardError {
         match self {
             ReadUdp(_) => "Failed receiving the first UDP datagram".fmt(f),
             ConnectUdp(_) => "Failed to connect UDP socket to peer".fmt(f),
-            ConnectTcp(error) => error.fmt(f),
+            ConnectTcp(_) => "Failed to connect to TCP forward address".fmt(f),
         }
     }
 }
@@ -77,29 +73,31 @@ impl std::error::Error for ForwardError {
 
 /// Struct allowing listening on UDP and forwarding the traffic over TCP.
 pub struct Udp2Tcp {
-    tcp_socket: TcpForwardSocket,
+    tcp_socket: TcpSocket,
     udp_socket: UdpSocket,
     tcp_forward_addr: SocketAddr,
     tcp_options: crate::TcpOptions,
 }
 
 impl Udp2Tcp {
-    /// Connects to the given TCP address (unless `TcpOptions::lazy_connect` is true) and binds to
-    /// the given UDP address.
+    /// Creates a TCP socket and binds to the given UDP address.
     /// Just calling this constructor won't forward any traffic over the sockets (see `run`).
     pub async fn new(
         udp_listen_addr: SocketAddr,
         tcp_forward_addr: SocketAddr,
         tcp_options: crate::TcpOptions,
-    ) -> Result<Self, ConnectError> {
-        let mut tcp_socket = TcpForwardSocket::new(tcp_forward_addr, &tcp_options)?;
-        if !tcp_options.lazy_connect {
-            tcp_socket = TcpForwardSocket::Stream(tcp_socket.connect().await?);
+    ) -> Result<Self, Error> {
+        let tcp_socket = match &tcp_forward_addr {
+            SocketAddr::V4(..) => TcpSocket::new_v4(),
+            SocketAddr::V6(..) => TcpSocket::new_v6(),
         }
+        .map_err(Error::CreateTcpSocket)?;
+
+        crate::tcp_options::apply(&tcp_socket, &tcp_options).map_err(Error::ApplyTcpOptions)?;
 
         let udp_socket = UdpSocket::bind(udp_listen_addr)
             .await
-            .map_err(ConnectError::BindUdp)?;
+            .map_err(Error::BindUdp)?;
         match udp_socket.local_addr() {
             Ok(addr) => log::info!("Listening on {}/UDP", addr),
             Err(e) => log::error!("Unable to get UDP local addr: {}", e),
@@ -127,7 +125,8 @@ impl Udp2Tcp {
         self.tcp_socket.as_raw_fd()
     }
 
-    /// Runs the forwarding until the TCP socket is closed, or an error occur.
+    /// Connects to the TCP address and runs the forwarding until the TCP socket is closed, or
+    /// an error occur.
     pub async fn run(self) -> Result<(), ForwardError> {
         // Wait for the first datagram, to get the UDP peer_addr to connect to.
         let mut tmp_buffer = crate::forward_traffic::datagram_buffer();
@@ -138,11 +137,13 @@ impl Udp2Tcp {
             .map_err(ForwardError::ReadUdp)?;
         log::info!("Incoming connection from {}/UDP", Redact(udp_peer_addr));
 
+        log::info!("Connecting to {}/TCP", self.tcp_forward_addr);
         let tcp_stream = self
             .tcp_socket
-            .connect()
+            .connect(self.tcp_forward_addr)
             .await
             .map_err(ForwardError::ConnectTcp)?;
+        log::info!("Connected to {}/TCP", self.tcp_forward_addr);
 
         // Connect the UDP socket to whoever sent the first datagram. This is where
         // all the returned traffic will be sent to.
@@ -166,52 +167,3 @@ impl Udp2Tcp {
         Ok(())
     }
 }
-
-enum TcpForwardSocket {
-    Socket((TcpSocket, SocketAddr)),
-    Stream(TcpStream),
-}
-
-impl TcpForwardSocket {
-    fn new(
-        tcp_forward_addr: SocketAddr,
-        options: &crate::TcpOptions,
-    ) -> Result<Self, ConnectError> {
-        let socket = match &tcp_forward_addr {
-            SocketAddr::V4(..) => TcpSocket::new_v4(),
-            SocketAddr::V6(..) => TcpSocket::new_v6(),
-        }
-        .map_err(ConnectError::CreateTcpSocket)?;
-
-        crate::tcp_options::apply(&socket, options).map_err(ConnectError::ApplyTcpOptions)?;
-
-        Ok(TcpForwardSocket::Socket((socket, tcp_forward_addr)))
-    }
-
-    async fn connect(self) -> Result<TcpStream, ConnectError> {
-        match self {
-            TcpForwardSocket::Socket((socket, addr)) => {
-                log::info!("Connecting to {}/TCP", addr);
-
-                let tcp_stream = socket
-                    .connect(addr)
-                    .await
-                    .map_err(ConnectError::ConnectTcp)?;
-
-                log::info!("Connected to {}/TCP", addr);
-                Ok(tcp_stream)
-            }
-            TcpForwardSocket::Stream(stream) => Ok(stream),
-        }
-    }
-}
-
-#[cfg(unix)]
-impl AsRawFd for TcpForwardSocket {
-    fn as_raw_fd(&self) -> RawFd {
-        match self {
-            TcpForwardSocket::Socket((sock, _)) => sock.as_raw_fd(),
-            TcpForwardSocket::Stream(stream) => stream.as_raw_fd(),
-        }
-    }
-}

tests/udp2tcp.rs

@@ -114,11 +114,7 @@ async fn pong_two_separate() -> Result<(), Box<dyn std::error::Error>> {
 /// Spawns a Udp2Tcp instance and connects to boths ends of it.
 /// Returns the UDP and TCP sockets that goes through it.
 async fn setup_udp2tcp() -> Result<
-    (
-        UdpSocket,
-        TcpStream,
-        JoinHandle<Result<(), udp2tcp::ForwardError>>,
-    ),
+    (UdpSocket, TcpStream, JoinHandle<Result<(), udp2tcp::Error>>),
     Box<dyn std::error::Error>,
 > {
     let tcp_listener = TcpListener::bind("127.0.0.1:0").await?;
@@ -130,7 +126,6 @@ async fn setup_udp2tcp() -> Result<
         TcpOptions::default(),
     )
     .await?;
-    let tcp_stream = tcp_listener.accept().await?.0;
 
     let udp_listen_addr = udp2tcp.local_udp_addr().unwrap();
     let join_handle = tokio::spawn(udp2tcp.run());
@@ -138,5 +133,11 @@ async fn setup_udp2tcp() -> Result<
     let udp_socket = UdpSocket::bind("127.0.0.1:0").await?;
     udp_socket.connect(udp_listen_addr).await?;
 
+    // Send empty datagram to connect TCP socket
+    udp_socket.send(&[]).await?;
+    let mut tcp_stream = tcp_listener.accept().await?.0;
+    let mut buf = [0; 1024];
+    tcp_stream.read(&mut buf).await?;
+
     Ok((udp_socket, tcp_stream, join_handle))
 }