From db90a18e5e4423726e19eb4c848ab526b4cc9618 Mon Sep 17 00:00:00 2001
From: Brian Martin <brayniac@gmail.com>
Date: Sun, 14 Feb 2016 12:00:30 -0800
Subject: [PATCH] add FramedTransport

- add a FramedTransport using a Vec<u8> tx buffer
- switch RwTransport to use `T` notation
---
 lib/rs/Cargo.toml           |  1 +
 lib/rs/src/lib.rs           |  1 +
 lib/rs/src/transport/mod.rs | 73 +++++++++++++++++++++++++++++++++++--
 3 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/lib/rs/Cargo.toml b/lib/rs/Cargo.toml
index 3cdabf68..caed8752 100644
--- a/lib/rs/Cargo.toml
+++ b/lib/rs/Cargo.toml
@@ -13,4 +13,5 @@ license = "MIT"
 podio = "0.1"
 log = "0"
 ordered-float = "0"
+byteorder = "0.4.2"
 
diff --git a/lib/rs/src/lib.rs b/lib/rs/src/lib.rs
index 8b000f30..af34785f 100644
--- a/lib/rs/src/lib.rs
+++ b/lib/rs/src/lib.rs
@@ -1,6 +1,7 @@
 #![recursion_limit="1024"]
 extern crate podio;
 extern crate ordered_float;
+extern crate byteorder;
 
 #[macro_use]
 extern crate log;
diff --git a/lib/rs/src/transport/mod.rs b/lib/rs/src/transport/mod.rs
index 91a3a264..4b02c42f 100644
--- a/lib/rs/src/transport/mod.rs
+++ b/lib/rs/src/transport/mod.rs
@@ -19,6 +19,8 @@
 
 use std::io::{self, Read, Write};
 
+use byteorder::{ByteOrder, BigEndian};
+
 pub mod server;
 
 pub trait Transport: Write + Read { }
@@ -26,15 +28,15 @@ pub trait Transport: Write + Read { }
 impl<'t, T> Transport for &'t mut T where T: Transport {}
 impl<'t> Transport for &'t mut Transport {}
 
-pub struct RwTransport<Rw>(pub Rw);
+pub struct RwTransport<T>(pub T);
 
-impl<R: Read> Read for RwTransport<R> {
+impl<T: Read> Read for RwTransport<T> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         self.0.read(buf)
     }
 }
 
-impl<W: Write> Write for RwTransport<W> {
+impl<T: Write> Write for RwTransport<T> {
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
         self.0.write(buf)
     }
@@ -44,5 +46,68 @@ impl<W: Write> Write for RwTransport<W> {
     }
 }
 
-impl<Rw: Read + Write> Transport for RwTransport<Rw> { }
+impl<T: Read + Write> Transport for RwTransport<T> { }
+
+pub struct FramedTransport<T> {
+    pub transport: T,
+    tx_buffer: Vec<u8>,
+    flushed: bool,
+}
+
+impl<T> FramedTransport<T> {
+    pub fn new(transport: T) -> FramedTransport<T> {
+        FramedTransport {
+            transport: transport,
+            tx_buffer: Vec::<u8>::new(),
+            flushed: true,
+        }
+    }
+
+    fn flush(&mut self) -> &Self {
+        self.tx_buffer.clear();
+        self.flushed = true;
+        self
+    }
+}
+
+impl<T: Read> Read for FramedTransport<T> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+
+        loop {
+            if self.flushed {
+                // read the size of the frame
+                let mut b = [0; 4];
+                if let Ok(_) = self.transport.read_exact(&mut b) {
+                    self.flushed = false;
+                }
+            } else {
+                break;
+            }
+        }
+
+        self.transport.read(buf)
+    }
+}
+
+impl<T: Write> Write for FramedTransport<T> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.tx_buffer.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        // write all buffered data to sink
+        let mut size = [0; 4];
+        BigEndian::write_i32(&mut size, self.tx_buffer.len() as i32);
+        try!(self.transport.write(&size));
+        try!(self.transport.write(&self.tx_buffer));
+
+        // flush self
+        self.flush();
+
+        // flush the sink
+        self.transport.flush()
+    }
+}
 
+impl<T: Read + Write> Transport for FramedTransport<T> { }