From 35451a05de2760c6d2e125dbdc8bb57f37baa65d Mon Sep 17 00:00:00 2001
From: Jedrzej Stuczynski <jedrzej.stuczynski@gmail.com>
Date: Thu, 23 Jan 2020 16:55:11 +0000
Subject: [PATCH] And by real traffic stream + necessary adjustments

---
 nym-client/src/client/real_traffic_stream.rs | 81 +++++++++++---------
 1 file changed, 45 insertions(+), 36 deletions(-)

diff --git a/nym-client/src/client/real_traffic_stream.rs b/nym-client/src/client/real_traffic_stream.rs
index f89257747b..66a80a2ebe 100644
--- a/nym-client/src/client/real_traffic_stream.rs
+++ b/nym-client/src/client/real_traffic_stream.rs
@@ -1,38 +1,38 @@
 use crate::client::mix_traffic::MixMessage;
+use crate::client::topology_control::TopologyInnerRef;
 use crate::client::{InputMessage, MESSAGE_SENDING_AVERAGE_DELAY};
-use directory_client::presence::Topology;
 use futures::channel::mpsc;
 use futures::task::{Context, Poll};
 use futures::{Future, Stream, StreamExt};
-use log::{debug, info, trace};
+use log::{info, trace, warn};
 use sphinx::route::Destination;
-use sphinx::SphinxPacket;
-use std::net::SocketAddr;
 use std::pin::Pin;
 use std::time::Duration;
 use tokio::time;
+use topology::NymTopology;
 
 // have a rather low value for test sake
 const AVERAGE_PACKET_DELAY: f64 = 0.1;
 
-pub(crate) struct OutQueueControl {
+pub(crate) struct OutQueueControl<T: NymTopology> {
     delay: time::Delay,
     mix_tx: mpsc::UnboundedSender<MixMessage>,
     input_rx: mpsc::UnboundedReceiver<InputMessage>,
     our_info: Destination,
-
-    // due to pinning, DerefMut trait, futures, etc its way easier to
-    // just have concrete implementation here rather than generic NymTopology
-    // considering that it will be replaced with refreshing topology within few days anyway
-    topology: Topology,
+    topology_ctrl_ref: TopologyInnerRef<T>,
 }
 
-impl Stream for OutQueueControl {
-    type Item = (SocketAddr, SphinxPacket);
+pub(crate) enum StreamMessage {
+    Cover,
+    Real(InputMessage),
+}
+
+impl<T: NymTopology> Stream for OutQueueControl<T> {
+    type Item = StreamMessage;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         // it is not yet time to return a message
-        if Future::poll(Pin::new(&mut self.delay), cx).is_pending() {
+        if Pin::new(&mut self.delay).poll(cx).is_pending() {
             return Poll::Pending;
         };
 
@@ -49,41 +49,26 @@ impl Stream for OutQueueControl {
         self.delay.reset(next);
 
         // decide what kind of message to send
-        match Stream::poll_next(Pin::new(&mut self.input_rx), cx) {
+        match Pin::new(&mut self.input_rx).poll_next(cx) {
             // in the case our real message channel stream was closed, we should also indicate we are closed
             // (and whoever is using the stream should panic)
             Poll::Ready(None) => Poll::Ready(None),
 
             // if there's an actual message - return it
-            Poll::Ready(Some(real_message)) => {
-                trace!("real message");
-                Poll::Ready(Some(mix_client::packet::encapsulate_message(
-                    real_message.0,
-                    real_message.1,
-                    &self.topology,
-                    AVERAGE_PACKET_DELAY,
-                )))
-            }
+            Poll::Ready(Some(real_message)) => Poll::Ready(Some(StreamMessage::Real(real_message))),
 
             // otherwise construct a dummy one
-            _ => {
-                trace!("loop cover message");
-                Poll::Ready(Some(mix_client::packet::loop_cover_message(
-                    self.our_info.address,
-                    self.our_info.identifier,
-                    &self.topology,
-                )))
-            }
+            Poll::Pending => Poll::Ready(Some(StreamMessage::Cover)),
         }
     }
 }
 
-impl OutQueueControl {
+impl<T: NymTopology> OutQueueControl<T> {
     pub(crate) fn new(
         mix_tx: mpsc::UnboundedSender<MixMessage>,
         input_rx: mpsc::UnboundedReceiver<InputMessage>,
         our_info: Destination,
-        topology: Topology,
+        topology: TopologyInnerRef<T>,
     ) -> Self {
         let initial_delay = time::delay_for(Duration::from_secs_f64(MESSAGE_SENDING_AVERAGE_DELAY));
         OutQueueControl {
@@ -91,20 +76,44 @@ impl OutQueueControl {
             mix_tx,
             input_rx,
             our_info,
-            topology,
+            topology_ctrl_ref: topology,
         }
     }
 
     pub(crate) async fn run_out_queue_control(mut self) {
         info!("starting out queue controller");
         while let Some(next_message) = self.next().await {
-            debug!("created new message");
+            trace!("created new message");
+            let read_lock = self.topology_ctrl_ref.read().await;
+            let topology = read_lock.topology.as_ref();
+
+            if topology.is_none() {
+                warn!("No valid topology detected - won't send any loop cover message this time");
+                continue;
+            }
+
+            let topology = topology.unwrap();
+
+            let next_packet = match next_message {
+                StreamMessage::Cover => mix_client::packet::loop_cover_message(
+                    self.our_info.address,
+                    self.our_info.identifier,
+                    topology,
+                ),
+                StreamMessage::Real(real_message) => mix_client::packet::encapsulate_message(
+                    real_message.0,
+                    real_message.1,
+                    topology,
+                    AVERAGE_PACKET_DELAY,
+                ),
+            };
+
             // if this one fails, there's no retrying because it means that either:
             // - we run out of memory
             // - the receiver channel is closed
             // in either case there's no recovery and we can only panic
             self.mix_tx
-                .unbounded_send(MixMessage::new(next_message.0, next_message.1))
+                .unbounded_send(MixMessage::new(next_packet.0, next_packet.1))
                 .unwrap();
         }
     }