<commit_before>extern crate tokio;

use tokio::io;
use tokio::net::TcpListener;
use tokio::prelude::*;

fn main() {
    let addr = "127.0.0.1:6142".parse().unwrap();
    let listener = TcpListener::bind(&addr).unwrap();

    let server = listener.incoming().for_each(|socket| {
        println!("accepted socket; addr={:?}", socket.peer_addr().unwrap());

        // Process socket here.

        Ok(())
    }).map_err(|err| {
        // All tasks must have an `Error` type of `()`
        // This forces error handling and helps avoid silent failures.
        // In our example, we are only going to log the error to STDOUT.
        println!("accept error = {:?}", err)
    });

    println!("server running on localhost:6142");
    tokio::run(server);
}
<commit_msg>add example Task and Executor<commit_after>extern crate tokio;

use tokio::io;
use tokio::net::TcpListener;
use tokio::prelude::*;

fn main() {
    let addr = "127.0.0.1:6142".parse().unwrap();
    let listener = TcpListener::bind(&addr).unwrap();

    let server = listener.incoming().for_each(|socket| {
        println!("accepted socket; addr={:?}", socket.peer_addr().unwrap());

        let connection = io::write_all(socket, "hello world\n").then(|res| {
            println!("wrote message; success={:?}", res.is_ok());
        });

        // spawn new task to process the socket
        tokio::spawn(connection);

        Ok(())
    }).map_err(|err| {
        // All tasks must have an `Error` type of `()`
        // This forces error handling and helps avoid silent failures.
        // In our example, we are only going to log the error to STDOUT.
        println!("accept error = {:?}", err)
    });

    println!("server running on localhost:6142");
    tokio::run(server);
}

// A task that polls a single widget and writes it to STDOUT
pub struct MyTask;

impl Future for MyTask {
    // The value this future will have when ready
    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Result<Async<()>, ()> {
        match poll_widget() {
            Async::Ready(widget) => {
                println!("widget={:?}", widget);
                Ok(Async::Ready(()))
            }
            Async::NotReady => {
                Ok(Async::NotReady)
            }
        }
    }
}

pub struct SpinExecutor {
    // the tasks an executor is responsible for in a double queue
    tasks: VecDeque<Box<Future<Item=(), Error=()>>>,
}

impl SpinExecutor {
    pub fn spawn<T>(&mut self, task: T)
        where T: Future<Item=(), Error=()> + 'static
    {
        self.tasks.push_back(Box::new(task))
    }

    pub fn run(&mut self) {
        // Of course, this would not be very efficient.
        // The executor spins in a busy loop and tries to poll all tasks even if the task will just return NotReady again.
        while let Some(mut task) = self.tasks.pop_front() {
            match task.poll().unwrap() {
                Async::Ready(_) => {}
                Async::NotReady => {
                    self.tasks.push_back(task);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use common::consts::get_version_string;
use platform::process::get_pid;
use platform::time::time_now;
use storage::Cache;
use storage::CacheError;
use storage::Key;
use storage::Value;
use tcp_transport::stats::TransportStats;

use super::cmd::Cmd;
use super::cmd::Delete;
use super::cmd::Get;
use super::cmd::Resp;
use super::cmd::Set;
use super::cmd::SetInstr;
use super::cmd::Stat;
use super::cmd::Touch;
use super::cmd::Value as CmdValue;

// For use to get an early exit from a function. The first parameter is a bool
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Option<Resp>. In case
// of None no return is made ie. execution continues. In case of Some(resp) we
// perform return resp;
//
// maybe_reply_stmt!(!noreply, match rv {
//     Ok(_) => None,   // do not return
//     Err(_) => Some(Resp::Error),
// });
//
// ->
//
// let opt = expr_provided;
// if !opt.is_none() {
//     if !noreply {
//         return opt.unwrap();
//     } else {
//         return Resp::Empty;
//     }
// }
macro_rules! maybe_reply_stmt {
    ( $cond:expr, $opt:expr ) => {
        {
            if !$opt.is_none() {
                if $cond {
                    return $opt.unwrap();
                } else {
                    return Resp::Empty;
                }
            };
        }
    };
}

// For use to conditionally produce a return value. The first parameter is used
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Resp.
//
// maybe_reply_expr!(!noreply, match rv {
//     Ok(_) => Resp::Stored,
//     Err(_) => Resp::Error,
// });
//
// ->
//
// let resp = expr_provided;
// if !noreply {
//     resp
// } else {
//     Resp::Empty
// }
macro_rules! maybe_reply_expr {
    ( $cond:expr, $resp:expr ) => {
        {
            if $cond {
                $resp
            } else {
                Resp::Empty
            }
        }
    };
}

struct DriverStats {
    cmd_get: u64,
    cmd_set: u64,
    cmd_flush: u64,
    cmd_touch: u64,
}

impl DriverStats {
    pub fn new() -> DriverStats {
        DriverStats {
            cmd_get: 0,
            cmd_set: 0,
            cmd_flush: 0,
            cmd_touch: 0,
        }
    }
}

pub struct Driver {
    cache: Cache,
    time_start: f64,

    stats: DriverStats,
    transport_stats: TransportStats, // this is a global snapshot
}

impl Driver {
    pub fn new(cache: Cache) -> Driver {
        Driver {
            cache: cache,
            stats: DriverStats::new(),
            time_start: time_now(),
            transport_stats: TransportStats::new(),
        }
    }

    fn set_exptime(&self, value: &mut Value, exptime: u32) {
        // If exptime is greater than zero we need to set it on the value
        if exptime > 0 {
            let tm;

            // Is it an interval greater than 30 days? Then it's a timestamp
            if exptime > 60 * 60 * 24 * 30 {
                tm = exptime as f64;

            } else {
                // Otherwise it's relative from now
                tm = time_now() + exptime as f64;
            }

            value.set_exptime(tm);
        }
    }

    fn do_add(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If so it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => {
                Some(Resp::NotStored)
            },
            Ok(false) => None,
            Err(_) => {
                Some(Resp::Error)
            }
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_append(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
            _ => None,
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Append the data we just received to the blob that is there
        value.item.extend(set.data);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_delete(&mut self, delete: Delete) -> Resp {
        let key = Key::new(delete.key.clone().into_bytes());

        let rv = self.cache.remove(&key);

        maybe_reply_expr!(!delete.noreply, match rv {
            Ok(_) => Resp::Deleted,
            Err(_) => Resp::NotFound,
        })
    }

    fn do_get(&mut self, get: Get) -> Resp {
        // Update stats
        self.stats.cmd_get += 1;

        let mut values = vec![];

        for key in get.keys {
            let key_str = key.clone();

            let key_st = Key::new(key.into_bytes());
            let rv = self.cache.get(&key_st);

            match rv {
                Ok(value) => {
                    let val_st = CmdValue {
                        key: key_str,
                        flags: value.flags,
                        data: value.item.clone(),
                    };
                    values.push(val_st);
                }
                // Keys that were not found are skipped, no error given
                Err(_) => (),
            }
        }

        Resp::Values(values)
    }

    fn do_prepend(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(_) => None,
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Prepend the data we just received to the blob that is there
        let mut new_item = Vec::with_capacity(set.data.len() + value.item.len());
        new_item.extend(set.data);
        new_item.extend(value.item);
        value.item = new_item;

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_replace(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If not it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_set(&mut self, set: Set) -> Resp {
        // Update stats
        self.stats.cmd_set += 1;

        let key = Key::new(set.key.into_bytes());
        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        match set.noreply {
            true => Resp::Empty,
            false => {
                match rv {
                    Ok(_) => Resp::Stored,
                    Err(_) => Resp::Error,
                }
            }
        }
    }

    fn do_stats(&self) -> Resp {
        let storage = self.cache.get_stats();

        let pid = get_pid().to_string();
        let uptime = ((time_now() - self.time_start) as u64).to_string();
        let time = (time_now() as u64).to_string();
        let version = get_version_string();
        let cmd_get = self.stats.cmd_get.to_string();
        let cmd_set = self.stats.cmd_set.to_string();
        let cmd_flush = self.stats.cmd_flush.to_string();
        let cmd_touch = self.stats.cmd_touch.to_string();
        let get_hits = storage.get_hits.to_string();
        let get_misses = storage.get_misses.to_string();
        let delete_hits = storage.delete_hits.to_string();
        let delete_misses = storage.delete_misses.to_string();
        let bytes_read = self.transport_stats.bytes_read.to_string();
        let bytes_written = self.transport_stats.bytes_written.to_string();
        let limit_maxbytes = self.cache.capacity.to_string();
        let bytes = storage.bytes.to_string();
        let curr_items = self.cache.len().to_string();
        let total_items = storage.total_items.to_string();
        let evictions = storage.evictions.to_string();
        let reclaimed = storage.reclaimed.to_string();

        let st_pid = Stat::new("pid", pid);
        let st_uptime = Stat::new("uptime", uptime);
        let st_time = Stat::new("time", time);
        let st_version = Stat::new("version", version);
        let st_cmd_get = Stat::new("cmd_get", cmd_get);
        let st_cmd_set = Stat::new("cmd_set", cmd_set);
        let st_cmd_flush = Stat::new("cmd_flush", cmd_flush);
        let st_cmd_touch = Stat::new("cmd_touch", cmd_touch);
        let st_get_hits = Stat::new("get_hits", get_hits);
        let st_get_misses = Stat::new("get_misses", get_misses);
        let st_delete_hits = Stat::new("delete_hits", delete_hits);
        let st_delete_misses = Stat::new("delete_misses", delete_misses);
        let st_bytes_read = Stat::new("bytes_read", bytes_read);
        let st_bytes_written = Stat::new("bytes_written", bytes_written);
        let st_limit_maxbytes = Stat::new("limit_maxbytes", limit_maxbytes);
        let st_bytes = Stat::new("bytes", bytes);
        let st_curr_items = Stat::new("curr_items", curr_items);
        let st_total_items = Stat::new("total_items", total_items);
        let st_evictions = Stat::new("evictions", evictions);
        let st_reclaimed = Stat::new("reclaimed", reclaimed);

        Resp::Stats(vec![st_pid,
                         st_uptime,
                         st_time,
                         st_version,
                         st_cmd_get,
                         st_cmd_set,
                         st_cmd_flush,
                         st_cmd_touch,
                         st_get_hits,
                         st_get_misses,
                         st_delete_hits,
                         st_delete_misses,
                         st_bytes_read,
                         st_bytes_written,
                         st_limit_maxbytes,
                         st_bytes,
                         st_curr_items,
                         st_total_items,
                         st_evictions,
                         st_reclaimed])
    }

    pub fn do_touch(&mut self, touch: Touch) -> Resp {
        let key = Key::new(touch.key.into_bytes());

        // See if the key is set
        let rv = self.cache.contains_key(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotFound),
            Err(_) => Some(Resp::Error),
        });

        // Load the value
        let rv = self.cache.remove(&key);
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(_) => None,
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        self.set_exptime(&mut value, touch.exptime);

        // Set it
        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!touch.noreply, match rv {
            Ok(_) => Resp::Touched,
            Err(_) => Resp::Error,
        })
    }

    pub fn do_version(&self) -> Resp {
        Resp::Version(get_version_string())
    }

    pub fn run(&mut self, cmd: Cmd) -> Resp {
        match cmd {
            Cmd::Delete(del) => self.do_delete(del),
            Cmd::Get(get) => self.do_get(get),
            Cmd::Quit => Resp::Empty,  // handled at transport level
            Cmd::Set(set) => match set.instr {
                SetInstr::Add => self.do_add(set),
                SetInstr::Append => self.do_append(set),
                SetInstr::Prepend => self.do_prepend(set),
                SetInstr::Replace => self.do_replace(set),
                SetInstr::Set => self.do_set(set),
                _ => Resp::Error,  // TODO
            },
            Cmd::Stats => self.do_stats(),
            Cmd::Touch(touch) => self.do_touch(touch),
            Cmd::Version => self.do_version(),
        }
    }

    pub fn update_transport_stats(&mut self, stats: TransportStats) {
        self.transport_stats = stats;
    }
}
<commit_msg>factor set to use macro<commit_after>use common::consts::get_version_string;
use platform::process::get_pid;
use platform::time::time_now;
use storage::Cache;
use storage::CacheError;
use storage::Key;
use storage::Value;
use tcp_transport::stats::TransportStats;

use super::cmd::Cmd;
use super::cmd::Delete;
use super::cmd::Get;
use super::cmd::Resp;
use super::cmd::Set;
use super::cmd::SetInstr;
use super::cmd::Stat;
use super::cmd::Touch;
use super::cmd::Value as CmdValue;

// For use to get an early exit from a function. The first parameter is a bool
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Option<Resp>. In case
// of None no return is made ie. execution continues. In case of Some(resp) we
// perform return resp;
//
// maybe_reply_stmt!(!noreply, match rv {
//     Ok(_) => None,   // do not return
//     Err(_) => Some(Resp::Error),
// });
//
// ->
//
// let opt = expr_provided;
// if !opt.is_none() {
//     if !noreply {
//         return opt.unwrap();
//     } else {
//         return Resp::Empty;
//     }
// }
macro_rules! maybe_reply_stmt {
    ( $cond:expr, $opt:expr ) => {
        {
            if !$opt.is_none() {
                if $cond {
                    return $opt.unwrap();
                } else {
                    return Resp::Empty;
                }
            };
        }
    };
}

// For use to conditionally produce a return value. The first parameter is used
// to indicate whether to omit responses (returns Resp::Empty instead). The
// second parameter is an expression that evaluates to Resp.
//
// maybe_reply_expr!(!noreply, match rv {
//     Ok(_) => Resp::Stored,
//     Err(_) => Resp::Error,
// });
//
// ->
//
// let resp = expr_provided;
// if !noreply {
//     resp
// } else {
//     Resp::Empty
// }
macro_rules! maybe_reply_expr {
    ( $cond:expr, $resp:expr ) => {
        {
            if $cond {
                $resp
            } else {
                Resp::Empty
            }
        }
    };
}

struct DriverStats {
    cmd_get: u64,
    cmd_set: u64,
    cmd_flush: u64,
    cmd_touch: u64,
}

impl DriverStats {
    pub fn new() -> DriverStats {
        DriverStats {
            cmd_get: 0,
            cmd_set: 0,
            cmd_flush: 0,
            cmd_touch: 0,
        }
    }
}

pub struct Driver {
    cache: Cache,
    time_start: f64,

    stats: DriverStats,
    transport_stats: TransportStats, // this is a global snapshot
}

impl Driver {
    pub fn new(cache: Cache) -> Driver {
        Driver {
            cache: cache,
            stats: DriverStats::new(),
            time_start: time_now(),
            transport_stats: TransportStats::new(),
        }
    }

    fn set_exptime(&self, value: &mut Value, exptime: u32) {
        // If exptime is greater than zero we need to set it on the value
        if exptime > 0 {
            let tm;

            // Is it an interval greater than 30 days? Then it's a timestamp
            if exptime > 60 * 60 * 24 * 30 {
                tm = exptime as f64;

            } else {
                // Otherwise it's relative from now
                tm = time_now() + exptime as f64;
            }

            value.set_exptime(tm);
        }
    }

    fn do_add(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If so it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => {
                Some(Resp::NotStored)
            },
            Ok(false) => None,
            Err(_) => {
                Some(Resp::Error)
            }
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_append(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
            _ => None,
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Append the data we just received to the blob that is there
        value.item.extend(set.data);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_delete(&mut self, delete: Delete) -> Resp {
        let key = Key::new(delete.key.clone().into_bytes());

        let rv = self.cache.remove(&key);

        maybe_reply_expr!(!delete.noreply, match rv {
            Ok(_) => Resp::Deleted,
            Err(_) => Resp::NotFound,
        })
    }

    fn do_get(&mut self, get: Get) -> Resp {
        // Update stats
        self.stats.cmd_get += 1;

        let mut values = vec![];

        for key in get.keys {
            let key_str = key.clone();

            let key_st = Key::new(key.into_bytes());
            let rv = self.cache.get(&key_st);

            match rv {
                Ok(value) => {
                    let val_st = CmdValue {
                        key: key_str,
                        flags: value.flags,
                        data: value.item.clone(),
                    };
                    values.push(val_st);
                }
                // Keys that were not found are skipped, no error given
                Err(_) => (),
            }
        }

        Resp::Values(values)
    }

    fn do_prepend(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Load the value
        let rv = self.cache.remove(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(_) => None,
            Err(CacheError::KeyNotFound) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        // Prepend the data we just received to the blob that is there
        let mut new_item = Vec::with_capacity(set.data.len() + value.item.len());
        new_item.extend(set.data);
        new_item.extend(value.item);
        value.item = new_item;

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_replace(&mut self, set: Set) -> Resp {
        let key = Key::new(set.key.into_bytes());

        // Do we store this item already? If not it's an early exit.
        let rv = self.cache.contains_key(&key);
        maybe_reply_stmt!(!set.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotStored),
            Err(_) => Some(Resp::Error),
        });

        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_set(&mut self, set: Set) -> Resp {
        // Update stats
        self.stats.cmd_set += 1;

        let key = Key::new(set.key.into_bytes());
        let mut value = Value::new(set.data);
        value.with_flags(set.flags);
        self.set_exptime(&mut value, set.exptime);

        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!set.noreply, match rv {
            Ok(_) => Resp::Stored,
            Err(_) => Resp::Error,
        })
    }

    fn do_stats(&self) -> Resp {
        let storage = self.cache.get_stats();

        let pid = get_pid().to_string();
        let uptime = ((time_now() - self.time_start) as u64).to_string();
        let time = (time_now() as u64).to_string();
        let version = get_version_string();
        let cmd_get = self.stats.cmd_get.to_string();
        let cmd_set = self.stats.cmd_set.to_string();
        let cmd_flush = self.stats.cmd_flush.to_string();
        let cmd_touch = self.stats.cmd_touch.to_string();
        let get_hits = storage.get_hits.to_string();
        let get_misses = storage.get_misses.to_string();
        let delete_hits = storage.delete_hits.to_string();
        let delete_misses = storage.delete_misses.to_string();
        let bytes_read = self.transport_stats.bytes_read.to_string();
        let bytes_written = self.transport_stats.bytes_written.to_string();
        let limit_maxbytes = self.cache.capacity.to_string();
        let bytes = storage.bytes.to_string();
        let curr_items = self.cache.len().to_string();
        let total_items = storage.total_items.to_string();
        let evictions = storage.evictions.to_string();
        let reclaimed = storage.reclaimed.to_string();

        let st_pid = Stat::new("pid", pid);
        let st_uptime = Stat::new("uptime", uptime);
        let st_time = Stat::new("time", time);
        let st_version = Stat::new("version", version);
        let st_cmd_get = Stat::new("cmd_get", cmd_get);
        let st_cmd_set = Stat::new("cmd_set", cmd_set);
        let st_cmd_flush = Stat::new("cmd_flush", cmd_flush);
        let st_cmd_touch = Stat::new("cmd_touch", cmd_touch);
        let st_get_hits = Stat::new("get_hits", get_hits);
        let st_get_misses = Stat::new("get_misses", get_misses);
        let st_delete_hits = Stat::new("delete_hits", delete_hits);
        let st_delete_misses = Stat::new("delete_misses", delete_misses);
        let st_bytes_read = Stat::new("bytes_read", bytes_read);
        let st_bytes_written = Stat::new("bytes_written", bytes_written);
        let st_limit_maxbytes = Stat::new("limit_maxbytes", limit_maxbytes);
        let st_bytes = Stat::new("bytes", bytes);
        let st_curr_items = Stat::new("curr_items", curr_items);
        let st_total_items = Stat::new("total_items", total_items);
        let st_evictions = Stat::new("evictions", evictions);
        let st_reclaimed = Stat::new("reclaimed", reclaimed);

        Resp::Stats(vec![st_pid,
                         st_uptime,
                         st_time,
                         st_version,
                         st_cmd_get,
                         st_cmd_set,
                         st_cmd_flush,
                         st_cmd_touch,
                         st_get_hits,
                         st_get_misses,
                         st_delete_hits,
                         st_delete_misses,
                         st_bytes_read,
                         st_bytes_written,
                         st_limit_maxbytes,
                         st_bytes,
                         st_curr_items,
                         st_total_items,
                         st_evictions,
                         st_reclaimed])
    }

    pub fn do_touch(&mut self, touch: Touch) -> Resp {
        let key = Key::new(touch.key.into_bytes());

        // See if the key is set
        let rv = self.cache.contains_key(&key);

        // If if it's not there we error out
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(true) => None,
            Ok(false) => Some(Resp::NotFound),
            Err(_) => Some(Resp::Error),
        });

        // Load the value
        let rv = self.cache.remove(&key);
        maybe_reply_stmt!(!touch.noreply, match rv {
            Ok(_) => None,
            Err(_) => Some(Resp::Error),
        });

        // Update the value
        let mut value = rv.unwrap();
        self.set_exptime(&mut value, touch.exptime);

        // Set it
        let rv = self.cache.set(key, value);

        maybe_reply_expr!(!touch.noreply, match rv {
            Ok(_) => Resp::Touched,
            Err(_) => Resp::Error,
        })
    }

    pub fn do_version(&self) -> Resp {
        Resp::Version(get_version_string())
    }

    pub fn run(&mut self, cmd: Cmd) -> Resp {
        match cmd {
            Cmd::Delete(del) => self.do_delete(del),
            Cmd::Get(get) => self.do_get(get),
            Cmd::Quit => Resp::Empty,  // handled at transport level
            Cmd::Set(set) => match set.instr {
                SetInstr::Add => self.do_add(set),
                SetInstr::Append => self.do_append(set),
                SetInstr::Prepend => self.do_prepend(set),
                SetInstr::Replace => self.do_replace(set),
                SetInstr::Set => self.do_set(set),
                _ => Resp::Error,  // TODO
            },
            Cmd::Stats => self.do_stats(),
            Cmd::Touch(touch) => self.do_touch(touch),
            Cmd::Version => self.do_version(),
        }
    }

    pub fn update_transport_stats(&mut self, stats: TransportStats) {
        self.transport_stats = stats;
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 `multipart` Crate Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Boundary parsing for `multipart` requests.

use ::safemem;

use super::buf_redux::BufReader;
use super::buf_redux::policy::MinBuffered;
use super::twoway;

use std::cmp;
use std::borrow::Borrow;

use std::io;
use std::io::prelude::*;

use self::State::*;

pub const MIN_BUF_SIZE: usize = 1024;

#[derive(Debug, PartialEq, Eq)]
enum State {
    Searching,
    BoundaryRead,
    AtEnd
}

/// A struct implementing `Read` and `BufRead` that will yield bytes until it sees a given sequence.
#[derive(Debug)]
pub struct BoundaryReader<R> {
    source: BufReader<R, MinBuffered>,
    boundary: Vec<u8>,
    search_idx: usize,
    state: State,
}

impl<R> BoundaryReader<R> where R: Read {
    #[doc(hidden)]
    pub fn from_reader<B: Into<Vec<u8>>>(reader: R, boundary: B) -> BoundaryReader<R> {
        let mut boundary = boundary.into();
        safemem::prepend(b"--", &mut boundary);
        let source = BufReader::new(reader).set_policy(MinBuffered(MIN_BUF_SIZE));

        BoundaryReader {
            source,
            boundary,
            search_idx: 0,
            state: Searching,
        }
    }

    fn read_to_boundary(&mut self) -> io::Result<&[u8]> {
        let buf = fill_buf_min(&mut self.source, self.search_idx + self.boundary.len())?;

        trace!("Buf: {:?}", String::from_utf8_lossy(buf));

        debug!("Before-loop Buf len: {} Search idx: {} State: {:?}",
               buf.len(), self.search_idx, self.state);

        if self.state == Searching && self.search_idx < buf.len() {
            let lookahead = &buf[self.search_idx..];

            debug!("Find boundary loop! Lookahead len: {}", lookahead.len());

            // Look for the boundary, or if it isn't found, stop near the end.
            match twoway::find_bytes(lookahead, &self.boundary) {
                Some(found_idx) => {
                    self.search_idx += found_idx;
                    self.state = BoundaryRead;
                },
                None => {
                    self.search_idx += lookahead.len().saturating_sub(self.boundary.len() + 2);
                }
            }
        }

        debug!("After-loop Buf len: {} Search idx: {} State: {:?}",
               buf.len(), self.search_idx, self.state);

        // don't modify search_idx so it always points to the start of the boundary
        let mut buf_len = self.search_idx;

        // back up the cursor to before the boundary's preceding CRLF
        if self.state != Searching && buf_len >= 2 {
            let two_bytes_before = &buf[buf_len - 2 .. buf_len];

            trace!("Two bytes before: {:?} ({:?}) (\"\\r\\n\": {:?})",
                   String::from_utf8_lossy(two_bytes_before), two_bytes_before, b"\r\n");

            if two_bytes_before == *b"\r\n" {
                debug!("Subtract two!");
                buf_len -= 2;
            }
        }

        let ret_buf = &buf[..buf_len];

        trace!("Returning buf: {:?}", String::from_utf8_lossy(ret_buf));

        Ok(ret_buf)
    }

    pub fn set_min_buf_size(&mut self, min_buf_size: usize) {
        // ensure the minimum buf size is at least enough to find a boundary with some extra
        let min_buf_size = cmp::max(self.boundary.len() * 2, min_buf_size);

        self.source.policy_mut().0 = min_buf_size;
    }

    #[doc(hidden)]
    pub fn consume_boundary(&mut self) -> io::Result<bool> {
        if self.state == AtEnd {
            return Ok(true);
        }

        while self.state == Searching {
            debug!("Boundary not found yet");

            let buf_len = self.read_to_boundary()?.len();

            debug!("Discarding {} bytes", buf_len);

            self.consume(buf_len);
        }

        let consume_amt = {
            let min_len = self.boundary.len() + 4;

            let buf = fill_buf_min(&mut self.source, min_len)?;

            if buf.len() < min_len {
                return Err(io::Error::new(io::ErrorKind::UnexpectedEof,
                                          "not enough bytes to verify boundary"));
            }

            // we have enough bytes to verify
            self.state = Searching;

            let mut consume_amt = self.search_idx + self.boundary.len();

            let last_two = &buf[consume_amt .. consume_amt + 2];

            match last_two {
                b"\r\n" => consume_amt += 2,
                b"--" => { consume_amt += 2; self.state = AtEnd },
                _ => debug!("Unexpected bytes following boundary: {:?}",
                            String::from_utf8_lossy(&last_two)),
            }

            consume_amt
        };

        trace!("Consuming {} bytes, remaining buf: {:?}",
               consume_amt,
               String::from_utf8_lossy(self.source.buffer()));

        self.source.consume(consume_amt);
        self.search_idx = 0;

        trace!("Consumed boundary (state: {:?}), remaining buf: {:?}", self.state,
               String::from_utf8_lossy(self.source.buffer()));

        Ok(self.state == AtEnd)
    }
}

#[cfg(feature = "bench")]
impl<'a> BoundaryReader<io::Cursor<&'a [u8]>> {
    fn new_with_bytes(bytes: &'a [u8], boundary: &str) -> Self {
        Self::from_reader(io::Cursor::new(bytes), boundary)
    }

    fn reset(&mut self) {
        // Dump buffer and reset cursor
        self.source.seek(io::SeekFrom::Start(0));
        self.state = Searching;
        self.search_idx = 0;
    }
}

impl<R> Borrow<R> for BoundaryReader<R> {
    fn borrow(&self) -> &R {
        self.source.get_ref()
    }
}

impl<R> Read for BoundaryReader<R> where R: Read {
    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
        let read = {
            let mut buf = self.read_to_boundary()?;
            // This shouldn't ever be an error so unwrapping is fine.
            buf.read(out).unwrap()
        };

        self.consume(read);
        Ok(read)
    }
}

impl<R> BufRead for BoundaryReader<R> where R: Read {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.read_to_boundary()
    }

    fn consume(&mut self, amt: usize) {
        let true_amt = cmp::min(amt, self.search_idx);

        debug!("Consume! amt: {} true amt: {}", amt, true_amt);

        self.source.consume(true_amt);
        self.search_idx -= true_amt;
    }
}

fn fill_buf_min<R: BufRead>(rdr: &mut R, min: usize) -> io::Result<&[u8]> {
    let mut last_len = 0;

    while {
        let buf = rdr.fill_buf()?;
        let cont = buf.len() > last_len && buf.len() < min;
        last_len = buf.len();
        cont
    } { /* nop */ }

    rdr.fill_buf()
}

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

    use std::io;
    use std::io::prelude::*;

    const BOUNDARY: &'static str = "boundary";
    const TEST_VAL: &'static str = "--boundary\r\n\
                                    dashed-value-1\r\n\
                                    --boundary\r\n\
                                    dashed-value-2\r\n\
                                    --boundary--";

    #[test]
    fn test_boundary() {
        let _logger = ::mock::log_on_panic();

        debug!("Testing boundary (no split)");

        let src = &mut TEST_VAL.as_bytes();
        let mut reader = BoundaryReader::from_reader(src, BOUNDARY);

        let mut buf = String::new();

        test_boundary_reader(&mut reader, &mut buf);
    }

    struct SplitReader<'a> {
        left: &'a [u8],
        right: &'a [u8],
    }

    impl<'a> SplitReader<'a> {
        fn split(data: &'a [u8], at: usize) -> SplitReader<'a> {
            let (left, right) = data.split_at(at);

            SplitReader {
                left: left,
                right: right,
            }
        }
    }

    impl<'a> Read for SplitReader<'a> {
        fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
            fn copy_bytes_partial(src: &mut &[u8], dst: &mut [u8]) -> usize {
                src.read(dst).unwrap()
            }

            let mut copy_amt = copy_bytes_partial(&mut self.left, dst);

            if copy_amt == 0 {
                copy_amt = copy_bytes_partial(&mut self.right, dst)
            };

            Ok(copy_amt)
        }
    }

    #[test]
    fn test_split_boundary() {
        let logger = ::mock::log_on_panic();

        debug!("Testing boundary (split)");

        let mut buf = String::new();

        // Substitute for `.step_by()` being unstable.
        for split_at in 0 .. TEST_VAL.len(){
            debug!("Testing split at: {}", split_at);

            let src = SplitReader::split(TEST_VAL.as_bytes(), split_at);
            let mut reader = BoundaryReader::from_reader(src, BOUNDARY);
            test_boundary_reader(&mut reader, &mut buf);
        }

        logger.clear();
    }

    fn test_boundary_reader<R: Read>(reader: &mut BoundaryReader<R>, buf: &mut String) {
        buf.clear();

        debug!("Read 1");
        let _ = reader.read_to_string(buf).unwrap();
        assert!(buf.is_empty(), "Buffer not empty: {:?}", buf);
        buf.clear();

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "dashed-value-1");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 3");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "dashed-value-2");
        buf.clear();

        debug!("Consume 3");
        reader.consume_boundary().unwrap();

        debug!("Read 4");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");
    }

    #[test]
    fn test_leading_crlf() {
        let mut body: &[u8] = b"\r\n\r\n--boundary\r\n\
                         asdf1234\
                         \r\n\r\n--boundary--";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\r\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2 (empty)");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");
    }

    #[test]
    fn test_trailing_crlf() {
        let mut body: &[u8] = b"--boundary\r\n\
                         asdf1234\
                         \r\n\r\n--boundary\r\n\
                         hjkl5678\r\n--boundary--";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");

        // Repro for https://github.com/abonander/multipart/issues/93
        // These two reads should produce the same buffer
        let buf1 = reader.read_to_boundary().unwrap().to_owned();
        let buf2 = reader.read_to_boundary().unwrap().to_owned();
        assert_eq!(buf1, buf2);

        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\r\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "hjkl5678");
        buf.clear();

        debug!("Consume 3");
        reader.consume_boundary().unwrap();

        debug!("Read 3 (empty)");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");
    }

    // https://github.com/abonander/multipart/issues/93#issuecomment-343610587
    #[test]
    fn test_trailing_lflf() {
        let mut body: &[u8] = b"--boundary\r\n\
                         asdf1234\
                         \n\n\r\n--boundary\r\n\
                         hjkl5678\r\n--boundary--";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");

        // same as above
        let buf1 = reader.read_to_boundary().unwrap().to_owned();
        let buf2 = reader.read_to_boundary().unwrap().to_owned();
        assert_eq!(buf1, buf2);

        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\n\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "hjkl5678");
        buf.clear();

        debug!("Consume 3");
        reader.consume_boundary().unwrap();

        debug!("Read 3 (empty)");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");
    }

    #[cfg(feature = "bench")]
    mod bench {
        extern crate test;
        use self::test::Bencher;

        use super::*;

        #[bench]
        fn bench_boundary_reader(b: &mut Bencher) {
            let mut reader = BoundaryReader::new_with_bytes(TEST_VAL.as_bytes(), BOUNDARY);
            let mut buf = String::with_capacity(256);

            b.iter(|| {
                reader.reset();
                test_boundary_reader(&mut reader, &mut buf);
            });
        }
    }
}
<commit_msg>fix reading of unterminated body<commit_after>// Copyright 2016 `multipart` Crate Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Boundary parsing for `multipart` requests.

use ::safemem;

use super::buf_redux::BufReader;
use super::buf_redux::policy::MinBuffered;
use super::twoway;

use std::cmp;
use std::borrow::Borrow;

use std::io;
use std::io::prelude::*;

use self::State::*;

pub const MIN_BUF_SIZE: usize = 1024;

#[derive(Debug, PartialEq, Eq)]
enum State {
    Searching,
    BoundaryRead,
    AtEnd
}

/// A struct implementing `Read` and `BufRead` that will yield bytes until it sees a given sequence.
#[derive(Debug)]
pub struct BoundaryReader<R> {
    source: BufReader<R, MinBuffered>,
    boundary: Vec<u8>,
    search_idx: usize,
    state: State,
}

impl<R> BoundaryReader<R> where R: Read {
    #[doc(hidden)]
    pub fn from_reader<B: Into<Vec<u8>>>(reader: R, boundary: B) -> BoundaryReader<R> {
        let mut boundary = boundary.into();
        safemem::prepend(b"--", &mut boundary);
        let source = BufReader::new(reader).set_policy(MinBuffered(MIN_BUF_SIZE));

        BoundaryReader {
            source,
            boundary,
            search_idx: 0,
            state: Searching,
        }
    }

    fn read_to_boundary(&mut self) -> io::Result<&[u8]> {
        if self.state == AtEnd { return Ok(&[])}

        let buf = self.source.fill_buf()?;

        if self.search_idx == 0 && buf.len() < self.boundary.len() {
            return Err(io::Error::new(io::ErrorKind::UnexpectedEof,
                                      "unexpected end of request body"));
        }

        trace!("Buf: {:?}", String::from_utf8_lossy(buf));

        debug!("Before-loop Buf len: {} Search idx: {} State: {:?}",
               buf.len(), self.search_idx, self.state);

        if self.state == Searching && self.search_idx < buf.len() {
            let lookahead = &buf[self.search_idx..];

            debug!("Find boundary loop! Lookahead len: {}", lookahead.len());

            // Look for the boundary, or if it isn't found, stop near the end.
            match twoway::find_bytes(lookahead, &self.boundary) {
                Some(found_idx) => {
                    self.search_idx += found_idx;
                    self.state = BoundaryRead;
                },
                None => {
                    self.search_idx += lookahead.len().saturating_sub(self.boundary.len() + 2);
                }
            }
        }

        debug!("After-loop Buf len: {} Search idx: {} State: {:?}",
               buf.len(), self.search_idx, self.state);

        // don't modify search_idx so it always points to the start of the boundary
        let mut buf_len = self.search_idx;

        // back up the cursor to before the boundary's preceding CRLF
        if self.state != Searching && buf_len >= 2 {
            let two_bytes_before = &buf[buf_len - 2 .. buf_len];

            trace!("Two bytes before: {:?} ({:?}) (\"\\r\\n\": {:?})",
                   String::from_utf8_lossy(two_bytes_before), two_bytes_before, b"\r\n");

            if two_bytes_before == *b"\r\n" {
                debug!("Subtract two!");
                buf_len -= 2;
            }
        }

        let ret_buf = &buf[..buf_len];

        trace!("Returning buf: {:?}", String::from_utf8_lossy(ret_buf));

        Ok(ret_buf)
    }

    pub fn set_min_buf_size(&mut self, min_buf_size: usize) {
        // ensure the minimum buf size is at least enough to find a boundary with some extra
        let min_buf_size = cmp::max(self.boundary.len() * 2, min_buf_size);

        self.source.policy_mut().0 = min_buf_size;
    }

    #[doc(hidden)]
    pub fn consume_boundary(&mut self) -> io::Result<bool> {
        if self.state == AtEnd {
            return Ok(true);
        }

        while self.state == Searching {
            debug!("Boundary not found yet");

            let buf_len = self.read_to_boundary()?.len();

            debug!("Discarding {} bytes", buf_len);

            self.consume(buf_len);
        }

        let consume_amt = {
            let min_len = self.boundary.len() + 4;

            let buf = self.source.fill_buf()?;

            if buf.len() < min_len {
                return Err(io::Error::new(io::ErrorKind::UnexpectedEof,
                                          "not enough bytes to verify boundary"));
            }

            // we have enough bytes to verify
            self.state = Searching;

            let mut consume_amt = self.search_idx + self.boundary.len();

            let last_two = &buf[consume_amt .. consume_amt + 2];

            match last_two {
                b"\r\n" => consume_amt += 2,
                b"--" => { consume_amt += 2; self.state = AtEnd },
                _ => debug!("Unexpected bytes following boundary: {:?}",
                            String::from_utf8_lossy(&last_two)),
            }

            consume_amt
        };

        trace!("Consuming {} bytes, remaining buf: {:?}",
               consume_amt,
               String::from_utf8_lossy(self.source.buffer()));

        self.source.consume(consume_amt);
        self.search_idx = 0;

        trace!("Consumed boundary (state: {:?}), remaining buf: {:?}", self.state,
               String::from_utf8_lossy(self.source.buffer()));

        Ok(self.state == AtEnd)
    }
}

#[cfg(feature = "bench")]
impl<'a> BoundaryReader<io::Cursor<&'a [u8]>> {
    fn new_with_bytes(bytes: &'a [u8], boundary: &str) -> Self {
        Self::from_reader(io::Cursor::new(bytes), boundary)
    }

    fn reset(&mut self) {
        // Dump buffer and reset cursor
        self.source.seek(io::SeekFrom::Start(0));
        self.state = Searching;
        self.search_idx = 0;
    }
}

impl<R> Borrow<R> for BoundaryReader<R> {
    fn borrow(&self) -> &R {
        self.source.get_ref()
    }
}

impl<R> Read for BoundaryReader<R> where R: Read {
    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
        let read = {
            let mut buf = self.read_to_boundary()?;
            // This shouldn't ever be an error so unwrapping is fine.
            buf.read(out).unwrap()
        };

        self.consume(read);
        Ok(read)
    }
}

impl<R> BufRead for BoundaryReader<R> where R: Read {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.read_to_boundary()
    }

    fn consume(&mut self, amt: usize) {
        let true_amt = cmp::min(amt, self.search_idx);

        debug!("Consume! amt: {} true amt: {}", amt, true_amt);

        self.source.consume(true_amt);
        self.search_idx -= true_amt;
    }
}

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

    use std::io;
    use std::io::prelude::*;

    const BOUNDARY: &'static str = "boundary";
    const TEST_VAL: &'static str = "--boundary\r\n\
                                    dashed-value-1\r\n\
                                    --boundary\r\n\
                                    dashed-value-2\r\n\
                                    --boundary--";

    #[test]
    fn test_boundary() {
        let _logger = ::mock::log_on_panic();

        debug!("Testing boundary (no split)");

        let src = &mut TEST_VAL.as_bytes();
        let mut reader = BoundaryReader::from_reader(src, BOUNDARY);

        let mut buf = String::new();

        test_boundary_reader(&mut reader, &mut buf);
    }

    struct SplitReader<'a> {
        left: &'a [u8],
        right: &'a [u8],
    }

    impl<'a> SplitReader<'a> {
        fn split(data: &'a [u8], at: usize) -> SplitReader<'a> {
            let (left, right) = data.split_at(at);

            SplitReader {
                left: left,
                right: right,
            }
        }
    }

    impl<'a> Read for SplitReader<'a> {
        fn read(&mut self, dst: &mut [u8]) -> io::Result<usize> {
            fn copy_bytes_partial(src: &mut &[u8], dst: &mut [u8]) -> usize {
                src.read(dst).unwrap()
            }

            let mut copy_amt = copy_bytes_partial(&mut self.left, dst);

            if copy_amt == 0 {
                copy_amt = copy_bytes_partial(&mut self.right, dst)
            };

            Ok(copy_amt)
        }
    }

    #[test]
    fn test_split_boundary() {
        let logger = ::mock::log_on_panic();

        debug!("Testing boundary (split)");

        let mut buf = String::new();

        // Substitute for `.step_by()` being unstable.
        for split_at in 0 .. TEST_VAL.len(){
            debug!("Testing split at: {}", split_at);

            let src = SplitReader::split(TEST_VAL.as_bytes(), split_at);
            let mut reader = BoundaryReader::from_reader(src, BOUNDARY);
            test_boundary_reader(&mut reader, &mut buf);
        }

        logger.clear();
    }

    fn test_boundary_reader<R: Read>(reader: &mut BoundaryReader<R>, buf: &mut String) {
        buf.clear();

        debug!("Read 1");
        let _ = reader.read_to_string(buf).unwrap();
        assert!(buf.is_empty(), "Buffer not empty: {:?}", buf);
        buf.clear();

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "dashed-value-1");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 3");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "dashed-value-2");
        buf.clear();

        debug!("Consume 3");
        reader.consume_boundary().unwrap();

        debug!("Read 4");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");
    }

    #[test]
    fn test_leading_crlf() {
        let logger = ::mock::log_on_panic();

        let mut body: &[u8] = b"\r\n\r\n--boundary\r\n\
                         asdf1234\
                         \r\n\r\n--boundary--";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\r\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2 (empty)");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");

        logger.clear()
    }

    #[test]
    fn test_trailing_crlf() {
        let logger = ::mock::log_on_panic();

        let mut body: &[u8] = b"--boundary\r\n\
                         asdf1234\
                         \r\n\r\n--boundary\r\n\
                         hjkl5678\r\n--boundary--";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");

        // Repro for https://github.com/abonander/multipart/issues/93
        // These two reads should produce the same buffer
        let buf1 = reader.read_to_boundary().unwrap().to_owned();
        let buf2 = reader.read_to_boundary().unwrap().to_owned();
        assert_eq!(buf1, buf2);

        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\r\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "hjkl5678");
        buf.clear();

        debug!("Consume 3");
        reader.consume_boundary().unwrap();

        debug!("Read 3 (empty)");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");

        logger.clear();
    }

    // https://github.com/abonander/multipart/issues/93#issuecomment-343610587
    #[test]
    fn test_trailing_lflf() {
        let logger = ::mock::log_on_panic();

        let mut body: &[u8] = b"--boundary\r\n\
                         asdf1234\
                         \n\n\r\n--boundary\r\n\
                         hjkl5678\r\n--boundary--";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");

        // same as above
        let buf1 = reader.read_to_boundary().unwrap().to_owned();
        let buf2 = reader.read_to_boundary().unwrap().to_owned();
        assert_eq!(buf1, buf2);

        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\n\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "hjkl5678");
        buf.clear();

        debug!("Consume 3");
        reader.consume_boundary().unwrap();

        debug!("Read 3 (empty)");
        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "");

        logger.clear();
    }

    // https://github.com/abonander/multipart/issues/104
    #[test]
    fn test_unterminated_body() {
        let logger = ::mock::log_on_panic();

        let mut body: &[u8] = b"--boundary\r\n\
                         asdf1234\
                         \n\n\r\n--boundary\r\n\
                         hjkl5678";

        let ref mut buf = String::new();
        let mut reader = BoundaryReader::from_reader(&mut body, BOUNDARY);

        debug!("Consume 1");
        reader.consume_boundary().unwrap();

        debug!("Read 1");

        // same as above
        let buf1 = reader.read_to_boundary().unwrap().to_owned();
        let buf2 = reader.read_to_boundary().unwrap().to_owned();
        assert_eq!(buf1, buf2);

        let _ = reader.read_to_string(buf).unwrap();
        assert_eq!(buf, "asdf1234\n\n");
        buf.clear();

        debug!("Consume 2");
        reader.consume_boundary().unwrap();

        debug!("Read 2");
        let _ = reader.read_to_string(buf).unwrap_err();

        logger.clear();
    }

    #[cfg(feature = "bench")]
    mod bench {
        extern crate test;
        use self::test::Bencher;

        use super::*;

        #[bench]
        fn bench_boundary_reader(b: &mut Bencher) {
            let mut reader = BoundaryReader::new_with_bytes(TEST_VAL.as_bytes(), BOUNDARY);
            let mut buf = String::with_capacity(256);

            b.iter(|| {
                reader.reset();
                test_boundary_reader(&mut reader, &mut buf);
            });
        }
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::fmt;
use std::result::Result;
use std::path::{Path, PathBuf};
use std::convert::From;
use std::convert::Into;

use regex::Regex;

#[derive(Debug)]
#[derive(Clone)]
#[derive(PartialEq)]
#[derive(Eq)]
// #[derive(Display)]
pub enum FileIDType {
    NONE,
    UUID,
}

impl Into<String> for FileIDType {

    fn into(self) -> String {
        let s = match self {
            FileIDType::UUID => "UUID",
            FileIDType::NONE => "",
        };

        String::from(s)
    }

}

impl<'a> From<&'a str> for FileIDType {

    fn from(s: &'a str) -> FileIDType {
        match s {
            "UUID"  => FileIDType::UUID,
            _       => FileIDType::NONE,
        }
    }

}

impl From<String> for FileIDType {

    fn from(s: String) -> FileIDType {
        FileIDType::from(&s[..])
    }

}

#[derive(Clone)]
pub struct FileID {
    id: Option<String>,
    id_type: FileIDType,
}

impl FileID {

    pub fn new(id_type: FileIDType, id: String) -> FileID {
        FileID {
            id: Some(id),
            id_type: id_type,
        }
    }

    pub fn is_valid(&self) -> bool {
        self.id.is_some()
    }

    pub fn get_type(&self) -> FileIDType {
        self.id_type.clone()
    }

}

impl Debug for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Display for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Into<String> for FileID {

    fn into(self) -> String {
        if let Some(id) = self.id {
            id.clone()
        } else {
            String::from("INVALID")
        }
    }

}

impl Into<FileIDType> for FileID {

    fn into(self) -> FileIDType {
        self.id_type.clone()
    }
}

impl From<String> for FileID {

    fn from(s: String) -> FileID {
        FileID::from(&s)
    }

}

impl<'a> From<&'a String> for FileID {

    fn from(string: &'a String) -> FileID {

        let regex = Regex::new(r"([:alnum:]*)-([:upper:]*)-([A-Za-z0-9-_]*)\.(.*)").unwrap();
        let s = string.split("/").last().unwrap_or("");

        debug!("Regex build: {:?}", regex);
        debug!("Matching string: '{}'", s);
        regex.captures(s).and_then(|capts| {
            // first one is the whole string, index 1-N are the matches.
            if capts.len() != 5 {
                debug!("Matches, but not expected number of groups");
                return None;
            }
            debug!("Matches: {}", capts.len());

            let modname     = capts.at(1).unwrap();
            let hashname    = capts.at(2).unwrap();
            let mut hash    = capts.at(3).unwrap();

            debug!("Destructure FilePath to ID:");
            debug!("                  FilePath: {:?}", s);
            debug!("               Module Name: {:?}", modname);
            debug!("                 Hash Name: {:?}", hashname);
            debug!("                      Hash: {:?}", hash);

            let idtype = FileIDType::from(hashname);
            match idtype {
                FileIDType::NONE => hash = "INVALID",
                _ => {},
            }

            Some(FileID::new(idtype, String::from(hash)))
        }).unwrap_or({
            debug!("Did not match");
            FileID {
                id_type: FileIDType::NONE,
                id: None,
            }
        })
    }

}

impl From<PathBuf> for FileID {

    fn from(s: PathBuf) -> FileID {
        unimplemented!()
    }

}

impl<'a> From<&'a PathBuf> for FileID {

    fn from(pb: &'a PathBuf) -> FileID {
        let s = pb.to_str().unwrap_or("");
        FileID::from(String::from(s))
    }

}

pub struct FileIDError {
    summary: String,
    descrip: String,
}

impl FileIDError {

    pub fn new(s: String, d: String) -> FileIDError {
        FileIDError {
            summary: s,
            descrip: d,
        }
    }

}

impl<'a> Error for FileIDError {

    fn description(&self) -> &str {
        &self.summary[..]
    }

    fn cause(&self) -> Option<&Error> {
        None
    }

}

impl<'a> Debug for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'\n{}", self.summary, self.descrip);
        Ok(())
    }

}

impl<'a> Display for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'", self.summary);
        Ok(())
    }

}

pub type FileIDResult = Result<FileID, FileIDError>;

#[cfg(test)]
mod test {

    use super::{FileID, FileIDType};

    #[test]
    fn file_id_from_string() {
        setup_logger();

        let s1 = String::from("/home/user/testmodule-UUID-some-id.imag");
        let s2 = String::from("/home/user/testmodule-UUID-some-id.extension.imag");
        let s3 = String::from("/home/user/testmodule-NOHASH-some-id.imag");

        let id1 = FileID::from(s1);
        let id2 = FileID::from(s2);
        let id3 = FileID::from(s3);

        println!("Id 1 : {:?}", id1);
        println!("Id 2 : {:?}", id2);
        println!("Id 3 : {:?}", id3);

        assert_eq!(FileIDType::UUID, id1.get_type());
        assert_eq!(FileIDType::UUID, id2.get_type());
        assert_eq!(FileIDType::NONE, id3.get_type());

        let f1 : String = id1.into();
        let f2 : String = id2.into();
        let f3 : String = id3.into();

        assert_eq!(String::from("some-id"), f1);
        assert_eq!(String::from("some-id"), f2);
        assert_eq!(String::from("INVALID"), f3);
    }

    fn setup_logger() {
        extern crate log;
        use log::{LogLevelFilter, set_logger};
        use runtime::ImagLogger;

        log::set_logger(|max_log_lvl| {
            let lvl = LogLevelFilter::Debug;
            max_log_lvl.set(lvl);
            Box::new(ImagLogger::new(lvl.to_log_level().unwrap()))
        });
        debug!("Init logger for test");
    }

}

<commit_msg>Fix: FileID::from() should also succeed if we pass a ID as string<commit_after>use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::fmt;
use std::result::Result;
use std::path::{Path, PathBuf};
use std::convert::From;
use std::convert::Into;

use regex::Regex;

#[derive(Debug)]
#[derive(Clone)]
#[derive(PartialEq)]
#[derive(Eq)]
// #[derive(Display)]
pub enum FileIDType {
    NONE,
    UUID,
}

impl Into<String> for FileIDType {

    fn into(self) -> String {
        let s = match self {
            FileIDType::UUID => "UUID",
            FileIDType::NONE => "",
        };

        String::from(s)
    }

}

impl<'a> From<&'a str> for FileIDType {

    fn from(s: &'a str) -> FileIDType {
        match s {
            "UUID"  => FileIDType::UUID,
            _       => FileIDType::NONE,
        }
    }

}

impl From<String> for FileIDType {

    fn from(s: String) -> FileIDType {
        FileIDType::from(&s[..])
    }

}

#[derive(Clone)]
pub struct FileID {
    id: Option<String>,
    id_type: FileIDType,
}

impl FileID {

    pub fn new(id_type: FileIDType, id: String) -> FileID {
        FileID {
            id: Some(id),
            id_type: id_type,
        }
    }

    pub fn is_valid(&self) -> bool {
        self.id.is_some()
    }

    pub fn get_type(&self) -> FileIDType {
        self.id_type.clone()
    }

}

impl Debug for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Display for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Into<String> for FileID {

    fn into(self) -> String {
        if let Some(id) = self.id {
            id.clone()
        } else {
            String::from("INVALID")
        }
    }

}

impl Into<FileIDType> for FileID {

    fn into(self) -> FileIDType {
        self.id_type.clone()
    }
}

impl From<String> for FileID {

    fn from(s: String) -> FileID {
        FileID::from(&s)
    }

}

impl<'a> From<&'a String> for FileID {

    fn from(string: &'a String) -> FileID {
        // we assume that it is an path
        let regex = Regex::new(r"([:alnum:]*)-([:upper:]*)-([A-Za-z0-9-_]*)\.(.*)").unwrap();
        let s = string.split("/").last().unwrap_or("");

        debug!("Regex build: {:?}", regex);
        debug!("Matching string: '{}'", s);
        regex.captures(s).and_then(|capts| {
            // first one is the whole string, index 1-N are the matches.
            if capts.len() != 5 {
                debug!("Matches, but not expected number of groups");
                return None;
            }
            debug!("Matches: {}", capts.len());

            let modname     = capts.at(1).unwrap();
            let hashname    = capts.at(2).unwrap();
            let mut hash    = capts.at(3).unwrap();

            debug!("Destructure FilePath to ID:");
            debug!("                  FilePath: {:?}", s);
            debug!("               Module Name: {:?}", modname);
            debug!("                 Hash Name: {:?}", hashname);
            debug!("                      Hash: {:?}", hash);

            let idtype = FileIDType::from(hashname);
            match idtype {
                FileIDType::NONE => hash = "INVALID",
                _ => {},
            }

            Some(FileID::new(idtype, String::from(hash)))
        }).unwrap_or({
            debug!("Did not match");
            debug!("It is no path, actually. So we assume it is an ID already");
            FileID {
                id_type: FileIDType::NONE,
                id: Some(string.clone()),
            }
        })
    }

}

impl From<PathBuf> for FileID {

    fn from(s: PathBuf) -> FileID {
        unimplemented!()
    }

}

impl<'a> From<&'a PathBuf> for FileID {

    fn from(pb: &'a PathBuf) -> FileID {
        let s = pb.to_str().unwrap_or("");
        FileID::from(String::from(s))
    }

}

pub struct FileIDError {
    summary: String,
    descrip: String,
}

impl FileIDError {

    pub fn new(s: String, d: String) -> FileIDError {
        FileIDError {
            summary: s,
            descrip: d,
        }
    }

}

impl<'a> Error for FileIDError {

    fn description(&self) -> &str {
        &self.summary[..]
    }

    fn cause(&self) -> Option<&Error> {
        None
    }

}

impl<'a> Debug for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'\n{}", self.summary, self.descrip);
        Ok(())
    }

}

impl<'a> Display for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'", self.summary);
        Ok(())
    }

}

pub type FileIDResult = Result<FileID, FileIDError>;

#[cfg(test)]
mod test {

    use super::{FileID, FileIDType};

    #[test]
    fn file_id_from_string() {
        setup_logger();

        let s1 = String::from("/home/user/testmodule-UUID-some-id.imag");
        let s2 = String::from("/home/user/testmodule-UUID-some-id.extension.imag");
        let s3 = String::from("/home/user/testmodule-NOHASH-some-id.imag");

        let id1 = FileID::from(s1);
        let id2 = FileID::from(s2);
        let id3 = FileID::from(s3);

        println!("Id 1 : {:?}", id1);
        println!("Id 2 : {:?}", id2);
        println!("Id 3 : {:?}", id3);

        assert_eq!(FileIDType::UUID, id1.get_type());
        assert_eq!(FileIDType::UUID, id2.get_type());
        assert_eq!(FileIDType::NONE, id3.get_type());

        let f1 : String = id1.into();
        let f2 : String = id2.into();
        let f3 : String = id3.into();

        assert_eq!(String::from("some-id"), f1);
        assert_eq!(String::from("some-id"), f2);
        assert_eq!(String::from("INVALID"), f3);
    }

    fn setup_logger() {
        extern crate log;
        use log::{LogLevelFilter, set_logger};
        use runtime::ImagLogger;

        log::set_logger(|max_log_lvl| {
            let lvl = LogLevelFilter::Debug;
            max_log_lvl.set(lvl);
            Box::new(ImagLogger::new(lvl.to_log_level().unwrap()))
        });
        debug!("Init logger for test");
    }

}

<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

#[feature(macro_rules)];
#[feature(globs)];

#[crate_type = "bin"];
#[no_uv];

extern crate capnp;
extern crate native;

use capnp::message::{MessageReader, MessageBuilder};

pub mod common;

pub mod carsales_capnp;
pub mod carsales;

pub mod catrank_capnp;
pub mod catrank;

pub mod eval_capnp;
pub mod eval;

mod Uncompressed {
    use capnp;
    use std;

    pub fn write<T : std::io::Writer, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        capnp::serialize::write_message(writer, message);
    }

    pub fn write_buffered<T : std::io::Writer, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        capnp::serialize::write_message(writer, message);
    }
/*
    pub fn new_reader<U : std::io::Reader, T>(
        inputStream : &mut U,
        options : capnp::message::ReaderOptions,
        cont : |&mut capnp::message::MessageReader| -> T) -> T {
        capnp::serialize::InputStreamMessageReader::new(inputStream, options, cont)
    }
*/
    pub fn new_buffered_reader<R: capnp::io::BufferedInputStream>(
        inputStream : &mut R,
        options : capnp::message::ReaderOptions) -> capnp::serialize::OwnedSpaceMessageReader {
        capnp::serialize::new_reader(inputStream, options).unwrap()
    }
}

mod Packed {
    use capnp;
    use std;
    use capnp::serialize_packed::{write_packed_message, write_packed_message_unbuffered};

    pub fn write<T : std::io::Writer, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        write_packed_message_unbuffered(writer, message);
    }

    pub fn write_buffered<T : capnp::io::BufferedOutputStream, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        write_packed_message(writer, message);
    }
/*
    pub fn new_reader<U : std::io::Reader, T>(
        inputStream : &mut U,
        options : capnp::message::ReaderOptions,
        cont : |&mut capnp::message::MessageReader| -> T) -> T {
        capnp::serialize::InputStreamMessageReader::new(
            &mut capnp::serialize_packed::PackedInputStream{
                inner : &mut capnp::io::BufferedInputStreamWrapper::new(inputStream)
            },
            options, cont)
    }
*/
    pub fn new_buffered_reader<R:capnp::io::BufferedInputStream>(
        inputStream : &mut R,
        options : capnp::message::ReaderOptions) -> capnp::serialize::OwnedSpaceMessageReader {
        capnp::serialize_packed::new_reader(inputStream, options).unwrap()
    }

}

static SCRATCH_SIZE : uint = 128 * 1024;

pub struct NoScratch;

impl NoScratch {
    fn new_builder(&mut self, _idx : uint) -> capnp::message::MallocMessageBuilder {
        capnp::message::MallocMessageBuilder::new_default()
    }
}

pub struct UseScratch {
    scratch_space : ~[capnp::common::Word]
}

impl UseScratch {
    pub fn new() -> UseScratch {
        UseScratch {
            scratch_space : capnp::common::allocate_zeroed_words(SCRATCH_SIZE * 6)
        }
    }

    fn new_builder<'a>(&mut self, idx : uint) -> capnp::message::ScratchSpaceMallocMessageBuilder<'a> {
        assert!(idx < 6);
        unsafe {
            capnp::message::ScratchSpaceMallocMessageBuilder::new_default(
                std::cast::transmute(
                    std::raw::Slice { data : self.scratch_space.unsafe_ref(idx * SCRATCH_SIZE),
                                      len : SCRATCH_SIZE }))
        }
    }
}

macro_rules! pass_by_object(
    ( $testcase:ident, $reuse:ident, $iters:expr ) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = $reuse.new_builder(0);
                let mut messageRes = $reuse.new_builder(1);

                let request = messageReq.init_root::<$testcase::RequestBuilder>();
                let response = messageRes.init_root::<$testcase::ResponseBuilder>();
                let expected = $testcase::setup_request(&mut rng, request);

                let requestReader = request.as_reader();
                $testcase::handle_request(requestReader, response);

                let responseReader = response.as_reader();
                if !$testcase::check_response(responseReader, expected) {
                    fail!("Incorrect response.");
                }
            }
        });
    )

macro_rules! pass_by_bytes(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr ) => ({
            let mut requestBytes : ~[u8] = std::vec::from_elem(SCRATCH_SIZE * 8, 0u8);
            let mut responseBytes : ~[u8] = std::vec::from_elem(SCRATCH_SIZE * 8, 0u8);
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = $reuse.new_builder(0);
                let mut messageRes = $reuse.new_builder(1);

                let request = messageReq.init_root::<$testcase::RequestBuilder>();
                let response = messageRes.init_root::<$testcase::ResponseBuilder>();
                let expected = $testcase::setup_request(&mut rng, request);

                {
                    let mut writer = capnp::io::ArrayOutputStream::new(requestBytes);
                    $compression::write_buffered(&mut writer, &messageReq)
                }

                let messageReader = $compression::new_buffered_reader(
                    &mut capnp::io::ArrayInputStream::new(requestBytes),
                    capnp::message::DefaultReaderOptions);

                let requestReader : $testcase::RequestReader = messageReader.get_root();
                $testcase::handle_request(requestReader, response);

                {
                    let mut writer = capnp::io::ArrayOutputStream::new(responseBytes);
                    $compression::write_buffered(&mut writer, &messageRes)
                }

                let messageReader = $compression::new_buffered_reader(
                    &mut capnp::io::ArrayInputStream::new(responseBytes),
                    capnp::message::DefaultReaderOptions);

                let responseReader : $testcase::ResponseReader = messageReader.get_root();
                if !$testcase::check_response(responseReader, expected) {
                    fail!("Incorrect response.");
                }
            }
        });
    )

macro_rules! server(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr, $input:expr, $output:expr) => ({
            let mut outBuffered = capnp::io::BufferedOutputStreamWrapper::new(&mut $output);
            let mut inBuffered = capnp::io::BufferedInputStreamWrapper::new(&mut $input);
            for _ in range(0, $iters) {
                let mut messageRes = $reuse.new_builder(0);

                let response = messageRes.init_root::<$testcase::ResponseBuilder>();
                let messageReader = $compression::new_buffered_reader(
                    &mut inBuffered,
                    capnp::message::DefaultReaderOptions);
                let requestReader : $testcase::RequestReader = messageReader.get_root();
                $testcase::handle_request(requestReader, response);

                $compression::write_buffered(&mut outBuffered, &messageRes);
                outBuffered.flush().unwrap();
            }
        });
    )

macro_rules! sync_client(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr) => ({
            let mut outStream = std::io::stdout();
            let mut inStream = std::io::stdin();
            let mut inBuffered = capnp::io::BufferedInputStreamWrapper::new(&mut inStream);
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = $reuse.new_builder(0);

                let request = messageReq.init_root::<$testcase::RequestBuilder>();

                let expected = $testcase::setup_request(&mut rng, request);
                $compression::write(&mut outStream, &messageReq);

                let messageReader = $compression::new_buffered_reader(
                    &mut inBuffered,
                    capnp::message::DefaultReaderOptions);
                let responseReader : $testcase::ResponseReader = messageReader.get_root();
                assert!($testcase::check_response(responseReader, expected));

            }
        });
    )

macro_rules! pass_by_pipe(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr) => ({
            use std::io::process;

            let mut args = std::os::args();
            args[2] = ~"client";

            let io = [process::CreatePipe(true, false), // stdin
                      process::CreatePipe(false, true), // stdout
                      process::Ignored];

            let config = process::ProcessConfig {
                program: args[0].as_slice(),
                args: args.slice(1, args.len()),
                io : io,
                .. process::ProcessConfig::new()
            };
            match process::Process::new(config) {
                Ok(ref mut p) => {
                    p.io.pop();
                    let mut childStdOut = p.io.pop().unwrap().unwrap();
                    let mut childStdIn = p.io.pop().unwrap().unwrap();

                    server!($testcase, $reuse, $compression, $iters, childStdOut, childStdIn);
                    println!("{}", p.wait());
                }
                Err(e) => {
                    println!("could not start process: {}", e);
                }
            }
        });
    )

macro_rules! do_testcase(
    ( $testcase:ident, $mode:expr, $reuse:ident, $compression:ident, $iters:expr ) => ({
            match $mode {
                ~"object" => pass_by_object!($testcase, $reuse, $iters),
                ~"bytes" => pass_by_bytes!($testcase, $reuse, $compression, $iters),
                ~"client" => sync_client!($testcase, $reuse, $compression, $iters),
                ~"server" => {
                    let mut input = std::io::stdin();
                    let mut output = std::io::stdout();
                    server!($testcase, $reuse, $compression, $iters, input, output)
                }
                ~"pipe" => pass_by_pipe!($testcase, $reuse, $compression, $iters),
                s => fail!("unrecognized mode: {}", s)
            }
        });
    )

macro_rules! do_testcase1(
    ( $testcase:expr, $mode:expr, $reuse:ident, $compression:ident, $iters:expr) => ({
            match $testcase {
                ~"carsales" => do_testcase!(carsales, $mode, $reuse, $compression, $iters),
                ~"catrank" => do_testcase!(catrank, $mode, $reuse, $compression, $iters),
                ~"eval" => do_testcase!(eval, $mode, $reuse, $compression, $iters),
                s => fail!("unrecognized test case: {}", s)
            }
        });
    )

macro_rules! do_testcase2(
    ( $testcase:expr, $mode:expr, $reuse:expr, $compression:ident, $iters:expr) => ({
            match $reuse {
                ~"no-reuse" => {
                    let mut scratch = NoScratch;
                    do_testcase1!($testcase, $mode, scratch, $compression, $iters)
                }
                ~"reuse" => {
                    let mut scratch = UseScratch::new();
                    do_testcase1!($testcase, $mode, scratch, $compression, $iters)
                }
                s => fail!("unrecognized reuse option: {}", s)
            }
        });
    )

#[start]
pub fn start (argc : int, argv: **u8) -> int {
    native::start(argc, argv, proc () {
            main();
        })
}

pub fn main() {
    let args = std::os::args();

    if args.len() != 6 {
        println!("USAGE: {} CASE MODE REUSE COMPRESSION ITERATION_COUNT", args[0]);
        std::os::set_exit_status(1);
        return;
    }

    let iters = match from_str::<u64>(args[5]) {
        Some (n) => n,
        None => {
            println!("Could not parse a u64 from: {}", args[5]);
            std::os::set_exit_status(1);
            return;
        }
    };

    match args[4] {
        ~"none" => do_testcase2!(args[1], args[2],  args[3], Uncompressed, iters),
        ~"packed" => do_testcase2!(args[1], args[2], args[3], Packed, iters),
        s => fail!("unrecognized compression: {}", s)
    }
}
<commit_msg>update for new std::io::process API<commit_after>/*
 * Copyright (c) 2013-2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

#[feature(macro_rules)];
#[feature(globs)];

#[crate_type = "bin"];
#[no_uv];

extern crate capnp;
extern crate native;

use capnp::message::{MessageReader, MessageBuilder};

pub mod common;

pub mod carsales_capnp;
pub mod carsales;

pub mod catrank_capnp;
pub mod catrank;

pub mod eval_capnp;
pub mod eval;

mod Uncompressed {
    use capnp;
    use std;

    pub fn write<T : std::io::Writer, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        capnp::serialize::write_message(writer, message);
    }

    pub fn write_buffered<T : std::io::Writer, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        capnp::serialize::write_message(writer, message);
    }

    pub fn new_buffered_reader<R: capnp::io::BufferedInputStream>(
        inputStream : &mut R,
        options : capnp::message::ReaderOptions) -> capnp::serialize::OwnedSpaceMessageReader {
        capnp::serialize::new_reader(inputStream, options).unwrap()
    }
}

mod Packed {
    use capnp;
    use std;
    use capnp::serialize_packed::{write_packed_message, write_packed_message_unbuffered};

    pub fn write<T : std::io::Writer, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        write_packed_message_unbuffered(writer, message);
    }

    pub fn write_buffered<T : capnp::io::BufferedOutputStream, U : capnp::message::MessageBuilder>(
        writer: &mut T,
        message: &U) {
        write_packed_message(writer, message);
    }

    pub fn new_buffered_reader<R:capnp::io::BufferedInputStream>(
        inputStream : &mut R,
        options : capnp::message::ReaderOptions) -> capnp::serialize::OwnedSpaceMessageReader {
        capnp::serialize_packed::new_reader(inputStream, options).unwrap()
    }

}

static SCRATCH_SIZE : uint = 128 * 1024;

pub struct NoScratch;

impl NoScratch {
    fn new_builder(&mut self, _idx : uint) -> capnp::message::MallocMessageBuilder {
        capnp::message::MallocMessageBuilder::new_default()
    }
}

pub struct UseScratch {
    scratch_space : ~[capnp::common::Word]
}

impl UseScratch {
    pub fn new() -> UseScratch {
        UseScratch {
            scratch_space : capnp::common::allocate_zeroed_words(SCRATCH_SIZE * 6)
        }
    }

    fn new_builder<'a>(&mut self, idx : uint) -> capnp::message::ScratchSpaceMallocMessageBuilder<'a> {
        assert!(idx < 6);
        unsafe {
            capnp::message::ScratchSpaceMallocMessageBuilder::new_default(
                std::cast::transmute(
                    std::raw::Slice { data : self.scratch_space.unsafe_ref(idx * SCRATCH_SIZE),
                                      len : SCRATCH_SIZE }))
        }
    }
}

macro_rules! pass_by_object(
    ( $testcase:ident, $reuse:ident, $iters:expr ) => ({
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = $reuse.new_builder(0);
                let mut messageRes = $reuse.new_builder(1);

                let request = messageReq.init_root::<$testcase::RequestBuilder>();
                let response = messageRes.init_root::<$testcase::ResponseBuilder>();
                let expected = $testcase::setup_request(&mut rng, request);

                let requestReader = request.as_reader();
                $testcase::handle_request(requestReader, response);

                let responseReader = response.as_reader();
                if !$testcase::check_response(responseReader, expected) {
                    fail!("Incorrect response.");
                }
            }
        });
    )

macro_rules! pass_by_bytes(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr ) => ({
            let mut requestBytes : ~[u8] = std::vec::from_elem(SCRATCH_SIZE * 8, 0u8);
            let mut responseBytes : ~[u8] = std::vec::from_elem(SCRATCH_SIZE * 8, 0u8);
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = $reuse.new_builder(0);
                let mut messageRes = $reuse.new_builder(1);

                let request = messageReq.init_root::<$testcase::RequestBuilder>();
                let response = messageRes.init_root::<$testcase::ResponseBuilder>();
                let expected = $testcase::setup_request(&mut rng, request);

                {
                    let mut writer = capnp::io::ArrayOutputStream::new(requestBytes);
                    $compression::write_buffered(&mut writer, &messageReq)
                }

                let messageReader = $compression::new_buffered_reader(
                    &mut capnp::io::ArrayInputStream::new(requestBytes),
                    capnp::message::DefaultReaderOptions);

                let requestReader : $testcase::RequestReader = messageReader.get_root();
                $testcase::handle_request(requestReader, response);

                {
                    let mut writer = capnp::io::ArrayOutputStream::new(responseBytes);
                    $compression::write_buffered(&mut writer, &messageRes)
                }

                let messageReader = $compression::new_buffered_reader(
                    &mut capnp::io::ArrayInputStream::new(responseBytes),
                    capnp::message::DefaultReaderOptions);

                let responseReader : $testcase::ResponseReader = messageReader.get_root();
                if !$testcase::check_response(responseReader, expected) {
                    fail!("Incorrect response.");
                }
            }
        });
    )

macro_rules! server(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr, $input:expr, $output:expr) => ({
            let mut outBuffered = capnp::io::BufferedOutputStreamWrapper::new(&mut $output);
            let mut inBuffered = capnp::io::BufferedInputStreamWrapper::new(&mut $input);
            for _ in range(0, $iters) {
                let mut messageRes = $reuse.new_builder(0);

                let response = messageRes.init_root::<$testcase::ResponseBuilder>();
                let messageReader = $compression::new_buffered_reader(
                    &mut inBuffered,
                    capnp::message::DefaultReaderOptions);
                let requestReader : $testcase::RequestReader = messageReader.get_root();
                $testcase::handle_request(requestReader, response);

                $compression::write_buffered(&mut outBuffered, &messageRes);
                outBuffered.flush().unwrap();
            }
        });
    )

macro_rules! sync_client(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr) => ({
            let mut outStream = std::io::stdout();
            let mut inStream = std::io::stdin();
            let mut inBuffered = capnp::io::BufferedInputStreamWrapper::new(&mut inStream);
            let mut rng = common::FastRand::new();
            for _ in range(0, $iters) {
                let mut messageReq = $reuse.new_builder(0);

                let request = messageReq.init_root::<$testcase::RequestBuilder>();

                let expected = $testcase::setup_request(&mut rng, request);
                $compression::write(&mut outStream, &messageReq);

                let messageReader = $compression::new_buffered_reader(
                    &mut inBuffered,
                    capnp::message::DefaultReaderOptions);
                let responseReader : $testcase::ResponseReader = messageReader.get_root();
                assert!($testcase::check_response(responseReader, expected));

            }
        });
    )

macro_rules! pass_by_pipe(
    ( $testcase:ident, $reuse:ident, $compression:ident, $iters:expr) => ({
            use std::io::process;

            let mut args = std::os::args();
            args[2] = ~"client";

            let config = process::ProcessConfig {
                program: args[0].as_slice(),
                args: args.slice(1, args.len()),
                stdin : process::CreatePipe(true, false),
                stdout : process::CreatePipe(false, true),
                stderr : process::Ignored,
                .. process::ProcessConfig::new()
            };
            match process::Process::configure(config) {
                Ok(ref mut p) => {
                    let mut childStdOut = p.stdout.take().unwrap();
                    let mut childStdIn = p.stdin.take().unwrap();

                    server!($testcase, $reuse, $compression, $iters, childStdOut, childStdIn);
                    println!("{}", p.wait());
                }
                Err(e) => {
                    println!("could not start process: {}", e);
                }
            }
        });
    )

macro_rules! do_testcase(
    ( $testcase:ident, $mode:expr, $reuse:ident, $compression:ident, $iters:expr ) => ({
            match $mode {
                ~"object" => pass_by_object!($testcase, $reuse, $iters),
                ~"bytes" => pass_by_bytes!($testcase, $reuse, $compression, $iters),
                ~"client" => sync_client!($testcase, $reuse, $compression, $iters),
                ~"server" => {
                    let mut input = std::io::stdin();
                    let mut output = std::io::stdout();
                    server!($testcase, $reuse, $compression, $iters, input, output)
                }
                ~"pipe" => pass_by_pipe!($testcase, $reuse, $compression, $iters),
                s => fail!("unrecognized mode: {}", s)
            }
        });
    )

macro_rules! do_testcase1(
    ( $testcase:expr, $mode:expr, $reuse:ident, $compression:ident, $iters:expr) => ({
            match $testcase {
                ~"carsales" => do_testcase!(carsales, $mode, $reuse, $compression, $iters),
                ~"catrank" => do_testcase!(catrank, $mode, $reuse, $compression, $iters),
                ~"eval" => do_testcase!(eval, $mode, $reuse, $compression, $iters),
                s => fail!("unrecognized test case: {}", s)
            }
        });
    )

macro_rules! do_testcase2(
    ( $testcase:expr, $mode:expr, $reuse:expr, $compression:ident, $iters:expr) => ({
            match $reuse {
                ~"no-reuse" => {
                    let mut scratch = NoScratch;
                    do_testcase1!($testcase, $mode, scratch, $compression, $iters)
                }
                ~"reuse" => {
                    let mut scratch = UseScratch::new();
                    do_testcase1!($testcase, $mode, scratch, $compression, $iters)
                }
                s => fail!("unrecognized reuse option: {}", s)
            }
        });
    )

#[start]
pub fn start (argc : int, argv: **u8) -> int {
    native::start(argc, argv, proc () {
            main();
        })
}

pub fn main() {
    let args = std::os::args();

    if args.len() != 6 {
        println!("USAGE: {} CASE MODE REUSE COMPRESSION ITERATION_COUNT", args[0]);
        std::os::set_exit_status(1);
        return;
    }

    let iters = match from_str::<u64>(args[5]) {
        Some (n) => n,
        None => {
            println!("Could not parse a u64 from: {}", args[5]);
            std::os::set_exit_status(1);
            return;
        }
    };

    match args[4] {
        ~"none" => do_testcase2!(args[1], args[2],  args[3], Uncompressed, iters),
        ~"packed" => do_testcase2!(args[1], args[2], args[3], Packed, iters),
        s => fail!("unrecognized compression: {}", s)
    }
}
<|endoftext|>