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     closed:       bool,
  23 }
  24
  25 impl ConnectionReader {
  26     pub fn new(
  27         local_addr: SocketAddr,
  28         peer_addr: SocketAddr,
  29         read_stream: Box<dyn AsyncRead + Send + Sync + Unpin>,
  30     ) -> Self {
  31         Self {
  32             local_addr,
  33             peer_addr,
  34             read_stream,
  35             pending_read: None,
  36             closed: false,
  37         }
  38     }
  39
  40     pub fn local_addr(&self) -> SocketAddr {
  41         self.local_addr.clone()
  42     }
  43
  44     pub fn peer_addr(&self) -> SocketAddr {
  45         self.peer_addr.clone()
  46     }
  47
  48     pub fn is_closed(&self) -> bool {
  49         self.closed
  50     }
  51 }
  52
  53 impl Stream for ConnectionReader {
  54     type Item = Any;
  55
  56     fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
  57         let mut buffer = BytesMut::new();
  58         buffer.resize(BUFFER_SIZE, 0);
  59
  60         trace!("Starting new read loop for {}", self.local_addr);
  61         loop {
  62             trace!("Reading from the stream");
  63             match futures::executor::block_on(self.read_stream.read(&mut buffer)) {
  64                 Ok(mut bytes_read) => {
  65                     if bytes_read > 0 {
  66                         trace!("Read {} bytes from the network stream", bytes_read)
  67                     } else if bytes_read == 0 && self.pending_read.is_none() {
  68                         return Poll::Pending;
  69                     }
  70
  71                     if let Some(mut pending_buf) = self.pending_read.take() {
  72                         trace!("Prepending broken data ({} bytes) encountered from earlier read of network stream", pending_buf.len());
  73                         bytes_read += pending_buf.len();
  74
  75                         pending_buf.unsplit(buffer);
  76                         buffer = pending_buf;
  77                     }
  78
  79                     let mut bytes_read_u64: u64 = bytes_read.try_into().expect(
  80                         format!("Conversion from usize ({}) to u64 failed", bytes_read).as_str(),
  81                     );
  82                     while bytes_read_u64 > 0 {
  83                         trace!(
  84                             "{} bytes from network stream still unprocessed",
  85                             bytes_read_u64
  86                         );
  87
  88                         buffer.resize(bytes_read, 0);
  89
  90                         match ConnectionMessage::parse_from_bytes(buffer.as_ref()) {
  91                             Ok(mut data) => {
  92                                 let serialized_size = data.compute_size();
  93                                 trace!("Deserialized message of size {} bytes", serialized_size);
  94
  95                                 buffer.advance(serialized_size as usize);
  96
  97                                 let serialized_size_u64: u64 = serialized_size.try_into().expect(
  98                                     format!(
  99                                         "Conversion from usize ({}) to u64 failed",
 100                                         serialized_size
 101                                     )
 102                                     .as_str(),
 103                                 );
 104                                 bytes_read_u64 -= serialized_size_u64;
 105                                 trace!("{} bytes still unprocessed", bytes_read_u64);
 106
 107                                 trace!("Sending deserialized message downstream");
 108                                 return Poll::Ready(Some(data.take_payload()));
 109                             }
 110
 111                             Err(err) => {
 112                                 warn!(
 113                                     "Could not deserialize data from the received bytes: {:#?}",
 114                                     err
 115                                 );
 116
 117                                 self.pending_read = Some(buffer);
 118                                 buffer = BytesMut::new();
 119                                 break;
 120                             }
 121                         }
 122                     }
 123
 124                     buffer.resize(BUFFER_SIZE, 0);
 125                 }
 126
 127                 // Close the stream
 128                 Err(_err) => {
 129                     trace!("Closing the stream");
 130                     self.pending_read.take();
 131                     self.closed = true;
 132                     return Poll::Ready(None);
 133                 }
 134             }
 135         }
 136     }
 137 }