Explicit proxy runner + closing local connection if remote is over (#314)

* Updated built dependency

So that it wouldn't fail on nwindows

* Concurrent RW client-side

* Draft of a ProxyRunner

* Super experimental proxy runner for provider

Very much WIP

* Extracting shared code + cleanup

* Logging cleanup

* Removed redundant close request

common/utils/Cargo.toml

@@ -8,9 +8,12 @@ edition = "2018"
 
 [dependencies]
 bytes = "0.5"
-# no need for any features as `AsyncRead` is always available
-tokio = { version = "0.2", features = [] }
+# TODO: "time" feature is only required for the delay loop which is going to go away very soon!
+tokio = { version = "0.2", features = [ "tcp", "io-util", "sync", "macros", "time" ] }
+futures = "0.3"
+log = "0.4"
+simple-socks5-requests = { path = "../../service-providers/simple-socks5/simple-socks5-requests" }
 
 [dev-dependencies]
-tokio = { version = "0.2", features = ["rt-threaded", "macros"] }
+tokio = { version = "0.2", features = ["rt-threaded"] }
 tokio-test = "0.2"

common/utils/src/connection_controller.rs

@@ -0,0 +1,109 @@
+// Copyright 2020 Nym Technologies SA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use futures::channel::mpsc;
+use log::*;
+use simple_socks5_requests::ConnectionId;
+use std::collections::HashMap;
+use tokio::stream::StreamExt;
+
+/// A generic message produced after reading from a socket/connection. It includes data that was
+/// actually read alongside boolean indicating whether the connection got closed so that
+/// remote could act accordingly.
+#[derive(Debug)]
+pub struct ConnectionMessage {
+    pub payload: Vec<u8>,
+    pub socket_closed: bool,
+}
+
+/// Channel responsible for sending data that was received from mix network into particular connection.
+/// Data includes the actual payload that is to be written onto the connection
+/// alongside boolean indicating whether the remote connection was closed after producing this message,
+/// so that the local connection should also shut down.
+pub type ConnectionSender = mpsc::UnboundedSender<ConnectionMessage>;
+
+/// Receiver part of the [`ConnectionSender`]
+pub type ConnectionReceiver = mpsc::UnboundedReceiver<ConnectionMessage>;
+
+pub type ControllerSender = mpsc::UnboundedSender<ControllerCommand>;
+pub type ControllerReceiver = mpsc::UnboundedReceiver<ControllerCommand>;
+
+pub enum ControllerCommand {
+    Insert(ConnectionId, ConnectionSender),
+    Remove(ConnectionId),
+    Send(ConnectionId, Vec<u8>, bool),
+}
+
+/// Controller represents a way of managing multiple open connections that are used for socks5
+/// proxy.
+pub struct Controller {
+    active_connections: HashMap<ConnectionId, ConnectionSender>,
+    receiver: ControllerReceiver,
+}
+
+impl Controller {
+    pub fn new() -> (Self, ControllerSender) {
+        let (sender, receiver) = mpsc::unbounded();
+        (
+            Controller {
+                active_connections: HashMap::new(),
+                receiver,
+            },
+            sender,
+        )
+    }
+
+    fn insert_connection(&mut self, conn_id: ConnectionId, sender: ConnectionSender) {
+        if self.active_connections.insert(conn_id, sender).is_some() {
+            panic!("there is already an active request with the same id present - it's probably a bug!")
+        }
+    }
+
+    fn remove_connection(&mut self, conn_id: ConnectionId) {
+        debug!("Removing {} from controller", conn_id);
+        if self.active_connections.remove(&conn_id).is_none() {
+            error!(
+                "tried to remove non-existing connection with id: {:?}",
+                conn_id
+            )
+        }
+    }
+
+    fn send_to_connection(&mut self, conn_id: ConnectionId, payload: Vec<u8>, is_closed: bool) {
+        if let Some(sender) = self.active_connections.get_mut(&conn_id) {
+            sender
+                .unbounded_send(ConnectionMessage {
+                    payload,
+                    socket_closed: is_closed,
+                })
+                .unwrap()
+        } else {
+            error!("no connection exists with id: {:?}", conn_id);
+        }
+    }
+
+    pub async fn run(&mut self) {
+        while let Some(command) = self.receiver.next().await {
+            match command {
+                ControllerCommand::Send(conn_id, data, is_closed) => {
+                    self.send_to_connection(conn_id, data, is_closed)
+                }
+                ControllerCommand::Insert(conn_id, sender) => {
+                    self.insert_connection(conn_id, sender)
+                }
+                ControllerCommand::Remove(conn_id) => self.remove_connection(conn_id),
+            }
+        }
+    }
+}

common/utils/src/lib.rs

@@ -13,4 +13,6 @@
 // limitations under the License.
 
 pub mod available_reader;
+pub mod connection_controller;
+pub mod proxy_runner;
 pub mod read_delay_loop;

common/utils/src/proxy_runner.rs

@@ -0,0 +1,229 @@
+// Copyright 2020 Nym Technologies SA
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use super::read_delay_loop::try_read_data;
+use crate::connection_controller::ConnectionReceiver;
+use futures::channel::mpsc;
+use log::*;
+use simple_socks5_requests::ConnectionId;
+use std::sync::Arc;
+use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
+use tokio::net::TcpStream;
+use tokio::prelude::*;
+use tokio::stream::StreamExt;
+use tokio::sync::Notify;
+
+#[derive(Debug)]
+pub struct ProxyMessage {
+    pub data: Vec<u8>,
+    pub socket_closed: bool,
+}
+
+impl Into<ProxyMessage> for (Vec<u8>, bool) {
+    fn into(self) -> ProxyMessage {
+        ProxyMessage {
+            data: self.0,
+            socket_closed: self.1,
+        }
+    }
+}
+
+pub type MixProxySender<S> = mpsc::UnboundedSender<S>;
+
+#[derive(Debug)]
+pub struct ProxyRunner<S> {
+    /// receives data from the mix network and sends that into the socket
+    mix_receiver: Option<ConnectionReceiver>,
+
+    /// sends whatever was read from the socket into the mix network
+    mix_sender: MixProxySender<S>,
+
+    socket: Option<TcpStream>,
+    connection_id: ConnectionId,
+}
+
+impl<S> ProxyRunner<S>
+where
+    S: Send + 'static,
+{
+    pub fn new(
+        socket: TcpStream,
+        mix_receiver: ConnectionReceiver,
+        mix_sender: MixProxySender<S>,
+        connection_id: ConnectionId,
+    ) -> Self {
+        ProxyRunner {
+            mix_receiver: Some(mix_receiver),
+            mix_sender,
+            socket: Some(socket),
+            connection_id,
+        }
+    }
+
+    async fn run_inbound<F>(
+        mut reader: OwnedReadHalf,
+        notify_closed: Arc<Notify>,
+        connection_id: ConnectionId,
+        mix_sender: MixProxySender<S>,
+        adapter_fn: F,
+    ) -> OwnedReadHalf
+    where
+        F: Fn(ConnectionId, Vec<u8>, bool) -> S + Send + 'static,
+    {
+        // TODO: to be removed with sequence numbers...
+        let socket_read_timeout_duration = std::time::Duration::from_millis(500);
+
+        let address = reader.as_ref().peer_addr().unwrap().to_string();
+        loop {
+            tokio::select! {
+                _ = notify_closed.notified() => {
+                    // the remote socket is closed, so there's no point
+                    // in reading anything more because we won't be able to write to remote anyway!
+                    break
+                }
+                // try to read from local socket and push everything to mixnet to the remote
+                reading_result = try_read_data(socket_read_timeout_duration, &mut reader, &address) => {
+                    let (read_data, timed_out) = match reading_result {
+                        Ok(data) => data,
+                        Err(err) => {
+                            error!("failed to read request from the socket - {}", err);
+                            break;
+                        }
+                    };
+
+                    if read_data.is_empty() && timed_out {
+                        // no point in writing empty data on each timeout
+                        continue
+                    }
+
+                    info!(
+                        "Going to send {} bytes via mixnet to remote {}. Is local closed: {}",
+                        read_data.len(),
+                        connection_id,
+                        !timed_out
+                    );
+
+                    mix_sender.unbounded_send(adapter_fn(connection_id, read_data, !timed_out)).unwrap();
+
+                    if !timed_out {
+                        // technically we already informed it when we sent the message to mixnet above
+                        info!("The local socket is closed - won't receive any more data. Informing remote about that...");
+                        // no point in reading from mixnet if connection is closed!
+                        notify_closed.notify();
+                        break;
+                    }
+                }
+            }
+        }
+
+        reader
+    }
+
+    async fn run_outbound(
+        mut writer: OwnedWriteHalf,
+        notify_closed: Arc<Notify>,
+        mut mix_receiver: ConnectionReceiver,
+        connection_id: ConnectionId,
+    ) -> (OwnedWriteHalf, ConnectionReceiver) {
+        loop {
+            tokio::select! {
+                _ = notify_closed.notified() => {
+                    // no need to read from mixnet as we won't be able to send to socket
+                    // anyway
+                    break
+                }
+                mix_data = mix_receiver.next() => {
+                    if mix_data.is_none() {
+                        warn!("mix receiver is none so we already got removed somewhere. This isn't really a warning, but shouldn't happen to begin with, so please say if you see this message");
+                        // we already got closed
+                        // not sure if we HAVE TO notify the other task, but might as well
+                        notify_closed.notify();
+                        break
+                    }
+                    let connection_message = mix_data.unwrap();
+
+                    info!(
+                        "Going to write {} bytes received from mixnet to connection {}. Is remote closed: {}",
+                        connection_message.payload.len(),
+                        connection_id,
+                        connection_message.socket_closed
+                    );
+
+                    if let Err(err) = writer.write_all(&connection_message.payload).await {
+                        // the other half is probably going to blow up too (if not, this task also needs to notify the other one!!)
+                        error!("failed to write response back to the socket - {}", err);
+                        break;
+                    }
+                    if connection_message.socket_closed {
+                        info!("Remote socket got closed - closing the local socket too");
+                        notify_closed.notify();
+                        break
+                    }
+                }
+            }
+        }
+
+        (writer, mix_receiver)
+    }
+
+    // The `adapter_fn` is used to transform whatever was read into appropriate
+    // request/response as required by entity running particular side of the proxy.
+    pub async fn run<F>(mut self, adapter_fn: F) -> Self
+    where
+        F: Fn(ConnectionId, Vec<u8>, bool) -> S + Send + 'static,
+    {
+        let notify_closed = Arc::new(Notify::new());
+        let notify_clone = Arc::clone(&notify_closed);
+
+        let (read_half, write_half) = self.socket.take().unwrap().into_split();
+        let mix_receiver = self.mix_receiver.take().unwrap();
+
+        // should run until either inbound closes or is notified from outbound
+        let inbound_future = Self::run_inbound(
+            read_half,
+            notify_closed,
+            self.connection_id,
+            self.mix_sender.clone(),
+            adapter_fn,
+        );
+        let outbound_future =
+            Self::run_outbound(write_half, notify_clone, mix_receiver, self.connection_id);
+
+        // TODO: this shouldn't really have to spawn tasks inside "library" code, but
+        // if we used join directly, stuff would have been executed on the same thread
+        // (it's not bad, but an unnecessary slowdown)
+        let handle_inbound = tokio::spawn(inbound_future);
+        let handle_outbound = tokio::spawn(outbound_future);
+
+        let (inbound_result, outbound_result) =
+            futures::future::join(handle_inbound, handle_outbound).await;
+
+        if inbound_result.is_err() || outbound_result.is_err() {
+            panic!("TODO: some future error?")
+        }
+
+        let (write_half, mix_receiver) = outbound_result.unwrap();
+
+        self.socket = Some(write_half.reunite(inbound_result.unwrap()).unwrap());
+        self.mix_receiver = Some(mix_receiver);
+        self
+    }
+
+    pub fn into_inner(mut self) -> (TcpStream, ConnectionReceiver) {
+        (
+            self.socket.take().unwrap(),
+            self.mix_receiver.take().unwrap(),
+        )
+    }
+}

common/utils/src/read_delay_loop.rs

@@ -24,11 +24,12 @@ use std::io;
 use tokio::io::AsyncRead;
 use tokio::time::Duration;
 
+// It returns data alognside information whether it timed out while reading from the socket
 pub async fn try_read_data<R>(
     timeout: Duration,
     mut reader: R,
     address: &str,
-) -> io::Result<Vec<u8>>
+) -> io::Result<(Vec<u8>, bool)>
 where
     R: AsyncRead + Unpin,
 {
@@ -41,7 +42,7 @@ where
         tokio::select! {
             _ = &mut delay => {
                 println!("Timed out. returning {} bytes received from {}", data.len(), address);
-                return Ok(data) // we return all response data on timeout
+                return Ok((data, true)) // we return all response data on timeout
             }
             read_data = &mut available_reader => {
                 match read_data {
@@ -53,7 +54,7 @@ where
                             println!("Connection is closed! Returning {} bytes received from {}", data.len(), address);
                             // we return all we managed to read because
                             // we know no more stuff is coming
-                            return Ok(data)
+                            return Ok((data, false))
                         }
                         let now = tokio::time::Instant::now();
                         let next = now + timeout;