src/reader.rs

   1 use crate::schema::ConnectionMessage;
   2 use async_std::net::SocketAddr;
   3 use async_std::pin::Pin;
   4 use bytes::{Buf, BytesMut};
   5 use futures::task::{Context, Poll};
   6 use futures::{AsyncRead, AsyncReadExt, Stream};
   7 use log::*;
   8 use protobuf::Message;
   9 use std::convert::TryInto;
  10
  11 pub use futures::SinkExt;
  12 pub use futures::StreamExt;
  13 use protobuf::well_known_types::Any;
  14
  15 const BUFFER_SIZE: usize = 8192;
  16
  17 pub struct ConnectionReader {
  18     local_addr: SocketAddr,
  19     peer_addr: SocketAddr,
  20     read_stream: Box<dyn AsyncRead + Send + Sync + Unpin>,
  21     pending_read: Option<BytesMut>,
  22 }
  23
  24 impl ConnectionReader {
  25     pub fn new(
  26         local_addr: SocketAddr,
  27         peer_addr: SocketAddr,
  28         read_stream: Box<dyn AsyncRead + Send + Sync + Unpin>,
  29     ) -> Self {
  30         Self {
  31             local_addr,
  32             peer_addr,
  33             read_stream,
  34             pending_read: None,
  35         }
  36     }
  37
  38     pub fn local_addr(&self) -> SocketAddr {
  39         self.local_addr.clone()
  40     }
  41
  42     pub fn peer_addr(&self) -> SocketAddr {
  43         self.peer_addr.clone()
  44     }
  45 }
  46
  47 impl Stream for ConnectionReader {
  48     type Item = Any;
  49
  50     fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
  51         let mut buffer = BytesMut::new();
  52         buffer.resize(BUFFER_SIZE, 0);
  53
  54         trace!("Starting new read loop for {}", self.local_addr);
  55         loop {
  56             trace!("Reading from the stream");
  57             match futures::executor::block_on(self.read_stream.read(&mut buffer)) {
  58                 Ok(mut bytes_read) => {
  59                     if bytes_read > 0 {
  60                         trace!("Read {} bytes from the network stream", bytes_read)
  61                     }
  62
  63                     if let Some(mut pending_buf) = self.pending_read.take() {
  64                         trace!("Prepending broken data ({} bytes) encountered from earlier read of network stream", pending_buf.len());
  65                         bytes_read += pending_buf.len();
  66
  67                         pending_buf.unsplit(buffer);
  68                         buffer = pending_buf;
  69                     }
  70
  71                     let mut bytes_read_u64: u64 = bytes_read.try_into().expect(
  72                         format!("Conversion from usize ({}) to u64 failed", bytes_read).as_str(),
  73                     );
  74                     while bytes_read_u64 > 0 {
  75                         trace!(
  76                             "{} bytes from network stream still unprocessed",
  77                             bytes_read_u64
  78                         );
  79
  80                         buffer.resize(bytes_read, 0);
  81
  82                         match ConnectionMessage::parse_from_bytes(buffer.as_ref()) {
  83                             Ok(mut data) => {
  84                                 let serialized_size = data.compute_size();
  85                                 trace!("Deserialized message of size {} bytes", serialized_size);
  86
  87                                 buffer.advance(serialized_size as usize);
  88
  89                                 let serialized_size_u64: u64 = serialized_size.try_into().expect(
  90                                     format!(
  91                                         "Conversion from usize ({}) to u64 failed",
  92                                         serialized_size
  93                                     )
  94                                     .as_str(),
  95                                 );
  96                                 bytes_read_u64 -= serialized_size_u64;
  97                                 trace!("{} bytes still unprocessed", bytes_read_u64);
  98
  99                                 trace!("Sending deserialized message downstream");
 100                                 return Poll::Ready(Some(data.take_payload()));
 101                             }
 102
 103                             Err(err) => {
 104                                 warn!(
 105                                     "Could not deserialize data from the received bytes: {:#?}",
 106                                     err
 107                                 );
 108
 109                                 self.pending_read = Some(buffer);
 110                                 buffer = BytesMut::new();
 111                                 break;
 112                             }
 113                         }
 114                     }
 115
 116                     buffer.resize(BUFFER_SIZE, 0);
 117                 }
 118
 119                 Err(_err) => return Poll::Pending,
 120             }
 121         }
 122     }
 123 }