Refactor streams

Author: afilini

src/client.rs

@@ -3,8 +3,6 @@
 //! This module contains definitions of all the complex data structures that are returned by calls
 
 use std::collections::{BTreeMap, VecDeque};
-#[cfg(test)]
-use std::fs::File;
 use std::io::{self, BufRead, BufReader, Read, Write};
 use std::net::{TcpStream, ToSocketAddrs};
 
@@ -31,8 +29,6 @@ use socks::{Socks5Stream, ToTargetAddr};
 use stream::ClonableStream;
 
 use batch::Batch;
-#[cfg(test)]
-use test_stream::TestStream;
 use types::*;
 
 macro_rules! impl_batch_call {
@@ -72,8 +68,8 @@ pub struct Client<S>
 where
     S: Read + Write,
 {
-    stream: S,
-    buf_reader: BufReader<S>,
+    stream: ClonableStream<S>,
+    buf_reader: BufReader<ClonableStream<S>>,
 
     headers: VecDeque<HeaderNotification>,
     script_notifications: BTreeMap<ScriptHash, VecDeque<ScriptStatus>>,
@@ -82,7 +78,7 @@ where
     calls: usize,
 }
 
-impl<S> From<S> for Client<ClonableStream<S>>
+impl<S> From<S> for Client<S>
 where
     S: Read + Write,
 {
@@ -107,23 +103,14 @@ impl Client<ElectrumPlaintextStream> {
     /// Creates a new plaintext client and tries to connect to `socket_addr`.
     pub fn new<A: ToSocketAddrs>(socket_addr: A) -> Result<Self, Error> {
         let stream = TcpStream::connect(socket_addr)?;
-        let buf_reader = BufReader::new(stream.try_clone()?);
-
-        Ok(Self {
-            stream,
-            buf_reader,
-            headers: VecDeque::new(),
-            script_notifications: BTreeMap::new(),
 
-            #[cfg(feature = "debug-calls")]
-            calls: 0,
-        })
+        Ok(stream.into())
     }
 }
 
 #[cfg(feature = "use-openssl")]
 /// Transport type used to establish an OpenSSL TLS encrypted/authenticated connection with the server
-pub type ElectrumSslStream = ClonableStream<SslStream<TcpStream>>;
+pub type ElectrumSslStream = SslStream<TcpStream>;
 #[cfg(feature = "use-openssl")]
 impl Client<ElectrumSslStream> {
     /// Creates a new SSL client and tries to connect to `socket_addr`. Optionally, if `domain` is not
@@ -174,7 +161,7 @@ mod danger {
     not(feature = "use-openssl")
 ))]
 /// Transport type used to establish a Rustls TLS encrypted/authenticated connection with the server
-pub type ElectrumSslStream = ClonableStream<StreamOwned<ClientSession, TcpStream>>;
+pub type ElectrumSslStream = StreamOwned<ClientSession, TcpStream>;
 #[cfg(all(
     any(feature = "default", feature = "use-rustls"),
     not(feature = "use-openssl")
@@ -209,7 +196,7 @@ impl Client<ElectrumSslStream> {
 
 #[cfg(any(feature = "default", feature = "proxy"))]
 /// Transport type used to establish a connection to a server through a socks proxy
-pub type ElectrumProxyStream = ClonableStream<Socks5Stream>;
+pub type ElectrumProxyStream = Socks5Stream;
 #[cfg(any(feature = "default", feature = "proxy"))]
 impl Client<ElectrumProxyStream> {
     /// Creates a new socks client and tries to connect to `target_addr` using `proxy_addr` as an
@@ -226,21 +213,6 @@ impl Client<ElectrumProxyStream> {
     }
 }
 
-#[cfg(test)]
-impl Client<TestStream> {
-    pub fn new_test(file: File) -> Self {
-        let stream = TestStream::new_out();
-        let buf_reader = BufReader::new(TestStream::new_in(file));
-
-        Self {
-            stream,
-            buf_reader,
-            headers: VecDeque::new(),
-            script_notifications: BTreeMap::new(),
-        }
-    }
-}
-
 impl<S: Read + Write> Client<S> {
     fn call(&mut self, req: Request) -> Result<serde_json::Value, Error> {
         let mut raw = serde_json::to_vec(&req)?;
@@ -650,9 +622,16 @@ impl<S: Read + Write> Client<S> {
 mod test {
     use std::fs::File;
     use std::io::Read;
+    use test_stream::TestStream;
 
     use client::Client;
 
+    impl Client<TestStream> {
+        pub fn new_test(file: File) -> Self {
+            TestStream::new(file).into()
+        }
+    }
+
     macro_rules! impl_test_prelude {
         ( $testcase:expr ) => {{
             let data_in = File::open(format!("./test_data/{}.in", $testcase)).unwrap();
@@ -665,7 +644,7 @@ mod test {
             let mut data_out = File::open(format!("./test_data/{}.out", $testcase)).unwrap();
             let mut buffer = Vec::new();
             data_out.read_to_end(&mut buffer).unwrap();
-            let stream_buffer: Vec<u8> = $stream.into();
+            let stream_buffer = $stream.stream().lock().unwrap().buffer.clone();
 
             assert_eq!(
                 stream_buffer,

src/stream.rs

@@ -1,6 +1,7 @@
 use std::io::{self, Read, Write};
 use std::sync::{Arc, Mutex};
 
+#[derive(Debug)]
 pub struct ClonableStream<T: Read + Write>(Arc<Mutex<T>>);
 
 impl<T: Read + Write> Read for ClonableStream<T> {
@@ -30,3 +31,10 @@ impl<T: Read + Write> Clone for ClonableStream<T> {
         ClonableStream(Arc::clone(&self.0))
     }
 }
+
+#[cfg(test)]
+impl<T: Read + Write> ClonableStream<T> {
+    pub fn stream(&self) -> Arc<Mutex<T>> {
+        Arc::clone(&self.0)
+    }
+}

src/test_stream.rs

@@ -3,44 +3,31 @@ use std::io::{Read, Result, Write};
 use std::fs::File;
 
 pub struct TestStream {
-    file: Option<File>,
-    buffer: Option<Vec<u8>>,
+    pub file: File,
+    pub buffer: Vec<u8>,
 }
 
 impl TestStream {
-    pub fn new_in(file: File) -> Self {
+    pub fn new(file: File) -> Self {
         TestStream {
-            file: Some(file),
-            buffer: None,
-        }
-    }
-
-    pub fn new_out() -> Self {
-        TestStream {
-            file: None,
-            buffer: Some(Vec::new()),
+            file,
+            buffer: Vec::new(),
         }
     }
 }
 
 impl Read for TestStream {
     fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
-        self.file.as_ref().unwrap().read(buf)
+        self.file.read(buf)
     }
 }
 
 impl Write for TestStream {
     fn write(&mut self, buf: &[u8]) -> Result<usize> {
-        self.buffer.as_mut().unwrap().write(buf)
+        self.buffer.write(buf)
     }
 
     fn flush(&mut self) -> Result<()> {
-        self.buffer.as_mut().unwrap().flush()
-    }
-}
-
-impl Into<Vec<u8>> for TestStream {
-    fn into(self) -> Vec<u8> {
-        self.buffer.unwrap()
+        self.buffer.flush()
     }
 }