<commit_before>use num::BigUint;
use num::traits::{Zero, One, FromPrimitive};
use digit_slice::{DigitSlice, FromDigits};
use {Shape, Noun, NockError, NockResult};

/// Evaluate the nock `*[subject formula]`
pub fn nock_on(mut subject: Noun, mut formula: Noun) -> NockResult {
    loop {
        if let Shape::Cell(ops, tail) = formula.clone().get() {
            match ops.as_u32() {
                // Axis
                Some(0) => return get_axis(&tail, &subject),

                // Just
                Some(1) => return Ok(tail.clone()),

                // Fire
                Some(2) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            let p = try!(nock_on(subject.clone(),
                                                 (*b).clone()));
                            let q = try!(nock_on(subject, (*c).clone()));
                            subject = p;
                            formula = q;
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Depth
                Some(3) => {
                    let p = try!(nock_on(subject.clone(), (*tail).clone()));
                    return match p.get() {
                        Shape::Cell(_, _) => Ok(Noun::from(0u32)),
                        _ => Ok(Noun::from(1u32)),
                    };
                }

                // Bump
                Some(4) => {
                    let p = try!(nock_on(subject.clone(), (*tail).clone()));
                    return match p.get() {
                        Shape::Atom(ref x) => {
                            // TODO: Non-bignum optimization
                            Ok(Noun::from(BigUint::from_digits(x).unwrap() +
                                          BigUint::one()))
                        }
                        _ => Err(NockError),
                    };
                }

                // Same
                Some(5) => {
                    let p = try!(nock_on(subject.clone(), (*tail).clone()));
                    return match p.get() {
                        Shape::Cell(ref a, ref b) => {
                            if a == b {
                                // Yes.
                                return Ok(Noun::from(0u32));
                            } else {
                                // No.
                                return Ok(Noun::from(1u32));
                            }
                        }
                        _ => return Err(NockError),
                    };
                }

                // If
                Some(6) => {
                    if let Some((b, c, d)) = tail.get_122() {
                        let p = try!(nock_on(subject.clone(), (*b).clone()));
                        match p.get() {
                            Shape::Atom(ref x) => {
                                if x == &0u32.as_digits() {
                                    formula = (*c).clone();
                                } else if x == &1u32.as_digits() {
                                    formula = (*d).clone();
                                } else {
                                    return Err(NockError);
                                }
                            }
                            _ => return Err(NockError),
                        }
                        continue;
                    } else {
                        return Err(NockError);
                    }
                }

                // Compose
                Some(7) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            let p = try!(nock_on(subject.clone(),
                                                 (*b).clone()));
                            subject = p;
                            formula = (*c).clone();
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Push
                Some(8) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            let p = try!(nock_on(subject.clone(),
                                                 (*b).clone()));
                            subject = Noun::cell(p, subject);
                            formula = (*c).clone();
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Call
                Some(9) => {
                    match tail.get() {
                        Shape::Cell(ref axis, ref c) => {
                            // Construct core.
                            subject = try!(nock_on(subject.clone(),
                                                   (*c).clone()));
                            // Fetch from core using axis.
                            formula = try!(get_axis(axis, &subject));

                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Hint
                Some(10) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            if let Shape::Cell(_, ref q) = b.get() {
                                // As per spec, try to nock a cell-shaped
                                // hint just to see whether it'll crash.
                                let _ = try!(nock_on(subject.clone(), (*q).clone()));
                            }

                            formula = (*c).clone();
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Unhandled opcode
                Some(_) => {
                    return Err(NockError);
                }

                None => {
                    if let Shape::Cell(_, _) = ops.get() {
                        // Autocons
                        let a = try!(nock_on(subject.clone(), ops.clone()));
                        let b = try!(nock_on(subject, tail.clone()));
                        return Ok(Noun::cell(a, b));
                    } else {
                        return Err(NockError);
                    }
                }
            }
        } else {
            return Err(NockError);
        }
    }
}

/// Evaluate nock `/[axis subject]`
pub fn get_axis(axis: &Noun, subject: &Noun) -> NockResult {
    // TODO: Optimize for small atoms.
    fn fas(x: BigUint, n: &Noun) -> NockResult {
        let two = BigUint::from_u32(2).unwrap();
        let three = BigUint::from_u32(3).unwrap();
        if x == BigUint::zero() {
            return Err(NockError);
        }
        if x == BigUint::one() {
            return Ok(n.clone());
        }
        if let Shape::Cell(ref a, ref b) = n.get() {
            if x == two {
                return Ok((*a).clone());
            } else if x == three {
                return Ok((*b).clone());
            } else {
                let half = x.clone() >> 1;
                let p = try!(fas(half, n));

                if x % two.clone() == BigUint::zero() {
                    fas(two, &p)
                } else {
                    fas(three, &p)
                }
            }
        } else {
            return Err(NockError);
        }
    }

    match axis.get() {
        Shape::Atom(ref x) => fas(BigUint::from_digits(x).unwrap(), subject),
        _ => Err(NockError)
    }
}
<commit_msg>New axis routine<commit_after>use num::BigUint;
use num::traits::One;
use digit_slice::{DigitSlice, FromDigits};
use {Shape, Noun, NockError, NockResult};

/// Evaluate the nock `*[subject formula]`
pub fn nock_on(mut subject: Noun, mut formula: Noun) -> NockResult {
    loop {
        if let Shape::Cell(ops, tail) = formula.clone().get() {
            match ops.as_u32() {
                // Axis
                Some(0) => return get_axis(&tail, &subject),

                // Just
                Some(1) => return Ok(tail.clone()),

                // Fire
                Some(2) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            let p = try!(nock_on(subject.clone(),
                                                 (*b).clone()));
                            let q = try!(nock_on(subject, (*c).clone()));
                            subject = p;
                            formula = q;
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Depth
                Some(3) => {
                    let p = try!(nock_on(subject.clone(), (*tail).clone()));
                    return match p.get() {
                        Shape::Cell(_, _) => Ok(Noun::from(0u32)),
                        _ => Ok(Noun::from(1u32)),
                    };
                }

                // Bump
                Some(4) => {
                    let p = try!(nock_on(subject.clone(), (*tail).clone()));
                    return match p.get() {
                        Shape::Atom(ref x) => {
                            // TODO: Non-bignum optimization
                            Ok(Noun::from(BigUint::from_digits(x).unwrap() +
                                          BigUint::one()))
                        }
                        _ => Err(NockError),
                    };
                }

                // Same
                Some(5) => {
                    let p = try!(nock_on(subject.clone(), (*tail).clone()));
                    return match p.get() {
                        Shape::Cell(ref a, ref b) => {
                            if a == b {
                                // Yes.
                                return Ok(Noun::from(0u32));
                            } else {
                                // No.
                                return Ok(Noun::from(1u32));
                            }
                        }
                        _ => return Err(NockError),
                    };
                }

                // If
                Some(6) => {
                    if let Some((b, c, d)) = tail.get_122() {
                        let p = try!(nock_on(subject.clone(), (*b).clone()));
                        match p.get() {
                            Shape::Atom(ref x) => {
                                if x == &0u32.as_digits() {
                                    formula = (*c).clone();
                                } else if x == &1u32.as_digits() {
                                    formula = (*d).clone();
                                } else {
                                    return Err(NockError);
                                }
                            }
                            _ => return Err(NockError),
                        }
                        continue;
                    } else {
                        return Err(NockError);
                    }
                }

                // Compose
                Some(7) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            let p = try!(nock_on(subject.clone(),
                                                 (*b).clone()));
                            subject = p;
                            formula = (*c).clone();
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Push
                Some(8) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            let p = try!(nock_on(subject.clone(),
                                                 (*b).clone()));
                            subject = Noun::cell(p, subject);
                            formula = (*c).clone();
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Call
                Some(9) => {
                    match tail.get() {
                        Shape::Cell(ref axis, ref c) => {
                            // Construct core.
                            subject = try!(nock_on(subject.clone(),
                                                   (*c).clone()));
                            // Fetch from core using axis.
                            formula = try!(get_axis(axis, &subject));

                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Hint
                Some(10) => {
                    match tail.get() {
                        Shape::Cell(ref b, ref c) => {
                            if let Shape::Cell(_, ref q) = b.get() {
                                // As per spec, try to nock a cell-shaped
                                // hint just to see whether it'll crash.
                                let _ = try!(nock_on(subject.clone(), (*q).clone()));
                            }

                            formula = (*c).clone();
                            continue;
                        }
                        _ => return Err(NockError),
                    }
                }

                // Unhandled opcode
                Some(_) => {
                    return Err(NockError);
                }

                None => {
                    if let Shape::Cell(_, _) = ops.get() {
                        // Autocons
                        let a = try!(nock_on(subject.clone(), ops.clone()));
                        let b = try!(nock_on(subject, tail.clone()));
                        return Ok(Noun::cell(a, b));
                    } else {
                        return Err(NockError);
                    }
                }
            }
        } else {
            return Err(NockError);
        }
    }
}

/// Evaluate nock `/[axis subject]`
pub fn get_axis(axis: &Noun, subject: &Noun) -> NockResult {
    fn fas(x: &[u8], n: usize, mut subject: &Noun) -> NockResult {
        for i in (0..(n - 1)).rev() {
            if let Shape::Cell(ref a, ref b) = subject.get() {
                if bit(x, i) {
                    subject = b;
                } else {
                    subject = a;
                }
            } else {
                return Err(NockError);
            }
        }
        Ok((*subject).clone())
    }

    match axis.get() {
        Shape::Atom(ref x) => {
            let start = msb(x);
            fas(x, start, subject)
        }
        _ => Err(NockError)
    }
}

#[inline]
fn bit(data: &[u8], pos: usize) -> bool {
    data[pos / 8] & (1 << (pos % 8)) != 0
}

/// Return the bit position of the most significant bit, interpreting the data
/// as a little-endian integer.
fn msb(data: &[u8]) -> usize {
    let mut offset = 0;
    let mut ret = 0;
    for byte in data.iter() {
        let mut b = *byte;
        if b == 0 { continue; }
        let mut x = 0;
        while b != 0 {
            x += 1;
            b >>= 1;
        }
        ret = offset * 8 + x;
        offset += 1;
    }
    ret
}

#[cfg(test)]
mod tests {
    use super::msb;

    #[test]
    fn test_msb() {
        assert_eq!(0, msb(&vec![]));
        assert_eq!(0, msb(&vec![0]));
        assert_eq!(1, msb(&vec![1]));
        assert_eq!(4, msb(&vec![15]));
        assert_eq!(5, msb(&vec![16]));
        assert_eq!(13, msb(&vec![123, 16]));
        assert_eq!(13, msb(&vec![123, 16, 0]));
    }
}
<|endoftext|>


<commit_before><commit_msg>Optimize node item offset calculation<commit_after><|endoftext|>


<commit_before>use address::{Address, Addressable};
use messages;
use std::fmt::Debug;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket, IpAddr, Ipv4Addr};
use std::collections::HashMap;
use std::cmp;
use time;
use transaction::TransactionId;

const FAILED_TO_RESPOND_THRESHOLD: usize = 3;
const MINUTES_UNTIL_QUESTIONABLE: i64 = 15;

pub fn deserialize(message: &messages::protobufs::Node) -> Box<Node> {
    let ip = message.get_ip_address();
    let ip = Ipv4Addr::new(ip[0], ip[1], ip[2], ip[3]);
    let port = message.get_port() as u16;
    let address = Address::from_str(message.get_id());
    Box::new(UdpNode::new(address, (ip, port)))
}

pub trait Serialize {
    fn serialize(&self) -> messages::protobufs::Node;
}

#[derive(Debug, PartialEq)]
pub enum Status {
    Good,
    Questionable,
    Bad
}

pub trait Node : Addressable + Debug + Serialize + Send + Sync {
    fn is_bad(&self) -> bool;
    fn is_good(&self) -> bool;
    fn is_questionable(&self) -> bool;
    fn last_seen(&self) -> time::Tm;
    fn received_query(&mut self, transaction_id: TransactionId);
    fn received_response(&mut self, transaction_id: TransactionId);
    fn send(&self, message: Vec<u8>);
    fn sent_query(&mut self, transaction_id: TransactionId);
}

impl PartialEq for Node {
    fn eq(&self, other: &Self) -> bool {
        self.address() == other.address()
    }
}

#[derive(Debug)]
pub struct UdpNode {
    address: Address,
    socket_address: SocketAddr,
    pending_queries: HashMap<TransactionId, time::Tm>,
    has_ever_responded: bool,
    last_received_query: time::Tm,
    last_received_response: time::Tm,
}

impl UdpNode {
    pub fn new<S: ToSocketAddrs>(address: Address, socket_address: S) -> UdpNode {
        let socket_address = socket_address
            .to_socket_addrs()
            .unwrap()
            .next()
            .unwrap();
        UdpNode {
            address: address,
            socket_address: socket_address,
            pending_queries: HashMap::new(),
            has_ever_responded: false,
            last_received_query: time::empty_tm(),
            last_received_response: time::empty_tm()
        }
    }

    fn status(&self) -> Status {
        let time_since_last_seen = time::now_utc() - self.last_seen();

        if self.has_ever_responded &&
            time_since_last_seen < time::Duration::minutes(MINUTES_UNTIL_QUESTIONABLE) {
            Status::Good
        } else if self.pending_queries.len() < FAILED_TO_RESPOND_THRESHOLD {
            Status::Questionable
        } else {
            Status::Bad
        }
    }
}

impl Node for UdpNode {
    fn is_bad(&self) -> bool {
        self.status() == Status::Bad
    }

    fn is_good(&self) -> bool {
        self.status() == Status::Good
    }

    fn is_questionable(&self) -> bool {
        self.status() == Status::Questionable
    }

    fn last_seen(&self) -> time::Tm {
        cmp::max(self.last_received_query, self.last_received_response)
    }

    fn received_query(&mut self, _: TransactionId) {
        self.last_received_query = time::now_utc();
    }

    fn received_response(&mut self, transaction_id: TransactionId) {
        self.last_received_response = time::now_utc();
        if let Some(_queried_at) = self.pending_queries.remove(&transaction_id) {
            self.has_ever_responded = true;
            //let time_to_respond = time::now_utc() - queried_at;
            //debug!("{} took {:?} to respond", transaction_id, time_to_respond);
        } else {
            debug!("{:?} was not expecting response to {}", self.address, transaction_id);
            debug!("pending queries: {:?}", self.pending_queries.len());
        }
    }

    fn send(&self, message: Vec<u8>) {
        let socket = UdpSocket::bind("0.0.0.0:0").unwrap();
        socket.send_to(&message[..], self.socket_address).unwrap();
    }

    fn sent_query(&mut self, transaction_id: TransactionId) {
        self.pending_queries.insert(transaction_id, time::now_utc());
    }
}

impl Serialize for UdpNode {
    fn serialize(&self) -> messages::protobufs::Node {
        let mut message = messages::protobufs::Node::new();
        message.set_id(self.address.to_str());
        match self.socket_address.ip() {
            IpAddr::V4(ipv4_addr) => {
                message.set_ip_address(ipv4_addr.octets().iter().cloned().collect());
            }
            IpAddr::V6(_) => {
                // TODO ipv6 node support
            }
        }
        message.set_port(self.socket_address.port() as u32);
        message
    }
}

impl Addressable for UdpNode {
    fn address(&self) -> Address {
        self.address
    }
}

#[cfg(test)]
mod tests {
    use address::Address;
    use super::{Node, UdpNode, Serialize};

    #[test]
    fn test_received_response() {
        let address = Address::for_content("some string");
        let mut node = UdpNode::new(address, ("0.0.0.0", 9000));

        // When it's not expecting the response
        node.received_response(1);
        assert!(!node.has_ever_responded);

        // When it's expecting the response
        let last_seen_before = node.last_seen();
        let transaction_id = 2;
        node.sent_query(transaction_id);
        node.received_response(transaction_id);
        let last_seen_after = node.last_seen();
        assert!(last_seen_before < last_seen_after);
        assert!(node.has_ever_responded);
    }

    #[test]
    fn test_serialize() {
        use messages;
        let mut protobuf = messages::protobufs::Node::new();
        protobuf.set_id("8b45e4bd1c6acb88bebf6407d16205f567e62a3e".to_string());
        protobuf.set_ip_address(vec![192, 168, 1, 2]);
        protobuf.set_port(9000);
        let address = Address::for_content("some string");
        let node = UdpNode::new(address, ("192.168.1.2", 9000));
        assert_eq!(node.serialize(), protobuf);
    }
}
<commit_msg>Increase bad node threshold to 5 failed responses<commit_after>use address::{Address, Addressable};
use messages;
use std::fmt::Debug;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket, IpAddr, Ipv4Addr};
use std::collections::HashMap;
use std::cmp;
use time;
use transaction::TransactionId;

const FAILED_TO_RESPOND_THRESHOLD: usize = 5;
const MINUTES_UNTIL_QUESTIONABLE: i64 = 15;

pub fn deserialize(message: &messages::protobufs::Node) -> Box<Node> {
    let ip = message.get_ip_address();
    let ip = Ipv4Addr::new(ip[0], ip[1], ip[2], ip[3]);
    let port = message.get_port() as u16;
    let address = Address::from_str(message.get_id());
    Box::new(UdpNode::new(address, (ip, port)))
}

pub trait Serialize {
    fn serialize(&self) -> messages::protobufs::Node;
}

#[derive(Debug, PartialEq)]
pub enum Status {
    Good,
    Questionable,
    Bad
}

pub trait Node : Addressable + Debug + Serialize + Send + Sync {
    fn is_bad(&self) -> bool;
    fn is_good(&self) -> bool;
    fn is_questionable(&self) -> bool;
    fn last_seen(&self) -> time::Tm;
    fn received_query(&mut self, transaction_id: TransactionId);
    fn received_response(&mut self, transaction_id: TransactionId);
    fn send(&self, message: Vec<u8>);
    fn sent_query(&mut self, transaction_id: TransactionId);
}

impl PartialEq for Node {
    fn eq(&self, other: &Self) -> bool {
        self.address() == other.address()
    }
}

#[derive(Debug)]
pub struct UdpNode {
    address: Address,
    socket_address: SocketAddr,
    pending_queries: HashMap<TransactionId, time::Tm>,
    has_ever_responded: bool,
    last_received_query: time::Tm,
    last_received_response: time::Tm,
}

impl UdpNode {
    pub fn new<S: ToSocketAddrs>(address: Address, socket_address: S) -> UdpNode {
        let socket_address = socket_address
            .to_socket_addrs()
            .unwrap()
            .next()
            .unwrap();
        UdpNode {
            address: address,
            socket_address: socket_address,
            pending_queries: HashMap::new(),
            has_ever_responded: false,
            last_received_query: time::empty_tm(),
            last_received_response: time::empty_tm()
        }
    }

    fn status(&self) -> Status {
        let time_since_last_seen = time::now_utc() - self.last_seen();

        if self.has_ever_responded &&
            time_since_last_seen < time::Duration::minutes(MINUTES_UNTIL_QUESTIONABLE) {
            Status::Good
        } else if self.pending_queries.len() < FAILED_TO_RESPOND_THRESHOLD {
            Status::Questionable
        } else {
            Status::Bad
        }
    }
}

impl Node for UdpNode {
    fn is_bad(&self) -> bool {
        self.status() == Status::Bad
    }

    fn is_good(&self) -> bool {
        self.status() == Status::Good
    }

    fn is_questionable(&self) -> bool {
        self.status() == Status::Questionable
    }

    fn last_seen(&self) -> time::Tm {
        cmp::max(self.last_received_query, self.last_received_response)
    }

    fn received_query(&mut self, _: TransactionId) {
        self.last_received_query = time::now_utc();
    }

    fn received_response(&mut self, transaction_id: TransactionId) {
        self.last_received_response = time::now_utc();
        if let Some(_queried_at) = self.pending_queries.remove(&transaction_id) {
            self.has_ever_responded = true;
            //let time_to_respond = time::now_utc() - queried_at;
            //debug!("{} took {:?} to respond", transaction_id, time_to_respond);
        } else {
            debug!("{:?} was not expecting response to {}", self.address, transaction_id);
            debug!("pending queries: {:?}", self.pending_queries.len());
        }
    }

    fn send(&self, message: Vec<u8>) {
        let socket = UdpSocket::bind("0.0.0.0:0").unwrap();
        socket.send_to(&message[..], self.socket_address).unwrap();
    }

    fn sent_query(&mut self, transaction_id: TransactionId) {
        self.pending_queries.insert(transaction_id, time::now_utc());
    }
}

impl Serialize for UdpNode {
    fn serialize(&self) -> messages::protobufs::Node {
        let mut message = messages::protobufs::Node::new();
        message.set_id(self.address.to_str());
        match self.socket_address.ip() {
            IpAddr::V4(ipv4_addr) => {
                message.set_ip_address(ipv4_addr.octets().iter().cloned().collect());
            }
            IpAddr::V6(_) => {
                // TODO ipv6 node support
            }
        }
        message.set_port(self.socket_address.port() as u32);
        message
    }
}

impl Addressable for UdpNode {
    fn address(&self) -> Address {
        self.address
    }
}

#[cfg(test)]
mod tests {
    use address::Address;
    use super::{Node, UdpNode, Serialize};

    #[test]
    fn test_received_response() {
        let address = Address::for_content("some string");
        let mut node = UdpNode::new(address, ("0.0.0.0", 9000));

        // When it's not expecting the response
        node.received_response(1);
        assert!(!node.has_ever_responded);

        // When it's expecting the response
        let last_seen_before = node.last_seen();
        let transaction_id = 2;
        node.sent_query(transaction_id);
        node.received_response(transaction_id);
        let last_seen_after = node.last_seen();
        assert!(last_seen_before < last_seen_after);
        assert!(node.has_ever_responded);
    }

    #[test]
    fn test_serialize() {
        use messages;
        let mut protobuf = messages::protobufs::Node::new();
        protobuf.set_id("8b45e4bd1c6acb88bebf6407d16205f567e62a3e".to_string());
        protobuf.set_ip_address(vec![192, 168, 1, 2]);
        protobuf.set_port(9000);
        let address = Address::for_content("some string");
        let node = UdpNode::new(address, ("192.168.1.2", 9000));
        assert_eq!(node.serialize(), protobuf);
    }
}
<|endoftext|>


<commit_before>use ahash::AHashMap;
use jwalk::WalkDir;
use serde::Deserialize;
use std::io::ErrorKind;
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use structopt::StructOpt;
use watchman_client::prelude::*;

#[derive(StructOpt, Debug)]
#[structopt(about = "Audit Watchman's in-memory database with the filesystem")]
pub(crate) struct AuditCmd {
    #[structopt(name = "path", parse(from_os_str))]
    path: PathBuf,

    #[structopt(
        long = "timeout",
        about = "seconds to wait for Watchman query result",
        default_value = "120"
    )]
    timeout_secs: u64,
}

query_result_type! {
    struct AuditQueryResult {
        name: NameField,
        mode: ModeAndPermissionsField,
        size: SizeField,
        mtime: MTimeField,
        oclock: ObservedClockField,
        ino: InodeNumberField,
    }
}

// Skip cookies when diffing.
fn is_cookie<T: AsRef<Path>>(name: T) -> bool {
    return name
        .as_ref()
        .file_name()
        .and_then(|s| s.to_str())
        .map_or(false, |s| s.starts_with(".watchman-cookie-"));
}

impl AuditCmd {
    pub(crate) async fn run(&self) -> crate::Result<()> {
        let client = Connector::new().connect().await?;
        let resolved = Arc::new(
            client
                .resolve_root(CanonicalPath::canonicalize(&self.path)?)
                .await?,
        );

        if resolved.watcher() == "eden" {
            return Err(watchman_client::Error::Generic(format!(
                "{} is an EdenFS mount - no need to audit",
                resolved.project_root().display()
            )));
        }

        let config = client.get_config(&resolved).await?;
        let mut ignore_dirs = config.ignore_dirs.unwrap_or_default().clone();

        // TODO: This list is duplicated in the Watchman query below.
        ignore_dirs.push(".hg".into());
        ignore_dirs.push(".git".into());
        ignore_dirs.push(".svn".into());

        let filesystem_state_handle = {
            let resolved = resolved.clone();
            tokio::spawn(async move {
                let mut filesystem_state: AHashMap<PathBuf, std::fs::Metadata> = AHashMap::new();

                let start_crawl = Instant::now();

                // Allocate outside the loop to save time.
                let resolved_path = Arc::new(resolved.path());

                let resolved_path_copy = resolved_path.clone();
                let walk_dir = WalkDir::new(&*resolved_path)
                    .skip_hidden(false)
                    .process_read_dir(move |_depth, path, _read_dir_state, children| {
                        let resolved_path: &Path = resolved_path_copy.as_ref();
                        let from_root = match path.strip_prefix(resolved_path) {
                            Ok(from_root) => from_root,
                            Err(_) => {
                                return;
                            }
                        };

                        if from_root == Path::new("") {
                            children.retain(|child| match child {
                                Ok(child) => {
                                    ignore_dirs.iter().all(|i| i.as_os_str() != child.file_name)
                                }
                                Err(_) => true,
                            });
                        }
                    });

                for entry in walk_dir {
                    let entry = match entry {
                        Ok(entry) => entry,
                        Err(err) => {
                            eprintln!("error while traversing directory: {}", err);
                            continue;
                        }
                    };

                    let entry_path = entry.path();
                    let relpath = match entry_path.strip_prefix(&*resolved_path) {
                        Ok(relpath) => relpath,
                        Err(err) => {
                            eprintln!(
                                "unable to form relative path from {} to {}: {}",
                                resolved_path.display(),
                                entry_path.display(),
                                err
                            );
                            continue;
                        }
                    };

                    let metadata = match entry.metadata() {
                        Ok(metadata) => metadata,
                        Err(err) => {
                            if err.io_error().map(|e| e.kind() == ErrorKind::NotFound) != Some(true)
                            {
                                eprintln!(
                                    "error fetching metadata for {}: {}",
                                    entry_path.display(),
                                    err
                                );
                            }
                            continue;
                        }
                    };
                    filesystem_state.insert(relpath.to_path_buf(), metadata);
                }
                // Watchman doesn't return information about the root, so remove it here.
                filesystem_state.remove(&PathBuf::new());
                eprintln!("Crawled filesystem in {:?}", start_crawl.elapsed());

                filesystem_state
            })
        };

        let start_query = Instant::now();

        // Do not ignore fresh instance results: the goal is to validate the
        // correctness of query results against the filesystem state, no
        // matter what happened in the crawler.

        use Expr::*;
        let result = client
            .query::<AuditQueryResult>(
                &resolved,
                QueryRequestCommon {
                    expression: Some(All(vec![
                        Exists,
                        Not(Box::new(Any(vec![
                            DirName(DirNameTerm {
                                path: ".git".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".hg".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".svn".into(),
                                depth: None,
                            }),
                            Name(NameTerm {
                                paths: vec![".git".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".svn".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".hg".into()],
                                wholename: true,
                            }),
                        ]))),
                    ])),
                    sync_timeout: SyncTimeout::Duration(Duration::new(self.timeout_secs, 0)),
                    ..Default::default()
                },
            )
            .await;

        let result = match result {
            Ok(result) => result,
            Err(err) => {
                eprintln!("Error during Watchman query: {}", err);
                // Use a different error code for Watchman query errors, including timeouts, so they can be differentiated by audit logging.
                std::process::exit(2);
            }
        };

        eprintln!(
            "Queried Watchman in {:?} (is_fresh_instance = {}, clock = {:?})",
            start_query.elapsed(),
            result.is_fresh_instance,
            result.clock
        );

        let filesystem_state = filesystem_state_handle.await.unwrap();

        let diff_start = Instant::now();

        let watchman_files = match result.files {
            Some(files) => files,
            None => {
                return Err(watchman_client::Error::Generic(
                    "No files set in result {}".into(),
                ));
            }
        };

        let mut any_differences = false;
        let mut phantoms = vec![];
        let mut missing = vec![];

        let mut watchman_state: AHashMap<&Path, &AuditQueryResult> =
            AHashMap::with_capacity(watchman_files.len());
        for watchman_file in &watchman_files {
            let filename = &*watchman_file.name;
            // Skip cookies when diffing.
            if is_cookie(filename) {
                continue;
            }
            watchman_state.insert(filename, watchman_file);

            let metadata = match filesystem_state.get(filename) {
                Some(metadata) => metadata,
                None => {
                    phantoms.push(watchman_file);
                    continue;
                }
            };

            let mut diffs = Vec::new();

            if *watchman_file.mode != u64::from(metadata.permissions().mode()) {
                diffs.push(format!(
                    "watchman mode is {} vs. fs {}",
                    *watchman_file.mode,
                    metadata.permissions().mode()
                ));
            }

            if metadata.is_file() && *watchman_file.size != metadata.size() {
                diffs.push(format!(
                    "watchman size is {} vs. fs {}",
                    *watchman_file.size,
                    metadata.len()
                ));
            }

            if metadata.is_file() && *watchman_file.mtime != metadata.mtime() {
                diffs.push(format!(
                    "watchman mtime is {} vs. fs {}",
                    *watchman_file.mtime,
                    metadata.mtime()
                ));
            }

            if *watchman_file.ino != metadata.ino() {
                diffs.push(format!(
                    "watchman ino is {} vs. fs {}",
                    *watchman_file.ino,
                    metadata.ino()
                ));
            }

            if !diffs.is_empty() {
                println!(
                    "Conflicting information for {}:",
                    watchman_file.name.display()
                );
                for diff in diffs {
                    println!("  {}", diff);
                }
                println!("  oclock is {:?}", *watchman_file.oclock);
                any_differences = true;
            }
        }

        for (path, val) in &filesystem_state {
            if is_cookie(path) {
                continue;
            }
            if !watchman_state.contains_key(&path.as_path()) {
                missing.push((path, val));
            }
        }

        phantoms.sort_by(|x, y| x.name.cmp(&y.name));
        missing.sort_by(|x, y| x.0.cmp(&y.0));

        if !phantoms.is_empty() {
            println!(
                "There are {} items reported by watchman not on the filesystem:",
                phantoms.len()
            );
            for phantom in &phantoms {
                println!("  {}", phantom.name.display());
            }
            any_differences = true;
        }

        if !missing.is_empty() {
            println!(
                "There are {} items on the filesystem not reported by watchman:",
                missing.len()
            );
            for (path, _) in &missing {
                println!("  {}", path.display());
            }
            any_differences = true;
        }

        if any_differences {
            // This is dumb, but Rust doesn't have a standard way to return
            // nonzero exit codes yet.
            std::process::exit(1);
        }

        eprintln!("Diffed in {:#?}", diff_start.elapsed());

        Ok(())
    }
}
<commit_msg>audit: include the version number in the query result<commit_after>use ahash::AHashMap;
use jwalk::WalkDir;
use serde::Deserialize;
use std::io::ErrorKind;
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use structopt::StructOpt;
use watchman_client::prelude::*;

#[derive(StructOpt, Debug)]
#[structopt(about = "Audit Watchman's in-memory database with the filesystem")]
pub(crate) struct AuditCmd {
    #[structopt(name = "path", parse(from_os_str))]
    path: PathBuf,

    #[structopt(
        long = "timeout",
        about = "seconds to wait for Watchman query result",
        default_value = "120"
    )]
    timeout_secs: u64,
}

query_result_type! {
    struct AuditQueryResult {
        name: NameField,
        mode: ModeAndPermissionsField,
        size: SizeField,
        mtime: MTimeField,
        oclock: ObservedClockField,
        ino: InodeNumberField,
    }
}

// Skip cookies when diffing.
fn is_cookie<T: AsRef<Path>>(name: T) -> bool {
    return name
        .as_ref()
        .file_name()
        .and_then(|s| s.to_str())
        .map_or(false, |s| s.starts_with(".watchman-cookie-"));
}

impl AuditCmd {
    pub(crate) async fn run(&self) -> crate::Result<()> {
        let client = Connector::new().connect().await?;
        let resolved = Arc::new(
            client
                .resolve_root(CanonicalPath::canonicalize(&self.path)?)
                .await?,
        );

        if resolved.watcher() == "eden" {
            return Err(watchman_client::Error::Generic(format!(
                "{} is an EdenFS mount - no need to audit",
                resolved.project_root().display()
            )));
        }

        let config = client.get_config(&resolved).await?;
        let mut ignore_dirs = config.ignore_dirs.unwrap_or_default().clone();

        // TODO: This list is duplicated in the Watchman query below.
        ignore_dirs.push(".hg".into());
        ignore_dirs.push(".git".into());
        ignore_dirs.push(".svn".into());

        let filesystem_state_handle = {
            let resolved = resolved.clone();
            tokio::spawn(async move {
                let mut filesystem_state: AHashMap<PathBuf, std::fs::Metadata> = AHashMap::new();

                let start_crawl = Instant::now();

                // Allocate outside the loop to save time.
                let resolved_path = Arc::new(resolved.path());

                let resolved_path_copy = resolved_path.clone();
                let walk_dir = WalkDir::new(&*resolved_path)
                    .skip_hidden(false)
                    .process_read_dir(move |_depth, path, _read_dir_state, children| {
                        let resolved_path: &Path = resolved_path_copy.as_ref();
                        let from_root = match path.strip_prefix(resolved_path) {
                            Ok(from_root) => from_root,
                            Err(_) => {
                                return;
                            }
                        };

                        if from_root == Path::new("") {
                            children.retain(|child| match child {
                                Ok(child) => {
                                    ignore_dirs.iter().all(|i| i.as_os_str() != child.file_name)
                                }
                                Err(_) => true,
                            });
                        }
                    });

                for entry in walk_dir {
                    let entry = match entry {
                        Ok(entry) => entry,
                        Err(err) => {
                            eprintln!("error while traversing directory: {}", err);
                            continue;
                        }
                    };

                    let entry_path = entry.path();
                    let relpath = match entry_path.strip_prefix(&*resolved_path) {
                        Ok(relpath) => relpath,
                        Err(err) => {
                            eprintln!(
                                "unable to form relative path from {} to {}: {}",
                                resolved_path.display(),
                                entry_path.display(),
                                err
                            );
                            continue;
                        }
                    };

                    let metadata = match entry.metadata() {
                        Ok(metadata) => metadata,
                        Err(err) => {
                            if err.io_error().map(|e| e.kind() == ErrorKind::NotFound) != Some(true)
                            {
                                eprintln!(
                                    "error fetching metadata for {}: {}",
                                    entry_path.display(),
                                    err
                                );
                            }
                            continue;
                        }
                    };
                    filesystem_state.insert(relpath.to_path_buf(), metadata);
                }
                // Watchman doesn't return information about the root, so remove it here.
                filesystem_state.remove(&PathBuf::new());
                eprintln!("Crawled filesystem in {:?}", start_crawl.elapsed());

                filesystem_state
            })
        };

        let start_query = Instant::now();

        // Do not ignore fresh instance results: the goal is to validate the
        // correctness of query results against the filesystem state, no
        // matter what happened in the crawler.

        use Expr::*;
        let result = client
            .query::<AuditQueryResult>(
                &resolved,
                QueryRequestCommon {
                    expression: Some(All(vec![
                        Exists,
                        Not(Box::new(Any(vec![
                            DirName(DirNameTerm {
                                path: ".git".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".hg".into(),
                                depth: None,
                            }),
                            DirName(DirNameTerm {
                                path: ".svn".into(),
                                depth: None,
                            }),
                            Name(NameTerm {
                                paths: vec![".git".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".svn".into()],
                                wholename: true,
                            }),
                            Name(NameTerm {
                                paths: vec![".hg".into()],
                                wholename: true,
                            }),
                        ]))),
                    ])),
                    sync_timeout: SyncTimeout::Duration(Duration::new(self.timeout_secs, 0)),
                    ..Default::default()
                },
            )
            .await;

        let result = match result {
            Ok(result) => result,
            Err(err) => {
                eprintln!("Error during Watchman query: {}", err);
                // Use a different error code for Watchman query errors, including timeouts, so they can be differentiated by audit logging.
                std::process::exit(2);
            }
        };

        eprintln!(
            "Queried Watchman in {:?} (is_fresh_instance = {}, clock = {:?}, version = {})",
            start_query.elapsed(),
            result.is_fresh_instance,
            result.clock,
            result.version
        );

        let filesystem_state = filesystem_state_handle.await.unwrap();

        let diff_start = Instant::now();

        let watchman_files = match result.files {
            Some(files) => files,
            None => {
                return Err(watchman_client::Error::Generic(
                    "No files set in result {}".into(),
                ));
            }
        };

        let mut any_differences = false;
        let mut phantoms = vec![];
        let mut missing = vec![];

        let mut watchman_state: AHashMap<&Path, &AuditQueryResult> =
            AHashMap::with_capacity(watchman_files.len());
        for watchman_file in &watchman_files {
            let filename = &*watchman_file.name;
            // Skip cookies when diffing.
            if is_cookie(filename) {
                continue;
            }
            watchman_state.insert(filename, watchman_file);

            let metadata = match filesystem_state.get(filename) {
                Some(metadata) => metadata,
                None => {
                    phantoms.push(watchman_file);
                    continue;
                }
            };

            let mut diffs = Vec::new();

            if *watchman_file.mode != u64::from(metadata.permissions().mode()) {
                diffs.push(format!(
                    "watchman mode is {} vs. fs {}",
                    *watchman_file.mode,
                    metadata.permissions().mode()
                ));
            }

            if metadata.is_file() && *watchman_file.size != metadata.size() {
                diffs.push(format!(
                    "watchman size is {} vs. fs {}",
                    *watchman_file.size,
                    metadata.len()
                ));
            }

            if metadata.is_file() && *watchman_file.mtime != metadata.mtime() {
                diffs.push(format!(
                    "watchman mtime is {} vs. fs {}",
                    *watchman_file.mtime,
                    metadata.mtime()
                ));
            }

            if *watchman_file.ino != metadata.ino() {
                diffs.push(format!(
                    "watchman ino is {} vs. fs {}",
                    *watchman_file.ino,
                    metadata.ino()
                ));
            }

            if !diffs.is_empty() {
                println!(
                    "Conflicting information for {}:",
                    watchman_file.name.display()
                );
                for diff in diffs {
                    println!("  {}", diff);
                }
                println!("  oclock is {:?}", *watchman_file.oclock);
                any_differences = true;
            }
        }

        for (path, val) in &filesystem_state {
            if is_cookie(path) {
                continue;
            }
            if !watchman_state.contains_key(&path.as_path()) {
                missing.push((path, val));
            }
        }

        phantoms.sort_by(|x, y| x.name.cmp(&y.name));
        missing.sort_by(|x, y| x.0.cmp(&y.0));

        if !phantoms.is_empty() {
            println!(
                "There are {} items reported by watchman not on the filesystem:",
                phantoms.len()
            );
            for phantom in &phantoms {
                println!("  {}", phantom.name.display());
            }
            any_differences = true;
        }

        if !missing.is_empty() {
            println!(
                "There are {} items on the filesystem not reported by watchman:",
                missing.len()
            );
            for (path, _) in &missing {
                println!("  {}", path.display());
            }
            any_differences = true;
        }

        if any_differences {
            // This is dumb, but Rust doesn't have a standard way to return
            // nonzero exit codes yet.
            std::process::exit(1);
        }

        eprintln!("Diffed in {:#?}", diff_start.elapsed());

        Ok(())
    }
}
<|endoftext|>


<commit_before><commit_msg>add some comments for possible log integration strategy<commit_after><|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::config::Error;
use diem_secure_storage::{
    GitHubStorage, InMemoryStorage, Namespaced, OnDiskStorage, Storage, VaultStorage,
};
use serde::{Deserialize, Serialize};
use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
};

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum SecureBackend {
    GitHub(GitHubConfig),
    InMemoryStorage,
    Vault(VaultConfig),
    OnDiskStorage(OnDiskStorageConfig),
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct GitHubConfig {
    /// The owner or account that hosts a repository
    pub repository_owner: String,
    /// The repository where storage will mount
    pub repository: String,
    /// The branch containing storage, defaults to master
    pub branch: Option<String>,
    /// The authorization token for accessing the repository
    pub token: Token,
    /// A namespace is an optional portion of the path to a key stored within GitHubConfig. For
    /// example, a key, S, without a namespace would be available in S, with a namespace, N, it
    /// would be in N/S.
    pub namespace: Option<String>,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct VaultConfig {
    /// Optional SSL Certificate for the vault host, this is expected to be a full path.
    pub ca_certificate: Option<PathBuf>,
    /// A namespace is an optional portion of the path to a key stored within Vault. For example,
    /// a secret, S, without a namespace would be available in secret/data/S, with a namespace, N, it
    /// would be in secret/data/N/S.
    pub namespace: Option<String>,
    /// Vault leverages leases on many tokens, specify this to automatically have your lease
    /// renewed up to that many seconds more. If this is not specified, the lease will not
    /// automatically be renewed.
    pub renew_ttl_secs: Option<u32>,
    /// Vault's URL, note: only HTTP is currently supported.
    pub server: String,
    /// The authorization token for accessing secrets
    pub token: Token,
    /// Disable check-and-set when writing secrets to Vault
    pub disable_cas: Option<bool>,
    /// Timeout for new vault socket connections, in milliseconds.
    pub connection_timeout_ms: Option<u64>,
    /// Timeout for generic vault operations (e.g., reads and writes), in milliseconds.
    pub response_timeout_ms: Option<u64>,
}

impl VaultConfig {
    pub fn ca_certificate(&self) -> Result<String, Error> {
        let path = self
            .ca_certificate
            .as_ref()
            .ok_or(Error::Missing("ca_certificate"))?;
        read_file(path)
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct OnDiskStorageConfig {
    // Required path for on disk storage
    pub path: PathBuf,
    /// A namespace is an optional portion of the path to a key stored within OnDiskStorage. For
    /// example, a key, S, without a namespace would be available in S, with a namespace, N, it
    /// would be in N/S.
    pub namespace: Option<String>,
    #[serde(skip)]
    data_dir: PathBuf,
}

/// Tokens can either be directly within this config or stored somewhere on disk.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Token {
    FromConfig(String),
    /// This is an absolute path and not relative to data_dir
    FromDisk(PathBuf),
}

impl Token {
    pub fn read_token(&self) -> Result<String, Error> {
        match self {
            Token::FromDisk(path) => read_file(path),
            Token::FromConfig(token) => Ok(token.clone()),
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct TokenFromConfig {
    token: String,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct TokenFromDisk {
    path: PathBuf,
}

impl Default for OnDiskStorageConfig {
    fn default() -> Self {
        Self {
            namespace: None,
            path: PathBuf::from("secure_storage.json"),
            data_dir: PathBuf::from("/opt/diem/data"),
        }
    }
}

impl OnDiskStorageConfig {
    pub fn path(&self) -> PathBuf {
        if self.path.is_relative() {
            self.data_dir.join(&self.path)
        } else {
            self.path.clone()
        }
    }

    pub fn set_data_dir(&mut self, data_dir: PathBuf) {
        self.data_dir = data_dir;
    }
}

fn read_file(path: &Path) -> Result<String, Error> {
    let mut file =
        File::open(path).map_err(|e| Error::IO(path.to_str().unwrap().to_string(), e))?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .map_err(|e| Error::IO(path.to_str().unwrap().to_string(), e))?;
    Ok(contents)
}

impl From<&SecureBackend> for Storage {
    fn from(backend: &SecureBackend) -> Self {
        match backend {
            SecureBackend::GitHub(config) => {
                let storage = Storage::from(GitHubStorage::new(
                    config.repository_owner.clone(),
                    config.repository.clone(),
                    config
                        .branch
                        .as_ref()
                        .cloned()
                        .unwrap_or_else(|| "master".to_string()),
                    config.token.read_token().expect("Unable to read token"),
                ));
                if let Some(namespace) = &config.namespace {
                    Storage::from(Namespaced::new(namespace, Box::new(storage)))
                } else {
                    storage
                }
            }
            SecureBackend::InMemoryStorage => Storage::from(InMemoryStorage::new()),
            SecureBackend::OnDiskStorage(config) => {
                let storage = Storage::from(OnDiskStorage::new(config.path()));
                if let Some(namespace) = &config.namespace {
                    Storage::from(Namespaced::new(namespace, Box::new(storage)))
                } else {
                    storage
                }
            }
            SecureBackend::Vault(config) => {
                let storage = Storage::from(VaultStorage::new(
                    config.server.clone(),
                    config.token.read_token().expect("Unable to read token"),
                    config
                        .ca_certificate
                        .as_ref()
                        .map(|_| config.ca_certificate().unwrap()),
                    config.renew_ttl_secs,
                    config.disable_cas.map_or_else(|| true, |disable| !disable),
                    config.connection_timeout_ms,
                    config.response_timeout_ms,
                ));
                if let Some(namespace) = &config.namespace {
                    Storage::from(Namespaced::new(namespace, Box::new(storage)))
                } else {
                    storage
                }
            }
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    #[derive(Debug, Deserialize, PartialEq, Serialize)]
    struct Config {
        vault: VaultConfig,
    }

    #[test]
    fn test_token_config_parsing() {
        let from_config = Config {
            vault: VaultConfig {
                namespace: None,
                server: "127.0.0.1:8200".to_string(),
                ca_certificate: None,
                token: Token::FromConfig("test".to_string()),
                renew_ttl_secs: None,
                disable_cas: None,
                connection_timeout_ms: None,
                response_timeout_ms: None,
            },
        };

        let text_from_config = r#"
vault:
    server: "127.0.0.1:8200"
    token:
        from_config: "test"
        "#;

        let de_from_config: Config = serde_yaml::from_str(text_from_config).unwrap();
        assert_eq!(de_from_config, from_config);
        // Just assert that it can be serialized, not about to do string comparison
        serde_yaml::to_string(&from_config).unwrap();
    }

    #[test]
    fn test_vault_timeout_parsing() {
        let from_config = Config {
            vault: VaultConfig {
                namespace: None,
                server: "127.0.0.1:8200".to_string(),
                ca_certificate: None,
                token: Token::FromConfig("test".to_string()),
                renew_ttl_secs: None,
                disable_cas: None,
                connection_timeout_ms: Some(3000),
                response_timeout_ms: Some(5000),
            },
        };

        let text_from_config = r#"
vault:
    server: "127.0.0.1:8200"
    token:
        from_config: "test"
    connection_timeout_ms: 3000
    response_timeout_ms: 5000
        "#;

        let de_from_config: Config = serde_yaml::from_str(text_from_config).unwrap();
        assert_eq!(de_from_config, from_config);
        // Just assert that it can be serialized, no need to do string comparison
        serde_yaml::to_string(&from_config).unwrap();
    }

    #[test]
    fn test_token_disk_parsing() {
        let from_disk = Config {
            vault: VaultConfig {
                namespace: None,
                server: "127.0.0.1:8200".to_string(),
                ca_certificate: None,
                token: Token::FromDisk(PathBuf::from("/token")),
                renew_ttl_secs: None,
                disable_cas: None,
                connection_timeout_ms: None,
                response_timeout_ms: None,
            },
        };

        let text_from_disk = r#"
vault:
    server: "127.0.0.1:8200"
    token:
        from_disk: "/token"
        "#;

        let de_from_disk: Config = serde_yaml::from_str(text_from_disk).unwrap();
        assert_eq!(de_from_disk, from_disk);
        // Just assert that it can be serialized, not about to do string comparison
        serde_yaml::to_string(&from_disk).unwrap();
    }

    #[test]
    fn test_token_reading() {
        let temppath = diem_temppath::TempPath::new();
        temppath.create_as_file().unwrap();
        let mut file = File::create(temppath.path()).unwrap();
        file.write_all(b"disk_token").unwrap();

        let disk = Token::FromDisk(temppath.path().to_path_buf());
        assert_eq!("disk_token", disk.read_token().unwrap());

        let config = Token::FromConfig("config_token".to_string());
        assert_eq!("config_token", config.read_token().unwrap());
    }
}
<commit_msg>config: add helper for querying if a backend config has a namespace<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::config::Error;
use diem_secure_storage::{
    GitHubStorage, InMemoryStorage, Namespaced, OnDiskStorage, Storage, VaultStorage,
};
use serde::{Deserialize, Serialize};
use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
};

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum SecureBackend {
    GitHub(GitHubConfig),
    InMemoryStorage,
    Vault(VaultConfig),
    OnDiskStorage(OnDiskStorageConfig),
}

impl SecureBackend {
    pub fn namespace(&self) -> Option<&str> {
        match self {
            SecureBackend::GitHub(GitHubConfig { namespace, .. })
            | SecureBackend::Vault(VaultConfig { namespace, .. })
            | SecureBackend::OnDiskStorage(OnDiskStorageConfig { namespace, .. }) => {
                namespace.as_deref()
            }
            SecureBackend::InMemoryStorage => None,
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct GitHubConfig {
    /// The owner or account that hosts a repository
    pub repository_owner: String,
    /// The repository where storage will mount
    pub repository: String,
    /// The branch containing storage, defaults to master
    pub branch: Option<String>,
    /// The authorization token for accessing the repository
    pub token: Token,
    /// A namespace is an optional portion of the path to a key stored within GitHubConfig. For
    /// example, a key, S, without a namespace would be available in S, with a namespace, N, it
    /// would be in N/S.
    pub namespace: Option<String>,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct VaultConfig {
    /// Optional SSL Certificate for the vault host, this is expected to be a full path.
    pub ca_certificate: Option<PathBuf>,
    /// A namespace is an optional portion of the path to a key stored within Vault. For example,
    /// a secret, S, without a namespace would be available in secret/data/S, with a namespace, N, it
    /// would be in secret/data/N/S.
    pub namespace: Option<String>,
    /// Vault leverages leases on many tokens, specify this to automatically have your lease
    /// renewed up to that many seconds more. If this is not specified, the lease will not
    /// automatically be renewed.
    pub renew_ttl_secs: Option<u32>,
    /// Vault's URL, note: only HTTP is currently supported.
    pub server: String,
    /// The authorization token for accessing secrets
    pub token: Token,
    /// Disable check-and-set when writing secrets to Vault
    pub disable_cas: Option<bool>,
    /// Timeout for new vault socket connections, in milliseconds.
    pub connection_timeout_ms: Option<u64>,
    /// Timeout for generic vault operations (e.g., reads and writes), in milliseconds.
    pub response_timeout_ms: Option<u64>,
}

impl VaultConfig {
    pub fn ca_certificate(&self) -> Result<String, Error> {
        let path = self
            .ca_certificate
            .as_ref()
            .ok_or(Error::Missing("ca_certificate"))?;
        read_file(path)
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct OnDiskStorageConfig {
    // Required path for on disk storage
    pub path: PathBuf,
    /// A namespace is an optional portion of the path to a key stored within OnDiskStorage. For
    /// example, a key, S, without a namespace would be available in S, with a namespace, N, it
    /// would be in N/S.
    pub namespace: Option<String>,
    #[serde(skip)]
    data_dir: PathBuf,
}

/// Tokens can either be directly within this config or stored somewhere on disk.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Token {
    FromConfig(String),
    /// This is an absolute path and not relative to data_dir
    FromDisk(PathBuf),
}

impl Token {
    pub fn read_token(&self) -> Result<String, Error> {
        match self {
            Token::FromDisk(path) => read_file(path),
            Token::FromConfig(token) => Ok(token.clone()),
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct TokenFromConfig {
    token: String,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct TokenFromDisk {
    path: PathBuf,
}

impl Default for OnDiskStorageConfig {
    fn default() -> Self {
        Self {
            namespace: None,
            path: PathBuf::from("secure_storage.json"),
            data_dir: PathBuf::from("/opt/diem/data"),
        }
    }
}

impl OnDiskStorageConfig {
    pub fn path(&self) -> PathBuf {
        if self.path.is_relative() {
            self.data_dir.join(&self.path)
        } else {
            self.path.clone()
        }
    }

    pub fn set_data_dir(&mut self, data_dir: PathBuf) {
        self.data_dir = data_dir;
    }
}

fn read_file(path: &Path) -> Result<String, Error> {
    let mut file =
        File::open(path).map_err(|e| Error::IO(path.to_str().unwrap().to_string(), e))?;
    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .map_err(|e| Error::IO(path.to_str().unwrap().to_string(), e))?;
    Ok(contents)
}

impl From<&SecureBackend> for Storage {
    fn from(backend: &SecureBackend) -> Self {
        match backend {
            SecureBackend::GitHub(config) => {
                let storage = Storage::from(GitHubStorage::new(
                    config.repository_owner.clone(),
                    config.repository.clone(),
                    config
                        .branch
                        .as_ref()
                        .cloned()
                        .unwrap_or_else(|| "master".to_string()),
                    config.token.read_token().expect("Unable to read token"),
                ));
                if let Some(namespace) = &config.namespace {
                    Storage::from(Namespaced::new(namespace, Box::new(storage)))
                } else {
                    storage
                }
            }
            SecureBackend::InMemoryStorage => Storage::from(InMemoryStorage::new()),
            SecureBackend::OnDiskStorage(config) => {
                let storage = Storage::from(OnDiskStorage::new(config.path()));
                if let Some(namespace) = &config.namespace {
                    Storage::from(Namespaced::new(namespace, Box::new(storage)))
                } else {
                    storage
                }
            }
            SecureBackend::Vault(config) => {
                let storage = Storage::from(VaultStorage::new(
                    config.server.clone(),
                    config.token.read_token().expect("Unable to read token"),
                    config
                        .ca_certificate
                        .as_ref()
                        .map(|_| config.ca_certificate().unwrap()),
                    config.renew_ttl_secs,
                    config.disable_cas.map_or_else(|| true, |disable| !disable),
                    config.connection_timeout_ms,
                    config.response_timeout_ms,
                ));
                if let Some(namespace) = &config.namespace {
                    Storage::from(Namespaced::new(namespace, Box::new(storage)))
                } else {
                    storage
                }
            }
        }
    }
}
#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;

    #[derive(Debug, Deserialize, PartialEq, Serialize)]
    struct Config {
        vault: VaultConfig,
    }

    #[test]
    fn test_token_config_parsing() {
        let from_config = Config {
            vault: VaultConfig {
                namespace: None,
                server: "127.0.0.1:8200".to_string(),
                ca_certificate: None,
                token: Token::FromConfig("test".to_string()),
                renew_ttl_secs: None,
                disable_cas: None,
                connection_timeout_ms: None,
                response_timeout_ms: None,
            },
        };

        let text_from_config = r#"
vault:
    server: "127.0.0.1:8200"
    token:
        from_config: "test"
        "#;

        let de_from_config: Config = serde_yaml::from_str(text_from_config).unwrap();
        assert_eq!(de_from_config, from_config);
        // Just assert that it can be serialized, not about to do string comparison
        serde_yaml::to_string(&from_config).unwrap();
    }

    #[test]
    fn test_vault_timeout_parsing() {
        let from_config = Config {
            vault: VaultConfig {
                namespace: None,
                server: "127.0.0.1:8200".to_string(),
                ca_certificate: None,
                token: Token::FromConfig("test".to_string()),
                renew_ttl_secs: None,
                disable_cas: None,
                connection_timeout_ms: Some(3000),
                response_timeout_ms: Some(5000),
            },
        };

        let text_from_config = r#"
vault:
    server: "127.0.0.1:8200"
    token:
        from_config: "test"
    connection_timeout_ms: 3000
    response_timeout_ms: 5000
        "#;

        let de_from_config: Config = serde_yaml::from_str(text_from_config).unwrap();
        assert_eq!(de_from_config, from_config);
        // Just assert that it can be serialized, no need to do string comparison
        serde_yaml::to_string(&from_config).unwrap();
    }

    #[test]
    fn test_token_disk_parsing() {
        let from_disk = Config {
            vault: VaultConfig {
                namespace: None,
                server: "127.0.0.1:8200".to_string(),
                ca_certificate: None,
                token: Token::FromDisk(PathBuf::from("/token")),
                renew_ttl_secs: None,
                disable_cas: None,
                connection_timeout_ms: None,
                response_timeout_ms: None,
            },
        };

        let text_from_disk = r#"
vault:
    server: "127.0.0.1:8200"
    token:
        from_disk: "/token"
        "#;

        let de_from_disk: Config = serde_yaml::from_str(text_from_disk).unwrap();
        assert_eq!(de_from_disk, from_disk);
        // Just assert that it can be serialized, not about to do string comparison
        serde_yaml::to_string(&from_disk).unwrap();
    }

    #[test]
    fn test_token_reading() {
        let temppath = diem_temppath::TempPath::new();
        temppath.create_as_file().unwrap();
        let mut file = File::create(temppath.path()).unwrap();
        file.write_all(b"disk_token").unwrap();

        let disk = Token::FromDisk(temppath.path().to_path_buf());
        assert_eq!("disk_token", disk.read_token().unwrap());

        let config = Token::FromConfig("config_token".to_string());
        assert_eq!("config_token", config.read_token().unwrap());
    }
}
<|endoftext|>


<commit_before>extern crate bytes;
extern crate chrono;
extern crate clap;
extern crate env_logger;
extern crate futures;
extern crate futures_cpupool;
extern crate grpc;
extern crate log;
extern crate protobuf;

use futures::future::Future;

extern crate grpc_interop;
use grpc_interop::*;

use bytes::Bytes;

use grpc::*;

use chrono::*;
use clap::{App, Arg};

fn empty_unary(client: TestServiceClient) {
    client
        .empty_call(grpc::RequestOptions::new(), Empty::new())
        .wait_drop_metadata()
        .expect("failed to get EmptyUnary result");
    println!("{} EmptyUnary done", Local::now().to_rfc3339());
}

// TODO: file a bug on this failure.
fn large_unary(client: TestServiceClient) {
    let mut payload = Payload::new();
    payload.set_body(vec![0; 271828]);
    let mut request = SimpleRequest::new();
    request.set_payload(payload);
    request.set_response_size(314159);
    let response = client
        .unary_call(grpc::RequestOptions::new(), request)
        .wait_drop_metadata()
        .expect("expected full frame");
    assert!(response.get_payload().body.len() == 314159);
    println!("{} LargeUnary done", Local::now().to_rfc3339());
}

fn client_streaming(client: TestServiceClient) {

    let (mut req, resp) = client
        .streaming_input_call(
            grpc::RequestOptions::new(),
        ).wait().expect("expected response");

    for size in [27182, 8, 1828, 45904].iter() {
        let mut request = StreamingInputCallRequest::new();
        let mut payload = Payload::new();
        payload.set_body(vec![0; size.to_owned()]);
        request.set_payload(payload);
        req.block_wait().expect("block_wait");
        req.send_data(request).expect("send_data");
    }

    req.finish().expect("finish");

    let resp = resp.wait_drop_metadata().expect("wait_drop_metadata");

    assert_eq!(resp.aggregated_payload_size, 74922);
    println!("{} ClientStreaming done", Local::now().to_rfc3339());
}

// TODO: test fails with an assertion error, we never see a response from the client.
// fails with 'expected a response: Other("partial frame")'
fn server_streaming(client: TestServiceClient) {
    let mut req = StreamingOutputCallRequest::new();
    let mut params = Vec::new();
    for &size in [31415, 9, 2653, 58979].iter() {
        let mut rp = ResponseParameters::new();
        rp.set_size(size as i32);
        params.push(rp);
    }
    req.set_response_parameters(::protobuf::RepeatedField::from_vec(params));

    let response_stream = client
        .streaming_output_call(grpc::RequestOptions::new(), req)
        .wait_drop_metadata();

    let mut response_sizes = Vec::new();

    {
        // this scope is to satisfy the borrow checker.
        let bar = response_stream.map(|response| {
            response_sizes.push(
                response
                    .expect("expected a response")
                    .get_payload()
                    .body
                    .len(),
            );
        });
        assert!(bar.count() == 4);
    }

    assert!(response_sizes.len() == 4);
    assert!(response_sizes[0] == 31415);
    assert!(response_sizes[1] == 9);
    assert!(response_sizes[2] == 2653);
    assert!(response_sizes[3] == 58979);
    println!("{} ServerStreaming done", Local::now().to_rfc3339());
}

// TODO: this test needs to interleave requests with responses and check along the way.
// Need to find a way to model that with the blocking API.
fn ping_pong(client: TestServiceClient) {
    let (mut req, resp) = client
        .full_duplex_call(
            grpc::RequestOptions::new(),
        ).wait().expect("start request");

    let mut resp = resp.wait_drop_metadata();

    for &(size, body_len) in [(31415, 27182), (9, 8), (2653, 1828), (58979, 45904)].iter() {
        let mut req_m = StreamingOutputCallRequest::new();
        let mut params = ResponseParameters::new();
        params.set_size(size);
        req_m.set_response_parameters(::protobuf::RepeatedField::from_vec(vec![params]));
        let mut payload = Payload::new();
        payload.set_body(vec![0; body_len]);
        req_m.set_payload(payload);
        req.block_wait().expect("block_wait");
        req.send_data(req_m).expect("send_data");

        let resp_m = resp.next().expect("next").unwrap();
        assert_eq!(size as usize, resp_m.payload.get_ref().body.len());
    }

    req.finish().expect("finish");
    assert!(resp.next().is_none());

    println!("{} PingPong done", Local::now().to_rfc3339());
}

fn empty_stream(client: TestServiceClient) {
    let (mut req, resp) = client
        .full_duplex_call(grpc::RequestOptions::new()).wait().expect("wait");
    req.finish().expect("finish");
    let resp: Vec<_> = resp.wait_drop_metadata().collect();
    assert!(resp.len() == 0);
    println!("{} EmptyStream done", Local::now().to_rfc3339());
}

fn custom_metadata(client: TestServiceClient) {
    fn make_options() -> grpc::RequestOptions {
        // The client attaches custom metadata with the following keys and values:
        // key: "x-grpc-test-echo-initial", value: "test_initial_metadata_value"
        // key: "x-grpc-test-echo-initial", value: "test_initial_metadata_value"
        let mut options = grpc::RequestOptions::new();
        options.metadata.add(
            MetadataKey::from("x-grpc-test-echo-initial"),
            Bytes::from("test_initial_metadata_value"),
        );
        options.metadata.add(
            MetadataKey::from("x-grpc-test-echo-trailing-bin"),
            Bytes::from(&b"\xab\xab\xab"[..]),
        );
        options
    }

    fn assert_result_metadata(initial: grpc::Metadata, trailing: grpc::Metadata) {
        assert_eq!(
            Some(&b"test_initial_metadata_value"[..]),
            initial.get("x-grpc-test-echo-initial")
        );
        assert_eq!(
            Some(&b"\xab\xab\xab"[..]),
            trailing.get("x-grpc-test-echo-trailing-bin")
        );
    }

    {
        // to a UnaryCall with request:
        // {
        //   response_size: 314159
        //   payload:{
        //     body: 271828 bytes of zeros
        //   }
        // }
        let mut req = SimpleRequest::new();
        req.set_response_size(314159);
        let mut payload = Payload::new();
        payload.set_body(vec![0; 271828]);
        req.set_payload(payload);
        let (initial, _result, trailing) = client
            .unary_call(make_options(), req)
            .wait()
            .expect("UnaryCall");

        assert_result_metadata(initial, trailing);
    }

    {
        // to a FullDuplexCall with request:
        // {
        //   response_parameters:{
        //     size: 314159
        //   }
        //   payload:{
        //     body: 271828 bytes of zeros
        //   }
        // }
        let mut req = StreamingOutputCallRequest::new();
        {
            let mut rp = ResponseParameters::new();
            rp.set_size(314159);
            req.mut_response_parameters().push(rp);
        }
        {
            let mut p = Payload::new();
            p.set_body(vec![0; 271828]);
            req.set_payload(p);
        }

        let (mut req1, resp) = client
            .full_duplex_call(make_options()).wait().expect("start request");
        req1.send_data(req).expect("send_data");
        req1.finish().expect("finish");

        let (initial, _, trailing) = resp.collect().wait().expect("collect");

        assert_result_metadata(initial, trailing);
    }
}

// The flags we use are defined in the gRPC Interopability doc
// https://github.com/grpc/grpc/blob/master/doc/interop-test-descriptions.md
fn main() {
    env_logger::init();

    let options = App::new("gRPC interopability client")
        .version("0.1")
        .author("Steve Jenson <stevej@buoyant.io>")
        .about("Interoperability Test Client for grpc-rust")
        .arg(
            Arg::with_name("server_host")
                .long("server_host")
                .help("The server host to connect to. For example, \"localhost\" or \"127.0.0.1\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("server_host_override")
                .long("server_host_override")
                .help(
                    "The server host to claim to be connecting to, for use in TLS and HTTP/2 \
                     :authority header. If unspecified, the value of --server_host will be used",
                ).takes_value(true),
        ).arg(
            Arg::with_name("server_port")
                .long("server_port")
                .help("The server port to connect to. For example, \"8080\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("test_case")
                .long("test_case")
                .help("The name of the test case to execute. For example, \"empty_unary\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("use_tls") // boolean
                .long("use_tls")
                .help("Whether to use a plaintext or encrypted connection")
                .takes_value(true),
        ).arg(
            Arg::with_name("use_test_ca")
                .long("use_test_ca")
                .help("Whether to replace platform root CAs with ca.pem as the CA root")
                .takes_value(true),
        ).arg(
            Arg::with_name("default_service_account")
                .long("default_service_account")
                .help("Email of the GCE default service account.")
                .takes_value(true),
        ).arg(
            Arg::with_name("oauth_scope")
                .long("oauth_scope")
                .help("OAuth scope. For example, \"https://www.googleapis.com/auth/xapi.zoo\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("service_account_key_file")
                .long("service_account_key_file")
                .help(
                    "The path to the service account JSON key file generated from GCE developer \
                     console.",
                ).takes_value(true),
        ).get_matches();

    let hostname = options.value_of("server_host").unwrap_or("localhost");
    let serverport = options
        .value_of("server_port")
        .map(|s| s.parse().unwrap())
        .unwrap_or(DEFAULT_PORT);

    let client =
        TestServiceClient::new_plain(hostname, serverport, Default::default()).expect("init");

    let testcase = options.value_of("test_case").unwrap_or("");
    match testcase.as_ref() {
        "empty_unary" => empty_unary(client),
        "cacheable_unary" => panic!("cacheable_unary not done yet"),
        "large_unary" => large_unary(client),
        "client_compressed_unary" => panic!("client_compressed_unary not done yet"),
        "server_compressed_unary" => panic!("server_compressed_unary not done yet"),
        "client_streaming" => client_streaming(client),
        "client_compressed_streaming" => panic!("client_compressed_streaming not done yet"),
        "server_streaming" => server_streaming(client),
        "server_compressed_streaming" => panic!("server_compressed_streaming not done yet"),
        "ping_pong" => ping_pong(client),
        "empty_stream" => empty_stream(client),
        "compute_engine_creds" => panic!("compute_engine_creds not done yet"),
        "jwt_token_creds" => panic!("jwt_token_creds not done yet"),
        "oauth2_auth_token" => panic!("oauth2_auth_token not done yet"),
        "per_rpc_creds" => panic!("per_rpc_creds not done yet"),
        "custom_metadata" => custom_metadata(client),
        "status_code_and_message" => panic!("status_code_and_message not done yet"),
        "unimplemented_method" => panic!("unimplemented_method not done yet"),
        "unimplemented_service" => panic!("unimplemented_service not done yet"),
        "cancel_after_begin" => panic!("cancel_after_begin not done yet"),
        "cancel_after_first_response" => panic!("cancel_after_first_response not done yet"),
        "timeout_on_sleeping_server" => panic!("timeout_on_sleeping_server not done yet"),
        _ => panic!("no test_case specified"),
    }
}
<commit_msg>interop test: connect to 127.0.0.1 instead of localhost<commit_after>extern crate bytes;
extern crate chrono;
extern crate clap;
extern crate env_logger;
extern crate futures;
extern crate futures_cpupool;
extern crate grpc;
extern crate log;
extern crate protobuf;

use futures::future::Future;

extern crate grpc_interop;
use grpc_interop::*;

use bytes::Bytes;

use grpc::*;

use chrono::*;
use clap::{App, Arg};

fn empty_unary(client: TestServiceClient) {
    client
        .empty_call(grpc::RequestOptions::new(), Empty::new())
        .wait_drop_metadata()
        .expect("failed to get EmptyUnary result");
    println!("{} EmptyUnary done", Local::now().to_rfc3339());
}

// TODO: file a bug on this failure.
fn large_unary(client: TestServiceClient) {
    let mut payload = Payload::new();
    payload.set_body(vec![0; 271828]);
    let mut request = SimpleRequest::new();
    request.set_payload(payload);
    request.set_response_size(314159);
    let response = client
        .unary_call(grpc::RequestOptions::new(), request)
        .wait_drop_metadata()
        .expect("expected full frame");
    assert!(response.get_payload().body.len() == 314159);
    println!("{} LargeUnary done", Local::now().to_rfc3339());
}

fn client_streaming(client: TestServiceClient) {

    let (mut req, resp) = client
        .streaming_input_call(
            grpc::RequestOptions::new(),
        ).wait().expect("expected response");

    for size in [27182, 8, 1828, 45904].iter() {
        let mut request = StreamingInputCallRequest::new();
        let mut payload = Payload::new();
        payload.set_body(vec![0; size.to_owned()]);
        request.set_payload(payload);
        req.block_wait().expect("block_wait");
        req.send_data(request).expect("send_data");
    }

    req.finish().expect("finish");

    let resp = resp.wait_drop_metadata().expect("wait_drop_metadata");

    assert_eq!(resp.aggregated_payload_size, 74922);
    println!("{} ClientStreaming done", Local::now().to_rfc3339());
}

// TODO: test fails with an assertion error, we never see a response from the client.
// fails with 'expected a response: Other("partial frame")'
fn server_streaming(client: TestServiceClient) {
    let mut req = StreamingOutputCallRequest::new();
    let mut params = Vec::new();
    for &size in [31415, 9, 2653, 58979].iter() {
        let mut rp = ResponseParameters::new();
        rp.set_size(size as i32);
        params.push(rp);
    }
    req.set_response_parameters(::protobuf::RepeatedField::from_vec(params));

    let response_stream = client
        .streaming_output_call(grpc::RequestOptions::new(), req)
        .wait_drop_metadata();

    let mut response_sizes = Vec::new();

    {
        // this scope is to satisfy the borrow checker.
        let bar = response_stream.map(|response| {
            response_sizes.push(
                response
                    .expect("expected a response")
                    .get_payload()
                    .body
                    .len(),
            );
        });
        assert!(bar.count() == 4);
    }

    assert!(response_sizes.len() == 4);
    assert!(response_sizes[0] == 31415);
    assert!(response_sizes[1] == 9);
    assert!(response_sizes[2] == 2653);
    assert!(response_sizes[3] == 58979);
    println!("{} ServerStreaming done", Local::now().to_rfc3339());
}

// TODO: this test needs to interleave requests with responses and check along the way.
// Need to find a way to model that with the blocking API.
fn ping_pong(client: TestServiceClient) {
    let (mut req, resp) = client
        .full_duplex_call(
            grpc::RequestOptions::new(),
        ).wait().expect("start request");

    let mut resp = resp.wait_drop_metadata();

    for &(size, body_len) in [(31415, 27182), (9, 8), (2653, 1828), (58979, 45904)].iter() {
        let mut req_m = StreamingOutputCallRequest::new();
        let mut params = ResponseParameters::new();
        params.set_size(size);
        req_m.set_response_parameters(::protobuf::RepeatedField::from_vec(vec![params]));
        let mut payload = Payload::new();
        payload.set_body(vec![0; body_len]);
        req_m.set_payload(payload);
        req.block_wait().expect("block_wait");
        req.send_data(req_m).expect("send_data");

        let resp_m = resp.next().expect("next").unwrap();
        assert_eq!(size as usize, resp_m.payload.get_ref().body.len());
    }

    req.finish().expect("finish");
    assert!(resp.next().is_none());

    println!("{} PingPong done", Local::now().to_rfc3339());
}

fn empty_stream(client: TestServiceClient) {
    let (mut req, resp) = client
        .full_duplex_call(grpc::RequestOptions::new()).wait().expect("wait");
    req.finish().expect("finish");
    let resp: Vec<_> = resp.wait_drop_metadata().collect();
    assert!(resp.len() == 0);
    println!("{} EmptyStream done", Local::now().to_rfc3339());
}

fn custom_metadata(client: TestServiceClient) {
    fn make_options() -> grpc::RequestOptions {
        // The client attaches custom metadata with the following keys and values:
        // key: "x-grpc-test-echo-initial", value: "test_initial_metadata_value"
        // key: "x-grpc-test-echo-initial", value: "test_initial_metadata_value"
        let mut options = grpc::RequestOptions::new();
        options.metadata.add(
            MetadataKey::from("x-grpc-test-echo-initial"),
            Bytes::from("test_initial_metadata_value"),
        );
        options.metadata.add(
            MetadataKey::from("x-grpc-test-echo-trailing-bin"),
            Bytes::from(&b"\xab\xab\xab"[..]),
        );
        options
    }

    fn assert_result_metadata(initial: grpc::Metadata, trailing: grpc::Metadata) {
        assert_eq!(
            Some(&b"test_initial_metadata_value"[..]),
            initial.get("x-grpc-test-echo-initial")
        );
        assert_eq!(
            Some(&b"\xab\xab\xab"[..]),
            trailing.get("x-grpc-test-echo-trailing-bin")
        );
    }

    {
        // to a UnaryCall with request:
        // {
        //   response_size: 314159
        //   payload:{
        //     body: 271828 bytes of zeros
        //   }
        // }
        let mut req = SimpleRequest::new();
        req.set_response_size(314159);
        let mut payload = Payload::new();
        payload.set_body(vec![0; 271828]);
        req.set_payload(payload);
        let (initial, _result, trailing) = client
            .unary_call(make_options(), req)
            .wait()
            .expect("UnaryCall");

        assert_result_metadata(initial, trailing);
    }

    {
        // to a FullDuplexCall with request:
        // {
        //   response_parameters:{
        //     size: 314159
        //   }
        //   payload:{
        //     body: 271828 bytes of zeros
        //   }
        // }
        let mut req = StreamingOutputCallRequest::new();
        {
            let mut rp = ResponseParameters::new();
            rp.set_size(314159);
            req.mut_response_parameters().push(rp);
        }
        {
            let mut p = Payload::new();
            p.set_body(vec![0; 271828]);
            req.set_payload(p);
        }

        let (mut req1, resp) = client
            .full_duplex_call(make_options()).wait().expect("start request");
        req1.send_data(req).expect("send_data");
        req1.finish().expect("finish");

        let (initial, _, trailing) = resp.collect().wait().expect("collect");

        assert_result_metadata(initial, trailing);
    }
}

// The flags we use are defined in the gRPC Interopability doc
// https://github.com/grpc/grpc/blob/master/doc/interop-test-descriptions.md
fn main() {
    env_logger::init();

    let options = App::new("gRPC interopability client")
        .version("0.1")
        .author("Steve Jenson <stevej@buoyant.io>")
        .about("Interoperability Test Client for grpc-rust")
        .arg(
            Arg::with_name("server_host")
                .long("server_host")
                .help("The server host to connect to. For example, \"localhost\" or \"127.0.0.1\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("server_host_override")
                .long("server_host_override")
                .help(
                    "The server host to claim to be connecting to, for use in TLS and HTTP/2 \
                     :authority header. If unspecified, the value of --server_host will be used",
                ).takes_value(true),
        ).arg(
            Arg::with_name("server_port")
                .long("server_port")
                .help("The server port to connect to. For example, \"8080\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("test_case")
                .long("test_case")
                .help("The name of the test case to execute. For example, \"empty_unary\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("use_tls") // boolean
                .long("use_tls")
                .help("Whether to use a plaintext or encrypted connection")
                .takes_value(true),
        ).arg(
            Arg::with_name("use_test_ca")
                .long("use_test_ca")
                .help("Whether to replace platform root CAs with ca.pem as the CA root")
                .takes_value(true),
        ).arg(
            Arg::with_name("default_service_account")
                .long("default_service_account")
                .help("Email of the GCE default service account.")
                .takes_value(true),
        ).arg(
            Arg::with_name("oauth_scope")
                .long("oauth_scope")
                .help("OAuth scope. For example, \"https://www.googleapis.com/auth/xapi.zoo\"")
                .takes_value(true),
        ).arg(
            Arg::with_name("service_account_key_file")
                .long("service_account_key_file")
                .help(
                    "The path to the service account JSON key file generated from GCE developer \
                     console.",
                ).takes_value(true),
        ).get_matches();

    let hostname = options.value_of("server_host").unwrap_or("127.0.0.1");
    let serverport = options
        .value_of("server_port")
        .map(|s| s.parse().unwrap())
        .unwrap_or(DEFAULT_PORT);

    let client =
        TestServiceClient::new_plain(hostname, serverport, Default::default()).expect("init");

    let testcase = options.value_of("test_case").unwrap_or("");
    match testcase.as_ref() {
        "empty_unary" => empty_unary(client),
        "cacheable_unary" => panic!("cacheable_unary not done yet"),
        "large_unary" => large_unary(client),
        "client_compressed_unary" => panic!("client_compressed_unary not done yet"),
        "server_compressed_unary" => panic!("server_compressed_unary not done yet"),
        "client_streaming" => client_streaming(client),
        "client_compressed_streaming" => panic!("client_compressed_streaming not done yet"),
        "server_streaming" => server_streaming(client),
        "server_compressed_streaming" => panic!("server_compressed_streaming not done yet"),
        "ping_pong" => ping_pong(client),
        "empty_stream" => empty_stream(client),
        "compute_engine_creds" => panic!("compute_engine_creds not done yet"),
        "jwt_token_creds" => panic!("jwt_token_creds not done yet"),
        "oauth2_auth_token" => panic!("oauth2_auth_token not done yet"),
        "per_rpc_creds" => panic!("per_rpc_creds not done yet"),
        "custom_metadata" => custom_metadata(client),
        "status_code_and_message" => panic!("status_code_and_message not done yet"),
        "unimplemented_method" => panic!("unimplemented_method not done yet"),
        "unimplemented_service" => panic!("unimplemented_service not done yet"),
        "cancel_after_begin" => panic!("cancel_after_begin not done yet"),
        "cancel_after_first_response" => panic!("cancel_after_first_response not done yet"),
        "timeout_on_sleeping_server" => panic!("timeout_on_sleeping_server not done yet"),
        _ => panic!("no test_case specified"),
    }
}
<|endoftext|>