<commit_before>
/// PCAP Header
///
/// * magic number (0xA1B23C4D)
/// * major version number
/// * minor version number
/// * GMT to local correction
/// * accuracy of timestamps (typically ignored as 0)
/// * max length of captured packets, in octets
/// * data link type
///
#[derive(Debug)]
pub struct Header {
    pub magic_number: u32,
    pub version_major: u16,
    pub version_minor: u16,
    pub thiszone: i32,
    pub sigfigs: u32,
    pub snaplen: u32,
    pub network: Link,
}

/// Link types as defined in http://www.tcpdump.org/linktypes.html
#[derive(Debug)]
pub enum Link {
    Null,
    Ethernet,
    Unknown(u32),
}

impl From<u32> for Link {
    fn from(link: u32) -> Self {
        match link {
            0 => Link::Null,
            1 => Link::Ethernet,
            otherwise => Link::Unknown(otherwise),
        }
    }
}

/// Record entry in a packet capture
///
/// * timestamp seconds
/// * timestamp nanoseconds
/// * number of octets of packet saved in file
/// * actual length of packet
#[derive(Debug)]
pub struct Record {
    pub ts_sec: u32,
    pub ts_usec: u32,
    pub incl_len: u32,
    pub orig_len: u32,
    pub payload: Vec<u8>,
}

use std::io;
use std::io::{Read, BufReader};

struct Buffer<R>(BufReader<R>);

impl<R: Read> Buffer<R> {
    fn take(&mut self, length: usize) -> io::Result<Vec<u8>> {
        let mut buffer = vec![0; length];
        self.0.read_exact(&mut buffer)?;
        Ok(buffer)
    }

    fn take_i32(&mut self) -> io::Result<i32> {
        let buffer = self.take(4)?;
        Ok((buffer[0] as i32) + ((buffer[1] as i32) << 8) + ((buffer[2] as i32) << 16) +
           ((buffer[3] as i32) << 24))
    }

    fn take_u32(&mut self) -> io::Result<u32> {
        let buffer = self.take(4)?;
        Ok((buffer[0] as u32) + ((buffer[1] as u32) << 8) + ((buffer[2] as u32) << 16) +
           ((buffer[3] as u32) << 24))
    }

    fn take_u16(&mut self) -> io::Result<u16> {
        let buffer = self.take(2)?;
        Ok((buffer[0] as u16) + ((buffer[1] as u16) << 8))
    }
}

/// PacketCapture: container for pcap byte-stream
pub struct PacketCapture<R> {
    capture: Buffer<R>,
}

impl<R: Read> PacketCapture<R> {
    pub fn new(capture: R) -> PacketCapture<R> {
        PacketCapture { capture: Buffer(BufReader::new(capture)) }
    }
}

impl Header {
    fn parse<R: Read>(data: &mut Buffer<R>) -> io::Result<Header> {
        Ok(Header {
               magic_number: data.take_u32()?,
               version_major: data.take_u16()?,
               version_minor: data.take_u16()?,
               thiszone: data.take_i32()?,
               sigfigs: data.take_u32()?,
               snaplen: data.take_u32()?,
               network: Link::from(data.take_u32()?),
           })
    }
}

impl Record {
    fn parse<R: Read>(data: &mut Buffer<R>) -> io::Result<Record> {
        let ts_sec = data.take_u32()?;
        let ts_usec = data.take_u32()?;
        let incl_len = data.take_u32()?;
        let orig_len = data.take_u32()?;
        let payload = data.take(incl_len as usize)?;

        Ok(Record {
               ts_sec: ts_sec,
               ts_usec: ts_usec,
               incl_len: incl_len,
               orig_len: orig_len,
               payload: payload,
           })
    }
}

impl<R: Read> PacketCapture<R> {
    pub fn parse(mut self) -> io::Result<(Header, Records<R>)> {
        Ok((Header::parse(&mut self.capture)?, Records { capture: self.capture }))
    }
}

pub struct Records<R> {
    capture: Buffer<R>,
}

impl<R: Read> Iterator for Records<R> {
    type Item = Record;

    fn next(&mut self) -> Option<Self::Item> {
        Record::parse(&mut self.capture).ok()
    }
}
<commit_msg>PCAP 'buffer' constructors.<commit_after>
/// PCAP Header
///
/// * magic number (0xA1B23C4D)
/// * major version number
/// * minor version number
/// * GMT to local correction
/// * accuracy of timestamps (typically ignored as 0)
/// * max length of captured packets, in octets
/// * data link type
///
#[derive(Debug)]
pub struct Header {
    pub magic_number: u32,
    pub version_major: u16,
    pub version_minor: u16,
    pub thiszone: i32,
    pub sigfigs: u32,
    pub snaplen: u32,
    pub network: Link,
}

/// Link types as defined in http://www.tcpdump.org/linktypes.html
#[derive(Debug)]
pub enum Link {
    Null,
    Ethernet,
    Unknown(u32),
}

impl From<u32> for Link {
    fn from(link: u32) -> Self {
        match link {
            0 => Link::Null,
            1 => Link::Ethernet,
            otherwise => Link::Unknown(otherwise),
        }
    }
}

/// Record entry in a packet capture
///
/// * timestamp seconds
/// * timestamp nanoseconds
/// * number of octets of packet saved in file
/// * actual length of packet
#[derive(Debug)]
pub struct Record {
    pub ts_sec: u32,
    pub ts_usec: u32,
    pub incl_len: u32,
    pub orig_len: u32,
    pub payload: Vec<u8>,
}

use std::io;
use std::io::{Read, BufReader, Cursor};

struct Buffer<R>(BufReader<R>);

impl<R: Read> Buffer<R> {
    fn new(buffer: R) -> Self {
        Buffer(BufReader::new(buffer))
    }

    fn take(&mut self, length: usize) -> io::Result<Vec<u8>> {
        let mut buffer = vec![0; length];
        self.0.read_exact(&mut buffer)?;
        Ok(buffer)
    }

    fn take_i32(&mut self) -> io::Result<i32> {
        let buffer = self.take(4)?;
        Ok((buffer[0] as i32) + ((buffer[1] as i32) << 8) + ((buffer[2] as i32) << 16) +
           ((buffer[3] as i32) << 24))
    }

    fn take_u32(&mut self) -> io::Result<u32> {
        let buffer = self.take(4)?;
        Ok((buffer[0] as u32) + ((buffer[1] as u32) << 8) + ((buffer[2] as u32) << 16) +
           ((buffer[3] as u32) << 24))
    }

    fn take_u16(&mut self) -> io::Result<u16> {
        let buffer = self.take(2)?;
        Ok((buffer[0] as u16) + ((buffer[1] as u16) << 8))
    }
}

/// PacketCapture: container for pcap byte-stream
pub struct PacketCapture<R> {
    capture: Buffer<R>,
}

impl<R: Read> PacketCapture<R> {
    pub fn new(capture: R) -> PacketCapture<R> {
        PacketCapture { capture: Buffer::new(capture) }
    }
}

impl From<Vec<u8>> for PacketCapture<Cursor<Vec<u8>>> {
    fn from(buffer: Vec<u8>) -> Self {
        PacketCapture { capture: Buffer::new(Cursor::new(buffer)) }
    }
}

impl Header {
    fn parse<R: Read>(data: &mut Buffer<R>) -> io::Result<Header> {
        Ok(Header {
               magic_number: data.take_u32()?,
               version_major: data.take_u16()?,
               version_minor: data.take_u16()?,
               thiszone: data.take_i32()?,
               sigfigs: data.take_u32()?,
               snaplen: data.take_u32()?,
               network: Link::from(data.take_u32()?),
           })
    }
}

impl Record {
    fn parse<R: Read>(data: &mut Buffer<R>) -> io::Result<Record> {
        let ts_sec = data.take_u32()?;
        let ts_usec = data.take_u32()?;
        let incl_len = data.take_u32()?;
        let orig_len = data.take_u32()?;
        let payload = data.take(incl_len as usize)?;

        Ok(Record {
               ts_sec: ts_sec,
               ts_usec: ts_usec,
               incl_len: incl_len,
               orig_len: orig_len,
               payload: payload,
           })
    }
}

impl<R: Read> PacketCapture<R> {
    pub fn parse(mut self) -> io::Result<(Header, Records<R>)> {
        Ok((Header::parse(&mut self.capture)?, Records { capture: self.capture }))
    }
}

pub struct Records<R> {
    capture: Buffer<R>,
}

impl<R: Read> Iterator for Records<R> {
    type Item = Record;

    fn next(&mut self) -> Option<Self::Item> {
        Record::parse(&mut self.capture).ok()
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::{TcpStream, ToSocketAddrs};
use std::time::Duration;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Sender};
use std::thread::{self, JoinHandle};
use util::codec::rpc;
use util::make_std_tcp_conn;
use protobuf::{self, MessageStatic};
use super::Result;

pub trait TRpcClient: Sync + Send {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic;
    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()>;
}

#[derive(Debug)]
pub struct RpcClient {
    stream: Mutex<TcpStream>,
}

const SOCKET_READ_TIMEOUT: u64 = 3;
const SOCKET_WRITE_TIMEOUT: u64 = 3;

impl RpcClient {
    pub fn new(addr: A) -> Result<RpcClient> {
        let stream = try!(make_std_tcp_conn(addr));
        try!(stream.set_read_timeout(Some(Duration::from_secs(SOCKET_READ_TIMEOUT))));
        try!(stream.set_write_timeout(Some(Duration::from_secs(SOCKET_WRITE_TIMEOUT))));
        Ok(RpcClient { stream: Mutex::new(stream) })
    }

    fn post_msg<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<Vec<u8>> {
        // TODO: reconnect if error
        let mut stream = self.stream.lock().unwrap();
        try!(rpc::encode_msg(&mut *stream, msg_id, message));
        let (id, data) = try!(rpc::decode_data(&mut *stream));
        if id != msg_id {
            return Err(box_err!("pd response msg_id not match, want {}, got {}", msg_id, id));
        }

        Ok(data)
    }
}

impl TRpcClient for RpcClient {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let mut resp = P::new();
        let data = try!(self.post_msg(msg_id, message));
        try!(rpc::decode_body(&data, &mut resp));
        Ok(resp)
    }

    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()> {
        let _ = try!(self.post_msg(msg_id, message));
        Ok(())
    }
}

#[derive(Debug)]
enum Msg {
    Rpc(u64, Box<protobuf::Message + Send>),
    Quit,
}

pub struct Client<C: TRpcClient + 'static> {
    msg_id: AtomicUsize,
    rpc_client: Arc<C>,
    tx: Mutex<Sender<Msg>>,
    handle: Option<JoinHandle<()>>,
}

impl<C: TRpcClient + 'static> Client<C> {
    pub fn new(rpc_client: C) -> Result<Client<C>> {
        let shared = Arc::new(rpc_client);
        let (tx, rx) = mpsc::channel();
        let handle = {
            let rpc_client = shared.clone();
            try!(thread::Builder::new()
                     .name("pd-client background worker".to_owned())
                     .spawn(move || {
                         info!("pd-client background worker starting");
                         while let Msg::Rpc(msg_id, msg) = rx.recv().unwrap() {
                             if let Err(e) = rpc_client.post(msg_id, msg.as_ref()) {
                                 error!("pd post error: {}", e);
                             } else {
                                 debug!("pd post done: {}", msg_id);
                             }
                         }
                         info!("pd-client background worker closing");
                     }))
        };
        Ok(Client {
            msg_id: AtomicUsize::new(0),
            rpc_client: shared.clone(),
            tx: Mutex::new(tx),
            handle: Some(handle),
        })
    }

    fn close(&mut self) {
        if let Err(e) = self.clone_tx().send(Msg::Quit) {
            error!("failed to send quit message to pd-client background worker: {}",
                   e);
        }
        let handle = self.handle.take().unwrap();
        if handle.join().is_err() {
            error!("failed to wait pd-client background worker quit");
        }
        info!("pd-client closed.");
    }

    pub fn send<M, P>(&self, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let id = self.alloc_msg_id();
        self.rpc_client.send(id, message)
    }

    pub fn post<M: protobuf::Message + Send>(&self, message: M) -> Result<()> {
        let id = self.alloc_msg_id();
        if let Err(e) = self.clone_tx().send(Msg::Rpc(id, box message)) {
            return Err(box_err!(format!("SendError: {:?}", e)));
        }
        Ok(())
    }

    fn alloc_msg_id(&self) -> u64 {
        self.msg_id.fetch_add(1, Ordering::Relaxed) as u64
    }

    fn clone_tx(&self) -> Sender<Msg> {
        let tx = self.tx.lock().unwrap();
        tx.clone()
    }
}

impl<C: TRpcClient + 'static> Drop for Client<C> {
    fn drop(&mut self) {
        self.close()
    }
}

#[cfg(test)]
mod tests {
    use super::super::{Result, PdClient};
    use super::*;
    use protobuf;
    use bytes::ByteBuf;
    use util::codec::rpc;
    use kvproto::{metapb, pdpb};

    struct MockRpcClient;

    impl MockRpcClient {
        fn reply(&self, req: &pdpb::Request) -> pdpb::Response {
            assert!(req.has_header());
            let header = req.get_header();
            let (uuid, cluster_id) = (header.get_uuid(), header.get_cluster_id());

            let mut resp = pdpb::Response::new();
            let mut header = pdpb::ResponseHeader::new();
            header.set_uuid(uuid.to_vec());
            header.set_cluster_id(cluster_id);

            match req.get_cmd_type() {
                pdpb::CommandType::AllocId => {
                    let mut alloc = pdpb::AllocIdResponse::new();
                    alloc.set_id(42);
                    resp.set_alloc_id(alloc);
                }
                pdpb::CommandType::PutMeta => {
                    let meta = req.get_put_meta();
                    if meta.get_meta_type() == pdpb::MetaType::StoreType {
                        let store = meta.get_store();
                        assert_eq!(store.get_id(), 233);
                        assert_eq!(store.get_address(), "localhost");
                    }
                }
                _ => {
                    let mut err = pdpb::Error::new();
                    err.set_message(String::from("I've got a bad feeling about this"));
                    header.set_error(err);
                }
            }

            resp.set_header(header);
            resp.set_cmd_type(req.get_cmd_type());
            resp
        }
    }

    impl TRpcClient for MockRpcClient {
        fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

            where M: protobuf::Message,
                  P: protobuf::Message + Default
        {
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, message).unwrap();
            let mut req = pdpb::Request::new();
            rpc::decode_msg(&mut buf.flip(), &mut req).unwrap();

            let resp = self.reply(&req);
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, &resp).unwrap();
            let mut msg = P::default();
            rpc::decode_msg(&mut buf.flip(), &mut msg).unwrap();
            Ok(msg)
        }

        fn post<M: protobuf::Message + ?Sized>(&self, _: u64, _: &M) -> Result<()> {
            Ok(())
        }
    }

    #[test]
    fn test_pd_client() {
        let mut client = Client::new(MockRpcClient).unwrap();
        assert_eq!(client.alloc_id().unwrap(), 42u64);
        assert!(client.is_cluster_bootstrapped(1).is_err());

        let mut store = metapb::Store::new();
        store.set_id(233);
        store.set_address("localhost".to_owned());
        assert!(client.put_store(1, store).is_ok());
    }
}
<commit_msg>pd: support pd list.<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::TcpStream;
use std::time::Duration;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Sender};
use std::thread::{self, JoinHandle};
use util::codec::rpc;
use util::make_std_tcp_conn;
use protobuf::{self, MessageStatic};
use super::Result;

pub trait TRpcClient: Sync + Send {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic;
    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()>;
}

#[derive(Debug)]
struct RpcClientCore {
    addrs: Vec<String>,
    // Try to connect pd with round-robin.
    next_index: usize,
    stream: Option<TcpStream>,
}

fn send_msg<M: protobuf::Message + ?Sized>(stream: &mut TcpStream,
                                           msg_id: u64,
                                           message: &M)
                                           -> Result<(u64, Vec<u8>)> {
    try!(stream.set_write_timeout(Some(Duration::from_secs(SOCKET_WRITE_TIMEOUT))));
    try!(rpc::encode_msg(stream, msg_id, message));

    try!(stream.set_read_timeout(Some(Duration::from_secs(SOCKET_READ_TIMEOUT))));
    let (id, data) = try!(rpc::decode_data(stream));
    Ok((id, data))
}

impl RpcClientCore {
    fn new(dsn: &str) -> RpcClientCore {
        let addrs: Vec<String> = dsn.split(',').map(From::from).collect();
        RpcClientCore {
            addrs: addrs,
            next_index: 0,
            stream: None,
        }
    }

    fn try_connect(&mut self) -> Result<()> {
        let index = self.next_index;
        self.next_index = (self.next_index + 1) % self.addrs.len();

        let addr = self.addrs.get(index).unwrap();
        let stream = try!(make_std_tcp_conn(&**addr));
        self.stream = Some(stream);
        Ok(())
    }

    fn post_msg<M: protobuf::Message + ?Sized>(&mut self,
                                               msg_id: u64,
                                               message: &M)
                                               -> Result<Vec<u8>> {
        // If we post failed, we should retry other addresses.
        let n = self.addrs.len();
        for _ in 0..n {
            // If no stream, try connect first.
            if self.stream.is_none() {
                if let Err(e) = self.try_connect() {
                    warn!("connect pd failed {:?}", e);
                    continue;
                }
            }

            let mut stream = self.stream.take().unwrap();
            // We may send message to a not leader pd, retry.

            let (id, data) = match send_msg(&mut stream, msg_id, message) {
                Err(e) => {
                    warn!("send message to pd failed {:?}", e);
                    continue;
                }
                Ok((id, data)) => (id, data),
            };

            if id != msg_id {
                return Err(box_err!("pd response msg_id not match, want {}, got {}", msg_id, id));
            }

            self.stream = Some(stream);
            return Ok(data);
        }

        Err(box_err!("post message to pd failed"))
    }
}

#[derive(Debug)]
pub struct RpcClient {
    core: Mutex<RpcClientCore>,
}

const SOCKET_READ_TIMEOUT: u64 = 3;
const SOCKET_WRITE_TIMEOUT: u64 = 3;

impl RpcClient {
    pub fn new(dsn: &str) -> Result<RpcClient> {
        Ok(RpcClient { core: Mutex::new(RpcClientCore::new(dsn)) })
    }
}

impl TRpcClient for RpcClient {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let mut resp = P::new();
        let data = try!(self.core.lock().unwrap().post_msg(msg_id, message));
        try!(rpc::decode_body(&data, &mut resp));
        Ok(resp)
    }

    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()> {
        let _ = try!(self.core.lock().unwrap().post_msg(msg_id, message));
        Ok(())
    }
}

#[derive(Debug)]
enum Msg {
    Rpc(u64, Box<protobuf::Message + Send>),
    Quit,
}

pub struct Client<C: TRpcClient + 'static> {
    msg_id: AtomicUsize,
    rpc_client: Arc<C>,
    tx: Mutex<Sender<Msg>>,
    handle: Option<JoinHandle<()>>,
}

impl<C: TRpcClient + 'static> Client<C> {
    pub fn new(rpc_client: C) -> Result<Client<C>> {
        let shared = Arc::new(rpc_client);
        let (tx, rx) = mpsc::channel();
        let handle = {
            let rpc_client = shared.clone();
            try!(thread::Builder::new()
                     .name("pd-client background worker".to_owned())
                     .spawn(move || {
                         info!("pd-client background worker starting");
                         while let Msg::Rpc(msg_id, msg) = rx.recv().unwrap() {
                             if let Err(e) = rpc_client.post(msg_id, msg.as_ref()) {
                                 error!("pd post error: {}", e);
                             } else {
                                 debug!("pd post done: {}", msg_id);
                             }
                         }
                         info!("pd-client background worker closing");
                     }))
        };
        Ok(Client {
            msg_id: AtomicUsize::new(0),
            rpc_client: shared.clone(),
            tx: Mutex::new(tx),
            handle: Some(handle),
        })
    }

    fn close(&mut self) {
        if let Err(e) = self.clone_tx().send(Msg::Quit) {
            error!("failed to send quit message to pd-client background worker: {}",
                   e);
        }
        let handle = self.handle.take().unwrap();
        if handle.join().is_err() {
            error!("failed to wait pd-client background worker quit");
        }
        info!("pd-client closed.");
    }

    pub fn send<M, P>(&self, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let id = self.alloc_msg_id();
        self.rpc_client.send(id, message)
    }

    pub fn post<M: protobuf::Message + Send>(&self, message: M) -> Result<()> {
        let id = self.alloc_msg_id();
        if let Err(e) = self.clone_tx().send(Msg::Rpc(id, box message)) {
            return Err(box_err!(format!("SendError: {:?}", e)));
        }
        Ok(())
    }

    fn alloc_msg_id(&self) -> u64 {
        self.msg_id.fetch_add(1, Ordering::Relaxed) as u64
    }

    fn clone_tx(&self) -> Sender<Msg> {
        let tx = self.tx.lock().unwrap();
        tx.clone()
    }
}

impl<C: TRpcClient + 'static> Drop for Client<C> {
    fn drop(&mut self) {
        self.close()
    }
}

#[cfg(test)]
mod tests {
    use std::net::TcpListener;
    use std::thread;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};

    use super::super::{Result, PdClient};
    use super::*;
    use protobuf;
    use bytes::ByteBuf;
    use util::codec::rpc;
    use kvproto::{metapb, pdpb};
    use rand;
    use util::make_std_tcp_conn;

    struct MockRpcClient;

    impl MockRpcClient {
        fn reply(&self, req: &pdpb::Request) -> pdpb::Response {
            assert!(req.has_header());
            let header = req.get_header();
            let (uuid, cluster_id) = (header.get_uuid(), header.get_cluster_id());

            let mut resp = pdpb::Response::new();
            let mut header = pdpb::ResponseHeader::new();
            header.set_uuid(uuid.to_vec());
            header.set_cluster_id(cluster_id);

            match req.get_cmd_type() {
                pdpb::CommandType::AllocId => {
                    let mut alloc = pdpb::AllocIdResponse::new();
                    alloc.set_id(42);
                    resp.set_alloc_id(alloc);
                }
                pdpb::CommandType::PutMeta => {
                    let meta = req.get_put_meta();
                    if meta.get_meta_type() == pdpb::MetaType::StoreType {
                        let store = meta.get_store();
                        assert_eq!(store.get_id(), 233);
                        assert_eq!(store.get_address(), "localhost");
                    }
                }
                _ => {
                    let mut err = pdpb::Error::new();
                    err.set_message(String::from("I've got a bad feeling about this"));
                    header.set_error(err);
                }
            }

            resp.set_header(header);
            resp.set_cmd_type(req.get_cmd_type());
            resp
        }
    }

    impl TRpcClient for MockRpcClient {
        fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

            where M: protobuf::Message,
                  P: protobuf::Message + Default
        {
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, message).unwrap();
            let mut req = pdpb::Request::new();
            rpc::decode_msg(&mut buf.flip(), &mut req).unwrap();

            let resp = self.reply(&req);
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, &resp).unwrap();
            let mut msg = P::default();
            rpc::decode_msg(&mut buf.flip(), &mut msg).unwrap();
            Ok(msg)
        }

        fn post<M: protobuf::Message + ?Sized>(&self, _: u64, _: &M) -> Result<()> {
            Ok(())
        }
    }

    #[test]
    fn test_pd_client() {
        let mut client = Client::new(MockRpcClient).unwrap();
        assert_eq!(client.alloc_id().unwrap(), 42u64);
        assert!(client.is_cluster_bootstrapped(1).is_err());

        let mut store = metapb::Store::new();
        store.set_id(233);
        store.set_address("localhost".to_owned());
        assert!(client.put_store(1, store).is_ok());
    }

    fn start_pd_server(index: usize,
                       leader_index: Arc<AtomicUsize>,
                       quit: Arc<AtomicBool>)
                       -> (thread::JoinHandle<()>, String) {
        let l = TcpListener::bind("127.0.0.1:0").unwrap();

        let addr = format!("{}", l.local_addr().unwrap());

        let h = thread::spawn(move || {
            loop {
                let (mut stream, _) = l.accept().unwrap();

                if quit.load(Ordering::SeqCst) {
                    return;
                }

                stream.set_nodelay(true).unwrap();

                let leader = leader_index.load(Ordering::SeqCst);
                if leader != index {
                    continue;
                }

                let (id, data) = rpc::decode_data(&mut stream).unwrap();
                rpc::encode_data(&mut stream, id, &data).unwrap();
            }
        });

        (h, addr)
    }

    #[test]
    fn test_rpc_client() {
        let leader = Arc::new(AtomicUsize::new(0));
        let quit = Arc::new(AtomicBool::new(false));
        let mut handlers = vec![];
        let mut addrs = vec![];

        let count = 3;
        for i in 0..count {
            let (h, addr) = start_pd_server(i, leader.clone(), quit.clone());
            handlers.push(h);
            addrs.push(addr);
        }

        let dsn = addrs.join(",");

        let msg = pdpb::Request::new();
        let client = RpcClient::new(&dsn).unwrap();

        for i in 0..10 {
            if let Err(_) = client.post(i as u64, &msg) {
                // If failed, next retry must be OK.
                client.post(i as u64, &msg).unwrap();
            }

            // select a leader randomly
            leader.store(rand::random::<usize>() % 3, Ordering::SeqCst);
        }

        quit.store(true, Ordering::SeqCst);

        // connect the server so that we can close the thread.
        for addr in addrs {
            make_std_tcp_conn(&*addr).unwrap();
        }

        for h in handlers {
            h.join().unwrap();
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

use util::tree::GornTree;

mod from_str;
// mod relabel;
pub mod negra;

/// Variable or terminal symbol in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub enum VarT<T> {
    /// `Var(i, j)` represents the `j`th component of the `i`th successor.
    /// Indexing starts from `0`.
    Var(usize, usize),
    T(T),
}

/// Composition function in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub struct Composition<T> {
    pub composition: Vec<Vec<VarT<T>>>,
}

impl<T> From<Vec<Vec<VarT<T>>>> for Composition<T> {
    fn from(encapsulated_value: Vec<Vec<VarT<T>>>) -> Self {
        Composition { composition: encapsulated_value }
    }
}

/// Rule of a weighted MCFG.
#[derive(Debug, PartialOrd, Ord, Clone)]
pub struct PMCFGRule<N, T, W> {
    pub head: N,
    pub tail: Vec<N>,
    pub composition: Composition<T>,
    pub weight: W,
}

impl<N, T, W> PMCFGRule<N, T, W>
    where T: Clone,
          W: Clone,
{
    pub fn map_nonterminals<F, M>(&self, f: F) -> PMCFGRule<M, T, W>
        where F: Fn(&N) -> M,
    {
        PMCFGRule {
            head: f(&self.head),
            tail: self.tail.iter().map(f).collect(),
            composition: self.composition.clone(),
            weight: self.weight.clone()
        }
    }
}

/// A weighted MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct PMCFG<N, T, W> {
    pub _dummy: PhantomData<T>,
    pub initial: Vec<N>,
    pub rules: Vec<PMCFGRule<N, T, W>>,
}

impl<N: Hash, T: Hash, W> Hash for PMCFGRule<N, T, W> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.tail.hash(state);
        self.composition.hash(state);
    }
}

impl<N: PartialEq, T: PartialEq, W> PartialEq for PMCFGRule<N, T, W> {
    fn eq(&self, other: &Self) -> bool {
        self.head == other.head && self.tail == other.tail && self.composition == other.composition
    }
}

impl<N: Eq, T: Eq, W> Eq for PMCFGRule<N, T, W> {}

impl<T: fmt::Display> fmt::Display for Composition<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter0 = self.composition.iter().peekable();
        let mut iter1;

        buffer.push_str("[");
        while let Some(proj) = iter0.next() {
            iter1 = proj.into_iter().peekable();
            buffer.push_str("[");
            while let Some(vart) = iter1.next() {
                buffer.push_str(format!("{}", vart).as_str());
                if iter1.peek().is_some() {
                    buffer.push_str(", ");
                }
            }
            buffer.push_str("]");
            if iter0.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]");

        write!(f, "{}", buffer)
    }
}

impl<T: fmt::Display> fmt::Display for VarT<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VarT::Var(i, j) => {
                write!(f, "Var {} {}", i, j)
            },
            VarT::T(ref x) => {
                write!(f, "T \"{}\"", x)
            },
        }
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFGRule<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
           let mut buffer = "".to_string();

           let mut iter = self.tail.iter().peekable();

           buffer.push_str("(");
           while let Some(nt) = iter.next() {
               buffer.push_str(format!("\"{}\"", nt).as_str());
               if iter.peek().is_some() {
                   buffer.push_str(", ");
               }
           }
           buffer.push_str(")");

        write!(f, "\"{}\" → {} {}  # {}", self.head, self.composition, buffer, self.weight)
        // write!(f, "\"{}\" → {}", self.head, self.composition)
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFG<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter = self.initial.iter().peekable();

        buffer.push_str("initial: [");
        while let Some(nt) = iter.next() {
            buffer.push_str(format!("\"{}\"", nt).as_str());
            if iter.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]\n\n");

        for r in &self.rules {
            buffer.push_str(format!("{}\n", r).as_str());
        }

        write!(f, "{}", buffer)
    }
}

pub fn evaluate<T: Clone + fmt::Debug>(term_map: &GornTree<Composition<T>>)
        -> Composition<T>
{
    evaluate_pos(term_map, vec![])
}

pub fn evaluate_pos<T>(term_map: &GornTree<Composition<T>>, address: Vec<usize>)
        -> Composition<T>
    where T: Clone + fmt::Debug,
{
    let unexpanded_composition = &term_map.get(&address).unwrap().composition;
    let mut expanded_nonterminals: BTreeMap<_, Vec<Vec<VarT<T>>>> = BTreeMap::new();
    let mut expanded_composition = Vec::new();

    for component in unexpanded_composition {
        let mut expanded_component = Vec::new();

        for variable in component {
            match variable {
                &VarT::Var(num_nonter, num_compon) => {
                    if !expanded_nonterminals.contains_key(&num_nonter) {
                        let mut child_address = address.clone();
                        child_address.push(num_nonter);
                        let nonter_compos = evaluate_pos(term_map, child_address).composition;
                        expanded_nonterminals.insert(num_nonter, nonter_compos);
                    }

                    let nonter_compos = expanded_nonterminals.get(&num_nonter).unwrap();

                    if let Some(compon) = nonter_compos.get(num_compon) {
                        for terminal in compon {
                            expanded_component.push(terminal.clone());
                        }
                    } else {
                        panic!("{:?}: use of {}-th component of nonterminal {} that has only {} components!",
                               unexpanded_composition, num_compon, num_nonter, nonter_compos.len());
                    }

                },
                &VarT::T(ref terminal) => {
                    expanded_component.push(VarT::T(terminal.clone()));
                },
            }
        }

        expanded_composition.push(expanded_component);
    }

    Composition::from(expanded_composition)
}

pub fn to_term<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>)
        -> (GornTree<Composition<T>>, GornTree<H>)
    where H: Clone, T: Clone,
{
    let mut term_map = GornTree::new();
    let mut head_map = GornTree::new();

    for (address, rule) in tree_map {
        let &PMCFGRule { ref head, tail: _, ref composition, weight: _ } = rule;
        term_map.insert(address.clone(), composition.clone());
        head_map.insert(address.clone(), head.clone());
    }

    (term_map, head_map)
}

#[cfg(test)]
mod tests {
    use super::*;
    use self::VarT::{Var, T};
    use std::str::FromStr;

    pub fn example_tree_map() -> GornTree<PMCFGRule<String, char, usize>> {
        let mut tree_map = GornTree::new();

        tree_map.insert(vec![], PMCFGRule::from_str(
            "S -> [[Var 0 0, Var 1 0, Var 0 1, Var 1 1]] (A, B) # 1"
        ).unwrap());
        tree_map.insert(vec![0], PMCFGRule::from_str(
            "A -> [[Var 1 0, Var 0 0], [Var 2 0, Var 0 1]] (A, a, c) # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0], PMCFGRule::from_str(
            "A -> [[Var 1 0, Var 0 0], [Var 2 0, Var 0 1]] (A, a, c) # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0, 0], PMCFGRule::from_str(
            "A -> [[], []] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0, 1], PMCFGRule::from_str(
            "a -> [[T a]] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0, 2], PMCFGRule::from_str(
            "c -> [[T c]] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 1], PMCFGRule::from_str(
            "a -> [[T a]] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 2], PMCFGRule::from_str(
            "c -> [[T c]] () # 1"
        ).unwrap());
        tree_map.insert(vec![1], PMCFGRule::from_str(
            "B -> [[Var 1 0, Var 0 0], [Var 2 0, Var 0 1]] (B, b, d) # 1"
        ).unwrap());
        tree_map.insert(vec![1, 0], PMCFGRule::from_str(
            "B -> [[], []] () # 1"
        ).unwrap());
        tree_map.insert(vec![1, 1], PMCFGRule::from_str(
            "b -> [[T b]] () # 1"
        ).unwrap());
        tree_map.insert(vec![1, 2], PMCFGRule::from_str(
            "d -> [[T d]] () # 1"
        ).unwrap());

        tree_map
    }

    #[test]
    fn test_evaluate() {
        let tree_map = example_tree_map();
        let mut term_map = GornTree::new();

        for (address, PMCFGRule { head: _, tail: _, composition, weight: _ }) in tree_map {
            term_map.insert(address, composition);
        }

        let expanded_compos = Composition::from(vec![
            vec![T('a'), T('a'), T('b'), T('c'), T('c'), T('d')]
        ]);

        assert_eq!(expanded_compos, evaluate(&term_map));
    }

    #[test]
    #[should_panic(expected =
        "[[Var(0, 0), Var(0, 1)]]: use of 1-th component of nonterminal 0 that has only 1 components!"
    )]
    fn test_evaluate_invalid_composition() {
        let mut term_map = GornTree::new();
        term_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), Var(0, 1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
            vec![T('a')]
        ]));

        evaluate(&term_map);
    }

    #[test]
    fn test_to_term() {
        let mut tree_map: GornTree<PMCFGRule<String, char, usize>> = GornTree::new();

        tree_map.insert(vec![], PMCFGRule::from_str(
            "A -> [[Var 0 0, T a, Var 0 1, T b]] (B) # 1"
        ).unwrap());
        tree_map.insert(vec![0], PMCFGRule::from_str(
            "B -> [[Var 1 0], [T c]] (C) # 1"
        ).unwrap());
        tree_map.insert(vec![0, 1], PMCFGRule::from_str(
            "C -> [[], []] () # 1"
        ).unwrap());

        let mut term_map = GornTree::new();
        term_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), T('a'), Var(0, 1), T('b')]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
            vec![Var(1, 0)],
            vec![T('c')]
        ]));
        term_map.insert(vec![0, 1], Composition::from(vec![
            vec![],
            vec![]
        ]));

        let mut head_map = GornTree::new();
        head_map.insert(vec![], String::from("A"));
        head_map.insert(vec![0], String::from("B"));
        head_map.insert(vec![0, 1], String::from("C"));

        assert_eq!((term_map, head_map), to_term(&tree_map));
    }
}
<commit_msg>src/pmcfg/mod.rs: tiny change to to_term()<commit_after>use std::collections::BTreeMap;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;

use util::tree::GornTree;

mod from_str;
// mod relabel;
pub mod negra;

/// Variable or terminal symbol in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub enum VarT<T> {
    /// `Var(i, j)` represents the `j`th component of the `i`th successor.
    /// Indexing starts from `0`.
    Var(usize, usize),
    T(T),
}

/// Composition function in an MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Hash)]
pub struct Composition<T> {
    pub composition: Vec<Vec<VarT<T>>>,
}

impl<T> From<Vec<Vec<VarT<T>>>> for Composition<T> {
    fn from(encapsulated_value: Vec<Vec<VarT<T>>>) -> Self {
        Composition { composition: encapsulated_value }
    }
}

/// Rule of a weighted MCFG.
#[derive(Debug, PartialOrd, Ord, Clone)]
pub struct PMCFGRule<N, T, W> {
    pub head: N,
    pub tail: Vec<N>,
    pub composition: Composition<T>,
    pub weight: W,
}

impl<N, T, W> PMCFGRule<N, T, W>
    where T: Clone,
          W: Clone,
{
    pub fn map_nonterminals<F, M>(&self, f: F) -> PMCFGRule<M, T, W>
        where F: Fn(&N) -> M,
    {
        PMCFGRule {
            head: f(&self.head),
            tail: self.tail.iter().map(f).collect(),
            composition: self.composition.clone(),
            weight: self.weight.clone()
        }
    }
}

/// A weighted MCFG.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
pub struct PMCFG<N, T, W> {
    pub _dummy: PhantomData<T>,
    pub initial: Vec<N>,
    pub rules: Vec<PMCFGRule<N, T, W>>,
}

impl<N: Hash, T: Hash, W> Hash for PMCFGRule<N, T, W> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.head.hash(state);
        self.tail.hash(state);
        self.composition.hash(state);
    }
}

impl<N: PartialEq, T: PartialEq, W> PartialEq for PMCFGRule<N, T, W> {
    fn eq(&self, other: &Self) -> bool {
        self.head == other.head && self.tail == other.tail && self.composition == other.composition
    }
}

impl<N: Eq, T: Eq, W> Eq for PMCFGRule<N, T, W> {}

impl<T: fmt::Display> fmt::Display for Composition<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter0 = self.composition.iter().peekable();
        let mut iter1;

        buffer.push_str("[");
        while let Some(proj) = iter0.next() {
            iter1 = proj.into_iter().peekable();
            buffer.push_str("[");
            while let Some(vart) = iter1.next() {
                buffer.push_str(format!("{}", vart).as_str());
                if iter1.peek().is_some() {
                    buffer.push_str(", ");
                }
            }
            buffer.push_str("]");
            if iter0.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]");

        write!(f, "{}", buffer)
    }
}

impl<T: fmt::Display> fmt::Display for VarT<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            VarT::Var(i, j) => {
                write!(f, "Var {} {}", i, j)
            },
            VarT::T(ref x) => {
                write!(f, "T \"{}\"", x)
            },
        }
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFGRule<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
           let mut buffer = "".to_string();

           let mut iter = self.tail.iter().peekable();

           buffer.push_str("(");
           while let Some(nt) = iter.next() {
               buffer.push_str(format!("\"{}\"", nt).as_str());
               if iter.peek().is_some() {
                   buffer.push_str(", ");
               }
           }
           buffer.push_str(")");

        write!(f, "\"{}\" → {} {}  # {}", self.head, self.composition, buffer, self.weight)
        // write!(f, "\"{}\" → {}", self.head, self.composition)
    }
}

impl<N: fmt::Display, T: fmt::Display, W: fmt::Display> fmt::Display for PMCFG<N, T, W> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut buffer = "".to_string();

        let mut iter = self.initial.iter().peekable();

        buffer.push_str("initial: [");
        while let Some(nt) = iter.next() {
            buffer.push_str(format!("\"{}\"", nt).as_str());
            if iter.peek().is_some() {
                buffer.push_str(", ");
            }
        }
        buffer.push_str("]\n\n");

        for r in &self.rules {
            buffer.push_str(format!("{}\n", r).as_str());
        }

        write!(f, "{}", buffer)
    }
}

pub fn evaluate<T: Clone + fmt::Debug>(term_map: &GornTree<Composition<T>>)
        -> Composition<T>
{
    evaluate_pos(term_map, vec![])
}

pub fn evaluate_pos<T>(term_map: &GornTree<Composition<T>>, address: Vec<usize>)
        -> Composition<T>
    where T: Clone + fmt::Debug,
{
    let unexpanded_composition = &term_map.get(&address).unwrap().composition;
    let mut expanded_nonterminals: BTreeMap<_, Vec<Vec<VarT<T>>>> = BTreeMap::new();
    let mut expanded_composition = Vec::new();

    for component in unexpanded_composition {
        let mut expanded_component = Vec::new();

        for variable in component {
            match variable {
                &VarT::Var(num_nonter, num_compon) => {
                    if !expanded_nonterminals.contains_key(&num_nonter) {
                        let mut child_address = address.clone();
                        child_address.push(num_nonter);
                        let nonter_compos = evaluate_pos(term_map, child_address).composition;
                        expanded_nonterminals.insert(num_nonter, nonter_compos);
                    }

                    let nonter_compos = expanded_nonterminals.get(&num_nonter).unwrap();

                    if let Some(compon) = nonter_compos.get(num_compon) {
                        for terminal in compon {
                            expanded_component.push(terminal.clone());
                        }
                    } else {
                        panic!("{:?}: use of {}-th component of nonterminal {} that has only {} components!",
                               unexpanded_composition, num_compon, num_nonter, nonter_compos.len());
                    }

                },
                &VarT::T(ref terminal) => {
                    expanded_component.push(VarT::T(terminal.clone()));
                },
            }
        }

        expanded_composition.push(expanded_component);
    }

    Composition::from(expanded_composition)
}

pub fn to_term<H, T, W>(tree_map: &GornTree<PMCFGRule<H, T, W>>)
        -> (GornTree<Composition<T>>, GornTree<H>)
    where H: Clone, T: Clone,
{
    let mut term_map = GornTree::new();
    let mut head_map = GornTree::new();

    for (address, &PMCFGRule { ref head, tail: _, ref composition, weight: _ }) in tree_map {
        term_map.insert(address.clone(), composition.clone());
        head_map.insert(address.clone(), head.clone());
    }

    (term_map, head_map)
}

#[cfg(test)]
mod tests {
    use super::*;
    use self::VarT::{Var, T};
    use std::str::FromStr;

    pub fn example_tree_map() -> GornTree<PMCFGRule<String, char, usize>> {
        let mut tree_map = GornTree::new();

        tree_map.insert(vec![], PMCFGRule::from_str(
            "S -> [[Var 0 0, Var 1 0, Var 0 1, Var 1 1]] (A, B) # 1"
        ).unwrap());
        tree_map.insert(vec![0], PMCFGRule::from_str(
            "A -> [[Var 1 0, Var 0 0], [Var 2 0, Var 0 1]] (A, a, c) # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0], PMCFGRule::from_str(
            "A -> [[Var 1 0, Var 0 0], [Var 2 0, Var 0 1]] (A, a, c) # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0, 0], PMCFGRule::from_str(
            "A -> [[], []] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0, 1], PMCFGRule::from_str(
            "a -> [[T a]] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 0, 2], PMCFGRule::from_str(
            "c -> [[T c]] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 1], PMCFGRule::from_str(
            "a -> [[T a]] () # 1"
        ).unwrap());
        tree_map.insert(vec![0, 2], PMCFGRule::from_str(
            "c -> [[T c]] () # 1"
        ).unwrap());
        tree_map.insert(vec![1], PMCFGRule::from_str(
            "B -> [[Var 1 0, Var 0 0], [Var 2 0, Var 0 1]] (B, b, d) # 1"
        ).unwrap());
        tree_map.insert(vec![1, 0], PMCFGRule::from_str(
            "B -> [[], []] () # 1"
        ).unwrap());
        tree_map.insert(vec![1, 1], PMCFGRule::from_str(
            "b -> [[T b]] () # 1"
        ).unwrap());
        tree_map.insert(vec![1, 2], PMCFGRule::from_str(
            "d -> [[T d]] () # 1"
        ).unwrap());

        tree_map
    }

    #[test]
    fn test_evaluate() {
        let tree_map = example_tree_map();
        let mut term_map = GornTree::new();

        for (address, PMCFGRule { head: _, tail: _, composition, weight: _ }) in tree_map {
            term_map.insert(address, composition);
        }

        let expanded_compos = Composition::from(vec![
            vec![T('a'), T('a'), T('b'), T('c'), T('c'), T('d')]
        ]);

        assert_eq!(expanded_compos, evaluate(&term_map));
    }

    #[test]
    #[should_panic(expected =
        "[[Var(0, 0), Var(0, 1)]]: use of 1-th component of nonterminal 0 that has only 1 components!"
    )]
    fn test_evaluate_invalid_composition() {
        let mut term_map = GornTree::new();
        term_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), Var(0, 1)]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
            vec![T('a')]
        ]));

        evaluate(&term_map);
    }

    #[test]
    fn test_to_term() {
        let mut tree_map: GornTree<PMCFGRule<String, char, usize>> = GornTree::new();

        tree_map.insert(vec![], PMCFGRule::from_str(
            "A -> [[Var 0 0, T a, Var 0 1, T b]] (B) # 1"
        ).unwrap());
        tree_map.insert(vec![0], PMCFGRule::from_str(
            "B -> [[Var 1 0], [T c]] (C) # 1"
        ).unwrap());
        tree_map.insert(vec![0, 1], PMCFGRule::from_str(
            "C -> [[], []] () # 1"
        ).unwrap());

        let mut term_map = GornTree::new();
        term_map.insert(vec![], Composition::from(vec![
            vec![Var(0, 0), T('a'), Var(0, 1), T('b')]
        ]));
        term_map.insert(vec![0], Composition::from(vec![
            vec![Var(1, 0)],
            vec![T('c')]
        ]));
        term_map.insert(vec![0, 1], Composition::from(vec![
            vec![],
            vec![]
        ]));

        let mut head_map = GornTree::new();
        head_map.insert(vec![], String::from("A"));
        head_map.insert(vec![0], String::from("B"));
        head_map.insert(vec![0, 1], String::from("C"));

        assert_eq!((term_map, head_map), to_term(&tree_map));
    }
}
<|endoftext|>


<commit_before>use super::web::{self, sqlite::create_connection_pool};
use crate::core::prelude::*;
use crate::infrastructure::osm;
use clap::{App, Arg, SubCommand};
use dotenv::dotenv;
use std::{env, process};

const DEFAULT_DB_URL: &str = "openfair.db";

fn update_event_locations<D: Db>(db: &mut D) -> Result<()> {
    let events = db.all_events()?;
    for mut e in events {
        if let Some(ref mut loc) = e.location {
            if let Some(ref addr) = loc.address {
                if let Some((lat, lng)) = web::api::geocoding::resolve_address_lat_lng(addr) {
                    loc.lat = lat;
                    loc.lng = lng;
                } else {
                    // Reset existing location if an address exists but is
                    // empty. This compensates for a bug in v0.4.3 where the
                    // check for an empty address was missing.
                    if addr.is_empty() {
                        info!("Resetting location of event {} with empty address", e.id);
                        loc.lat = 0.0;
                        loc.lng = 0.0;
                    }
                }
                if let Err(err) = db.update_event(&e) {
                    warn!("Failed to update location of event {}: {}", e.id, err);
                } else {
                    info!("Updated location of event {}", e.id);
                }
            }
        }
    }
    Ok(())
}

pub fn run() {
    dotenv().ok();
    let matches = App::new("openFairDB")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Kohlhase <mail@markus-kohlhase.de>")
        .arg(
            Arg::with_name("db-url")
                .long("db-url")
                .value_name("DATABASE_URL")
                .help("URL to the database"),
        )
        .arg(
            Arg::with_name("enable-cors")
                .long("enable-cors")
                .help("Allow requests from any origin"),
        )
        .arg(
            Arg::with_name("fix-event-address-location")
                .long("fix-event-address-location")
                .help("Update the location of ALL events by resolving their address"),
        )
        .subcommand(
            SubCommand::with_name("osm")
                .about("OpenStreetMap functionalities")
                .subcommand(
                    SubCommand::with_name("import")
                        .about("import entries from OSM (JSON file)")
                        .arg(
                            Arg::with_name("osm-file")
                                .value_name("OSM_FILE")
                                .help("JSON file with osm nodes"),
                        ),
                ),
        )
        .get_matches();

    let db_url = match matches.value_of("db-url") {
        Some(db_url) => db_url.into(),
        None => match env::var("DATABASE_URL") {
            Ok(url) => url,
            Err(_) => DEFAULT_DB_URL.to_string(),
        },
    };
    let pool = create_connection_pool(&db_url).unwrap();

    match matches.subcommand() {
        ("osm", Some(osm_matches)) => match osm_matches.subcommand() {
            ("import", Some(import_matches)) => {
                let osm_file = match import_matches.value_of("osm-file") {
                    Some(osm_file) => osm_file,
                    None => {
                        println!("{}", matches.usage());
                        process::exit(1)
                    }
                };
                if let Err(err) = osm::import_from_osm_file(&db_url, osm_file) {
                    println!("Could not import from '{}': {}", osm_file, err);
                    process::exit(1)
                }
            }
            _ => println!("{}", osm_matches.usage()),
        },
        _ => {
            if matches.is_present("fix-event-address-location") {
                info!("Updating all event locations...");
                update_event_locations(&mut *pool.get().unwrap()).unwrap();
            }
            web::run(pool, matches.is_present("enable-cors"));
        }
    }
}
<commit_msg>Don't reset location of events with empty address<commit_after>use super::web::{self, sqlite::create_connection_pool};
use crate::core::prelude::*;
use crate::infrastructure::osm;
use clap::{App, Arg, SubCommand};
use dotenv::dotenv;
use std::{env, process};

const DEFAULT_DB_URL: &str = "openfair.db";

fn update_event_locations<D: Db>(db: &mut D) -> Result<()> {
    let events = db.all_events()?;
    for mut e in events {
        if let Some(ref mut loc) = e.location {
            if let Some(ref addr) = loc.address {
                if let Some((lat, lng)) = web::api::geocoding::resolve_address_lat_lng(addr) {
                    loc.lat = lat;
                    loc.lng = lng;
                    if let Err(err) = db.update_event(&e) {
                        warn!("Failed to update location of event {}: {}", e.id, err);
                    } else {
                        info!("Updated location of event {}", e.id);
                    }
                }
            }
        }
    }
    Ok(())
}

pub fn run() {
    dotenv().ok();
    let matches = App::new("openFairDB")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Kohlhase <mail@markus-kohlhase.de>")
        .arg(
            Arg::with_name("db-url")
                .long("db-url")
                .value_name("DATABASE_URL")
                .help("URL to the database"),
        )
        .arg(
            Arg::with_name("enable-cors")
                .long("enable-cors")
                .help("Allow requests from any origin"),
        )
        .arg(
            Arg::with_name("fix-event-address-location")
                .long("fix-event-address-location")
                .help("Update the location of ALL events by resolving their address"),
        )
        .subcommand(
            SubCommand::with_name("osm")
                .about("OpenStreetMap functionalities")
                .subcommand(
                    SubCommand::with_name("import")
                        .about("import entries from OSM (JSON file)")
                        .arg(
                            Arg::with_name("osm-file")
                                .value_name("OSM_FILE")
                                .help("JSON file with osm nodes"),
                        ),
                ),
        )
        .get_matches();

    let db_url = match matches.value_of("db-url") {
        Some(db_url) => db_url.into(),
        None => match env::var("DATABASE_URL") {
            Ok(url) => url,
            Err(_) => DEFAULT_DB_URL.to_string(),
        },
    };
    let pool = create_connection_pool(&db_url).unwrap();

    match matches.subcommand() {
        ("osm", Some(osm_matches)) => match osm_matches.subcommand() {
            ("import", Some(import_matches)) => {
                let osm_file = match import_matches.value_of("osm-file") {
                    Some(osm_file) => osm_file,
                    None => {
                        println!("{}", matches.usage());
                        process::exit(1)
                    }
                };
                if let Err(err) = osm::import_from_osm_file(&db_url, osm_file) {
                    println!("Could not import from '{}': {}", osm_file, err);
                    process::exit(1)
                }
            }
            _ => println!("{}", osm_matches.usage()),
        },
        _ => {
            if matches.is_present("fix-event-address-location") {
                info!("Updating all event locations...");
                update_event_locations(&mut *pool.get().unwrap()).unwrap();
            }
            web::run(pool, matches.is_present("enable-cors"));
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Sam Payson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
// associated documentation files (the "Software"), to deal in the Software without restriction,
// including without limitation the rights to use, copy, modify, merge, publish, distribute,
// sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
// NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io;
use std::mem;

pub enum Primitive {
    UInt8(u8),
    UInt16(u16),
    UInt32(u32),
    UInt64(u64),

    Int8(i8),
    Int16(i16),
    Int32(i32),
    Int64(i64),

    Fixed32(u32),
    Fixed64(u64),

    Float32(f32),
    Float64(f64),

    Bool(bool),

    Bytes(Vec<u8>),
    String(String),

    Enum(i64),
}

pub fn write_uvarint<W>(w: &mut W, mut x: u64) -> io::Result<usize>
    where W: io::Write {

    // A 64-bit varint can be at most 10 bytes long.
    let mut buf = [0u8; 10];
    let mut idx = 0;

    while x > 0x7F {
        buf[idx] = 0x80 | (x & 0x7F) as u8;
        x = x >> 7;
        idx += 1;
    }

    buf[idx] = x as u8;

    try!(w.write_all(&buf[idx..]));

    Ok((idx + 1) as usize)
}

pub fn write_varint<W>(w: &mut W, x: i64) -> io::Result<usize>
    where W: io::Write {
    let ux: u64 = if x < 0 {
        !((x as u64) << 1)
    } else {
        (x as u64) << 1
    };

    write_uvarint(w, ux)
}

pub fn write_le_u8<W>(w: &mut W, x: u8) -> io::Result<usize>
    where W: io::Write {

    let buf = [x];
    try!(w.write_all(&buf));
    Ok(1)
}

pub fn write_le_u16<W>(w: &mut W, x: u16) -> io::Result<usize>
    where W: io::Write {

    let buf = [
        ((x >> 0) & 0xFF) as u8,
        ((x >> 8) & 0xFF) as u8,
    ];

    try!(w.write_all(&buf));
    Ok(2)
}

pub fn write_le_u32<W>(w: &mut W, x: u32) -> io::Result<usize>
    where W: io::Write {

    let buf = [
        ((x >>  0) & 0xFF) as u8,
        ((x >>  8) & 0xFF) as u8,
        ((x >> 16) & 0xFF) as u8,
        ((x >> 24) & 0xFF) as u8,
    ];

    try!(w.write_all(&buf));
    Ok(4)
}

pub fn write_le_u64<W>(w: &mut W, x: u64) -> io::Result<usize>
    where W: io::Write {

    let buf = [
        ((x >>  0) & 0xFF) as u8,
        ((x >>  8) & 0xFF) as u8,
        ((x >> 16) & 0xFF) as u8,
        ((x >> 24) & 0xFF) as u8,
        ((x >> 32) & 0xFF) as u8,
        ((x >> 40) & 0xFF) as u8,
        ((x >> 48) & 0xFF) as u8,
        ((x >> 56) & 0xFF) as u8,
    ];

    try!(w.write_all(&buf));
    Ok(8)
}

pub fn write_le_i8<W>(w: &mut W, x: i8) -> io::Result<usize>
    where W: io::Write {

    write_le_u8(w, x as u8)
}

pub fn write_le_i16<W>(w: &mut W, x: i16) -> io::Result<usize>
    where W: io::Write {

    write_le_u16(w, x as u16)
}

pub fn write_le_i32<W>(w: &mut W, x: i32) -> io::Result<usize>
    where W: io::Write {

    write_le_u32(w, x as u32)
}

pub fn write_le_i64<W>(w: &mut W, x: i64) -> io::Result<usize>
    where W: io::Write {

    write_le_u64(w, x as u64)
}

pub fn write_le_f32<W>(w: &mut W, x: f32) -> io::Result<usize>
    where W: io::Write {

    write_le_u32(w, unsafe { mem::transmute(x) })
}

pub fn write_le_f64<W>(w: &mut W, x: f64) -> io::Result<usize>
    where W: io::Write {

    write_le_u64(w, unsafe { mem::transmute(x) })
}
<commit_msg>Revised write_varint to be a bit cleaner.<commit_after>// Copyright (c) 2015, Sam Payson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
// associated documentation files (the "Software"), to deal in the Software without restriction,
// including without limitation the rights to use, copy, modify, merge, publish, distribute,
// sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
// NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

use std::io;
use std::mem;

pub enum Primitive {
    UInt8(u8),
    UInt16(u16),
    UInt32(u32),
    UInt64(u64),

    Int8(i8),
    Int16(i16),
    Int32(i32),
    Int64(i64),

    Fixed32(u32),
    Fixed64(u64),

    Float32(f32),
    Float64(f64),

    Bool(bool),

    Bytes(Vec<u8>),
    String(String),

    Enum(i64),
}

pub fn write_uvarint<W>(w: &mut W, mut x: u64) -> io::Result<usize>
    where W: io::Write {

    // A 64-bit varint can be at most 10 bytes long.
    let mut buf = [0u8; 10];
    let mut idx = 0;

    while x > 0x7F {
        buf[idx] = 0x80 | (x & 0x7F) as u8;
        x = x >> 7;
        idx += 1;
    }

    buf[idx] = x as u8;

    try!(w.write_all(&buf[idx..]));

    Ok((idx + 1) as usize)
}

pub fn write_varint<W>(w: &mut W, x: i64) -> io::Result<usize>
    where W: io::Write {

    let ux = (x as u64) << 1;

    write_uvarint(w, if x < 0 { !ux } else { ux })
}

pub fn write_le_u8<W>(w: &mut W, x: u8) -> io::Result<usize>
    where W: io::Write {

    let buf = [x];
    try!(w.write_all(&buf));
    Ok(1)
}

pub fn write_le_u16<W>(w: &mut W, x: u16) -> io::Result<usize>
    where W: io::Write {

    let buf = [
        ((x >> 0) & 0xFF) as u8,
        ((x >> 8) & 0xFF) as u8,
    ];

    try!(w.write_all(&buf));
    Ok(2)
}

pub fn write_le_u32<W>(w: &mut W, x: u32) -> io::Result<usize>
    where W: io::Write {

    let buf = [
        ((x >>  0) & 0xFF) as u8,
        ((x >>  8) & 0xFF) as u8,
        ((x >> 16) & 0xFF) as u8,
        ((x >> 24) & 0xFF) as u8,
    ];

    try!(w.write_all(&buf));
    Ok(4)
}

pub fn write_le_u64<W>(w: &mut W, x: u64) -> io::Result<usize>
    where W: io::Write {

    let buf = [
        ((x >>  0) & 0xFF) as u8,
        ((x >>  8) & 0xFF) as u8,
        ((x >> 16) & 0xFF) as u8,
        ((x >> 24) & 0xFF) as u8,
        ((x >> 32) & 0xFF) as u8,
        ((x >> 40) & 0xFF) as u8,
        ((x >> 48) & 0xFF) as u8,
        ((x >> 56) & 0xFF) as u8,
    ];

    try!(w.write_all(&buf));
    Ok(8)
}

pub fn write_le_i8<W>(w: &mut W, x: i8) -> io::Result<usize>
    where W: io::Write {

    write_le_u8(w, x as u8)
}

pub fn write_le_i16<W>(w: &mut W, x: i16) -> io::Result<usize>
    where W: io::Write {

    write_le_u16(w, x as u16)
}

pub fn write_le_i32<W>(w: &mut W, x: i32) -> io::Result<usize>
    where W: io::Write {

    write_le_u32(w, x as u32)
}

pub fn write_le_i64<W>(w: &mut W, x: i64) -> io::Result<usize>
    where W: io::Write {

    write_le_u64(w, x as u64)
}

pub fn write_le_f32<W>(w: &mut W, x: f32) -> io::Result<usize>
    where W: io::Write {

    write_le_u32(w, unsafe { mem::transmute(x) })
}

pub fn write_le_f64<W>(w: &mut W, x: f64) -> io::Result<usize>
    where W: io::Write {

    write_le_u64(w, unsafe { mem::transmute(x) })
}
<|endoftext|>


<commit_before>//! MC Server List Ping protocol.
//!
//! Reference: http://wiki.vg/Server_List_Ping

use std::fs::File;
use std::io::ErrorKind::InvalidInput;
use std::io::prelude::*;
use std::io;
use std::net::TcpStream;
use std::ops::Sub; // Sub for Timespec
use std::path::Path;

use packet::{PacketRead, PacketWrite, Protocol};

use rustc_serialize::base64::{ToBase64, STANDARD};
use rustc_serialize::json;
use time;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Description {
    pub text: String,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Players {
    pub max: i32,
    pub online: i32,
    pub sample: Option<Vec<Sample>>,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Sample {
    pub name: String,
    pub id: String,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub name: String,
    pub protocol: i32,
}

/// Response sent to clients as JSON.
#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Response {
    // FIXME(toqueteos): This is ChatJson
    pub description: String,
    pub favicon: Option<String>,
    pub players: Players,
    pub version: Version,
}

impl Protocol for Response {
    type Clean = Response;

    fn proto_len(value: &Response) -> usize {
        <String as Protocol>::proto_len(&json::encode(&value).unwrap())
    }
    fn proto_encode(value: &Response, dst: &mut Write) -> io::Result<()> {
        try!(<String as Protocol>::proto_encode(&json::encode(&value).unwrap(), dst));
        Ok(())
    }
    fn proto_decode(src: &mut Read) -> io::Result<Response> {
        let s = try!(<String as Protocol>::proto_decode(src));
        println!("Response proto_decode {}", s);
        json::decode(&s).map_err(|_| io::Error::new(InvalidInput, "found invalid JSON", None))
    }
}

// FIXME(toqueteos): This is yelling to be a method of a Server struct or
// something more useful. We need the Handshake's `next_state` field in order
// to perform login for a player.
/// Server-side Server List response
pub fn response(mut stream: &mut TcpStream) -> io::Result<()> {
    use packet::status::serverbound::Packet::{self, StatusRequest};
    use packet::status::clientbound::StatusResponse;

    // C->S: Status Request packet
    match try!(Packet::read(stream)) {
        StatusRequest(_) => {
            // S->C: Status Response packet
            let mut file = try!(File::open(&Path::new("assets/favicon.png")));
            let mut contents = Vec::new();
            try!(file.read_to_end(&mut contents));
            let favicon = contents.as_slice().to_base64(STANDARD);
            // FIXME(toqueteos): Micro-optimization? We could totally drop JSON
            // encoding and just replace player values (online & max) with format! all
            // other values are static.
            let resp = Response{
                version: Version{
                    name: "1.8.3".to_string(),
                    protocol: 47,
                },
                players: Players{
                    // FIXME(toqueteos): This is value should be a internal counter of server
                    online: 0,
                    // FIXME(toqueteos): This is value read from server.properties file
                    max: 20,
                    sample: None
                },
                description: "With custom favicons! Woot :D".to_string(),
                favicon: Some(format!("data:image/png;base64,{:}", favicon)),
            };
            try!(StatusResponse { response: resp }.write(stream));
            Ok(())
        }
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting C->S StatusRequest packet, got {:?}", wrong_packet))))
    }
}

/// Server-side pong response, optional
pub fn pong(mut stream: &mut TcpStream) -> io::Result<()> {
    use packet::status::clientbound::Pong;
    use packet::status::serverbound::Packet::{self, Ping};

    // C->S: Ping packet
    match try!(Packet::read(stream)) {
        Ping(ping) => {
            // S->C: Pong packet
            try!(Pong { time: ping.time }.write(stream));
            Ok(())
        }
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting C->S Ping packet, got {:?}", wrong_packet))))
    }
}

/// Client-side Server List request
pub fn request(mut stream: &mut TcpStream) -> io::Result<Response> {
    use packet::status::serverbound::StatusRequest;
    use packet::status::clientbound::Packet::{self, StatusResponse};

    // C->S: Status Request packet
    try!(StatusRequest.write(stream));

    // S->C: Status Response packet
    match try!(Packet::read(stream)) {
        StatusResponse(resp) => Ok(resp.response),
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting S->C StatusResponse packet, got {:?}", wrong_packet))))
    }
}

/// Client-side ping request, optional
pub fn ping(mut stream: &mut TcpStream) -> io::Result<i64> {
    use packet::status::clientbound::Packet::{self, Pong};
    use packet::status::serverbound::Ping;

    // C->S: Ping packet
    let start = time::get_time();
    try!(Ping { time: start.sec }.write(stream));

    // S->C: Pong packet
    match try!(Packet::read(stream)) {
        Pong(_) => {
            let end = time::get_time();
            let elapsed = end.sub(start).num_milliseconds();
            Ok(elapsed)
        }
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting S->C Pong packet, got {:?}", wrong_packet))))
    }
}

#[cfg(test)]
mod tests {
    #[test]
    #[cfg(vanilla_server_required)]
    fn client_server_list_ping() {
        // NOTE(toqueteos): There's just one test right now and it doesn't run
        // by default so rustc sees all of these imports as unused, that's why
        // they are in here.
        use super::*;

        use std::io::prelude::*;
        use std::net::TcpStream;

        use packet::handshake::Handshake;
        use packet::{PacketWrite, NextState};

        let mut stream = TcpStream::connect("127.0.0.1:25565").unwrap();
        Handshake {
            proto_version: 47,
            server_address: "127.0.0.1".to_string(),
            server_port: 25565,
            next_state: NextState::Status
        }.write(&mut stream).unwrap();
        request(&mut stream).unwrap();
        ping(&mut stream).unwrap();
    }
}
<commit_msg>Add request-ping reverse slp test.<commit_after>//! MC Server List Ping protocol.
//!
//! Reference: http://wiki.vg/Server_List_Ping

use std::fs::File;
use std::io::ErrorKind::InvalidInput;
use std::io::prelude::*;
use std::io;
use std::net::TcpStream;
use std::ops::Sub; // Sub for Timespec
use std::path::Path;

use packet::{PacketRead, PacketWrite, Protocol};

use rustc_serialize::base64::{ToBase64, STANDARD};
use rustc_serialize::json;
use time;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Description {
    pub text: String,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Players {
    pub max: i32,
    pub online: i32,
    pub sample: Option<Vec<Sample>>,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Sample {
    pub name: String,
    pub id: String,
}

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub name: String,
    pub protocol: i32,
}

/// Response sent to clients as JSON.
#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Response {
    // FIXME(toqueteos): This is ChatJson
    pub description: String,
    pub favicon: Option<String>,
    pub players: Players,
    pub version: Version,
}

impl Protocol for Response {
    type Clean = Response;

    fn proto_len(value: &Response) -> usize {
        <String as Protocol>::proto_len(&json::encode(&value).unwrap())
    }
    fn proto_encode(value: &Response, dst: &mut Write) -> io::Result<()> {
        try!(<String as Protocol>::proto_encode(&json::encode(&value).unwrap(), dst));
        Ok(())
    }
    fn proto_decode(src: &mut Read) -> io::Result<Response> {
        let s = try!(<String as Protocol>::proto_decode(src));
        println!("Response proto_decode {}", s);
        json::decode(&s).map_err(|_| io::Error::new(InvalidInput, "found invalid JSON", None))
    }
}

// FIXME(toqueteos): This is yelling to be a method of a Server struct or
// something more useful. We need the Handshake's `next_state` field in order
// to perform login for a player.
/// Server-side Server List response
pub fn response(mut stream: &mut TcpStream) -> io::Result<()> {
    use packet::status::serverbound::Packet::{self, StatusRequest};
    use packet::status::clientbound::StatusResponse;

    // C->S: Status Request packet
    match try!(Packet::read(stream)) {
        StatusRequest(_) => {
            // S->C: Status Response packet
            let mut file = try!(File::open(&Path::new("assets/favicon.png")));
            let mut contents = Vec::new();
            try!(file.read_to_end(&mut contents));
            let favicon = contents.as_slice().to_base64(STANDARD);
            // FIXME(toqueteos): Micro-optimization? We could totally drop JSON
            // encoding and just replace player values (online & max) with format! all
            // other values are static.
            let resp = Response{
                version: Version{
                    name: "1.8.3".to_string(),
                    protocol: 47,
                },
                players: Players{
                    // FIXME(toqueteos): This is value should be a internal counter of server
                    online: 0,
                    // FIXME(toqueteos): This is value read from server.properties file
                    max: 20,
                    sample: None
                },
                description: "With custom favicons! Woot :D".to_string(),
                favicon: Some(format!("data:image/png;base64,{:}", favicon)),
            };
            try!(StatusResponse { response: resp }.write(stream));
            Ok(())
        }
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting C->S StatusRequest packet, got {:?}", wrong_packet))))
    }
}

/// Server-side pong response, optional
pub fn pong(mut stream: &mut TcpStream) -> io::Result<()> {
    use packet::status::clientbound::Pong;
    use packet::status::serverbound::Packet::{self, Ping};

    // C->S: Ping packet
    match try!(Packet::read(stream)) {
        Ping(ping) => {
            // S->C: Pong packet
            try!(Pong { time: ping.time }.write(stream));
            Ok(())
        }
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting C->S Ping packet, got {:?}", wrong_packet))))
    }
}

/// Client-side Server List request
pub fn request(mut stream: &mut TcpStream) -> io::Result<Response> {
    use packet::status::serverbound::StatusRequest;
    use packet::status::clientbound::Packet::{self, StatusResponse};

    // C->S: Status Request packet
    try!(StatusRequest.write(stream));

    // S->C: Status Response packet
    match try!(Packet::read(stream)) {
        StatusResponse(resp) => Ok(resp.response),
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting S->C StatusResponse packet, got {:?}", wrong_packet))))
    }
}

/// Client-side ping request, optional
pub fn ping(mut stream: &mut TcpStream) -> io::Result<i64> {
    use packet::status::clientbound::Packet::{self, Pong};
    use packet::status::serverbound::Ping;

    // C->S: Ping packet
    let start = time::get_time();
    try!(Ping { time: start.sec }.write(stream));

    // S->C: Pong packet
    match try!(Packet::read(stream)) {
        Pong(_) => {
            let end = time::get_time();
            let elapsed = end.sub(start).num_milliseconds();
            Ok(elapsed)
        }
        wrong_packet => Err(io::Error::new(InvalidInput, "invalid packet read", Some(format!("expecting S->C Pong packet, got {:?}", wrong_packet))))
    }
}

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

    use std::io::prelude::*;
    use std::net::TcpStream;

    use packet::handshake::Handshake;
    use packet::{PacketWrite, NextState};

    #[test]
    #[cfg(vanilla_server_required)]
    fn client_server_list_ping() {
        let mut stream = TcpStream::connect("127.0.0.1:25565").unwrap();
        Handshake {
            proto_version: 47,
            server_address: "127.0.0.1".to_string(),
            server_port: 25565,
            next_state: NextState::Status
        }.write(&mut stream).unwrap();
        let response = request(&mut stream).unwrap();
        println!("request {:?}", response);
        let elapsed = ping(&mut stream).unwrap();
        println!("ping {}", elapsed);
    }

    #[test]
    #[should_panic]
    #[cfg(vanilla_server_required)]
    fn client_slp_reversed() {
        let mut stream = TcpStream::connect("127.0.0.1:25565").unwrap();
        let elapsed = ping(&mut stream).unwrap();
        println!("ping {}", elapsed);
        Handshake {
            proto_version: 47,
            server_address: "127.0.0.1".to_string(),
            server_port: 25565,
            next_state: NextState::Status
        }.write(&mut stream).unwrap();
        let response = request(&mut stream).unwrap();
        println!("request {:?}", response);
    }
}
<|endoftext|>


<commit_before>//! the [GPOS] table
//!
//! [GPOS]: https://docs.microsoft.com/en-us/typography/opentype/spec/gpos

include!("../../generated/generated_gpos.rs");

use std::collections::HashSet;

use super::value_record::ValueRecord;
use super::{
    ChainedSequenceContext, ClassDef, CoverageTable, Device, FeatureList, FeatureVariations,
    Lookup, LookupList, LookupType, ScriptList, SequenceContext,
};

#[cfg(all(test, feature = "parsing"))]
#[path = "../tests/gpos.rs"]
mod tests;

/// A GPOS lookup list table.
type PositionLookupList = LookupList<PositionLookup>;

table_newtype!(
    PositionSequenceContext,
    SequenceContext,
    read_fonts::layout::SequenceContext<'a>
);

table_newtype!(
    PositionChainContext,
    ChainedSequenceContext,
    read_fonts::layout::ChainedSequenceContext<'a>
);

impl Gpos {
    fn compute_version(&self) -> MajorMinor {
        if self.feature_variations_offset.get().is_none() {
            MajorMinor::VERSION_1_0
        } else {
            MajorMinor::VERSION_1_1
        }
    }
}

lookup_type!(SinglePos, 1);
lookup_type!(PairPos, 2);
lookup_type!(CursivePosFormat1, 3);
lookup_type!(MarkBasePosFormat1, 4);
lookup_type!(MarkLigPosFormat1, 5);
lookup_type!(MarkMarkPosFormat1, 6);
lookup_type!(PositionSequenceContext, 7);
lookup_type!(PositionChainContext, 8);
lookup_type!(ExtensionSubtable, 9);

impl<T: LookupType + FontWrite> FontWrite for ExtensionPosFormat1<T> {
    fn write_into(&self, writer: &mut TableWriter) {
        1u16.write_into(writer);
        T::TYPE.write_into(writer);
        self.extension_offset.write_into(writer);
    }
}

// these can't have auto impls because the traits don't support generics
#[cfg(feature = "parsing")]
impl<'a> FontRead<'a> for PositionLookup {
    fn read(data: FontData<'a>) -> Result<Self, ReadError> {
        read_fonts::layout::gpos::PositionLookup::read(data).map(|x| x.to_owned_table())
    }
}

#[cfg(feature = "parsing")]
impl<'a> FontRead<'a> for PositionLookupList {
    fn read(data: FontData<'a>) -> Result<Self, ReadError> {
        read_fonts::layout::gpos::PositionLookupList::read(data).map(|x| x.to_owned_table())
    }
}

impl SinglePosFormat1 {
    fn compute_value_format(&self) -> ValueFormat {
        self.value_record.format()
    }
}

impl SinglePosFormat2 {
    fn compute_value_format(&self) -> ValueFormat {
        self.value_records
            .first()
            .map(ValueRecord::format)
            .unwrap_or(ValueFormat::empty())
    }
}

impl PairPosFormat1 {
    fn compute_value_format1(&self) -> ValueFormat {
        self.pair_set_offsets
            .first()
            .and_then(|off| off.get())
            .and_then(|pairset| pairset.pair_value_records.first())
            .map(|rec| rec.value_record1.format())
            .unwrap_or(ValueFormat::empty())
    }

    fn compute_value_format2(&self) -> ValueFormat {
        self.pair_set_offsets
            .first()
            .and_then(|off| off.get())
            .and_then(|pairset| pairset.pair_value_records.first())
            .map(|rec| rec.value_record2.format())
            .unwrap_or(ValueFormat::empty())
    }
}

impl PairPosFormat2 {
    fn compute_value_format1(&self) -> ValueFormat {
        self.class1_records
            .first()
            .and_then(|rec| rec.class2_records.first())
            .map(|rec| rec.value_record1.format())
            .unwrap_or(ValueFormat::empty())
    }

    fn compute_value_format2(&self) -> ValueFormat {
        self.class1_records
            .first()
            .and_then(|rec| rec.class2_records.first())
            .map(|rec| rec.value_record2.format())
            .unwrap_or(ValueFormat::empty())
    }

    fn compute_class1_count(&self) -> u16 {
        self.class_def1_offset
            .get()
            .map(|cls| cls.class_count())
            .unwrap_or_default()
    }

    fn compute_class2_count(&self) -> u16 {
        self.class_def2_offset
            .get()
            .map(|cls| cls.class_count())
            .unwrap_or_default()
    }
}

impl MarkBasePosFormat1 {
    fn compute_mark_class_count(&self) -> u16 {
        mark_array_class_count(self.mark_array_offset.get())
    }
}

impl MarkMarkPosFormat1 {
    fn compute_mark_class_count(&self) -> u16 {
        mark_array_class_count(self.mark1_array_offset.get())
    }
}

impl MarkLigPosFormat1 {
    fn compute_mark_class_count(&self) -> u16 {
        mark_array_class_count(self.mark_array_offset.get())
    }
}

fn mark_array_class_count(mark_array: Option<&MarkArray>) -> u16 {
    mark_array.map(MarkArray::class_count).unwrap_or_default()
}

impl MarkArray {
    fn class_count(&self) -> u16 {
        self.mark_records
            .iter()
            .map(|rec| rec.mark_class)
            .collect::<HashSet<_>>()
            .len() as u16
    }
}
<commit_msg>Add convenience init methods to AnchorTable<commit_after>//! the [GPOS] table
//!
//! [GPOS]: https://docs.microsoft.com/en-us/typography/opentype/spec/gpos

include!("../../generated/generated_gpos.rs");

use std::collections::HashSet;

use super::value_record::ValueRecord;
use super::{
    ChainedSequenceContext, ClassDef, CoverageTable, Device, FeatureList, FeatureVariations,
    Lookup, LookupList, LookupType, ScriptList, SequenceContext,
};

#[cfg(all(test, feature = "parsing"))]
#[path = "../tests/gpos.rs"]
mod tests;

/// A GPOS lookup list table.
type PositionLookupList = LookupList<PositionLookup>;

table_newtype!(
    PositionSequenceContext,
    SequenceContext,
    read_fonts::layout::SequenceContext<'a>
);

table_newtype!(
    PositionChainContext,
    ChainedSequenceContext,
    read_fonts::layout::ChainedSequenceContext<'a>
);

impl Gpos {
    fn compute_version(&self) -> MajorMinor {
        if self.feature_variations_offset.get().is_none() {
            MajorMinor::VERSION_1_0
        } else {
            MajorMinor::VERSION_1_1
        }
    }
}

lookup_type!(SinglePos, 1);
lookup_type!(PairPos, 2);
lookup_type!(CursivePosFormat1, 3);
lookup_type!(MarkBasePosFormat1, 4);
lookup_type!(MarkLigPosFormat1, 5);
lookup_type!(MarkMarkPosFormat1, 6);
lookup_type!(PositionSequenceContext, 7);
lookup_type!(PositionChainContext, 8);
lookup_type!(ExtensionSubtable, 9);

impl<T: LookupType + FontWrite> FontWrite for ExtensionPosFormat1<T> {
    fn write_into(&self, writer: &mut TableWriter) {
        1u16.write_into(writer);
        T::TYPE.write_into(writer);
        self.extension_offset.write_into(writer);
    }
}

// these can't have auto impls because the traits don't support generics
#[cfg(feature = "parsing")]
impl<'a> FontRead<'a> for PositionLookup {
    fn read(data: FontData<'a>) -> Result<Self, ReadError> {
        read_fonts::layout::gpos::PositionLookup::read(data).map(|x| x.to_owned_table())
    }
}

#[cfg(feature = "parsing")]
impl<'a> FontRead<'a> for PositionLookupList {
    fn read(data: FontData<'a>) -> Result<Self, ReadError> {
        read_fonts::layout::gpos::PositionLookupList::read(data).map(|x| x.to_owned_table())
    }
}

impl AnchorTable {
    /// Create a new [`AnchorFormat1`] table.
    pub fn format_1(x_coordinate: i16, y_coordinate: i16) -> Self {
        Self::Format1(AnchorFormat1 {
            x_coordinate,
            y_coordinate,
        })
    }

    /// Create a new [`AnchorFormat2`] table.
    pub fn format_2(x_coordinate: i16, y_coordinate: i16, anchor_point: u16) -> Self {
        Self::Format2(AnchorFormat2 {
            x_coordinate,
            y_coordinate,
            anchor_point,
        })
    }
}

impl SinglePosFormat1 {
    fn compute_value_format(&self) -> ValueFormat {
        self.value_record.format()
    }
}

impl SinglePosFormat2 {
    fn compute_value_format(&self) -> ValueFormat {
        self.value_records
            .first()
            .map(ValueRecord::format)
            .unwrap_or(ValueFormat::empty())
    }
}

impl PairPosFormat1 {
    fn compute_value_format1(&self) -> ValueFormat {
        self.pair_set_offsets
            .first()
            .and_then(|off| off.get())
            .and_then(|pairset| pairset.pair_value_records.first())
            .map(|rec| rec.value_record1.format())
            .unwrap_or(ValueFormat::empty())
    }

    fn compute_value_format2(&self) -> ValueFormat {
        self.pair_set_offsets
            .first()
            .and_then(|off| off.get())
            .and_then(|pairset| pairset.pair_value_records.first())
            .map(|rec| rec.value_record2.format())
            .unwrap_or(ValueFormat::empty())
    }
}

impl PairPosFormat2 {
    fn compute_value_format1(&self) -> ValueFormat {
        self.class1_records
            .first()
            .and_then(|rec| rec.class2_records.first())
            .map(|rec| rec.value_record1.format())
            .unwrap_or(ValueFormat::empty())
    }

    fn compute_value_format2(&self) -> ValueFormat {
        self.class1_records
            .first()
            .and_then(|rec| rec.class2_records.first())
            .map(|rec| rec.value_record2.format())
            .unwrap_or(ValueFormat::empty())
    }

    fn compute_class1_count(&self) -> u16 {
        self.class_def1_offset
            .get()
            .map(|cls| cls.class_count())
            .unwrap_or_default()
    }

    fn compute_class2_count(&self) -> u16 {
        self.class_def2_offset
            .get()
            .map(|cls| cls.class_count())
            .unwrap_or_default()
    }
}

impl MarkBasePosFormat1 {
    fn compute_mark_class_count(&self) -> u16 {
        mark_array_class_count(self.mark_array_offset.get())
    }
}

impl MarkMarkPosFormat1 {
    fn compute_mark_class_count(&self) -> u16 {
        mark_array_class_count(self.mark1_array_offset.get())
    }
}

impl MarkLigPosFormat1 {
    fn compute_mark_class_count(&self) -> u16 {
        mark_array_class_count(self.mark_array_offset.get())
    }
}

fn mark_array_class_count(mark_array: Option<&MarkArray>) -> u16 {
    mark_array.map(MarkArray::class_count).unwrap_or_default()
}

impl MarkArray {
    fn class_count(&self) -> u16 {
        self.mark_records
            .iter()
            .map(|rec| rec.mark_class)
            .collect::<HashSet<_>>()
            .len() as u16
    }
}
<|endoftext|>