<commit_before>use std::thread;
use std::thread::JoinHandle;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{Ordering, AtomicBool, AtomicUsize,
  ATOMIC_USIZE_INIT};
use std::ops::{Deref, Drop};
use std::fmt::{Display, Formatter, Error};
use std::collections::HashMap;
use std::net::{UdpSocket, SocketAddr, ToSocketAddrs};

use serialize;
use message::Message;
use payload::{Service, Payload, Device};
use net2::{UdpBuilder, UdpSocketExt};

/// udp broadcast ip address and lifx default port.
///
static BROADCAST_IP: &'static str = "255.255.255.255:56700";

/// sequence number counter used to confirm acks.
///
static SEQUENCE_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

/// returns the next sequence number (global, shared counter).
///
fn next_sequence() -> u8 {
  SEQUENCE_COUNTER.fetch_add(1, Ordering::SeqCst) as u8
}

/// sends a message to the specified address.
///
fn send_msg<S : Deref<Target = UdpSocket>, A : ToSocketAddrs>(
  socket: &S,
  addr: A,
  payload: Payload,
  ack_required: bool,
  target: u64
) -> Result<u8, String>
{
  let seq = next_sequence();
  let msg = Message::new(payload, ack_required, target, seq);
  let encoded = try!(serialize::encode(&msg).or(err!("failed to encode")));
  let bytes = try!(
    socket.send_to(&encoded[..], addr).or(err!("failed to send message")));

  if bytes == encoded.len() {
    Ok(seq)
  } else {
    err!("wrong number of bytes written")
  }
}

/// a bulb is a LiFX device where the service is Udp.
///
#[derive(Clone)]
pub struct Bulb<A : ToSocketAddrs> {
  label: Option<String>,
  ip: A,
  port: u32,
  target: u64,
  socket: Arc<UdpSocket>
}

impl<A> Bulb<A> where A : ToSocketAddrs {
  pub fn send_msg(&self, payload: Payload, ack_required: bool) -> Result<u8, String> {
    send_msg(&self.socket, &self.ip, payload, ack_required, self.target)
  }
}

impl<A> Display for Bulb<A> where A : ToSocketAddrs + Display {
  fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
    write!(
      f,
      "'{:?}' ({}:{})",
      self.label,
      self.ip,
      self.port)
  }
}

/// the client handles device messages from from any lifx bulb.
///
pub struct Client {
  closed: Arc<AtomicBool>,
  socket: Arc<UdpSocket>,
  devices: Arc<RwLock<HashMap<u64, Bulb<SocketAddr>>>>,
}

impl Client {
  /// creates a new client that will read responses from any lifx bulb.
  ///
  pub fn new<A : ToSocketAddrs>(addr: A) -> Result<Client, String> {
    let closed = Arc::new(AtomicBool::new(false));
    let devices = Arc::new(RwLock::new(HashMap::new()));
    let udp_builder = try!(
      UdpBuilder::new_v4().or(err!("failed to create builder")));
    let udp_socket = Arc::new(try!(
      udp_builder.bind(addr).or(err!("failed to bind to addr"))));

    let client = Client {
      closed: closed,
      socket: udp_socket,
      devices: devices
    };

    Ok(client)
  }

  /// listens for certain messages, and updates the client object accordingly
  ///
  pub fn listen(&self) -> JoinHandle<()> {
    let socket = self.socket.clone();
    let closed = self.closed.clone();
    let devices = self.devices.clone();

    thread::spawn(move || {
      let mut buf = [0; 256];

      while !closed.load(Ordering::SeqCst) {
        let (amt, src) = socket.recv_from(&mut buf[..]).unwrap();
        let resp = serialize::decode::<Message>(&buf[..amt]).unwrap();

        match *resp.payload() {
          Payload::Device(Device::StateService(Service::Udp, port)) =>
            {
              println!("Received device with port: {}", port);

              devices.write().unwrap().insert(
                resp.target(),
                Bulb {
                  label: None,
                  ip: src,
                  port: port,
                  target: resp.target(),
                  socket: socket.clone()
                });

              println!("Devices:");

              for d in devices.read().unwrap().values() {
                println!("  Devices: {}", d);
              }
            }
          _ =>
            ()
        }
      }
    })
  }

  /// broadcasts messages to the client at a set interval. use `listen` to
  /// have the client process certain messages.
  ///
  pub fn discover(&self, wait: u32) -> JoinHandle<()> {
    use Device::*;

    let socket = self.socket.clone();
    let closed = self.closed.clone();

    thread::spawn(move || {
      while !closed.load(Ordering::SeqCst) {
        let _ = socket.set_broadcast(true);
        let _ = send_msg(&socket, BROADCAST_IP, Payload::Device(GetService), false, 0);
        let _ = socket.set_broadcast(false);

        thread::sleep_ms(wait);
      }
    })
  }

  /// sends a message to the specified address.
  ///
  pub fn send_msg<A : ToSocketAddrs>(
    &self,
    addr: A,
    payload: Payload,
    ack_required: bool,
    target: u64
  ) -> Result<u8, String>
  {
    send_msg(&self.socket, addr, payload, ack_required, target)
  }

  /// closes a client. it will no longer receive responses from the socket.
  ///
  #[inline(always)]
  pub fn close(&self) {
    self.closed.store(true, Ordering::SeqCst)
  }

  /// checks if a client is closed.
  ///
  #[inline(always)]
  pub fn is_closed(&self) -> bool {
    self.closed.load(Ordering::SeqCst)
  }
}

impl Drop for Client {
  fn drop(&mut self) {
    self.close();
  }
}

#[test]
fn test_sequence_counter_overflow() {
  use std::u8;

  assert_eq!(0, next_sequence());
  for _ in 1..(u8::MAX as usize) + 1 { next_sequence(); }
  assert_eq!(0, next_sequence());
  for _ in 1..(u8::MAX as usize) + 1 { next_sequence(); }
  assert_eq!(0, next_sequence());
}<commit_msg>Added retrieving device label to client<commit_after>use std::thread;
use std::thread::JoinHandle;
use std::sync::{Arc, RwLock};
use std::sync::atomic::{Ordering, AtomicBool, AtomicUsize,
  ATOMIC_USIZE_INIT};
use std::ops::{Deref, Drop};
use std::fmt::{Display, Formatter, Error};
use std::collections::HashMap;
use std::net::{UdpSocket, SocketAddr, ToSocketAddrs};

use serialize;
use message::Message;
use payload::{Service, Payload, Device};
use net2::{UdpBuilder, UdpSocketExt};

/// udp broadcast ip address and lifx default port.
///
static BROADCAST_IP: &'static str = "255.255.255.255:56700";

/// sequence number counter used to confirm acks.
///
static SEQUENCE_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

/// returns the next sequence number (global, shared counter).
///
fn next_sequence() -> u8 {
  SEQUENCE_COUNTER.fetch_add(1, Ordering::SeqCst) as u8
}

/// sends a message to the specified address.
///
fn send_msg<S : Deref<Target = UdpSocket>, A : ToSocketAddrs>(
  socket: &S,
  addr: A,
  payload: Payload,
  ack_required: bool,
  target: u64
) -> Result<u8, String>
{
  let seq = next_sequence();
  let msg = Message::new(payload, ack_required, target, seq);
  let encoded = try!(serialize::encode(&msg).or(err!("failed to encode")));
  let bytes = try!(
    socket.send_to(&encoded[..], addr).or(err!("failed to send message")));

  println!("    Sending: {:?}", msg);

  if bytes == encoded.len() {
    Ok(seq)
  } else {
    err!("wrong number of bytes written")
  }
}

/// a bulb is a LiFX device where the service is Udp.
///
#[derive(Clone)]
pub struct Bulb<A : ToSocketAddrs> {
  label: Option<String>,
  ip: A,
  port: u32,
  target: u64,
  socket: Arc<UdpSocket>
}

impl<A> Bulb<A> where A : ToSocketAddrs {
  pub fn send_msg(&self, payload: Payload, ack_required: bool) -> Result<u8, String> {
    send_msg(&self.socket, &self.ip, payload, ack_required, self.target)
  }
}

impl<A> Display for Bulb<A> where A : ToSocketAddrs + Display {
  fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
    write!(
      f,
      "'{:?}' ({}:{} {})",
      self.label,
      self.ip,
      self.port,
      self.target)
  }
}

/// the client handles device messages from from any lifx bulb.
///
pub struct Client {
  closed: Arc<AtomicBool>,
  socket: Arc<UdpSocket>,
  devices: Arc<RwLock<HashMap<u64, Bulb<SocketAddr>>>>,
}

impl Client {
  /// creates a new client that will read responses from any lifx bulb.
  ///
  pub fn new<A : ToSocketAddrs>(addr: A) -> Result<Client, String> {
    let closed = Arc::new(AtomicBool::new(false));
    let devices = Arc::new(RwLock::new(HashMap::new()));
    let udp_builder = try!(
      UdpBuilder::new_v4().or(err!("failed to create builder")));
    let udp_socket = Arc::new(try!(
      udp_builder.bind(addr).or(err!("failed to bind to addr"))));

    let client = Client {
      closed: closed,
      socket: udp_socket,
      devices: devices
    };

    Ok(client)
  }

  /// listens for certain messages, and updates the client object accordingly
  ///
  pub fn listen(&self) -> JoinHandle<()> {
    let socket = self.socket.clone();
    let closed = self.closed.clone();
    let devices = self.devices.clone();

    thread::spawn(move || {
      let mut buf = [0; 256];

      while !closed.load(Ordering::SeqCst) {
        let (amt, src) = socket.recv_from(&mut buf[..]).unwrap();
        let resp = serialize::decode::<Message>(&buf[..amt]).unwrap();

        match *resp.payload() {
          Payload::Device(Device::StateService(Service::Udp, port)) =>
            {
              println!("Received device with port: {}", port);

              devices.write().unwrap().entry(resp.target()).or_insert(
                Bulb {
                  label: None,
                  ip: src,
                  port: port,
                  target: resp.target(),
                  socket: socket.clone()
                });

              println!("Devices:");

              for d in devices.read().unwrap().values() {
                println!("  Devices: {}", d);

                let _ = d.send_msg(Payload::Device(Device::GetLabel), false);

                thread::sleep_ms(200);
              }
            }
          Payload::Device(Device::StateLabel(ref label)) =>
            {
              println!("Received device label: '{:?}' for {}", label, resp.target());

              // TODO: Move the label with some UNSAFE code?
              if let Some(bulb) = devices.write().unwrap().get_mut(&resp.target()) {
                bulb.label = Some(label.clone());
              }
            }
          _ =>
            ()
        }
      }
    })
  }

  /// broadcasts messages to the client at a set interval. use `listen` to
  /// have the client process certain messages.
  ///
  pub fn discover(&self, wait: u32) -> JoinHandle<()> {
    use Device::*;

    let socket = self.socket.clone();
    let closed = self.closed.clone();

    thread::spawn(move || {
      while !closed.load(Ordering::SeqCst) {
        // TODO: Technically, there should be a LOCK on the socket here. Messages
        // should not be able to be sent between the time the socket is set to
        // broadcast.
        let _ = socket.set_broadcast(true);
        let _ = send_msg(&socket, BROADCAST_IP, Payload::Device(GetService), false, 0);
        let _ = socket.set_broadcast(false);

        thread::sleep_ms(wait);
      }
    })
  }

  /// sends a message to the specified address.
  ///
  pub fn send_msg<A : ToSocketAddrs>(
    &self,
    addr: A,
    payload: Payload,
    ack_required: bool,
    target: u64
  ) -> Result<u8, String>
  {
    send_msg(&self.socket, addr, payload, ack_required, target)
  }

  /// closes a client. it will no longer receive responses from the socket.
  ///
  #[inline(always)]
  pub fn close(&self) {
    self.closed.store(true, Ordering::SeqCst)
  }

  /// checks if a client is closed.
  ///
  #[inline(always)]
  pub fn is_closed(&self) -> bool {
    self.closed.load(Ordering::SeqCst)
  }
}

impl Drop for Client {
  fn drop(&mut self) {
    self.close();
  }
}

#[test]
fn test_sequence_counter_overflow() {
  use std::u8;

  assert_eq!(0, next_sequence());
  for _ in 1..(u8::MAX as usize) + 1 { next_sequence(); }
  assert_eq!(0, next_sequence());
  for _ in 1..(u8::MAX as usize) + 1 { next_sequence(); }
  assert_eq!(0, next_sequence());
}<|endoftext|>


<commit_before>//! The `Client` allows users of the `raft` library to connect to remote `Server` instances and
//! issue commands to be applied to the `StateMachine`.

use std::collections::HashSet;
use std::fmt;
use std::io::Write;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::str::FromStr;

use bufstream::BufStream;
use capnp::{serialize, MessageReader, ReaderOptions};

use messages_capnp::{client_response, proposal_response};
use messages;
use ClientId;
use Result;
use RaftError;

/// The representation of a Client connection to the cluster.
pub struct Client {
    /// The `Uuid` of the client, should be unique in the cluster.
    pub id: ClientId,
    /// The current connection to the current leader.
    /// If it is none it may mean that there is no estabished leader or that there has been
    /// a disconnection.
    leader_connection: Option<BufStream<TcpStream>>,
    /// A lookup for the cluster's nodes.
    cluster: HashSet<SocketAddr>,
}

impl Client {

    /// Creates a new client.
    pub fn new(cluster: HashSet<SocketAddr>) -> Client {
        Client {
            id: ClientId::new(),
            leader_connection: None,
            cluster: cluster,
        }
    }

    /// Proposes an entry to be appended to the replicated log. This will only
    /// return once the entry has been durably committed.
    /// Returns `Error` when the entire cluster has an unknown leader. Try proposing again later.
    pub fn propose(&mut self, entry: &[u8]) -> Result<()> {
        scoped_trace!("{:?}: propose", self);
        let mut message = messages::proposal_request(entry);

        let mut members = self.cluster.iter().cloned();

        loop {
            let mut connection = match self.leader_connection.take() {
                Some(cxn) => {
                    scoped_debug!("had existing connection {:?}", cxn.get_ref().peer_addr());
                    cxn
                },
                None => {
                    let leader = try!(members.next().ok_or(RaftError::LeaderSearchExhausted));
                    scoped_debug!("connecting to potential leader {}", leader);
                    // Send the preamble.
                    let preamble = messages::client_connection_preamble(self.id);
                    let mut stream = BufStream::new(try!(TcpStream::connect(leader)));
                    try!(serialize::write_message(&mut stream, &*preamble));
                    stream
                }
            };
            try!(serialize::write_message(&mut connection, &mut message));
            try!(connection.flush());
            let response = try!(serialize::read_message(&mut connection, ReaderOptions::new()));
            match try!(response.get_root::<client_response::Reader>()).which().unwrap() {
                client_response::Which::Proposal(Ok(response)) => {
                    match response.which().unwrap() {
                        proposal_response::Which::Success(()) => {
                            scoped_debug!("received response Success");
                            self.leader_connection = Some(connection);
                            return Ok(())
                        },
                        proposal_response::Which::UnknownLeader(()) => {
                            scoped_debug!("received response UnknownLeader");
                            ()
                        },
                        proposal_response::Which::NotLeader(leader) => {
                            scoped_debug!("received response NotLeader");
                            let leader_str = try!(leader);
                            if !self.cluster.contains(&try!(SocketAddr::from_str(leader_str))) {
                                scoped_debug!("cluster violation detected");
                                return Err(RaftError::ClusterViolation.into())
                            }
                            let mut connection: TcpStream = try!(TcpStream::connect(leader_str));
                            let preamble = messages::client_connection_preamble(self.id);
                            try!(serialize::write_message(&mut connection, &*preamble));
                            self.leader_connection = Some(BufStream::new(connection));
                        }
                    }
                },
                _ => panic!("Unexpected message type"), // TODO: return a proper error
            }
        }
    }
}

impl fmt::Debug for Client {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Client({})", self.id)
    }
}

#[cfg(test)]
mod test {
    extern crate env_logger;

    use std::collections::HashSet;
    use std::io::{Read, Write};
    use std::net::{TcpStream, SocketAddr, TcpListener};
    use std::str::FromStr;
    use std::thread;

    use uuid::Uuid;
    use capnp::{serialize, ReaderOptions};
    use capnp::message::MessageReader;
    use bufstream::BufStream;

    use {Client, ClientId, messages, Result};
    use messages_capnp::{connection_preamble, client_request};

    fn expect_preamble(connection: &mut TcpStream, client_id: Uuid) -> Result<bool> {
        let message = try!(serialize::read_message(connection, ReaderOptions::new()));
        let preamble = try!(message.get_root::<connection_preamble::Reader>());
        // Test to make sure preamble has the right id.
        if let connection_preamble::id::Which::Client(Ok(id)) = try!(preamble.get_id().which()) {
            Ok(Uuid::from_bytes(id).unwrap() == client_id)
        } else {
            Ok(false)
        }
    }

    fn expect_proposal(connection: &mut TcpStream, value: &[u8]) -> Result<bool> {
        let message = try!(serialize::read_message(connection, ReaderOptions::new()));
        let request = try!(message.get_root::<client_request::Reader>());
        // Test to make sure request has the right value.
        if let client_request::Which::Proposal(Ok(proposal)) = try!(request.which()) {
            Ok(proposal.get_entry().unwrap() == value)
        } else {
            Ok(false)
        }
    }

    #[test]
    fn test_proposal_success() {
        setup_test!("test_proposal_success");
        // Start up the cluster and get what we need.
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            let (mut connection, _)  = test_server.accept().unwrap();

            // Proposal should be fine, no errors.
            scoped_debug!("Should get preamble and proposal. Responds Success");
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();
            // Send response! (success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();
        });

        // Propose. It's a marriage made in heaven! :)
        // Should be ok
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());

        child.join().unwrap();
    }

    #[test]
    fn test_proposal_unknown_leader() {
        setup_test!("test_proposal_unknown_leader");
        // Start up the cluster and get what we need.
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            let (mut connection, _)  = test_server.accept().unwrap();

            // Proposal should report unknown leader, and have the client return error.
            scoped_debug!("Should get proposal. Responds UnknownLeader");
            expect_proposal(&mut connection, to_propose).unwrap();
            // Send response! (unknown leader!) Client should drop connection.
            let response = messages::proposal_response_unknown_leader();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();
        });

        // Propose. It's a marriage made in heaven! :)
        assert!(client.propose(to_propose).is_err());

        child.join().unwrap();
    }

    #[test]
    fn test_proposal_not_leader() {
        setup_test!("test_proposal_not_leader");
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let second_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let second_addr = second_server.local_addr().unwrap();
        cluster.insert(second_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the first proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            // Proposal should report NotLeader. Client should choose the server we direct it to.
            scoped_debug!("Should get preamble and proposal. Responds NotLeader.");
            let (mut connection, _)  = test_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send response! (not leader!)
            let response = messages::proposal_response_not_leader(&format!("{}", second_addr));
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // Test that it seeks out other server and proposes.
            scoped_debug!("Second server should get preamble and proposal. Responds Success.");
            let (mut connection, _)  = second_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send final response! (Success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();
        });

        // Workaround to set up rigged selection of servers.
        client.leader_connection = {
            let preamble = messages::client_connection_preamble(client.id);
            let mut stream = BufStream::new(TcpStream::connect(test_addr).unwrap());
            serialize::write_message(&mut stream, &*preamble).unwrap();
            Some(stream)
        };

        // Should be ok, change leader connection.
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());

        child.join().unwrap();
    }

    /// This test makes sure that the client cannot be redirected to a leader which exists outside
    /// the cluster. This is a necessary test since it would introduce error into the cluster.
    #[test]
    fn test_proposal_leader_not_in_cluster() {
        setup_test!("test_proposal_leader_not_in_cluster");
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let second_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let second_addr = second_server.local_addr().unwrap();
        // cluster.insert(second_addr); <--- NOT in cluster.

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the first proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            // Proposal should report NotLeader. Client should choose the server we direct it to.
            scoped_debug!("Should get preamble and proposal. Responds NotLeader.");
            let (mut connection, _)  = test_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send response! (not leader!)
            let response = messages::proposal_response_not_leader(&format!("{}", second_addr));
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush();

            // No more...
        });

        // Workaround to set up rigged selection of servers.
        client.leader_connection = {
            let preamble = messages::client_connection_preamble(client.id);
            let mut stream = BufStream::new(TcpStream::connect(test_addr).unwrap());
            serialize::write_message(&mut stream, &*preamble).unwrap();
            Some(stream)
        };

        // Should be err, change leader connectio but to wrong ip..
        assert!(client.propose(to_propose).is_err());
        assert!(client.leader_connection.is_none());

        child.join().unwrap();
    }
}
<commit_msg>Fix client.rs merge conflicts<commit_after>//! The `Client` allows users of the `raft` library to connect to remote `Server` instances and
//! issue commands to be applied to the `StateMachine`.

use std::collections::HashSet;
use std::fmt;
use std::io::Write;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::str::FromStr;

use bufstream::BufStream;
use capnp::{serialize, MessageReader, ReaderOptions};

use messages_capnp::{client_response, proposal_response};
use messages;
use ClientId;
use Result;
use RaftError;

/// The representation of a Client connection to the cluster.
pub struct Client {
    /// The `Uuid` of the client, should be unique in the cluster.
    pub id: ClientId,
    /// The current connection to the current leader.
    /// If it is none it may mean that there is no estabished leader or that there has been
    /// a disconnection.
    leader_connection: Option<BufStream<TcpStream>>,
    /// A lookup for the cluster's nodes.
    cluster: HashSet<SocketAddr>,
}

impl Client {

    /// Creates a new client.
    pub fn new(cluster: HashSet<SocketAddr>) -> Client {
        Client {
            id: ClientId::new(),
            leader_connection: None,
            cluster: cluster,
        }
    }

    /// Proposes an entry to be appended to the replicated log. This will only
    /// return once the entry has been durably committed.
    /// Returns `Error` when the entire cluster has an unknown leader. Try proposing again later.
    pub fn propose(&mut self, entry: &[u8]) -> Result<()> {
        scoped_trace!("{:?}: propose", self);
        let mut message = messages::proposal_request(entry);

        let mut members = self.cluster.iter().cloned();

        loop {
            let mut connection = match self.leader_connection.take() {
                Some(cxn) => {
                    scoped_debug!("had existing connection {:?}", cxn.get_ref().peer_addr());
                    cxn
                },
                None => {
                    let leader = try!(members.next().ok_or(RaftError::LeaderSearchExhausted));
                    scoped_debug!("connecting to potential leader {}", leader);
                    // Send the preamble.
                    let preamble = messages::client_connection_preamble(self.id);
                    let mut stream = BufStream::new(try!(TcpStream::connect(leader)));
                    try!(serialize::write_message(&mut stream, &*preamble));
                    stream
                }
            };
            try!(serialize::write_message(&mut connection, &mut message));
            try!(connection.flush());
            let response = try!(serialize::read_message(&mut connection, ReaderOptions::new()));
            match try!(response.get_root::<client_response::Reader>()).which().unwrap() {
                client_response::Which::Proposal(Ok(response)) => {
                    match response.which().unwrap() {
                        proposal_response::Which::Success(()) => {
                            scoped_debug!("received response Success");
                            self.leader_connection = Some(connection);
                            return Ok(())
                        },
                        proposal_response::Which::UnknownLeader(()) => {
                            scoped_debug!("received response UnknownLeader");
                            ()
                        },
                        proposal_response::Which::NotLeader(leader) => {
                            scoped_debug!("received response NotLeader");
                            let leader_str = try!(leader);
                            if !self.cluster.contains(&try!(SocketAddr::from_str(leader_str))) {
                                scoped_debug!("cluster violation detected");
                                return Err(RaftError::ClusterViolation.into())
                            }
                            let mut connection: TcpStream = try!(TcpStream::connect(leader_str));
                            let preamble = messages::client_connection_preamble(self.id);
                            try!(serialize::write_message(&mut connection, &*preamble));
                            self.leader_connection = Some(BufStream::new(connection));
                        }
                    }
                },
                _ => panic!("Unexpected message type"), // TODO: return a proper error
            }
        }
    }
}

impl fmt::Debug for Client {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Client({})", self.id)
    }
}

#[cfg(test)]
mod test {
    extern crate env_logger;

    use std::collections::HashSet;
    use std::io::Write;
    use std::net::{TcpStream, TcpListener};
    use std::thread;

    use uuid::Uuid;
    use capnp::{serialize, ReaderOptions};
    use capnp::message::MessageReader;
    use bufstream::BufStream;

    use {Client, messages, Result};
    use messages_capnp::{connection_preamble, client_request};

    fn expect_preamble(connection: &mut TcpStream, client_id: Uuid) -> Result<bool> {
        let message = try!(serialize::read_message(connection, ReaderOptions::new()));
        let preamble = try!(message.get_root::<connection_preamble::Reader>());
        // Test to make sure preamble has the right id.
        if let connection_preamble::id::Which::Client(Ok(id)) = try!(preamble.get_id().which()) {
            Ok(Uuid::from_bytes(id).unwrap() == client_id)
        } else {
            Ok(false)
        }
    }

    fn expect_proposal(connection: &mut TcpStream, value: &[u8]) -> Result<bool> {
        let message = try!(serialize::read_message(connection, ReaderOptions::new()));
        let request = try!(message.get_root::<client_request::Reader>());
        // Test to make sure request has the right value.
        if let client_request::Which::Proposal(Ok(proposal)) = try!(request.which()) {
            Ok(proposal.get_entry().unwrap() == value)
        } else {
            Ok(false)
        }
    }

    #[test]
    fn test_proposal_success() {
        setup_test!("test_proposal_success");
        // Start up the cluster and get what we need.
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            let (mut connection, _)  = test_server.accept().unwrap();

            // Proposal should be fine, no errors.
            scoped_debug!("Should get preamble and proposal. Responds Success");
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();
            // Send response! (success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush().unwrap();
        });

        // Propose. It's a marriage made in heaven! :)
        // Should be ok
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());

        child.join().unwrap();
    }

    #[test]
    fn test_proposal_unknown_leader() {
        setup_test!("test_proposal_unknown_leader");
        // Start up the cluster and get what we need.
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let mut client = Client::new(cluster);
        let to_propose = b"Bears";

        // The client connects on the proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            let (mut connection, _)  = test_server.accept().unwrap();

            // Proposal should report unknown leader, and have the client return error.
            scoped_debug!("Should get proposal. Responds UnknownLeader");
            expect_proposal(&mut connection, to_propose).unwrap();
            // Send response! (unknown leader!) Client should drop connection.
            let response = messages::proposal_response_unknown_leader();
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush().unwrap();
        });

        // Propose. It's a marriage made in heaven! :)
        assert!(client.propose(to_propose).is_err());

        child.join().unwrap();
    }

    #[test]
    fn test_proposal_not_leader() {
        setup_test!("test_proposal_not_leader");
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let second_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let second_addr = second_server.local_addr().unwrap();
        cluster.insert(second_addr);

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the first proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            // Proposal should report NotLeader. Client should choose the server we direct it to.
            scoped_debug!("Should get preamble and proposal. Responds NotLeader.");
            let (mut connection, _)  = test_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send response! (not leader!)
            let response = messages::proposal_response_not_leader(&second_addr);
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush().unwrap();

            // Test that it seeks out other server and proposes.
            scoped_debug!("Second server should get preamble and proposal. Responds Success.");
            let (mut connection, _)  = second_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send final response! (Success!)
            let response = messages::proposal_response_success();
            serialize::write_message(&mut connection, &*response).unwrap();
        });

        // Workaround to set up rigged selection of servers.
        client.leader_connection = {
            let preamble = messages::client_connection_preamble(client.id);
            let mut stream = BufStream::new(TcpStream::connect(test_addr).unwrap());
            serialize::write_message(&mut stream, &*preamble).unwrap();
            Some(stream)
        };

        // Should be ok, change leader connection.
        client.propose(to_propose).unwrap();
        assert!(client.leader_connection.is_some());

        child.join().unwrap();
    }

    /// This test makes sure that the client cannot be redirected to a leader which exists outside
    /// the cluster. This is a necessary test since it would introduce error into the cluster.
    #[test]
    fn test_proposal_leader_not_in_cluster() {
        setup_test!("test_proposal_leader_not_in_cluster");
        let mut cluster = HashSet::new();
        let test_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let test_addr = test_server.local_addr().unwrap();
        cluster.insert(test_addr);

        let second_server = TcpListener::bind("127.0.0.1:0").unwrap();
        let second_addr = second_server.local_addr().unwrap();
        // cluster.insert(second_addr); <--- NOT in cluster.

        let mut client = Client::new(cluster);
        let client_id = client.id.0.clone();
        let to_propose = b"Bears";

        // The client connects on the first proposal.
        // Wait for it.
        let child = thread::spawn(move || {
            // Proposal should report NotLeader. Client should choose the server we direct it to.
            scoped_debug!("Should get preamble and proposal. Responds NotLeader.");
            let (mut connection, _)  = test_server.accept().unwrap();
            expect_preamble(&mut connection, client_id).unwrap();
            expect_proposal(&mut connection, to_propose).unwrap();

            // Send response! (not leader!)
            let response = messages::proposal_response_not_leader(&second_addr);
            serialize::write_message(&mut connection, &*response).unwrap();
            connection.flush().unwrap();

            // No more...
        });

        // Workaround to set up rigged selection of servers.
        client.leader_connection = {
            let preamble = messages::client_connection_preamble(client.id);
            let mut stream = BufStream::new(TcpStream::connect(test_addr).unwrap());
            serialize::write_message(&mut stream, &*preamble).unwrap();
            Some(stream)
        };

        // Should be err, change leader connectio but to wrong ip..
        assert!(client.propose(to_propose).is_err());
        assert!(client.leader_connection.is_none());

        child.join().unwrap();
    }
}
<|endoftext|>


<commit_before>use hyper::client::Client;
use hyper::client::response::Response;
use hyper::error::Error;
use hyper::header::Headers;
use hyper::method::Method;

/// Trait that represents some methods to send a specific request
pub trait GetResponse {
    /// Given a specific URL, get the response from the target server
    fn get_http_response(&self, url: &str) -> Result<Response, Error>;

    /// Given a specific URL and an header, get the response from the target server
    fn get_http_response_using_headers(&self,
                                       url: &str,
                                       header: Headers)
                                       -> Result<Response, Error>;
}

impl GetResponse for Client {
    fn get_http_response(&self, url: &str) -> Result<Response, Error> {
        self.get_http_response_using_headers(url, Headers::new())
    }

    fn get_http_response_using_headers(&self,
                                       url: &str,
                                       custom_header: Headers)
                                       -> Result<Response, Error> {
        self.request(Method::Get, url).headers(custom_header).send()
    }
}<commit_msg>Add a method trait to get a header response from a server request, without any content body<commit_after>use hyper::client::Client;
use hyper::client::response::Response;
use hyper::error::Error;
use hyper::header::Headers;
use hyper::method::Method;

/// Trait that represents some methods to send a specific request
pub trait GetResponse {
    /// Given a specific URL, get the header without the content body (useful to not waste time,
    /// ressources and informations)
    fn get_head_response(&self, url: &str) -> Result<Response, Error>;

    /// Given a specific URL, get the response from the target server
    fn get_http_response(&self, url: &str) -> Result<Response, Error>;

    /// Given a specific URL and an header, get the response from the target server
    fn get_http_response_using_headers(&self,
                                       url: &str,
                                       header: Headers)
                                       -> Result<Response, Error>;
}

impl GetResponse for Client {
    fn get_head_response(&self, url: &str) -> Result<Response, Error> {
        self.request(Method::Head, url).send()
    }

    fn get_http_response(&self, url: &str) -> Result<Response, Error> {
        self.get_http_response_using_headers(url, Headers::new())
    }

    fn get_http_response_using_headers(&self,
                                       url: &str,
                                       custom_header: Headers)
                                       -> Result<Response, Error> {
        self.request(Method::Get, url).headers(custom_header).send()
    }
}
<|endoftext|>


<commit_before>//! This module is UNSTABLE, the reason is - recently added.

use std::convert::From;
use std::io::{self, Read, BufRead};
use std::str::{from_utf8, Utf8Error};

use super::{read_marker};
use super::{
    ReadError,
    MarkerReadError,
};
use super::{BigEndianRead};

use super::super::init::Marker;
use super::super::value::ValueRef;

trait USizeCast {
    fn from(v: Self) -> Option<usize> where Self: Sized;
}

impl USizeCast for u8 {
    fn from(v: u8) -> Option<usize> {
        // Impossible to panic, since u8 always fits in usize.
        Some(v as usize)
    }
}

impl USizeCast for u16 {
    fn from(v: u16) -> Option<usize> {
        // TODO: This can overflow on 8-bit systems.
        Some(v as usize)
    }
}

impl USizeCast for u32 {
    fn from(v: u32) -> Option<usize> {
        // TODO: This can overflow on 8- and 16-bit systems.
        Some(v as usize)
    }
}

// TODO: Display trait.
#[derive(Debug)]
pub enum Error<'r> {
    /// Failed to fill the internal reader buffer.
    ///
    /// RMP tries to obtain the buffer at the beginning of read operation using `read_buf` function.
    ///
    /// According to the Rust documentation, `fill_buf` function will return an I/O error if the
    /// underlying reader was read, but returned an error.
    InvalidBufferFill(io::Error),
    /// Failed to read the type marker value.
    InvalidMarkerRead(ReadError),
    /// Failed to read string/array/map size.
    InvalidLengthRead(ReadError),
    /// Failed to read packed non-marker data.
    InvalidDataRead(ReadError),
    /// Failed to cast the length read to machine size.
    InvalidLengthSize,
    /// Failed to interpret a byte slice as a UTF-8 string.
    ///
    /// Contains untouched bytearray with the underlying decoding error.
    InvalidUtf8(&'r [u8], Utf8Error),
    /// Failed to read ext type.
    InvalidExtTypeRead(ReadError),
}

impl<'r> From<MarkerReadError> for Error<'r> {
    fn from(err: MarkerReadError) -> Error<'r> {
        Error::InvalidMarkerRead(From::from(err))
    }
}

fn read_length<R, D>(rd: &mut R) -> Result<D, ReadError>
    where R: Read,
          D: BigEndianRead
{
    D::read(rd).map_err(From::from)
}

fn read_str(buf: &[u8], len: usize) -> Result<&str, Error> {
    let buf = try!(read_bin(buf, len));

    // Try to decode sliced buffer as UTF-8.
    let res = try!(from_utf8(buf).map_err(|err| Error::InvalidUtf8(buf, err)));

    Ok(res)
}

fn read_bin(buf: &[u8], len: usize) -> Result<&[u8], Error> {
    if len > buf.len() {
        return Err(Error::InvalidDataRead(ReadError::UnexpectedEOF));
    }

    // Take a slice.
    let buf = &buf[..len];

    Ok(buf)
}

// Helper function that reads a single byte from the given `Read` and interpret it as an Ext type.
fn read_ext_type<R>(rd: &mut R) -> Result<i8, ReadError>
    where R: Read
{
    i8::read(rd).map_err(From::from)
}

fn read_ext(mut buf: &[u8], len: usize) -> Result<(i8, &[u8]), Error> {
    let ty  = try!(read_ext_type(&mut buf).map_err(|err| Error::InvalidExtTypeRead(err)));
    let buf = try!(read_bin(buf, len));

    Ok((ty, buf))
}

#[inline]
fn read_str_value<U>(buf: &[u8], len: U) -> Result<ValueRef, Error>
    where U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_str(buf, len));

    Ok(ValueRef::String(res))
}

#[inline]
fn read_bin_value<U>(buf: &[u8], len: U) -> Result<ValueRef, Error>
    where U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_bin(buf, len));

    Ok(ValueRef::Binary(res))
}

#[inline]
fn read_ext_value<U>(mut buf: &[u8], len: U) -> Result<ValueRef, Error>
    where U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let (ty, buf) = try!(read_ext(&mut buf, len));

    Ok(ValueRef::Ext(ty, buf))
}

fn read_map(buf: &[u8], len: usize) -> Result<(Vec<(ValueRef, ValueRef)>, usize), Error>
{
    let mut vec = Vec::with_capacity(len);

    let buf = &buf[..];
    let mut pos = 0;

    for _ in 0..len {
        let kbuf = &buf[pos..];
        let (key, num) = try!(read_value_ref_impl(&kbuf));
        pos += num;

        let vbuf = &buf[pos..];
        let (val, num) = try!(read_value_ref_impl(&vbuf));
        pos += num;

        vec.push((key, val));
    }

    Ok((vec, pos as usize))
}

// TODO: Valid for [u8], Vec<u8>... all wrappers, where get() -> &[u8]
fn read_value_ref_impl(buf: &[u8]) -> Result<(ValueRef, usize), Error> {
    let mut buf = buf;
    let mut pos = 0usize;

    // Reading the marker involves either 1 byte read or nothing. On success consumes strictly
    // 1 byte from the `buf`, not from the `rd`.
    let marker = try!(read_marker(&mut buf));
    pos += 1;

    let val = match marker {
        Marker::FixedString(len) => {
            pos += len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Str8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 1 + len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Str16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 2 + len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Str32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 4 + len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Bin8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 1 + len as usize;
            try!(read_bin_value(buf, len))
        }
        Marker::Bin16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 2 + len as usize;
            try!(read_bin_value(buf, len))
        }
        Marker::Bin32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 4 + len as usize;
            try!(read_bin_value(buf, len))
        }
        Marker::FixedMap(len) => {
            let len = len as usize;
            let (map, bytes) = try!(read_map(&mut buf, len));
            pos += bytes;
            ValueRef::Map(map)
        }
        Marker::FixExt1 => {
            let len: u8 = 1;
            pos += 2;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt2 => {
            let len: u8 = 2;
            pos += 3;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt4 => {
            let len: u8 = 4;
            pos += 5;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt8 => {
            let len: u8 = 8;
            pos += 9;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt16 => {
            let len: u8 = 16;
            pos += 17;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 2 + len as usize;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 3 + len as usize;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 5 + len as usize;
            try!(read_ext_value(&mut buf, len))
        }
        _ => unimplemented!(),
    };

    Ok((val, pos as usize))
}

// NOTE: Consumes nothing from the given `BufRead` both on success and fail.
pub fn read_value_ref<R>(rd: &mut R) -> Result<ValueRef, Error>
    where R: BufRead
{
    let mut buf = try!(rd.fill_buf().map_err(|err| Error::InvalidBufferFill(err)));

    // Reading the marker involves either 1 byte read or nothing. On success consumes strictly
    // 1 byte from the `buf`, not from the `rd`.
    let marker = try!(read_marker(&mut buf));

    let val = match marker {
        Marker::FixedString(len) => {
            try!(read_str_value(buf, len))
        }
        Marker::Str8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(buf, len))
        }
        Marker::Str16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(buf, len))
        }
        Marker::Str32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(buf, len))
        }
        Marker::Bin8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(buf, len))
        }
        Marker::Bin16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(buf, len))
        }
        Marker::Bin32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(buf, len))
        }
        Marker::FixedMap(len) => {
            let len = len as usize;
            let (map, _) = try!(read_map(&mut buf, len));
            ValueRef::Map(map)
        }
        Marker::FixExt1 => {
            let len: u8 = 1;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt2 => {
            let len: u8 = 2;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt4 => {
            let len: u8 = 4;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt8 => {
            let len: u8 = 8;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt16 => {
            let len: u8 = 16;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(&mut buf, len))
        }
        _ => unimplemented!(),
    };

    Ok(val)
}

// TODO: Concepts: less code; more tests; fast refactoring after each test; no more than 5 min for
// each red-yellow-green.
<commit_msg>ValueRef: map - code clean.<commit_after>//! This module is UNSTABLE, the reason is - recently added.

use std::convert::From;
use std::io::{self, Read, BufRead};
use std::str::{from_utf8, Utf8Error};

use super::{read_marker};
use super::{
    ReadError,
    MarkerReadError,
};
use super::{BigEndianRead};

use super::super::init::Marker;
use super::super::value::ValueRef;

trait USizeCast {
    fn from(v: Self) -> Option<usize> where Self: Sized;
}

impl USizeCast for u8 {
    fn from(v: u8) -> Option<usize> {
        // Impossible to panic, since u8 always fits in usize.
        Some(v as usize)
    }
}

impl USizeCast for u16 {
    fn from(v: u16) -> Option<usize> {
        // TODO: This can overflow on 8-bit systems.
        Some(v as usize)
    }
}

impl USizeCast for u32 {
    fn from(v: u32) -> Option<usize> {
        // TODO: This can overflow on 8- and 16-bit systems.
        Some(v as usize)
    }
}

// TODO: Display trait.
#[derive(Debug)]
pub enum Error<'r> {
    /// Failed to fill the internal reader buffer.
    ///
    /// RMP tries to obtain the buffer at the beginning of read operation using `read_buf` function.
    ///
    /// According to the Rust documentation, `fill_buf` function will return an I/O error if the
    /// underlying reader was read, but returned an error.
    InvalidBufferFill(io::Error),
    /// Failed to read the type marker value.
    InvalidMarkerRead(ReadError),
    /// Failed to read string/array/map size.
    InvalidLengthRead(ReadError),
    /// Failed to read packed non-marker data.
    InvalidDataRead(ReadError),
    /// Failed to cast the length read to machine size.
    InvalidLengthSize,
    /// Failed to interpret a byte slice as a UTF-8 string.
    ///
    /// Contains untouched bytearray with the underlying decoding error.
    InvalidUtf8(&'r [u8], Utf8Error),
    /// Failed to read ext type.
    InvalidExtTypeRead(ReadError),
}

impl<'r> From<MarkerReadError> for Error<'r> {
    fn from(err: MarkerReadError) -> Error<'r> {
        Error::InvalidMarkerRead(From::from(err))
    }
}

fn read_length<R, D>(rd: &mut R) -> Result<D, ReadError>
    where R: Read,
          D: BigEndianRead
{
    D::read(rd).map_err(From::from)
}

fn read_str(buf: &[u8], len: usize) -> Result<&str, Error> {
    let buf = try!(read_bin(buf, len));

    // Try to decode sliced buffer as UTF-8.
    let res = try!(from_utf8(buf).map_err(|err| Error::InvalidUtf8(buf, err)));

    Ok(res)
}

fn read_bin(buf: &[u8], len: usize) -> Result<&[u8], Error> {
    if len > buf.len() {
        return Err(Error::InvalidDataRead(ReadError::UnexpectedEOF));
    }

    // Take a slice.
    let buf = &buf[..len];

    Ok(buf)
}

// Helper function that reads a single byte from the given `Read` and interpret it as an Ext type.
fn read_ext_type<R>(rd: &mut R) -> Result<i8, ReadError>
    where R: Read
{
    i8::read(rd).map_err(From::from)
}

fn read_ext(mut buf: &[u8], len: usize) -> Result<(i8, &[u8]), Error> {
    let ty  = try!(read_ext_type(&mut buf).map_err(|err| Error::InvalidExtTypeRead(err)));
    let buf = try!(read_bin(buf, len));

    Ok((ty, buf))
}

#[inline]
fn read_str_value<U>(buf: &[u8], len: U) -> Result<ValueRef, Error>
    where U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_str(buf, len));

    Ok(ValueRef::String(res))
}

#[inline]
fn read_bin_value<U>(buf: &[u8], len: U) -> Result<ValueRef, Error>
    where U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let res = try!(read_bin(buf, len));

    Ok(ValueRef::Binary(res))
}

#[inline]
fn read_ext_value<U>(mut buf: &[u8], len: U) -> Result<ValueRef, Error>
    where U: USizeCast
{
    let len = try!(U::from(len).ok_or(Error::InvalidLengthSize));
    let (ty, buf) = try!(read_ext(&mut buf, len));

    Ok(ValueRef::Ext(ty, buf))
}

fn read_map(buf: &[u8], len: usize) -> Result<(Vec<(ValueRef, ValueRef)>, usize), Error>
{
    let mut vec = Vec::with_capacity(len);
    let mut pos = 0;

    for _ in 0..len {
        let (key, num) = try!(read_value_ref_impl(&buf[pos..]));
        pos += num;

        let (val, num) = try!(read_value_ref_impl(&buf[pos..]));
        pos += num;

        vec.push((key, val));
    }

    Ok((vec, pos))
}

// TODO: Valid for [u8], Vec<u8>... all wrappers, where get() -> &[u8]
fn read_value_ref_impl(buf: &[u8]) -> Result<(ValueRef, usize), Error> {
    let mut buf = buf;
    let mut pos = 0usize;

    // Reading the marker involves either 1 byte read or nothing. On success consumes strictly
    // 1 byte from the `buf`, not from the `rd`.
    let marker = try!(read_marker(&mut buf));
    pos += 1;

    let val = match marker {
        Marker::FixedString(len) => {
            pos += len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Str8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 1 + len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Str16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 2 + len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Str32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 4 + len as usize;
            try!(read_str_value(buf, len))
        }
        Marker::Bin8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 1 + len as usize;
            try!(read_bin_value(buf, len))
        }
        Marker::Bin16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 2 + len as usize;
            try!(read_bin_value(buf, len))
        }
        Marker::Bin32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 4 + len as usize;
            try!(read_bin_value(buf, len))
        }
        Marker::FixedMap(len) => {
            let len = len as usize;
            let (map, bytes) = try!(read_map(&mut buf, len));
            pos += bytes;
            ValueRef::Map(map)
        }
        Marker::FixExt1 => {
            let len: u8 = 1;
            pos += 2;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt2 => {
            let len: u8 = 2;
            pos += 3;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt4 => {
            let len: u8 = 4;
            pos += 5;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt8 => {
            let len: u8 = 8;
            pos += 9;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt16 => {
            let len: u8 = 16;
            pos += 17;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 2 + len as usize;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 3 + len as usize;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            pos += 5 + len as usize;
            try!(read_ext_value(&mut buf, len))
        }
        _ => unimplemented!(),
    };

    Ok((val, pos as usize))
}

// NOTE: Consumes nothing from the given `BufRead` both on success and fail.
pub fn read_value_ref<R>(rd: &mut R) -> Result<ValueRef, Error>
    where R: BufRead
{
    let mut buf = try!(rd.fill_buf().map_err(|err| Error::InvalidBufferFill(err)));

    // Reading the marker involves either 1 byte read or nothing. On success consumes strictly
    // 1 byte from the `buf`, not from the `rd`.
    let marker = try!(read_marker(&mut buf));

    let val = match marker {
        Marker::FixedString(len) => {
            try!(read_str_value(buf, len))
        }
        Marker::Str8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(buf, len))
        }
        Marker::Str16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(buf, len))
        }
        Marker::Str32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_str_value(buf, len))
        }
        Marker::Bin8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(buf, len))
        }
        Marker::Bin16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(buf, len))
        }
        Marker::Bin32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_bin_value(buf, len))
        }
        Marker::FixedMap(len) => {
            let len = len as usize;
            let (map, _) = try!(read_map(&mut buf, len));
            ValueRef::Map(map)
        }
        Marker::FixExt1 => {
            let len: u8 = 1;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt2 => {
            let len: u8 = 2;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt4 => {
            let len: u8 = 4;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt8 => {
            let len: u8 = 8;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::FixExt16 => {
            let len: u8 = 16;
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext8 => {
            let len: u8 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext16 => {
            let len: u16 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(&mut buf, len))
        }
        Marker::Ext32 => {
            let len: u32 = try!(read_length(&mut buf).map_err(|err| Error::InvalidLengthRead(err)));
            try!(read_ext_value(&mut buf, len))
        }
        _ => unimplemented!(),
    };

    Ok(val)
}

// TODO: Concepts: less code; more tests; fast refactoring after each test; no more than 5 min for
// each red-yellow-green.
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};

use xdg;

fn xdg_dir() -> xdg::BaseDirectories {
    let prefix = Path::new("afl.rs")
        .join(afl_rustc_version())
        .join(pkg_version());
    xdg::BaseDirectories::with_prefix(prefix).unwrap()
}

const SHORT_COMMIT_HASH_LEN: usize = 7;

pub fn afl_rustc_version() -> String {
    let version_meta = rustc_version::version_meta().unwrap();
    let mut ret = String::from("rustc-");
    ret.push_str(&version_meta.semver.to_string());
    if let Some(commit_hash) = version_meta.commit_hash {
        ret.push('-');
        ret.push_str(&commit_hash[..SHORT_COMMIT_HASH_LEN]);
    }
    ret
}

fn pkg_version() -> String {
    let mut ret = String::from("afl.rs-");

    let version = env!("CARGO_PKG_VERSION");
    assert!(!version.is_empty());

    ret.push_str(version);
    ret
}

pub fn afl_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl").unwrap()
}

pub fn afl_llvm_rt_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl-llvm-rt").unwrap()
}

#[allow(dead_code)]
pub fn object_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.o")
}

#[allow(dead_code)]
pub fn archive_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.a")
}
<commit_msg>Use local directories when building for docs.rs (#185)<commit_after>use std::path::{Path, PathBuf};

use xdg;

fn xdg_dir() -> xdg::BaseDirectories {
    let prefix = Path::new("afl.rs")
        .join(afl_rustc_version())
        .join(pkg_version());
    xdg::BaseDirectories::with_prefix(prefix).unwrap()
}

const SHORT_COMMIT_HASH_LEN: usize = 7;

pub fn afl_rustc_version() -> String {
    let version_meta = rustc_version::version_meta().unwrap();
    let mut ret = String::from("rustc-");
    ret.push_str(&version_meta.semver.to_string());
    if let Some(commit_hash) = version_meta.commit_hash {
        ret.push('-');
        ret.push_str(&commit_hash[..SHORT_COMMIT_HASH_LEN]);
    }
    ret
}

fn pkg_version() -> String {
    let mut ret = String::from("afl.rs-");

    let version = env!("CARGO_PKG_VERSION");
    assert!(!version.is_empty());

    ret.push_str(version);
    ret
}

// Place directories inside the crate when building for docs.rs.
// (Modifying system paths are forbidden.)

#[cfg(docsrs)]
pub fn afl_dir() -> PathBuf {
    let path = PathBuf::from("./afl");
    std::fs::create_dir_all(&path).unwrap();
    path
}

#[cfg(not(docsrs))]
pub fn afl_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl").unwrap()
}

#[cfg(docsrs)]
pub fn afl_llvm_rt_dir() -> PathBuf {
    let path = PathBuf::from("./afl-llvm-rt");
    std::fs::create_dir_all(&path).unwrap();
    path
}

#[cfg(not(docsrs))]
pub fn afl_llvm_rt_dir() -> PathBuf {
    xdg_dir().create_data_directory("afl-llvm-rt").unwrap()
}

#[allow(dead_code)]
pub fn object_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.o")
}

#[allow(dead_code)]
pub fn archive_file_path() -> PathBuf {
    afl_llvm_rt_dir().join("libafl-llvm-rt.a")
}
<|endoftext|>


<commit_before>use rustc_serialize::json::{Json, ToJson};
use rustc_serialize::{Encodable, Encoder};
use std::collections::BTreeMap;
use std::error::Error;

use error::{WebDriverResult, WebDriverError, ErrorStatus};

pub static ELEMENT_KEY: &'static str = "element-6066-11e4-a52e-4f735466cecf";

#[derive(RustcEncodable, PartialEq, Clone, Debug)]
pub struct Date(u64);

impl Date {
    pub fn new(timestamp: u64) -> Date {
        Date(timestamp)
    }
}

impl ToJson for Date {
    fn to_json(&self) -> Json {
        let &Date(x) = self;
        x.to_json()
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum Nullable<T: ToJson> {
    Value(T),
    Null
}

impl<T: ToJson> Nullable<T> {
     pub fn is_null(&self) -> bool {
        match *self {
            Nullable::Value(_) => false,
            Nullable::Null => true
        }
    }

     pub fn is_value(&self) -> bool {
        match *self {
            Nullable::Value(_) => true,
            Nullable::Null => false
        }
    }
}

impl<T: ToJson> Nullable<T> {
    //This is not very pretty
    pub fn from_json<F: FnOnce(&Json) -> WebDriverResult<T>>(value: &Json, f: F) -> WebDriverResult<Nullable<T>> {
        if value.is_null() {
            Ok(Nullable::Null)
        } else {
            Ok(Nullable::Value(try!(f(value))))
        }
    }
}

impl<T: ToJson> ToJson for Nullable<T> {
    fn to_json(&self) -> Json {
        match *self {
            Nullable::Value(ref x) => x.to_json(),
            Nullable::Null => Json::Null
        }
    }
}

impl<T: ToJson> Encodable for Nullable<T> {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        match *self {
            Nullable::Value(ref x) => x.to_json().encode(s),
            Nullable::Null => s.emit_option_none()
        }
    }
}

#[derive(PartialEq)]
pub struct WebElement {
    pub id: String
}

impl WebElement {
    pub fn new(id: String) -> WebElement {
        WebElement {
            id: id
        }
    }

    pub fn from_json(data: &Json) -> WebDriverResult<WebElement> {
        let object = try_opt!(data.as_object(),
                              ErrorStatus::InvalidArgument,
                              "Could not convert webelement to object");
        let id_value = try_opt!(object.get(ELEMENT_KEY),
                                ErrorStatus::InvalidArgument,
                                "Could not find webelement key");

        let id = try_opt!(id_value.as_string(),
                          ErrorStatus::InvalidArgument,
                          "Could not convert web element to string").to_string();

        Ok(WebElement::new(id))
    }
}

impl ToJson for WebElement {
    fn to_json(&self) -> Json {
        let mut data = BTreeMap::new();
        data.insert(ELEMENT_KEY.to_string(), self.id.to_json());
        Json::Object(data)
    }
}

#[derive(PartialEq)]
pub enum FrameId {
    Short(u16),
    Element(WebElement),
    Null
}

impl FrameId {
    pub fn from_json(data: &Json) -> WebDriverResult<FrameId> {
        match data {
            &Json::U64(x) => {
                if x > u16::max_value() as u64 || x < u16::min_value() as u64 {
                    return Err(WebDriverError::new(ErrorStatus::NoSuchFrame,
                                                   "frame id out of range"))
                };
                Ok(FrameId::Short(x as u16))
            },
            &Json::Null => Ok(FrameId::Null),
            &Json::Object(_) => Ok(FrameId::Element(
                try!(WebElement::from_json(data)))),
            _ => Err(WebDriverError::new(ErrorStatus::NoSuchFrame,
                                         "frame id has unexpected type"))
        }
    }
}

impl ToJson for FrameId {
    fn to_json(&self) -> Json {
        match *self {
            FrameId::Short(x) => {
                Json::U64(x as u64)
            },
            FrameId::Element(ref x) => {
                Json::String(x.id.clone())
            },
            FrameId::Null => {
                Json::Null
            }
        }
    }
}

#[derive(PartialEq)]
pub enum LocatorStrategy {
    CSSSelector,
    LinkText,
    PartialLinkText,
    XPath
}

impl LocatorStrategy {
    pub fn from_json(body: &Json) -> WebDriverResult<LocatorStrategy> {
        match try_opt!(body.as_string(),
                       ErrorStatus::InvalidArgument,
                       "Cound not convert strategy to string") {
            "css selector" => Ok(LocatorStrategy::CSSSelector),
            "link text" => Ok(LocatorStrategy::LinkText),
            "partial link text" => Ok(LocatorStrategy::PartialLinkText),
            "xpath" => Ok(LocatorStrategy::XPath),
            x => Err(WebDriverError::new(ErrorStatus::InvalidArgument,
                                         &format!("Unknown locator strategy {}", x)[..]))
        }
    }
}

impl ToJson for LocatorStrategy {
    fn to_json(&self) -> Json {
        Json::String(match *self {
            LocatorStrategy::CSSSelector => "css selector",
            LocatorStrategy::LinkText => "link text",
            LocatorStrategy::PartialLinkText => "partial link text",
            LocatorStrategy::XPath => "xpath"
        }.to_string())
    }
}
<commit_msg>derive Debug trait for WebElement and FrameId<commit_after>use rustc_serialize::{Encodable, Encoder};
use rustc_serialize::json::{Json, ToJson};
use std::collections::BTreeMap;
use std::error::Error;

use error::{WebDriverResult, WebDriverError, ErrorStatus};

pub static ELEMENT_KEY: &'static str = "element-6066-11e4-a52e-4f735466cecf";

#[derive(RustcEncodable, PartialEq, Clone, Debug)]
pub struct Date(u64);

impl Date {
    pub fn new(timestamp: u64) -> Date {
        Date(timestamp)
    }
}

impl ToJson for Date {
    fn to_json(&self) -> Json {
        let &Date(x) = self;
        x.to_json()
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum Nullable<T: ToJson> {
    Value(T),
    Null
}

impl<T: ToJson> Nullable<T> {
     pub fn is_null(&self) -> bool {
        match *self {
            Nullable::Value(_) => false,
            Nullable::Null => true
        }
    }

     pub fn is_value(&self) -> bool {
        match *self {
            Nullable::Value(_) => true,
            Nullable::Null => false
        }
    }
}

impl<T: ToJson> Nullable<T> {
    //This is not very pretty
    pub fn from_json<F: FnOnce(&Json) -> WebDriverResult<T>>(value: &Json, f: F) -> WebDriverResult<Nullable<T>> {
        if value.is_null() {
            Ok(Nullable::Null)
        } else {
            Ok(Nullable::Value(try!(f(value))))
        }
    }
}

impl<T: ToJson> ToJson for Nullable<T> {
    fn to_json(&self) -> Json {
        match *self {
            Nullable::Value(ref x) => x.to_json(),
            Nullable::Null => Json::Null
        }
    }
}

impl<T: ToJson> Encodable for Nullable<T> {
    fn encode<S: Encoder>(&self, s: &mut S) -> Result<(), S::Error> {
        match *self {
            Nullable::Value(ref x) => x.to_json().encode(s),
            Nullable::Null => s.emit_option_none()
        }
    }
}

#[derive(PartialEq, Debug)]
pub struct WebElement {
    pub id: String
}

impl WebElement {
    pub fn new(id: String) -> WebElement {
        WebElement {
            id: id
        }
    }

    pub fn from_json(data: &Json) -> WebDriverResult<WebElement> {
        let object = try_opt!(data.as_object(),
                              ErrorStatus::InvalidArgument,
                              "Could not convert webelement to object");
        let id_value = try_opt!(object.get(ELEMENT_KEY),
                                ErrorStatus::InvalidArgument,
                                "Could not find webelement key");

        let id = try_opt!(id_value.as_string(),
                          ErrorStatus::InvalidArgument,
                          "Could not convert web element to string").to_string();

        Ok(WebElement::new(id))
    }
}

impl ToJson for WebElement {
    fn to_json(&self) -> Json {
        let mut data = BTreeMap::new();
        data.insert(ELEMENT_KEY.to_string(), self.id.to_json());
        Json::Object(data)
    }
}

#[derive(PartialEq, Debug)]
pub enum FrameId {
    Short(u16),
    Element(WebElement),
    Null
}

impl FrameId {
    pub fn from_json(data: &Json) -> WebDriverResult<FrameId> {
        match data {
            &Json::U64(x) => {
                if x > u16::max_value() as u64 || x < u16::min_value() as u64 {
                    return Err(WebDriverError::new(ErrorStatus::NoSuchFrame,
                                                   "frame id out of range"))
                };
                Ok(FrameId::Short(x as u16))
            },
            &Json::Null => Ok(FrameId::Null),
            &Json::Object(_) => Ok(FrameId::Element(
                try!(WebElement::from_json(data)))),
            _ => Err(WebDriverError::new(ErrorStatus::NoSuchFrame,
                                         "frame id has unexpected type"))
        }
    }
}

impl ToJson for FrameId {
    fn to_json(&self) -> Json {
        match *self {
            FrameId::Short(x) => {
                Json::U64(x as u64)
            },
            FrameId::Element(ref x) => {
                Json::String(x.id.clone())
            },
            FrameId::Null => {
                Json::Null
            }
        }
    }
}

#[derive(PartialEq)]
pub enum LocatorStrategy {
    CSSSelector,
    LinkText,
    PartialLinkText,
    XPath
}

impl LocatorStrategy {
    pub fn from_json(body: &Json) -> WebDriverResult<LocatorStrategy> {
        match try_opt!(body.as_string(),
                       ErrorStatus::InvalidArgument,
                       "Cound not convert strategy to string") {
            "css selector" => Ok(LocatorStrategy::CSSSelector),
            "link text" => Ok(LocatorStrategy::LinkText),
            "partial link text" => Ok(LocatorStrategy::PartialLinkText),
            "xpath" => Ok(LocatorStrategy::XPath),
            x => Err(WebDriverError::new(ErrorStatus::InvalidArgument,
                                         &format!("Unknown locator strategy {}", x)[..]))
        }
    }
}

impl ToJson for LocatorStrategy {
    fn to_json(&self) -> Json {
        Json::String(match *self {
            LocatorStrategy::CSSSelector => "css selector",
            LocatorStrategy::LinkText => "link text",
            LocatorStrategy::PartialLinkText => "partial link text",
            LocatorStrategy::XPath => "xpath"
        }.to_string())
    }
}
<|endoftext|>


<commit_before>use nom::{alphanumeric, digit, eof, is_alphanumeric, line_ending, multispace, space};
use nom::{IResult, Err, ErrorKind, Needed};
use parser::{Column, Table};
use std::str;
use std::str::FromStr;

#[derive(Clone, Debug, PartialEq)]
pub enum Operator {
    Not,
    And,
    Or,
    Like,
    NotLike,
    Equal,
    NotEqual,
    Greater,
    GreaterOrEqual,
    Less,
    LessOrEqual,
}

#[derive(Clone, Debug, PartialEq)]
pub enum FieldExpression {
    All,
    Seq(Vec<Column>),
}

impl Default for FieldExpression {
    fn default() -> FieldExpression {
        FieldExpression::All
    }
}

#[inline]
pub fn is_sql_identifier(chr: u8) -> bool {
    // XXX(malte): dot should not be in here once we have proper alias handling
    is_alphanumeric(chr) || chr == '_' as u8
}

#[inline]
pub fn is_fp_number(chr: u8) -> bool {
    is_alphanumeric(chr) || chr == '.' as u8
}

/// Parses a SQL identifier (alphanumeric and "_").
named!(pub fp_number<&[u8], &[u8]>,
// f64::from_str(&format!("{}.{}", integral, fractional)).unwrap()
//      integral: map_res!(take_while1!(is_digit), str::from_utf8) ~
//      fractional: map_res!(take_while!(is_digit), str::from_utf8),
    take_while1!(is_fp_number)
);

/// Parses a SQL column identifier in the table.column format
named!(pub column_identifier<&[u8], Column>,
    chain!(
        table: opt!(
            chain!(
                tbl_name: map_res!(sql_identifier, str::from_utf8) ~
                tag!("."),
                || { tbl_name }
            )
        ) ~
        column: map_res!(sql_identifier, str::from_utf8),
        || {
            Column {
                name: String::from(column),
                table: match table {
                    None => None,
                    Some(t) => Some(String::from(t)),
                },
            }
        }
    )
);

/// Parses a SQL identifier (alphanumeric and "_").
named!(pub sql_identifier<&[u8], &[u8]>,
    take_while1!(is_sql_identifier)
);

/// Parse an unsigned integer.
named!(pub unsigned_number<&[u8], u64>,
    map_res!(
        map_res!(digit, str::from_utf8),
        FromStr::from_str
    )
);

/// Parse a terminator that ends a SQL statement.
named!(pub statement_terminator,
    delimited!(opt!(multispace),
               alt_complete!(tag!(";") | line_ending | eof),
               opt!(multispace)
    )
);

/// Parse binary comparison operators
named!(pub binary_comparison_operator<&[u8], Operator>,
    alt_complete!(
           map!(caseless_tag!("like"), |_| Operator::Like)
         | map!(caseless_tag!("not_like"), |_| Operator::NotLike)
         | map!(caseless_tag!("<>"), |_| Operator::NotEqual)
         | map!(caseless_tag!(">="), |_| Operator::GreaterOrEqual)
         | map!(caseless_tag!("<="), |_| Operator::LessOrEqual)
         | map!(caseless_tag!("="), |_| Operator::Equal)
         | map!(caseless_tag!("<"), |_| Operator::Less)
         | map!(caseless_tag!(">"), |_| Operator::Greater)
    )
);

/// Parse logical operators
named!(pub binary_logical_operator<&[u8], Operator>,
    alt_complete!(
           map!(caseless_tag!("and"), |_| Operator::And)
         | map!(caseless_tag!("or"), |_| Operator::Or)
    )
);

/// Parse unary comparison operators
named!(pub unary_comparison_operator<&[u8], &str>,
    map_res!(
        alt_complete!(
               tag_s!(b"ISNULL")
             | tag_s!(b"NOT")
             | tag_s!(b"-") // ??? (number neg)
        ),
        str::from_utf8
    )
);

/// Parse unary comparison operators
named!(pub unary_negation_operator<&[u8], Operator>,
    alt_complete!(
          map!(caseless_tag!("not"), |_| Operator::Not)
        | map!(caseless_tag!("!"), |_| Operator::Not)
    )
);

/// Parse rule for a comma-separated list.
named!(pub field_list<&[u8], Vec<Column> >,
       many0!(
           chain!(
               fieldname: column_identifier ~
               opt!(
                   chain!(
                       multispace? ~
                       tag!(",") ~
                       multispace?,
                       ||{}
                   )
               ),
               || { fieldname }
           )
       )
);

/// Parse list of columns/fields.
/// XXX(malte): add support for named table notation
named!(pub field_expr<&[u8], FieldExpression>,
       alt_complete!(
           tag!("*") => { |_| FieldExpression::All }
         | map!(field_list, |v| FieldExpression::Seq(v))
       )
);

/// Parse list of table names.
/// XXX(malte): add support for aliases
named!(pub table_list<&[u8], Vec<Table> >,
       many0!(
           chain!(
               table: table_reference ~
               opt!(
                   chain!(
                       multispace? ~
                       tag!(",") ~
                       multispace?,
                       || {}
                   )
               ),
               || { table }
           )
       )
);

/// Parse a list of values (e.g., for INSERT syntax).
/// XXX(malte): proper value type
named!(pub value_list<&[u8], Vec<&str> >,
       many0!(
           map_res!(
               chain!(
                   val: alt_complete!(
                         tag_s!(b"?")
                       | tag_s!(b"CURRENT_TIMESTAMP")
                       | delimited!(tag!("'"), alphanumeric, tag!("'"))
                       | fp_number
                   ) ~
                   opt!(
                       chain!(
                           multispace? ~
                           tag!(",") ~
                           multispace?,
                           ||{}
                       )
                   ),
                   || { val }
               ),
               str::from_utf8
           )
       )
);

/// Parse a reference to a named table
/// XXX(malte): add support for schema.table notation
named!(pub table_reference<&[u8], Table>,
    chain!(
        table: map_res!(sql_identifier, str::from_utf8) ~
        alias: opt!(
            chain!(
                space ~
                caseless_tag!("as") ~
                space ~
                alias: map_res!(sql_identifier, str::from_utf8),
                || { String::from(alias) }
            )
        ),
        || {
            Table {
                name: String::from(table),
                alias: alias,
            }
        }
    )
);

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

    #[test]
    fn sql_identifiers() {
        let id1 = b"foo";
        let id2 = b"f_o_o";
        let id3 = b"foo12";
        let id4 = b":fo oo";

        assert!(sql_identifier(id1).is_done());
        assert!(sql_identifier(id2).is_done());
        assert!(sql_identifier(id3).is_done());
        assert!(sql_identifier(id4).is_err());
    }
}
<commit_msg>Order operators by length<commit_after>use nom::{alphanumeric, digit, eof, is_alphanumeric, line_ending, multispace, space};
use nom::{IResult, Err, ErrorKind, Needed};
use parser::{Column, Table};
use std::str;
use std::str::FromStr;

#[derive(Clone, Debug, PartialEq)]
pub enum Operator {
    Not,
    And,
    Or,
    Like,
    NotLike,
    Equal,
    NotEqual,
    Greater,
    GreaterOrEqual,
    Less,
    LessOrEqual,
}

#[derive(Clone, Debug, PartialEq)]
pub enum FieldExpression {
    All,
    Seq(Vec<Column>),
}

impl Default for FieldExpression {
    fn default() -> FieldExpression {
        FieldExpression::All
    }
}

#[inline]
pub fn is_sql_identifier(chr: u8) -> bool {
    // XXX(malte): dot should not be in here once we have proper alias handling
    is_alphanumeric(chr) || chr == '_' as u8
}

#[inline]
pub fn is_fp_number(chr: u8) -> bool {
    is_alphanumeric(chr) || chr == '.' as u8
}

/// Parses a SQL identifier (alphanumeric and "_").
named!(pub fp_number<&[u8], &[u8]>,
// f64::from_str(&format!("{}.{}", integral, fractional)).unwrap()
//      integral: map_res!(take_while1!(is_digit), str::from_utf8) ~
//      fractional: map_res!(take_while!(is_digit), str::from_utf8),
    take_while1!(is_fp_number)
);

/// Parses a SQL column identifier in the table.column format
named!(pub column_identifier<&[u8], Column>,
    chain!(
        table: opt!(
            chain!(
                tbl_name: map_res!(sql_identifier, str::from_utf8) ~
                tag!("."),
                || { tbl_name }
            )
        ) ~
        column: map_res!(sql_identifier, str::from_utf8),
        || {
            Column {
                name: String::from(column),
                table: match table {
                    None => None,
                    Some(t) => Some(String::from(t)),
                },
            }
        }
    )
);

/// Parses a SQL identifier (alphanumeric and "_").
named!(pub sql_identifier<&[u8], &[u8]>,
    take_while1!(is_sql_identifier)
);

/// Parse an unsigned integer.
named!(pub unsigned_number<&[u8], u64>,
    map_res!(
        map_res!(digit, str::from_utf8),
        FromStr::from_str
    )
);

/// Parse a terminator that ends a SQL statement.
named!(pub statement_terminator,
    delimited!(opt!(multispace),
               alt_complete!(tag!(";") | line_ending | eof),
               opt!(multispace)
    )
);

/// Parse binary comparison operators
named!(pub binary_comparison_operator<&[u8], Operator>,
    alt_complete!(
           map!(caseless_tag!("not_like"), |_| Operator::NotLike)
         | map!(caseless_tag!("like"), |_| Operator::Like)
         | map!(caseless_tag!("<>"), |_| Operator::NotEqual)
         | map!(caseless_tag!(">="), |_| Operator::GreaterOrEqual)
         | map!(caseless_tag!("<="), |_| Operator::LessOrEqual)
         | map!(caseless_tag!("="), |_| Operator::Equal)
         | map!(caseless_tag!("<"), |_| Operator::Less)
         | map!(caseless_tag!(">"), |_| Operator::Greater)
    )
);

/// Parse logical operators
named!(pub binary_logical_operator<&[u8], Operator>,
    alt_complete!(
           map!(caseless_tag!("and"), |_| Operator::And)
         | map!(caseless_tag!("or"), |_| Operator::Or)
    )
);

/// Parse unary comparison operators
named!(pub unary_comparison_operator<&[u8], &str>,
    map_res!(
        alt_complete!(
               tag_s!(b"ISNULL")
             | tag_s!(b"NOT")
             | tag_s!(b"-") // ??? (number neg)
        ),
        str::from_utf8
    )
);

/// Parse unary comparison operators
named!(pub unary_negation_operator<&[u8], Operator>,
    alt_complete!(
          map!(caseless_tag!("not"), |_| Operator::Not)
        | map!(caseless_tag!("!"), |_| Operator::Not)
    )
);

/// Parse rule for a comma-separated list.
named!(pub field_list<&[u8], Vec<Column> >,
       many0!(
           chain!(
               fieldname: column_identifier ~
               opt!(
                   chain!(
                       multispace? ~
                       tag!(",") ~
                       multispace?,
                       ||{}
                   )
               ),
               || { fieldname }
           )
       )
);

/// Parse list of columns/fields.
/// XXX(malte): add support for named table notation
named!(pub field_expr<&[u8], FieldExpression>,
       alt_complete!(
           tag!("*") => { |_| FieldExpression::All }
         | map!(field_list, |v| FieldExpression::Seq(v))
       )
);

/// Parse list of table names.
/// XXX(malte): add support for aliases
named!(pub table_list<&[u8], Vec<Table> >,
       many0!(
           chain!(
               table: table_reference ~
               opt!(
                   chain!(
                       multispace? ~
                       tag!(",") ~
                       multispace?,
                       || {}
                   )
               ),
               || { table }
           )
       )
);

/// Parse a list of values (e.g., for INSERT syntax).
/// XXX(malte): proper value type
named!(pub value_list<&[u8], Vec<&str> >,
       many0!(
           map_res!(
               chain!(
                   val: alt_complete!(
                         tag_s!(b"?")
                       | tag_s!(b"CURRENT_TIMESTAMP")
                       | delimited!(tag!("'"), alphanumeric, tag!("'"))
                       | fp_number
                   ) ~
                   opt!(
                       chain!(
                           multispace? ~
                           tag!(",") ~
                           multispace?,
                           ||{}
                       )
                   ),
                   || { val }
               ),
               str::from_utf8
           )
       )
);

/// Parse a reference to a named table
/// XXX(malte): add support for schema.table notation
named!(pub table_reference<&[u8], Table>,
    chain!(
        table: map_res!(sql_identifier, str::from_utf8) ~
        alias: opt!(
            chain!(
                space ~
                caseless_tag!("as") ~
                space ~
                alias: map_res!(sql_identifier, str::from_utf8),
                || { String::from(alias) }
            )
        ),
        || {
            Table {
                name: String::from(table),
                alias: alias,
            }
        }
    )
);

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

    #[test]
    fn sql_identifiers() {
        let id1 = b"foo";
        let id2 = b"f_o_o";
        let id3 = b"foo12";
        let id4 = b":fo oo";

        assert!(sql_identifier(id1).is_done());
        assert!(sql_identifier(id2).is_done());
        assert!(sql_identifier(id3).is_done());
        assert!(sql_identifier(id4).is_err());
    }
}
<|endoftext|>


<commit_before>// In heavy WIP
use std::fs::File;
use std::io::Read;
use std::path::Path;
use toml;

use error::PackageError;
use repository::RepositoryList;

/// Represents the config file.
#[derive(Deserialize)]
pub struct Config {
    pub buffer_size: Option<u64>,
    pub config_path: Option<String>,
    pub download_path: Option<String>,
    pub repository: Option<RepositoryList>,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            buffer_size: Some(65536),
            config_path: Some(String::from("~/.wiz/config/")),
            download_path: Some(String::from("~/.wiz/downloads/")),
            repository: Some(RepositoryList(Vec::new())),
        }
    }
}

impl Config {
    /// Read the config file from the config path specified in `path`.
    /// If the config file is read & parsed properly, it should return
    /// a `Config`.
    pub fn read_from(path: &Path) -> Result<Self, PackageError> {
        // Read the config file into a string
        let mut config = File::open(path)?;
        let mut content = String::new();
        config.read_to_string(&mut content)?;

        // Try parsing the string, and convert it into a `Config`.
        let config = content.parse::<toml::Value>()?;
        let config = config.try_into::<Self>()?;
        Ok(config)
    }

    /// Read the config file from the default config path.
    /// If the config file is read & parsed properly, it should return
    /// a `Config`.
    pub fn read_from_default() -> Result<Self, PackageError> {
        if let Self { config_path: Some(path), .. } = Self::default() {
            Self::read_from(Path::new(&path))
        } else {
            // This should never, ever, happen.
            Err(PackageError::Parsing(
                String::from("Unable to get config_path from the default cfg.")
            ))
        }
    }

}<commit_msg>Function for config file's None-gap filling<commit_after>// In heavy WIP
use std::fs::File;
use std::io::Read;
use std::path::Path;
use toml;

use error::PackageError;
use repository::RepositoryList;

/// Represents the config file.
#[derive(Deserialize)]
pub struct Config {
    pub buffer_size: Option<u64>,
    pub config_path: Option<String>,
    pub download_path: Option<String>,
    pub repository: Option<RepositoryList>,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            buffer_size: Some(65536),
            config_path: Some(String::from("~/.wiz/config/")),
            download_path: Some(String::from("~/.wiz/downloads/")),
            repository: Some(RepositoryList(Vec::new())),
        }
    }
}

/// Internal macro, to ease the implementation of `fill_with_default`.
macro_rules! set_on_none {
    ($dest:ident, $key:ident, $value:expr) => (
        match $dest {
            Self { $key: ref mut x, .. } => {
                match x {
                    &mut Some(_) => {},
                    &mut None => *x = $value,
                }
            }
        }

        // Wait for RFC-2086 to be implemented first.
        /*
            [allow(irrefutable_let_pattern)]
            if let Config { $key: mut x, .. } = $dest {
                if let None = x { x = $value }
            }
        */
    );
}

impl Config {
    /// Read the config file from the config path specified in `path`.
    /// If the config file is read & parsed properly, it should return
    /// a `Config`.
    pub fn read_from(path: &Path) -> Result<Self, PackageError> {
        // Read the config file into a string
        let mut config = File::open(path)?;
        let mut content = String::new();
        config.read_to_string(&mut content)?;

        // Try parsing the string, and convert it into a `Config`.
        let config = content.parse::<toml::Value>()?;
        let config = config.try_into::<Self>()?;
        Ok(config)
    }

    /// Read the config file from the default config path.
    /// If the config file is read & parsed properly, it should return
    /// a `Config`.
    pub fn read_from_default() -> Result<Self, PackageError> {
        if let Self { config_path: Some(path), .. } = Self::default() {
            Self::read_from(Path::new(&path))
        } else {
            // This should never, ever, happen.
            Err(PackageError::Parsing(
                String::from("Unable to get config_path from the default cfg.")
            ))
        }
    }

    pub fn fill_with_default(mut self) -> Self {
        let Self {
            buffer_size: buf_size,
            config_path: conf_path,
            download_path: dl_path,
            repository: repo,
        } = Self::default();

        set_on_none!(self, buffer_size, buf_size);
        set_on_none!(self, config_path, conf_path);
        set_on_none!(self, download_path, dl_path);
        set_on_none!(self, repository, repo);

        self
    }
}
<|endoftext|>


<commit_before>extern crate xdg;

pub fn config_dir(name: &str) -> Path {
    let mut path: Path = xdg::get_config_home();
    path.push(name);
    path
}
<commit_msg>config: update for new dependencies<commit_after>extern crate xdg_basedir;
use std::path::PathBuf;
use self::xdg_basedir::error::Error;

pub fn config_dir(name: &str) -> Result<PathBuf, Error> {
    let mut path: PathBuf = try!(xdg_basedir::get_config_home());
    path.push(name);
    Ok(path)
}
<|endoftext|>