Merge pull request #380 from tox-rs/friend_request

Handle friend requests inside onion client

Author: Roman Proskuryakov

src/toxcore/dht/server/mod.rs

@@ -118,7 +118,7 @@ pub struct Server {
     /// Tx split of a channel to send packets to this peer via UDP socket.
     pub tx: Tx,
     /// Sink to send friend's `SocketAddr` when it gets known.
-    friend_saddr_sink: Arc<RwLock<Option<mpsc::UnboundedSender<PackedNode>>>>,
+    friend_saddr_sink: Arc<RwLock<OptionalSink<mpsc::UnboundedSender<PackedNode>>>>,
     /// Struct that stores and manages requests IDs and timeouts.
     request_queue: Arc<RwLock<RequestQueue<PublicKey>>>,
     /// Close nodes list which contains nodes close to own DHT `PublicKey`.
@@ -759,13 +759,9 @@ impl Server {
         for friend in friends.values_mut() {
             friend.try_add_to_close(node);
         }
-        if let Some(ref friend_saddr_sink) = *self.friend_saddr_sink.read() {
-            if friends.contains_key(&node.pk) {
-                Either::A(send_to(friend_saddr_sink, node)
-                    .map_err(|e| e.context(HandlePacketErrorKind::FriendSaddr).into()))
-            } else {
-                Either::B(future::ok(()))
-            }
+        if friends.contains_key(&node.pk) {
+            Either::A(send_to(&*self.friend_saddr_sink.read(), node)
+                .map_err(|e| e.context(HandlePacketErrorKind::FriendSaddr).into()))
         } else {
             Either::B(future::ok(()))
         }
@@ -1467,7 +1463,7 @@ impl Server {
 
     /// Set sink to send friend's `SocketAddr` when it gets known.
     pub fn set_friend_saddr_sink(&self, friend_saddr_sink: mpsc::UnboundedSender<PackedNode>) {
-        *self.friend_saddr_sink.write() = Some(friend_saddr_sink);
+        self.friend_saddr_sink.write().set(friend_saddr_sink);
     }
 
     /// Get `PrecomputedKey`s cache.

src/toxcore/io_tokio.rs

@@ -3,7 +3,7 @@
 use std::fmt::Debug;
 use std::io::{Error as IoError};
 
-use futures::{Future, Sink, Stream};
+use futures::{Async, AsyncSink, Future, Poll, Sink, StartSend, Stream};
 
 /// A convenience typedef around a `Future` whose error component is `io::Error`
 pub type IoFuture<T> = Box<Future<Item = T, Error = IoError> + Send>;
@@ -32,3 +32,51 @@ pub fn send_all_to<T: Send + 'static, S, Tx, E: Debug>(tx: &Tx, s: S) -> impl Fu
         .map(|_tx| ()) // ignore tx because it was cloned
 }
 
+/// `Sink` type that can either be empty or store inner `Sink`. When it's empty
+/// it drops all sent elements.
+#[derive(Debug, Clone)]
+pub struct OptionalSink<S: Sink>(Option<S>);
+
+impl<S: Sink> OptionalSink<S> {
+    /// Create new `OptionalSink`.
+    pub fn new() -> Self {
+        OptionalSink(None)
+    }
+
+    /// Set a sink to `OptionalSink`.
+    pub fn set(&mut self, sink: S) {
+        self.0 = Some(sink);
+    }
+}
+
+impl<S: Sink> Default for OptionalSink<S> {
+    fn default() -> Self {
+        OptionalSink::new()
+    }
+}
+
+impl<S: Sink> Sink for OptionalSink<S> {
+    type SinkItem = S::SinkItem;
+    type SinkError = S::SinkError;
+    fn start_send(&mut self, item: S::SinkItem) -> StartSend<S::SinkItem, S::SinkError> {
+        if let Some(ref mut sink) = self.0 {
+            sink.start_send(item)
+        } else {
+            Ok(AsyncSink::Ready)
+        }
+    }
+    fn poll_complete(&mut self) -> Poll<(), S::SinkError> {
+        if let Some(ref mut sink) = self.0 {
+            sink.poll_complete()
+        } else {
+            Ok(Async::Ready(()))
+        }
+    }
+    fn close(&mut self) -> Poll<(), S::SinkError> {
+        if let Some(ref mut sink) = self.0 {
+            sink.close()
+        } else {
+            Ok(Async::Ready(()))
+        }
+    }
+}

src/toxcore/net_crypto/mod.rs

@@ -128,10 +128,10 @@ pub struct NetCrypto {
     /// packet. If key from `Cookie` is not equal to saved key inside
     /// `CryptoConnection` then `NetCrypto` module will send message to this
     /// sink.
-    dht_pk_tx: Arc<RwLock<Option<DhtPkTx>>>,
+    dht_pk_tx: Arc<RwLock<OptionalSink<DhtPkTx>>>,
     /// Sink to send a connection status when it becomes connected or
     /// disconnected. The key is a long term key of the connection.
-    connection_status_tx: Arc<RwLock<Option<ConnectionStatusTx>>>,
+    connection_status_tx: Arc<RwLock<OptionalSink<ConnectionStatusTx>>>,
     /// Sink to send lossless packets. The key is a long term public key of the
     /// peer that sent this packet.
     lossless_tx: LosslessTx,
@@ -229,11 +229,7 @@ impl NetCrypto {
     /// or disconnected.
     fn send_connection_status(&self, connection: &CryptoConnection, status: bool) -> impl Future<Item = (), Error = mpsc::SendError<(PublicKey, bool)>> {
         if connection.is_established() != status {
-            if let Some(ref connection_status_tx) = *self.connection_status_tx.read() {
-                Either::A(send_to(connection_status_tx, (connection.peer_real_pk, status)))
-            } else {
-                Either::B(future::ok(()))
-            }
+            Either::A(send_to(&*self.connection_status_tx.read(), (connection.peer_real_pk, status)))
         } else {
             Either::B(future::ok(()))
         }
@@ -452,12 +448,8 @@ impl NetCrypto {
             return Box::new(future::err(HandlePacketError::from(HandlePacketErrorKind::InvalidRealPk)))
         }
         if cookie.dht_pk != connection.peer_dht_pk {
-            let dht_pk_future = if let Some(ref dht_pk_tx) = *self.dht_pk_tx.read() {
-                Either::A(send_to(dht_pk_tx, (connection.peer_real_pk, cookie.dht_pk))
-                    .map_err(|e| e.context(HandlePacketErrorKind::SendToDhtpk).into()))
-            } else {
-                Either::B(future::ok(()))
-            };
+            let dht_pk_future = send_to(&*self.dht_pk_tx.read(), (connection.peer_real_pk, cookie.dht_pk))
+                .map_err(|e| e.context(HandlePacketErrorKind::SendToDhtpk).into());
             return Box::new(dht_pk_future
                 .and_then(|()| future::err(HandlePacketError::from(HandlePacketErrorKind::InvalidDhtPk)))
             )
@@ -561,12 +553,8 @@ impl NetCrypto {
         let connection = Arc::new(RwLock::new(connection));
         connections.insert(cookie.real_pk, connection);
 
-        let dht_pk_future = if let Some(ref dht_pk_tx) = *self.dht_pk_tx.read() {
-            Either::A(send_to(dht_pk_tx, (cookie.real_pk, cookie.dht_pk))
-                .map_err(|e| e.context(HandlePacketErrorKind::SendToDhtpk).into()))
-        } else {
-            Either::B(future::ok(()))
-        };
+        let dht_pk_future = send_to(&*self.dht_pk_tx.read(), (cookie.real_pk, cookie.dht_pk))
+            .map_err(|e| e.context(HandlePacketErrorKind::SendToDhtpk).into());
 
         Either::B(kill_future.join(dht_pk_future).map(|_| ()))
     }
@@ -986,13 +974,13 @@ impl NetCrypto {
 
     /// Set sink to send DHT `PublicKey` when it gets known.
     pub fn set_dht_pk_sink(&self, dht_pk_tx: DhtPkTx) {
-        *self.dht_pk_tx.write() = Some(dht_pk_tx);
+        self.dht_pk_tx.write().set(dht_pk_tx);
     }
 
     /// Set sink to send a connection status when it becomes connected or
     /// disconnected.
     pub fn set_connection_status_sink(&self, connection_status_tx: ConnectionStatusTx) {
-        *self.connection_status_tx.write() = Some(connection_status_tx);
+        self.connection_status_tx.write().set(connection_status_tx);
     }
 }
 

src/toxcore/onion/client/errors.rs

@@ -64,6 +64,9 @@ error_kind! {
         #[doc = "Failed to handle DHT `PublicKey` announce."]
         #[fail(display = "Failed to handle DHT PublicKey announce")]
         DhtPkAnnounce,
+        #[doc = "Failed to send a friend request."]
+        #[fail(display = "Failed to send a friend request")]
+        FriendRequest,
     }
 }
 

src/toxcore/onion/client/mod.rs

@@ -40,6 +40,10 @@ use crate::toxcore::time::*;
 /// key is a DHT key.
 type DhtPkTx = mpsc::UnboundedSender<(PublicKey, PublicKey)>;
 
+/// Shorthand for the transmit half of the message channel for sending friend
+/// requests when we get them. The key is a long term key.
+type FriendRequestTx = mpsc::UnboundedSender<(PublicKey, FriendRequest)>;
+
 /// Number of friend's close nodes to store.
 const MAX_ONION_FRIEND_NODES: u8 = 8;
 
@@ -302,7 +306,10 @@ struct OnionClientState {
     friends: HashMap<PublicKey, OnionFriend>,
     /// Sink to send DHT `PublicKey` when it gets known. The first key is a long
     /// term key, the second key is a DHT key.
-    dht_pk_tx: Option<DhtPkTx>,
+    dht_pk_tx: OptionalSink<DhtPkTx>,
+    /// Sink to send friend requests when we get them. The key is a long term
+    /// key.
+    friend_request_tx: OptionalSink<FriendRequestTx>,
 }
 
 impl OnionClientState {
@@ -312,7 +319,8 @@ impl OnionClientState {
             announce_list: Kbucket::new(MAX_ONION_ANNOUNCE_NODES),
             announce_requests: RequestQueue::new(ANNOUNCE_TIMEOUT),
             friends: HashMap::new(),
-            dht_pk_tx: None,
+            dht_pk_tx: OptionalSink::new(),
+            friend_request_tx: OptionalSink::new(),
         }
     }
 }
@@ -360,7 +368,7 @@ impl OnionClient {
 
     /// Set sink to send DHT `PublicKey` when it gets known.
     pub fn set_dht_pk_sink(&self, dht_pk_tx: DhtPkTx) {
-        self.state.lock().dht_pk_tx = Some(dht_pk_tx);
+        self.state.lock().dht_pk_tx.set(dht_pk_tx);
     }
 
     /// Check if a node was pinged recently.
@@ -511,11 +519,7 @@ impl OnionClient {
         friend.dht_pk = Some(dht_pk_announce.dht_pk);
         friend.last_seen = Some(clock_now());
 
-        let dht_pk_future = if let Some(ref dht_pk_tx) = state.dht_pk_tx {
-            Either::A(send_to(dht_pk_tx, (friend_pk, dht_pk_announce.dht_pk)))
-        } else {
-            Either::B(future::ok(()))
-        };
+        let dht_pk_future = send_to(&state.dht_pk_tx, (friend_pk, dht_pk_announce.dht_pk));
 
         let futures = dht_pk_announce.nodes.into_iter().map(|node| match node.ip_port.protocol {
             ProtocolType::UDP => {
@@ -545,13 +549,15 @@ impl OnionClient {
             Ok(payload) => payload,
             Err(e) => return Either::A(future::err(e.context(HandleDataResponseErrorKind::InvalidInnerPayload).into()))
         };
-        match iner_payload {
+        let future = match iner_payload {
             OnionDataResponseInnerPayload::DhtPkAnnounce(dht_pk_announce) =>
-                Either::B(self.handle_dht_pk_announce(payload.real_pk, dht_pk_announce)
+                Either::A(self.handle_dht_pk_announce(payload.real_pk, dht_pk_announce)
                     .map_err(|e| e.context(HandleDataResponseErrorKind::DhtPkAnnounce).into())),
-            OnionDataResponseInnerPayload::FriendRequest(_) =>
-                Either::A(future::ok(()))
-        }
+            OnionDataResponseInnerPayload::FriendRequest(friend_request) =>
+                Either::B(send_to(&self.state.lock().friend_request_tx, (payload.real_pk, friend_request))
+                    .map_err(|e| e.context(HandleDataResponseErrorKind::FriendRequest).into()))
+        };
+        Either::B(future)
     }
 
     /// Add new node to random nodes pool to use them to build random paths.
@@ -991,7 +997,8 @@ mod tests {
             announce_list: Kbucket::new(8),
             announce_requests: RequestQueue::new(Duration::from_secs(42)),
             friends: HashMap::new(),
-            dht_pk_tx: None,
+            dht_pk_tx: OptionalSink::new(),
+            friend_request_tx: OptionalSink::new(),
         };
 
         let _onion_client_state_c = onion_client_state.clone();