Add client implementation

5 files changed, 363 insertions, 5 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 585e3b6..1ae1fb6 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -21,6 +21,9 @@ log = "0.4"
 simple_logger = "0.5"
 yaml-rust = "0.3"
 ctrlc = { version = "3.1", features = ["termination"] }
+clap = "2"
+chrono = "0.4"
+hex = "0.3"
 
 [dev-dependencies]
 criterion = "0.2.0"
@@ -28,3 +31,4 @@ criterion = "0.2.0"
 [[bench]]
 name = "messages"
 harness = false
+
diff --git a/src/bin/client.rs b/src/bin/client.rs
new file mode 100644
index 0000000..d5bd3d4
--- /dev/null
+++ b/src/bin/client.rs
@@ -0,0 +1,227 @@
+#[macro_use]
+extern crate clap;
+extern crate roughenough;
+extern crate ring;
+extern crate time;
+extern crate chrono;
+extern crate byteorder;
+extern crate hex;
+
+use ring::rand;
+use ring::rand::SecureRandom;
+use ring::digest;
+
+use byteorder::{LittleEndian, ReadBytesExt};
+
+use chrono::TimeZone;
+use chrono::offset::Utc;
+
+use std::iter::Iterator;
+use std::collections::HashMap;
+use std::net::{UdpSocket, ToSocketAddrs};
+
+use roughenough::{RtMessage, Tag, TREE_NODE_TWEAK, TREE_LEAF_TWEAK, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT};
+use roughenough::sign::Verifier;
+use clap::{Arg, App};
+
+fn create_nonce() -> [u8; 64] {
+    let rng = rand::SystemRandom::new();
+    let mut nonce = [0u8; 64];
+    rng.fill(&mut nonce).unwrap();
+
+    nonce
+}
+
+fn make_request(nonce: &[u8]) -> Vec<u8> {
+    let mut msg = RtMessage::new(1);
+    msg.add_field(Tag::NONC, nonce).unwrap();
+    msg.pad_to_kilobyte();
+
+    msg.encode().unwrap()
+}
+
+fn receive_response(sock: &mut UdpSocket) -> RtMessage {
+    let mut buf = [0; 744];
+    let resp_len = sock.recv_from(&mut buf).unwrap().0;
+    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
+}
+
+
+struct ResponseHandler {
+    pub_key: Option<Vec<u8>>,
+    msg: HashMap<Tag, Vec<u8>>,
+    srep: HashMap<Tag, Vec<u8>>,
+    cert: HashMap<Tag, Vec<u8>>,
+    dele: HashMap<Tag, Vec<u8>>,
+    nonce: [u8; 64]
+}
+
+impl ResponseHandler {
+    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
+        let msg = response.into_hash_map();
+        let srep = RtMessage::from_bytes(&msg[&Tag::SREP]).unwrap().into_hash_map();
+        let cert = RtMessage::from_bytes(&msg[&Tag::CERT]).unwrap().into_hash_map();
+
+        let dele = RtMessage::from_bytes(&cert[&Tag::DELE]).unwrap().into_hash_map();
+
+        ResponseHandler {
+            pub_key,
+            msg,
+            srep,
+            cert,
+            dele,
+            nonce
+        }
+    }
+
+    pub fn extract_time(&self) -> (u64, u32) {
+        let midpoint = self.srep[&Tag::MIDP].as_slice().read_u64::<LittleEndian>().unwrap();
+        let radius = self.srep[&Tag::RADI].as_slice().read_u32::<LittleEndian>().unwrap();
+
+        if self.pub_key.is_some() {
+            self.validate_dele();
+            self.validate_srep();
+            self.validate_merkle();
+            self.validate_midpoint(midpoint);
+        }
+
+        (midpoint, radius)
+    }
+
+    fn validate_dele(&self) {
+        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
+        full_cert.extend(&self.cert[&Tag::DELE]);
+
+        assert!(self.validate_sig(self.pub_key.as_ref().unwrap(), &self.cert[&Tag::SIG], &full_cert),
+                "Invalid signature on DELE tag!");
+    }
+
+    fn validate_srep(&self) {
+        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
+        full_srep.extend(&self.msg[&Tag::SREP]);
+
+        assert!(self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
+                "Invalid signature on SREP tag!");
+    }
+
+    fn validate_merkle(&self) {
+        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP]).unwrap().into_hash_map();
+        let mut index = self.msg[&Tag::INDX].as_slice().read_u32::<LittleEndian>().unwrap();
+        let paths = &self.msg[&Tag::PATH];
+
+        let mut hash = sha_512(TREE_LEAF_TWEAK, &self.nonce);
+
+        assert_eq!(paths.len() % 64, 0);
+
+        for path in paths.chunks(64) {
+            let mut ctx = digest::Context::new(&digest::SHA512);
+            ctx.update(TREE_NODE_TWEAK);
+
+            if index & 1 == 0 {
+                // Left
+                ctx.update(&hash);
+                ctx.update(path);
+            } else {
+                // Right
+                ctx.update(path);
+                ctx.update(&hash);
+            }
+            hash = Vec::from(ctx.finish().as_ref());
+
+            index >>= 1;
+        }
+
+        assert_eq!(hash, srep[&Tag::ROOT], "Nonce not in merkle tree!");
+
+    }
+
+    fn validate_midpoint(&self, midpoint: u64) {
+        let mint = self.dele[&Tag::MINT].as_slice().read_u64::<LittleEndian>().unwrap();
+        let maxt = self.dele[&Tag::MAXT].as_slice().read_u64::<LittleEndian>().unwrap();
+
+        assert!(midpoint >= mint, "Response midpoint {} lies before delegation span ({}, {})");
+        assert!(midpoint <= maxt, "Response midpoint {} lies after delegation span ({}, {})");
+    }
+
+    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
+        let mut verifier = Verifier::new(public_key);
+        verifier.update(data);
+        verifier.verify(sig)
+    }
+}
+
+fn sha_512(prefix: &[u8], data: &[u8]) -> Vec<u8> {
+    let mut ctx = digest::Context::new(&digest::SHA512);
+    ctx.update(prefix);
+    ctx.update(data);
+    Vec::from(ctx.finish().as_ref())
+}
+
+fn main() {
+    let matches = App::new("roughenough client")
+                      .version("0.1.1")
+                      .arg(Arg::with_name("host")
+                           .required(true)
+                           .help("The Roughtime server to connect to")
+                           .takes_value(true))
+                      .arg(Arg::with_name("port")
+                           .required(true)
+                           .help("The Roughtime server port to connect to")
+                           .takes_value(true))
+                      .arg(Arg::with_name("public-key")
+                           .short("p")
+                           .long("public-key")
+                           .takes_value(true)
+                           .help("The server public key used to validate responses. If unset, no validation will be performed"))
+                      .arg(Arg::with_name("time-format")
+                           .short("f")
+                           .long("time-format")
+                           .takes_value(true)
+                           .help("The strftime format string used to print the time recieved from the server")
+                           .default_value("%b %d %Y %H:%M:%S")
+                        )
+                      .arg(Arg::with_name("num-requests")
+                           .short("n")
+                           .long("num-requests")
+                           .takes_value(true)
+                           .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling")
+                           .default_value("1")
+                        )
+                      .get_matches();
+
+    let host = matches.value_of("host").unwrap();
+    let port = value_t_or_exit!(matches.value_of("port"), u16);
+    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
+    let pub_key = matches.value_of("public-key").map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
+    let time_format = matches.value_of("time-format").unwrap();
+
+    println!("Requesting time from: {:?}:{:?}", host, port);
+
+    let addrs: Vec<_> = (host, port).to_socket_addrs().unwrap().collect();
+
+    let mut requests = Vec::with_capacity(num_requests);
+
+    for _ in 0..num_requests {
+        let nonce = create_nonce();
+        let mut socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
+        let request = make_request(&nonce);
+
+        requests.push((nonce, request, socket));
+    }
+
+    for &mut (_, ref request, ref mut socket) in requests.iter_mut() {
+        socket.send_to(request, addrs.as_slice()).unwrap();
+    }
+
+    for (nonce, _, mut socket) in requests {
+        let resp = receive_response(&mut socket);
+
+        let (midpoint, radius) = ResponseHandler::new(pub_key.clone(), resp, nonce).extract_time();
+
+        let seconds = midpoint / 10_u64.pow(6);
+        let spec = Utc.timestamp(seconds as i64, ((midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3)) as u32);
+        let out = spec.format(time_format).to_string();
+
+        println!("Recieved time from server: midpoint={:?}, radius={:?}", out, radius);
+    }
+}
diff --git a/src/error.rs b/src/error.rs
index 9b63f66..decc0bc 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -22,6 +22,8 @@ pub enum Error {
     /// The associated tag was added to an `RtMessage` in non-increasing order.
     TagNotStrictlyIncreasing(Tag),
 
+    InvalidTag(Box<[u8]>),
+
     /// Encoding failed. The associated `std::io::Error` should provide more information.
     EncodingFailure(std::io::Error),
 
diff --git a/src/message.rs b/src/message.rs
index 6b093a4..8a2e1b6 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::io::Write;
-use byteorder::{LittleEndian, WriteBytesExt};
+use std::io::{Read, Write, Cursor};
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+use std::iter::once;
+use std::collections::HashMap;
 
 use tag::Tag;
 use error::Error;
@@ -21,7 +23,7 @@ use error::Error;
 ///
 /// A Roughtime protocol message; a map of u32 tags to arbitrary byte-strings.
 ///
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct RtMessage {
     tags: Vec<Tag>,
     values: Vec<Vec<u8>>,
@@ -34,13 +36,77 @@ impl RtMessage {
     ///
     /// * `num_fields` - Reserve space for this many fields.
     ///
-    pub fn new(num_fields: u8) -> Self {
+    pub fn new(num_fields: u32) -> Self {
         RtMessage {
             tags: Vec::with_capacity(num_fields as usize),
             values: Vec::with_capacity(num_fields as usize),
         }
     }
 
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        let mut msg = Cursor::new(bytes);
+
+        let num_tags = msg.read_u32::<LittleEndian>()?;
+        let mut rt_msg = RtMessage::new(num_tags);
+
+        if num_tags == 1 {
+            let pos = msg.position() as usize;
+            let tag = Tag::from_wire(&bytes[pos..pos+4])?;
+            msg.set_position((pos + 4) as u64);
+
+            let mut value = Vec::new();
+            
+            msg.read_to_end(&mut value).unwrap();
+            rt_msg.add_field(tag, &value)?;
+            return Ok(rt_msg)
+        }
+
+        let mut offsets = Vec::with_capacity((num_tags - 1) as usize);
+        let mut tags = Vec::with_capacity(num_tags as usize);
+
+        for _ in 0..num_tags - 1 {
+            let offset =  msg.read_u32::<LittleEndian>()?;
+            if offset % 4 != 0 {
+                panic!("Invalid offset {:?} in message {:?}", offset, bytes);
+            }
+            offsets.push(offset as usize);
+        }
+
+
+        let mut buf = [0; 4];
+        for _ in 0..num_tags {
+            msg.read_exact(&mut buf).unwrap();
+            let tag = Tag::from_wire(&buf)?;
+
+            if let Some(last_tag) = tags.last() {
+                if tag <= *last_tag {
+                    return Err(Error::TagNotStrictlyIncreasing(tag))
+                }
+            }
+            tags.push(tag);
+        }
+
+        // All offsets are relative to the end of the header,
+        // which is our current position
+        let header_end = msg.position() as usize;
+        // Compute the end of the last value,
+        // as an offset from the end of the header
+        let msg_end = bytes.len() - header_end;
+
+        assert_eq!(offsets.len(), tags.len() - 1);
+
+        for (tag, (value_start, value_end)) in tags.into_iter().zip(
+            once(&0).chain(offsets.iter()).zip(
+                offsets.iter().chain(once(&msg_end))
+            )
+        ) {
+
+            let value = bytes[(header_end + value_start)..(header_end + value_end)].to_vec();
+            rt_msg.add_field(tag, &value)?;
+        }
+        Ok(rt_msg)
+    }
+
     /// Add a field to this `RtMessage`
     ///
     /// ## Arguments
@@ -69,6 +135,18 @@ impl RtMessage {
         self.tags.len() as u32
     }
 
+    pub fn tags(&self) -> &[Tag] {
+        &self.tags
+    }
+
+    pub fn values(&self) -> &[Vec<u8>] {
+        &self.values
+    }
+
+    pub fn into_hash_map(self) -> HashMap<Tag, Vec<u8>> {
+        self.tags.into_iter().zip(self.values.into_iter()).collect()
+    }
+
     /// Encode this message into its on-the-wire representation.
     pub fn encode(&self) -> Result<Vec<u8>, Error> {
         let num_tags = self.tags.len();
@@ -112,6 +190,31 @@ impl RtMessage {
 
         4 + tags_size + offsets_size + values_size
     }
+
+    /// Adds a PAD tag to the end of this message, with a length
+    /// set such that the final encoded size of this message is 1KB
+    /// 
+    /// If the encoded size of this message is already >= 1KB,
+    /// this method does nothing
+    pub fn pad_to_kilobyte(&mut self) {
+        let size = self.encoded_size();
+        if size >= 1024 {
+            return;
+        }
+
+        let mut padding_needed = 1024 - size;
+        if self.tags.len() == 1 {
+            // If we currently only have one tag, adding a padding tag will cause
+            // a 32-bit offset values to be written
+            padding_needed -= 4;
+        }
+        padding_needed -= Tag::PAD.wire_value().len();
+        let padding = vec![0; padding_needed];
+
+        self.add_field(Tag::PAD, &padding).unwrap();
+
+        assert_eq!(self.encoded_size(), 1024);
+    }
 }
 
 #[cfg(test)]
diff --git a/src/tag.rs b/src/tag.rs
index 05e5da9..4ec02b0 100644
--- a/src/tag.rs
+++ b/src/tag.rs
@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use error::Error;
+
 /// An unsigned 32-bit value (key) that maps to a byte-string (value).
-#[derive(Debug, PartialEq, PartialOrd)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Hash, Clone)]
 pub enum Tag {
     // Enforcement of the "tags in strictly increasing order" rule is done using the
     // little-endian encoding of the ASCII tag value; e.g. 'SIG\x00' is 0x00474953 and
@@ -56,4 +58,24 @@ impl Tag {
             Tag::SREP => b"SREP",
         }
     }
+
+    pub fn from_wire(bytes: &[u8]) -> Result<Self, Error> {
+        match bytes {
+            b"CERT" => Ok(Tag::CERT),
+            b"DELE" => Ok(Tag::DELE),
+            b"INDX" => Ok(Tag::INDX),
+            b"MAXT" => Ok(Tag::MAXT),
+            b"MIDP" => Ok(Tag::MIDP),
+            b"MINT" => Ok(Tag::MINT),
+            b"NONC" => Ok(Tag::NONC),
+            b"PAD\xff" => Ok(Tag::PAD),
+            b"PATH" => Ok(Tag::PATH),
+            b"PUBK" => Ok(Tag::PUBK),
+            b"RADI" => Ok(Tag::RADI),
+            b"ROOT" => Ok(Tag::ROOT),
+            b"SIG\x00" => Ok(Tag::SIG),
+            b"SREP" => Ok(Tag::SREP),
+            _ => Err(Error::InvalidTag(Box::from(bytes)))
+        }
+    }
 }