<commit_before>use std::{io, path, mem, fs};
use std::io::Read;
use std::time::Duration;
use libc::{c_ulong, c_ushort, c_uint, c_long, c_schar, c_char};
use libc::statvfs;
use data::*;
use super::common::*;
use super::unix;
use nom::{digit, not_line_ending, space, is_space};
use std::str;
use std::path::Path;

fn read_file(path: &str) -> io::Result<String> {
    let mut s = String::new();
    fs::File::open(path)
        .and_then(|mut f| f.read_to_string(&mut s))
        .map(|_| s)
}

fn value_from_file<T: str::FromStr>(path: &str) -> io::Result<T> {
    try!(read_file(path))
        .trim_right_matches("\n")
        .parse()
        .and_then(|n| Ok(n))
        .or_else(|_| {
            Err(io::Error::new(io::ErrorKind::Other,
                               format!("File: \"{}\" doesn't contain an int value", &path)))
        })
}

fn capacity(charge_full: i32, charge_now: i32) -> f32 {
    charge_now as f32 / charge_full as f32
}

fn time(on_ac: bool, charge_full: i32, charge_now: i32, current_now: i32) -> Duration {
    if current_now != 0 {
        if on_ac {
            // Charge time
            Duration::from_secs((charge_full - charge_now).abs() as u64 * 3600u64 / current_now as u64)
        }
        else {
            // Discharge time
            Duration::from_secs(charge_now as u64 * 3600u64 / current_now as u64)
        }
    } else {
        Duration::new(0, 0)
    }
}

/// Parse an unsigned integer out of a string, surrounded by whitespace
named!(
    usize_s<usize>,
    ws!(map_res!(
        map_res!(digit, str::from_utf8),
        str::FromStr::from_str
    ))
);

/// Parse `cpuX`, where X is a number
named!(proc_stat_cpu_prefix<()>, do_parse!(tag!("cpu") >> digit >> ()));

/// Parse a `/proc/stat` CPU line into a `CpuTime` struct
named!(
    proc_stat_cpu_time<CpuTime>,
    do_parse!(
        ws!(proc_stat_cpu_prefix) >>
        user: usize_s >>
        nice: usize_s >>
        system: usize_s >>
        idle: usize_s >>
        iowait: usize_s >>
        irq: usize_s >>
            (CpuTime {
                 user: user,
                 nice: nice,
                 system: system,
                 idle: idle,
                 interrupt: irq,
                 other: iowait,
             })
    )
);

/// Parse the top CPU load aggregate line of `/proc/stat`
named!(proc_stat_cpu_aggregate<()>, do_parse!(tag!("cpu") >> space >> ()));

/// Parse `/proc/stat` to extract per-CPU loads
named!(
    proc_stat_cpu_times<Vec<CpuTime>>,
    do_parse!(
        ws!(flat_map!(not_line_ending, proc_stat_cpu_aggregate)) >>
        result: many1!(ws!(flat_map!(not_line_ending, proc_stat_cpu_time))) >>
        (result)
    )
);

/// Get the current per-CPU `CpuTime` statistics
fn cpu_time() -> io::Result<Vec<CpuTime>> {
    read_file("/proc/stat").and_then(|data| {
        proc_stat_cpu_times(data.as_bytes()).to_result().map_err(
            |err| {
                io::Error::new(io::ErrorKind::InvalidData, err)
            },
        )
    })
}

/// Parse a `/proc/meminfo` line into (key, ByteSize)
named!(
    proc_meminfo_line<(String, ByteSize)>,
    complete!(do_parse!(
        key: flat_map!(take_until!(":"), parse_to!(String)) >>
        tag!(":") >>
        value: usize_s >>
        ws!(tag!("kB")) >>
        ((key, ByteSize::kib(value)))
    ))
);

/// Optionally parse a `/proc/meminfo` line`
named!(
    proc_meminfo_line_opt<Option<(String, ByteSize)>>,
    opt!(proc_meminfo_line)
);

/// Parse `/proc/meminfo` into a hashmap
named!(
    proc_meminfo<BTreeMap<String, ByteSize>>,
    fold_many0!(
        ws!(flat_map!(not_line_ending, proc_meminfo_line_opt)),
        BTreeMap::new(),
        |mut map: BTreeMap<String, ByteSize>, opt| {
            if let Some((key, val)) = opt {
                map.insert(key, val);
            }
            map
        }
    )
);

/// Get memory statistics
fn memory_stats() -> io::Result<BTreeMap<String, ByteSize>> {
    read_file("/proc/meminfo").and_then(|data| {
        proc_meminfo(data.as_bytes()).to_result().map_err(|err| {
            io::Error::new(io::ErrorKind::InvalidData, err)
        })
    })
}

/// Parse a single word
named!(word_s<String>, flat_map!(
    map_res!(take_till!(is_space), str::from_utf8),
    parse_to!(String)
));

/// `/proc/mounts` data
struct ProcMountsData {
    source: String,
    target: String,
    fstype: String,
}

/// Parse a `/proc/mounts` line to get a mountpoint
named!(
    proc_mounts_line<ProcMountsData>,
    ws!(do_parse!(
        source: word_s >>
        target: word_s >>
        fstype: word_s >>
        (ProcMountsData {
            source: source,
            target: target,
            fstype: fstype,
        })
    ))
);

/// Parse `/proc/mounts` to get a list of mountpoints
named!(
    proc_mounts<Vec<ProcMountsData>>,
    many1!(ws!(flat_map!(not_line_ending, proc_mounts_line)))
);

/// Stat a mountpoint to gather filesystem statistics
fn stat_mount(mount: ProcMountsData) -> io::Result<Filesystem> {
    let mut info: statvfs = unsafe { mem::zeroed() };
    let result = unsafe { statvfs(mount.target.as_ptr() as *const c_char, &mut info) };
    match result {
        0 => Ok(Filesystem {
            files: info.f_files as usize,
            free: ByteSize::b(info.f_bfree as usize * info.f_bsize as usize),
            avail: ByteSize::b(info.f_bavail as usize * info.f_bsize as usize),
            total: ByteSize::b(info.f_blocks as usize * info.f_bsize as usize),
            name_max: info.f_namemax as usize,
            fs_type: mount.fstype,
            fs_mounted_from: mount.source,
            fs_mounted_on: mount.target,
        }),
        _ => Err(io::Error::last_os_error()),
    }
}

pub struct PlatformImpl;

/// An implementation of `Platform` for Linux.
/// See `Platform` for documentation.
impl Platform for PlatformImpl {
    #[inline(always)]
    fn new() -> Self {
        PlatformImpl
    }

    fn cpu_load(&self) -> io::Result<DelayedMeasurement<Vec<CPULoad>>> {
        cpu_time().map(|times| {
            DelayedMeasurement::new(Box::new(move || {
                cpu_time().map(|delay_times| {
                    delay_times
                        .iter()
                        .zip(times.iter())
                        .map(|(now, prev)| (*now - prev).to_cpuload())
                        .collect::<Vec<_>>()
                })
            }))
        })
    }

    fn load_average(&self) -> io::Result<LoadAverage> {
        unix::load_average()
    }

    fn memory(&self) -> io::Result<Memory> {
        memory_stats()
            .or_else(|_| {
                // If there's no procfs, e.g. in a chroot without mounting it or something
                let mut meminfo = BTreeMap::new();
                let mut info: sysinfo = unsafe { mem::zeroed() };
                unsafe { sysinfo(&mut info) };
                let unit = info.mem_unit as usize;
                meminfo.insert(
                    "MemTotal".to_owned(),
                    ByteSize::b(info.totalram as usize * unit),
                );
                meminfo.insert(
                    "MemFree".to_owned(),
                    ByteSize::b(info.freeram as usize * unit),
                );
                meminfo.insert(
                    "Shmem".to_owned(),
                    ByteSize::b(info.sharedram as usize * unit),
                );
                meminfo.insert(
                    "Buffers".to_owned(),
                    ByteSize::b(info.bufferram as usize * unit),
                );
                Ok(meminfo)
            })
            .map(|meminfo| {
                Memory {
                    total: meminfo.get("MemTotal").map(|x| x.clone()).unwrap_or(
                        ByteSize::b(0),
                    ),
                    free: meminfo.get("MemFree").map(|x| x.clone()).unwrap_or(
                        ByteSize::b(0),
                    ) +
                        meminfo.get("Buffers").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ) +
                        meminfo.get("Cached").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ) +
                        meminfo.get("SReclaimable").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ) -
                        meminfo.get("Shmem").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ),
                    platform_memory: PlatformMemory { meminfo: meminfo },
                }
            })
    }

    fn uptime(&self) -> io::Result<Duration> {
        let mut info: sysinfo = unsafe { mem::zeroed() };
        unsafe { sysinfo(&mut info) };
        Ok(Duration::from_secs(info.uptime as u64))
    }

    fn battery_life(&self) -> io::Result<BatteryLife> {
        let dir = "/sys/class/power_supply";
        let entries = try!(fs::read_dir(&dir));
        let mut full = 0;
        let mut now = 0;
        let mut current = 0;
        for e in entries {
            let p = e.unwrap().path();
            let s = p.to_str().unwrap();
            let f = p.file_name().unwrap().to_str().unwrap();
            if f.len() > 3 {
                if f.split_at(3).0 == "BAT" {
                    full += try!(
                        value_from_file::<i32>(&(s.to_string() + "/energy_full"))
                            .or_else(|_| value_from_file::<i32>(&(s.to_string() + "/charge_full")))
                    );
                    now += try!(
                        value_from_file::<i32>(&(s.to_string() + "/energy_now"))
                            .or_else(|_| value_from_file::<i32>(&(s.to_string() + "/charge_now")))
                    );
                    current += try!(
                        value_from_file::<i32>(&(s.to_string() + "/power_now"))
                            .or_else(|_| value_from_file::<i32>(&(s.to_string() + "/current_now")))
                    );
                }
            }
        }
        if full != 0 {
            let on_ac =
                match self.on_ac_power() {
                    Ok(true) => true,
                    _ => false,
                };
            Ok(BatteryLife {
                remaining_capacity: capacity(full, now),
                remaining_time: time(on_ac, full, now, current),
            })
        } else {
            Err(io::Error::new(io::ErrorKind::Other, "Missing battery information"))
        }
    }

    fn on_ac_power(&self) -> io::Result<bool> {
        value_from_file::<i32>("/sys/class/power_supply/AC/online").map(|v| v == 1)
    }

    fn mounts(&self) -> io::Result<Vec<Filesystem>> {
        read_file("/proc/mounts")
            .and_then(|data| {
                proc_mounts(data.as_bytes()).to_result().map_err(|err| {
                    io::Error::new(io::ErrorKind::InvalidData, err)
                })
            })
            .map(|mounts| {
                mounts
                    .into_iter()
                    .filter_map(|mount| stat_mount(mount).ok())
                    .collect()
            })
    }

    fn mount_at<P: AsRef<path::Path>>(&self, path: P) -> io::Result<Filesystem> {
        read_file("/proc/mounts")
            .and_then(|data| {
                proc_mounts(data.as_bytes()).to_result().map_err(|err| {
                    io::Error::new(io::ErrorKind::InvalidData, err)
                })
            })
            .and_then(|mounts| {
                mounts
                    .into_iter()
                    .find(|mount| Path::new(&mount.target) == path.as_ref())
                    .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "No such mount"))
            })
            .and_then(stat_mount)
    }

    fn networks(&self) -> io::Result<BTreeMap<String, Network>> {
        unix::networks()
    }

    fn network_stats(&self, interface: &str) -> io::Result<NetworkStats> {
        let path_root: String = ("/sys/class/net/".to_string() + interface) + "/statistics/";
        let stats_file = |file: &str| (&path_root).to_string() + file;

        let rx_bytes: usize = try!(value_from_file::<usize>(&stats_file("rx_bytes")));
        let tx_bytes: usize = try!(value_from_file::<usize>(&stats_file("tx_bytes")));
        let rx_packets: usize = try!(value_from_file::<usize>(&stats_file("rx_packets")));
        let tx_packets: usize = try!(value_from_file::<usize>(&stats_file("tx_packets")));
        let rx_errors: usize = try!(value_from_file::<usize>(&stats_file("rx_errors")));
        let tx_errors: usize = try!(value_from_file::<usize>(&stats_file("tx_errors")));

        Ok(NetworkStats {
            rx_bytes,
            tx_bytes,
            rx_packets,
            tx_packets,
            rx_errors,
            tx_errors,
        })
    }

    fn cpu_temp(&self) -> io::Result<f32> {
        read_file("/sys/class/thermal/thermal_zone0/temp")
            .and_then(|data| match data.parse::<f32>() {
                Ok(x) => Ok(x),
                Err(_) => Err(io::Error::new(io::ErrorKind::Other, "Could not parse float")),
            })
            .map(|num| num / 1000.0)
    }
}

#[repr(C)]
#[derive(Debug)]
struct sysinfo {
    uptime: c_long,
    loads: [c_ulong; 3],
    totalram: c_ulong,
    freeram: c_ulong,
    sharedram: c_ulong,
    bufferram: c_ulong,
    totalswap: c_ulong,
    freeswap: c_ulong,
    procs: c_ushort,
    totalhigh: c_ulong,
    freehigh: c_ulong,
    mem_unit: c_uint,
    padding: [c_schar; 8],
}

#[link(name = "c")]
extern "C" {
    fn sysinfo(info: *mut sysinfo);
}
<commit_msg>Expose swap stats on failed meminfo<commit_after>use std::{io, path, mem, fs};
use std::io::Read;
use std::time::Duration;
use libc::{c_ulong, c_ushort, c_uint, c_long, c_schar, c_char};
use libc::statvfs;
use data::*;
use super::common::*;
use super::unix;
use nom::{digit, not_line_ending, space, is_space};
use std::str;
use std::path::Path;

fn read_file(path: &str) -> io::Result<String> {
    let mut s = String::new();
    fs::File::open(path)
        .and_then(|mut f| f.read_to_string(&mut s))
        .map(|_| s)
}

fn value_from_file<T: str::FromStr>(path: &str) -> io::Result<T> {
    try!(read_file(path))
        .trim_right_matches("\n")
        .parse()
        .and_then(|n| Ok(n))
        .or_else(|_| {
            Err(io::Error::new(io::ErrorKind::Other,
                               format!("File: \"{}\" doesn't contain an int value", &path)))
        })
}

fn capacity(charge_full: i32, charge_now: i32) -> f32 {
    charge_now as f32 / charge_full as f32
}

fn time(on_ac: bool, charge_full: i32, charge_now: i32, current_now: i32) -> Duration {
    if current_now != 0 {
        if on_ac {
            // Charge time
            Duration::from_secs((charge_full - charge_now).abs() as u64 * 3600u64 / current_now as u64)
        }
        else {
            // Discharge time
            Duration::from_secs(charge_now as u64 * 3600u64 / current_now as u64)
        }
    } else {
        Duration::new(0, 0)
    }
}

/// Parse an unsigned integer out of a string, surrounded by whitespace
named!(
    usize_s<usize>,
    ws!(map_res!(
        map_res!(digit, str::from_utf8),
        str::FromStr::from_str
    ))
);

/// Parse `cpuX`, where X is a number
named!(proc_stat_cpu_prefix<()>, do_parse!(tag!("cpu") >> digit >> ()));

/// Parse a `/proc/stat` CPU line into a `CpuTime` struct
named!(
    proc_stat_cpu_time<CpuTime>,
    do_parse!(
        ws!(proc_stat_cpu_prefix) >>
        user: usize_s >>
        nice: usize_s >>
        system: usize_s >>
        idle: usize_s >>
        iowait: usize_s >>
        irq: usize_s >>
            (CpuTime {
                 user: user,
                 nice: nice,
                 system: system,
                 idle: idle,
                 interrupt: irq,
                 other: iowait,
             })
    )
);

/// Parse the top CPU load aggregate line of `/proc/stat`
named!(proc_stat_cpu_aggregate<()>, do_parse!(tag!("cpu") >> space >> ()));

/// Parse `/proc/stat` to extract per-CPU loads
named!(
    proc_stat_cpu_times<Vec<CpuTime>>,
    do_parse!(
        ws!(flat_map!(not_line_ending, proc_stat_cpu_aggregate)) >>
        result: many1!(ws!(flat_map!(not_line_ending, proc_stat_cpu_time))) >>
        (result)
    )
);

/// Get the current per-CPU `CpuTime` statistics
fn cpu_time() -> io::Result<Vec<CpuTime>> {
    read_file("/proc/stat").and_then(|data| {
        proc_stat_cpu_times(data.as_bytes()).to_result().map_err(
            |err| {
                io::Error::new(io::ErrorKind::InvalidData, err)
            },
        )
    })
}

/// Parse a `/proc/meminfo` line into (key, ByteSize)
named!(
    proc_meminfo_line<(String, ByteSize)>,
    complete!(do_parse!(
        key: flat_map!(take_until!(":"), parse_to!(String)) >>
        tag!(":") >>
        value: usize_s >>
        ws!(tag!("kB")) >>
        ((key, ByteSize::kib(value)))
    ))
);

/// Optionally parse a `/proc/meminfo` line`
named!(
    proc_meminfo_line_opt<Option<(String, ByteSize)>>,
    opt!(proc_meminfo_line)
);

/// Parse `/proc/meminfo` into a hashmap
named!(
    proc_meminfo<BTreeMap<String, ByteSize>>,
    fold_many0!(
        ws!(flat_map!(not_line_ending, proc_meminfo_line_opt)),
        BTreeMap::new(),
        |mut map: BTreeMap<String, ByteSize>, opt| {
            if let Some((key, val)) = opt {
                map.insert(key, val);
            }
            map
        }
    )
);

/// Get memory statistics
fn memory_stats() -> io::Result<BTreeMap<String, ByteSize>> {
    read_file("/proc/meminfo").and_then(|data| {
        proc_meminfo(data.as_bytes()).to_result().map_err(|err| {
            io::Error::new(io::ErrorKind::InvalidData, err)
        })
    })
}

/// Parse a single word
named!(word_s<String>, flat_map!(
    map_res!(take_till!(is_space), str::from_utf8),
    parse_to!(String)
));

/// `/proc/mounts` data
struct ProcMountsData {
    source: String,
    target: String,
    fstype: String,
}

/// Parse a `/proc/mounts` line to get a mountpoint
named!(
    proc_mounts_line<ProcMountsData>,
    ws!(do_parse!(
        source: word_s >>
        target: word_s >>
        fstype: word_s >>
        (ProcMountsData {
            source: source,
            target: target,
            fstype: fstype,
        })
    ))
);

/// Parse `/proc/mounts` to get a list of mountpoints
named!(
    proc_mounts<Vec<ProcMountsData>>,
    many1!(ws!(flat_map!(not_line_ending, proc_mounts_line)))
);

/// Stat a mountpoint to gather filesystem statistics
fn stat_mount(mount: ProcMountsData) -> io::Result<Filesystem> {
    let mut info: statvfs = unsafe { mem::zeroed() };
    let result = unsafe { statvfs(mount.target.as_ptr() as *const c_char, &mut info) };
    match result {
        0 => Ok(Filesystem {
            files: info.f_files as usize,
            free: ByteSize::b(info.f_bfree as usize * info.f_bsize as usize),
            avail: ByteSize::b(info.f_bavail as usize * info.f_bsize as usize),
            total: ByteSize::b(info.f_blocks as usize * info.f_bsize as usize),
            name_max: info.f_namemax as usize,
            fs_type: mount.fstype,
            fs_mounted_from: mount.source,
            fs_mounted_on: mount.target,
        }),
        _ => Err(io::Error::last_os_error()),
    }
}

pub struct PlatformImpl;

/// An implementation of `Platform` for Linux.
/// See `Platform` for documentation.
impl Platform for PlatformImpl {
    #[inline(always)]
    fn new() -> Self {
        PlatformImpl
    }

    fn cpu_load(&self) -> io::Result<DelayedMeasurement<Vec<CPULoad>>> {
        cpu_time().map(|times| {
            DelayedMeasurement::new(Box::new(move || {
                cpu_time().map(|delay_times| {
                    delay_times
                        .iter()
                        .zip(times.iter())
                        .map(|(now, prev)| (*now - prev).to_cpuload())
                        .collect::<Vec<_>>()
                })
            }))
        })
    }

    fn load_average(&self) -> io::Result<LoadAverage> {
        unix::load_average()
    }

    fn memory(&self) -> io::Result<Memory> {
        memory_stats()
            .or_else(|_| {
                // If there's no procfs, e.g. in a chroot without mounting it or something
                let mut meminfo = BTreeMap::new();
                let mut info: sysinfo = unsafe { mem::zeroed() };
                unsafe { sysinfo(&mut info) };
                let unit = info.mem_unit as usize;
                meminfo.insert(
                    "MemTotal".to_owned(),
                    ByteSize::b(info.totalram as usize * unit),
                );
                meminfo.insert(
                    "MemFree".to_owned(),
                    ByteSize::b(info.freeram as usize * unit),
                );
                meminfo.insert(
                    "Shmem".to_owned(),
                    ByteSize::b(info.sharedram as usize * unit),
                );
                meminfo.insert(
                    "Buffers".to_owned(),
                    ByteSize::b(info.bufferram as usize * unit),
                );
                meminfo.insert(
                    "SwapTotal".to_owned(),
                    ByteSize::b(info.totalswap as usize * unit),
                );
                meminfo.insert(
                    "SwapFree".to_owned(),
                    ByteSize::b(info.freeswap as usize * unit),
                );
                Ok(meminfo)
            })
            .map(|meminfo| {
                Memory {
                    total: meminfo.get("MemTotal").map(|x| x.clone()).unwrap_or(
                        ByteSize::b(0),
                    ),
                    free: meminfo.get("MemFree").map(|x| x.clone()).unwrap_or(
                        ByteSize::b(0),
                    ) +
                        meminfo.get("Buffers").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ) +
                        meminfo.get("Cached").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ) +
                        meminfo.get("SReclaimable").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ) -
                        meminfo.get("Shmem").map(|x| x.clone()).unwrap_or(
                            ByteSize::b(0),
                        ),
                    platform_memory: PlatformMemory { meminfo: meminfo },
                }
            })
    }

    fn uptime(&self) -> io::Result<Duration> {
        let mut info: sysinfo = unsafe { mem::zeroed() };
        unsafe { sysinfo(&mut info) };
        Ok(Duration::from_secs(info.uptime as u64))
    }

    fn battery_life(&self) -> io::Result<BatteryLife> {
        let dir = "/sys/class/power_supply";
        let entries = try!(fs::read_dir(&dir));
        let mut full = 0;
        let mut now = 0;
        let mut current = 0;
        for e in entries {
            let p = e.unwrap().path();
            let s = p.to_str().unwrap();
            let f = p.file_name().unwrap().to_str().unwrap();
            if f.len() > 3 {
                if f.split_at(3).0 == "BAT" {
                    full += try!(
                        value_from_file::<i32>(&(s.to_string() + "/energy_full"))
                            .or_else(|_| value_from_file::<i32>(&(s.to_string() + "/charge_full")))
                    );
                    now += try!(
                        value_from_file::<i32>(&(s.to_string() + "/energy_now"))
                            .or_else(|_| value_from_file::<i32>(&(s.to_string() + "/charge_now")))
                    );
                    current += try!(
                        value_from_file::<i32>(&(s.to_string() + "/power_now"))
                            .or_else(|_| value_from_file::<i32>(&(s.to_string() + "/current_now")))
                    );
                }
            }
        }
        if full != 0 {
            let on_ac =
                match self.on_ac_power() {
                    Ok(true) => true,
                    _ => false,
                };
            Ok(BatteryLife {
                remaining_capacity: capacity(full, now),
                remaining_time: time(on_ac, full, now, current),
            })
        } else {
            Err(io::Error::new(io::ErrorKind::Other, "Missing battery information"))
        }
    }

    fn on_ac_power(&self) -> io::Result<bool> {
        value_from_file::<i32>("/sys/class/power_supply/AC/online").map(|v| v == 1)
    }

    fn mounts(&self) -> io::Result<Vec<Filesystem>> {
        read_file("/proc/mounts")
            .and_then(|data| {
                proc_mounts(data.as_bytes()).to_result().map_err(|err| {
                    io::Error::new(io::ErrorKind::InvalidData, err)
                })
            })
            .map(|mounts| {
                mounts
                    .into_iter()
                    .filter_map(|mount| stat_mount(mount).ok())
                    .collect()
            })
    }

    fn mount_at<P: AsRef<path::Path>>(&self, path: P) -> io::Result<Filesystem> {
        read_file("/proc/mounts")
            .and_then(|data| {
                proc_mounts(data.as_bytes()).to_result().map_err(|err| {
                    io::Error::new(io::ErrorKind::InvalidData, err)
                })
            })
            .and_then(|mounts| {
                mounts
                    .into_iter()
                    .find(|mount| Path::new(&mount.target) == path.as_ref())
                    .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "No such mount"))
            })
            .and_then(stat_mount)
    }

    fn networks(&self) -> io::Result<BTreeMap<String, Network>> {
        unix::networks()
    }

    fn network_stats(&self, interface: &str) -> io::Result<NetworkStats> {
        let path_root: String = ("/sys/class/net/".to_string() + interface) + "/statistics/";
        let stats_file = |file: &str| (&path_root).to_string() + file;

        let rx_bytes: usize = try!(value_from_file::<usize>(&stats_file("rx_bytes")));
        let tx_bytes: usize = try!(value_from_file::<usize>(&stats_file("tx_bytes")));
        let rx_packets: usize = try!(value_from_file::<usize>(&stats_file("rx_packets")));
        let tx_packets: usize = try!(value_from_file::<usize>(&stats_file("tx_packets")));
        let rx_errors: usize = try!(value_from_file::<usize>(&stats_file("rx_errors")));
        let tx_errors: usize = try!(value_from_file::<usize>(&stats_file("tx_errors")));

        Ok(NetworkStats {
            rx_bytes,
            tx_bytes,
            rx_packets,
            tx_packets,
            rx_errors,
            tx_errors,
        })
    }

    fn cpu_temp(&self) -> io::Result<f32> {
        read_file("/sys/class/thermal/thermal_zone0/temp")
            .and_then(|data| match data.parse::<f32>() {
                Ok(x) => Ok(x),
                Err(_) => Err(io::Error::new(io::ErrorKind::Other, "Could not parse float")),
            })
            .map(|num| num / 1000.0)
    }
}

#[repr(C)]
#[derive(Debug)]
struct sysinfo {
    uptime: c_long,
    loads: [c_ulong; 3],
    totalram: c_ulong,
    freeram: c_ulong,
    sharedram: c_ulong,
    bufferram: c_ulong,
    totalswap: c_ulong,
    freeswap: c_ulong,
    procs: c_ushort,
    totalhigh: c_ulong,
    freehigh: c_ulong,
    mem_unit: c_uint,
    padding: [c_schar; 8],
}

#[link(name = "c")]
extern "C" {
    fn sysinfo(info: *mut sysinfo);
}
<|endoftext|>


<commit_before>const NRBITS: u32 = 8;
const TYPEBITS: u32 = 8;

#[cfg(target_arch = "mips")]
mod consts {
    pub const NONE: u8 = 1;
    pub const READ: u8 = 2;
    pub const WRITE: u8 = 4;
    pub const SIZEBITS: u8 = 13;
    pub const DIRBITS: u8 = 3;
}
#[cfg(target_arch = "powerpc")]
mod consts {
    pub const NONE: u8 = 1;
    pub const READ: u8 = 2;
    pub const WRITE: u8 = 4;
    pub const SIZEBITS: u8 = 13;
    pub const DIRBITS: u8 = 3;
}
#[cfg(not(any(target_arch = "powerpc", target_arch = "mips")))]
mod consts {
    pub const NONE: u8 = 0;
    pub const READ: u8 = 2;
    pub const WRITE: u8 = 1;
    pub const SIZEBITS: u8 = 14;
    pub const DIRBITS: u8 = 2;
}

pub use self::consts::*;

pub const NRSHIFT: u32 = 0;
pub const TYPESHIFT: u32 = NRSHIFT + NRBITS as u32;
pub const SIZESHIFT: u32 = TYPESHIFT + SIZEBITS as u32;
pub const DIRSHIFT: u32 = SIZESHIFT + DIRBITS as u32;

pub const NRMASK: u32 = (1 << NRBITS) - 1;
pub const TYPEMASK: u32 = (1 << TYPEBITS) - 1;
pub const SIZEMASK: u32 = (1 << SIZEBITS) - 1;
pub const DIRMASK: u32 = (1 << DIRBITS) - 1;

/// Encode an ioctl command.
#[macro_export]
macro_rules! ioc {
    ($dir:expr, $ty:expr, $nr:expr, $sz:expr) => (
    ($dir as u32) << DIRBITS |
        ($ty as u32) << TYPEBITS |
        ($nr as u32) << NRSHIFT |
        ($sz as u32) << SIZESHIFT)
}

/// Encode an ioctl command that reads.
#[macro_export]
macro_rules! ior {
    ($ty:expr, $nr:expr, $sz:expr) => (ioc!(READ, $ty, $nr, $sz))
}

/// Encode an ioctl command that writes.
#[macro_export]
macro_rules! iow {
    ($ty:expr, $nr:expr, $sz:expr) => (ioc!(WRITE, $ty, $nr, $sz))
}

/// Encode an ioctl command that both reads and writes.
#[macro_export]
macro_rules! iorw {
    ($ty:expr, $nr:expr, $sz:expr) => (ioc!(READ|WRITE, $ty, $nr, $sz))
}

/// Declare a wrapper function around an ioctl.
#[macro_export]
macro_rules! ioctl {
    (read $name:ident with $ioty:expr, $nr:expr; $ty:ty) => (
        pub unsafe fn $name(fd: $crate::libc::c_int, val: *mut $ty) {
            $crate::ioctl(fd, ior!($ioty, $nr, ::std::mem::size_of::<$ty>()) as ::libc::c_ulong, val);
        }
    );
    (write $name:ident with $ioty:expr, $nr:expr; $ty:ty) => (
        pub unsafe fn $name(fd: c_int, val: *const $ty) {
            $crate::ioctl(fd, ior!($ioty, $nr, ::std::mem::size_of::<$ty>()) as ::libc::c_ulong, val);
        }
    );
}

#[inline(always)]
pub fn ioc_dir(nr: u32) -> u8 {
    ((nr >> DIRSHIFT) & DIRMASK) as u8
}

#[inline(always)]
pub fn ioc_type(nr: u32) -> u32 {
    (nr >> TYPESHIFT) & TYPEMASK
}

#[inline(always)]
pub fn ioc_nr(nr: u32) -> u32 {
    (nr >> NRSHIFT) & NRMASK
}

#[inline(always)]
pub fn ioc_size(nr: u32) -> u32 {
    ((nr >> SIZESHIFT) as u32) & SIZEMASK
}

pub const IN: u32 = (WRITE as u32) << DIRSHIFT;
pub const OUT: u32 = (READ as u32) << DIRSHIFT;
pub const INOUT: u32 = ((READ|WRITE) as u32) << DIRSHIFT;
pub const SIZE_MASK: u32 = SIZEMASK << SIZESHIFT;
<commit_msg>[feat] Add buffer versions of the things<commit_after>const NRBITS: u32 = 8;
const TYPEBITS: u32 = 8;

#[cfg(target_arch = "mips")]
mod consts {
    pub const NONE: u8 = 1;
    pub const READ: u8 = 2;
    pub const WRITE: u8 = 4;
    pub const SIZEBITS: u8 = 13;
    pub const DIRBITS: u8 = 3;
}
#[cfg(target_arch = "powerpc")]
mod consts {
    pub const NONE: u8 = 1;
    pub const READ: u8 = 2;
    pub const WRITE: u8 = 4;
    pub const SIZEBITS: u8 = 13;
    pub const DIRBITS: u8 = 3;
}
#[cfg(not(any(target_arch = "powerpc", target_arch = "mips")))]
mod consts {
    pub const NONE: u8 = 0;
    pub const READ: u8 = 2;
    pub const WRITE: u8 = 1;
    pub const SIZEBITS: u8 = 14;
    pub const DIRBITS: u8 = 2;
}

pub use self::consts::*;

pub const NRSHIFT: u32 = 0;
pub const TYPESHIFT: u32 = NRSHIFT + NRBITS as u32;
pub const SIZESHIFT: u32 = TYPESHIFT + SIZEBITS as u32;
pub const DIRSHIFT: u32 = SIZESHIFT + DIRBITS as u32;

pub const NRMASK: u32 = (1 << NRBITS) - 1;
pub const TYPEMASK: u32 = (1 << TYPEBITS) - 1;
pub const SIZEMASK: u32 = (1 << SIZEBITS) - 1;
pub const DIRMASK: u32 = (1 << DIRBITS) - 1;

/// Encode an ioctl command.
#[macro_export]
macro_rules! ioc {
    ($dir:expr, $ty:expr, $nr:expr, $sz:expr) => (
    ($dir as u32) << DIRBITS |
        ($ty as u32) << TYPEBITS |
        ($nr as u32) << NRSHIFT |
        ($sz as u32) << SIZESHIFT)
}

/// Encode an ioctl command that reads.
#[macro_export]
macro_rules! ior {
    ($ty:expr, $nr:expr, $sz:expr) => (ioc!(READ, $ty, $nr, $sz))
}

/// Encode an ioctl command that writes.
#[macro_export]
macro_rules! iow {
    ($ty:expr, $nr:expr, $sz:expr) => (ioc!(WRITE, $ty, $nr, $sz))
}

/// Encode an ioctl command that both reads and writes.
#[macro_export]
macro_rules! iorw {
    ($ty:expr, $nr:expr, $sz:expr) => (ioc!(READ|WRITE, $ty, $nr, $sz))
}

/// Declare a wrapper function around an ioctl.
#[macro_export]
macro_rules! ioctl {
    (read $name:ident with $ioty:expr, $nr:expr; $ty:ty) => (
        pub unsafe fn $name(fd: $crate::libc::c_int, val: *mut $ty) {
            $crate::ioctl(fd, ior!($ioty, $nr, ::std::mem::size_of::<$ty>()) as ::libc::c_ulong, val);
        }
    );
    (write $name:ident with $ioty:expr, $nr:expr; $ty:ty) => (
        pub unsafe fn $name(fd: c_int, val: *const $ty) {
            $crate::ioctl(fd, ior!($ioty, $nr, ::std::mem::size_of::<$ty>()) as ::libc::c_ulong, val);
        }
    );
    (read buf $name:ident with $ioty:expr, $nr:expr; $ty:ty) => (
        pub unsafe fn $name(fd: $crate::libc::c_int, val: *mut $ty, len: usize) {
            $crate::ioctl(fd, ior!($ioty, $nr, len) as ::libc::c_ulong, val);
        }
    );
    (write buf $name:ident with $ioty:expr, $nr:expr; $ty:ty) => (
        pub unsafe fn $name(fd: c_int, val: *const $ty, len: usize) {
            $crate::ioctl(fd, ior!($ioty, $nr, len) as ::libc::c_ulong, val);
        }
    );
}

#[inline(always)]
pub fn ioc_dir(nr: u32) -> u8 {
    ((nr >> DIRSHIFT) & DIRMASK) as u8
}

#[inline(always)]
pub fn ioc_type(nr: u32) -> u32 {
    (nr >> TYPESHIFT) & TYPEMASK
}

#[inline(always)]
pub fn ioc_nr(nr: u32) -> u32 {
    (nr >> NRSHIFT) & NRMASK
}

#[inline(always)]
pub fn ioc_size(nr: u32) -> u32 {
    ((nr >> SIZESHIFT) as u32) & SIZEMASK
}

pub const IN: u32 = (WRITE as u32) << DIRSHIFT;
pub const OUT: u32 = (READ as u32) << DIRSHIFT;
pub const INOUT: u32 = ((READ|WRITE) as u32) << DIRSHIFT;
pub const SIZE_MASK: u32 = SIZEMASK << SIZESHIFT;
<|endoftext|>


<commit_before>use std::fmt;
use std::collections::HashMap;
use std::io::{Read, Write, BufRead, BufReader};

use protocol::binary_packet::Opcode;
use client::Stats;
use error::MemcacheError;
use stream::Stream;
use value::{ToMemcacheValue, FromMemcacheValue};

#[derive(Default)]
pub struct Options {
    pub noreply: bool,
    pub exptime: u32,
    pub flags: u32,
}

enum StoreCommand {
    Set,
    Add,
    Replace,
    Append,
    Prepend,
}

impl fmt::Display for StoreCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StoreCommand::Set => write!(f, "set"),
            StoreCommand::Add => write!(f, "add"),
            StoreCommand::Replace => write!(f, "replace"),
            StoreCommand::Append => write!(f, "append"),
            StoreCommand::Prepend => write!(f, "prepend"),
        }
    }
}
pub struct AsciiProtocol<C: Read + Write + Sized> {
    pub reader: BufReader<C>,
}

impl AsciiProtocol<Stream> {
    fn store<V: ToMemcacheValue<Stream>>(
        &mut self,
        command: StoreCommand,
        key: &str,
        value: V,
        options: &Options,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }

        let mut header = format!(
            "{} {} {} {} {}",
            command,
            key,
            value.get_flags(),
            options.exptime,
            value.get_length()
        );
        if options.noreply {
            header += " noreply";
        }
        header += "\r\n";
        self.reader.get_mut().write(header.as_bytes())?;
        value.write_to(self.reader.get_mut())?;
        self.reader.get_mut().write(b"\r\n")?;
        self.reader.get_mut().flush()?;

        if options.noreply {
            return Ok(());
        }

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "STORED\r\n" {
            return Ok(());
        } else if s == "NOT_STORED\r\n" {
            return Ok(());
        } else {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
    }

    pub(super) fn version(&mut self) -> Result<String, MemcacheError> {
        self.reader.get_mut().write(b"version\r\n")?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if !s.starts_with("VERSION") {
            return Err(MemcacheError::ServerError(0));
        }
        let s = s.trim_left_matches("VERSION ");
        let s = s.trim_right_matches("\r\n");

        return Ok(s.to_string());
    }

    pub(super) fn flush(&mut self) -> Result<(), MemcacheError> {
        unimplemented!();
    }

    pub(super) fn flush_with_delay(&mut self, delay: u32) -> Result<(), MemcacheError> {
        unimplemented!();
    }

    pub(super) fn get<V: FromMemcacheValue>(&mut self, key: &str) -> Result<Option<V>, MemcacheError> {
        write!(self.reader.get_mut(), "get {}\r\n", key)?;

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;

        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s.starts_with("END") {
            return Ok(None);
        } else if !s.starts_with("VALUE") {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let header: Vec<_> = s.trim_right_matches("\r\n").split(" ").collect();
        if header.len() != 4 {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let key = header[1];
        if key != key {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        let flags = header[2].parse()?;
        let length = header[3].parse()?;

        let mut buffer = vec![0; length];
        self.reader.read_exact(buffer.as_mut_slice())?;

        // read the rest \r\n and END\r\n
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "END\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        return Ok(Some(FromMemcacheValue::from_memcache_value(buffer, flags)?));
    }

    pub(super) fn gets<V: FromMemcacheValue>(
        &mut self,
        keys: Vec<&str>,
    ) -> Result<HashMap<String, V>, MemcacheError> {
        write!(self.reader.get_mut(), "gets {}\r\n", keys.join(" "))?;

        let mut result: HashMap<String, V> = HashMap::new();
        loop {
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;

            if is_memcache_error(s.as_str()) {
                return Err(MemcacheError::from(s));
            } else if s.starts_with("END") {
                break;
            } else if !s.starts_with("VALUE") {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let header: Vec<_> = s.trim_right_matches("\r\n").split(" ").collect();
            if header.len() != 5 {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let key = header[1];
            if key != key {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
            let flags = header[2].parse()?;
            let length = header[3].parse()?;

            let mut buffer = vec![0; length];
            self.reader.read_exact(buffer.as_mut_slice())?;

            result.insert(key.to_string(), FromMemcacheValue::from_memcache_value(buffer, flags)?);

            // read the rest \r\n
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;
            if s != "\r\n" {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
        }

        return Ok(result);
    }

    pub(super) fn set<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options{exptime: expiration, ..Default::default()};
        return self.store(StoreCommand::Set, key, value, &options);
    }

    pub(super) fn add<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options{exptime: expiration, ..Default::default()};
        return self.store(StoreCommand::Add, key, value, &options);
    }

    pub(super) fn replace<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options{exptime: expiration, ..Default::default()};
        return self.store(StoreCommand::Replace, key, value, &options);
    }

    pub(super) fn append<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn prepend<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn delete(&mut self, key: &str) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn increment(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn decrement(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn touch(&mut self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn stats(&mut self) -> Result<Stats, MemcacheError> {
        unimplemented!();
    }
}

fn is_memcache_error(s: &str) -> bool {
    return s == "ERROR\r\n" || s.starts_with("CIENT_ERROR") || s.starts_with("SERVER_ERROR");
}
<commit_msg>:rotating_light: Remove using of deprecated functions in stdlib<commit_after>use std::fmt;
use std::collections::HashMap;
use std::io::{Read, Write, BufRead, BufReader};

use protocol::binary_packet::Opcode;
use client::Stats;
use error::MemcacheError;
use stream::Stream;
use value::{ToMemcacheValue, FromMemcacheValue};

#[derive(Default)]
pub struct Options {
    pub noreply: bool,
    pub exptime: u32,
    pub flags: u32,
}

enum StoreCommand {
    Set,
    Add,
    Replace,
    Append,
    Prepend,
}

impl fmt::Display for StoreCommand {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            StoreCommand::Set => write!(f, "set"),
            StoreCommand::Add => write!(f, "add"),
            StoreCommand::Replace => write!(f, "replace"),
            StoreCommand::Append => write!(f, "append"),
            StoreCommand::Prepend => write!(f, "prepend"),
        }
    }
}
pub struct AsciiProtocol<C: Read + Write + Sized> {
    pub reader: BufReader<C>,
}

impl AsciiProtocol<Stream> {
    fn store<V: ToMemcacheValue<Stream>>(
        &mut self,
        command: StoreCommand,
        key: &str,
        value: V,
        options: &Options,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }

        let mut header = format!(
            "{} {} {} {} {}",
            command,
            key,
            value.get_flags(),
            options.exptime,
            value.get_length()
        );
        if options.noreply {
            header += " noreply";
        }
        header += "\r\n";
        self.reader.get_mut().write(header.as_bytes())?;
        value.write_to(self.reader.get_mut())?;
        self.reader.get_mut().write(b"\r\n")?;
        self.reader.get_mut().flush()?;

        if options.noreply {
            return Ok(());
        }

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s == "STORED\r\n" {
            return Ok(());
        } else if s == "NOT_STORED\r\n" {
            return Ok(());
        } else {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
    }

    pub(super) fn version(&mut self) -> Result<String, MemcacheError> {
        self.reader.get_mut().write(b"version\r\n")?;
        self.reader.get_mut().flush()?;
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s);
        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if !s.starts_with("VERSION") {
            return Err(MemcacheError::ServerError(0));
        }
        let s = s.trim_start_matches("VERSION ");
        let s = s.trim_end_matches("\r\n");

        return Ok(s.to_string());
    }

    pub(super) fn flush(&mut self) -> Result<(), MemcacheError> {
        unimplemented!();
    }

    pub(super) fn flush_with_delay(&mut self, delay: u32) -> Result<(), MemcacheError> {
        unimplemented!();
    }

    pub(super) fn get<V: FromMemcacheValue>(&mut self, key: &str) -> Result<Option<V>, MemcacheError> {
        write!(self.reader.get_mut(), "get {}\r\n", key)?;

        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;

        if is_memcache_error(s.as_str()) {
            return Err(MemcacheError::from(s));
        } else if s.starts_with("END") {
            return Ok(None);
        } else if !s.starts_with("VALUE") {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let header: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
        if header.len() != 4 {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        let key = header[1];
        if key != key {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        let flags = header[2].parse()?;
        let length = header[3].parse()?;

        let mut buffer = vec![0; length];
        self.reader.read_exact(buffer.as_mut_slice())?;

        // read the rest \r\n and END\r\n
        let mut s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }
        s = String::new();
        let _ = self.reader.read_line(&mut s)?;
        if s != "END\r\n" {
            return Err(MemcacheError::ClientError("invalid server response".into()));
        }

        return Ok(Some(FromMemcacheValue::from_memcache_value(buffer, flags)?));
    }

    pub(super) fn gets<V: FromMemcacheValue>(
        &mut self,
        keys: Vec<&str>,
    ) -> Result<HashMap<String, V>, MemcacheError> {
        write!(self.reader.get_mut(), "gets {}\r\n", keys.join(" "))?;

        let mut result: HashMap<String, V> = HashMap::new();
        loop {
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;

            if is_memcache_error(s.as_str()) {
                return Err(MemcacheError::from(s));
            } else if s.starts_with("END") {
                break;
            } else if !s.starts_with("VALUE") {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let header: Vec<_> = s.trim_end_matches("\r\n").split(" ").collect();
            if header.len() != 5 {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }

            let key = header[1];
            if key != key {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
            let flags = header[2].parse()?;
            let length = header[3].parse()?;

            let mut buffer = vec![0; length];
            self.reader.read_exact(buffer.as_mut_slice())?;

            result.insert(key.to_string(), FromMemcacheValue::from_memcache_value(buffer, flags)?);

            // read the rest \r\n
            let mut s = String::new();
            let _ = self.reader.read_line(&mut s)?;
            if s != "\r\n" {
                return Err(MemcacheError::ClientError("invalid server response".into()));
            }
        }

        return Ok(result);
    }

    pub(super) fn set<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options{exptime: expiration, ..Default::default()};
        return self.store(StoreCommand::Set, key, value, &options);
    }

    pub(super) fn add<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options{exptime: expiration, ..Default::default()};
        return self.store(StoreCommand::Add, key, value, &options);
    }

    pub(super) fn replace<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
        expiration: u32,
    ) -> Result<(), MemcacheError> {
        let options = Options{exptime: expiration, ..Default::default()};
        return self.store(StoreCommand::Replace, key, value, &options);
    }

    pub(super) fn append<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn prepend<V: ToMemcacheValue<Stream>>(
        &mut self,
        key: &str,
        value: V,
    ) -> Result<(), MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn delete(&mut self, key: &str) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn increment(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn decrement(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn touch(&mut self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
        if key.len() > 250 {
            return Err(MemcacheError::ClientError(String::from("key is too long")));
        }
        unimplemented!();
    }

    pub(super) fn stats(&mut self) -> Result<Stats, MemcacheError> {
        unimplemented!();
    }
}

fn is_memcache_error(s: &str) -> bool {
    return s == "ERROR\r\n" || s.starts_with("CIENT_ERROR") || s.starts_with("SERVER_ERROR");
}
<|endoftext|>


<commit_before>use platform::process::get_pid;
use platform::time::sleep_secs;
use platform::time::time_now;
use storage::Cache;

use super::Driver;
use super::cmd::Cmd;
use super::cmd::Get;
use super::cmd::Resp;
use super::cmd::Set;
use super::cmd::SetInstr;
use super::cmd::Stat;
use super::cmd::Value;

#[test]
fn test_cmd_set_and_get_a_key() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let key_name = "x";
    let blob = vec![1, 2, 3];

    // Try to retrieve a key not set
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(0, resp.get_values().unwrap().len());

    // Set a key
    let set = Set::new(SetInstr::Set, key_name, 15, 0, blob.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(15, resp.get_first_value().unwrap().flags);
    assert_eq!(blob, resp.get_first_value().unwrap().data);

    // Set a key with noreply flag
    let set = Set::new(SetInstr::Set, "y", 15, 0, blob.clone(), true);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Empty);

    // Retrieve it
    let cmd = Cmd::Get(Get::one("y"));
    let resp = driver.run(cmd);
    assert_eq!(15, resp.get_first_value().unwrap().flags);
    assert_eq!(blob, resp.get_first_value().unwrap().data);
}

#[test]
fn test_cmd_set_and_get_multiple_keys() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let val1 = vec![1];
    let val3 = vec![3];

    // Set two keys
    let set = Set::new(SetInstr::Set, "a", 15, 0, val1.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    let set = Set::new(SetInstr::Set, "c", 17, 0, val3.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Try to retrieve three keys - get two
    let keys = vec!["a".to_string(), "b".to_string(), "c".to_string()];
    let cmd = Cmd::Get(Get::new(keys));
    let resp = driver.run(cmd);

    let values = resp.get_values().unwrap();
    let val1 = Value::new("a", 15, val1);
    let val3 = Value::new("c", 17, val3);
    assert_eq!(2, values.len());
    assert_eq!(val1, values[0]);
    assert_eq!(val3, values[1]);
}

#[test]
fn test_cmd_stats() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    // Set a key
    let set = Set::new(SetInstr::Set, "x", 0, 0, vec![8, 9], false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it
    let cmd = Cmd::Get(Get::one("x"));
    driver.run(cmd);

    // Run stats
    let cmd = Cmd::Stats;
    let resp = driver.run(cmd);

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

    let stats = resp.get_stats().unwrap();
    assert_eq!(*stats,
               (vec![st_pid,
                                st_uptime,
                                st_time,
                                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]));
}

// this is a slow test that relies on sleeps
#[ignore]
#[test]
fn test_cmd_relative_exptime() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let key_name = "x";
    let blob = vec![1, 2, 3];

    // Set a key with exptime of 1 second
    let set = Set::new(SetInstr::Set, key_name, 0, 1, blob.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it right away - succeeds
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(blob, resp.get_first_value().unwrap().data);

    // sleep 1.5 secs - long enough to expire key
    sleep_secs(1.5);

    // Retrieve the key again - it's gone
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Error);
}

// this is a slow test that relies on sleeps
#[ignore]
#[test]
fn test_cmd_absolute_exptime() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let key_name = "x";
    let blob = vec![1, 2, 3];
    let exp = time_now().round() as u32 + 1;

    // Set a key with exptime of 1 second
    let set = Set::new(SetInstr::Set, key_name, 0, exp, blob.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it right away - succeeds
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(blob, resp.get_first_value().unwrap().data);

    // sleep 2.5 secs - long enough to expire key
    sleep_secs(2.5);

    // Retrieve the key again - it's gone
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Error);
}
<commit_msg>fix two tests<commit_after>use platform::process::get_pid;
use platform::time::sleep_secs;
use platform::time::time_now;
use storage::Cache;

use super::Driver;
use super::cmd::Cmd;
use super::cmd::Get;
use super::cmd::Resp;
use super::cmd::Set;
use super::cmd::SetInstr;
use super::cmd::Stat;
use super::cmd::Value;

#[test]
fn test_cmd_set_and_get_a_key() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let key_name = "x";
    let blob = vec![1, 2, 3];

    // Try to retrieve a key not set
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(0, resp.get_values().unwrap().len());

    // Set a key
    let set = Set::new(SetInstr::Set, key_name, 15, 0, blob.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(15, resp.get_first_value().unwrap().flags);
    assert_eq!(blob, resp.get_first_value().unwrap().data);

    // Set a key with noreply flag
    let set = Set::new(SetInstr::Set, "y", 15, 0, blob.clone(), true);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Empty);

    // Retrieve it
    let cmd = Cmd::Get(Get::one("y"));
    let resp = driver.run(cmd);
    assert_eq!(15, resp.get_first_value().unwrap().flags);
    assert_eq!(blob, resp.get_first_value().unwrap().data);
}

#[test]
fn test_cmd_set_and_get_multiple_keys() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let val1 = vec![1];
    let val3 = vec![3];

    // Set two keys
    let set = Set::new(SetInstr::Set, "a", 15, 0, val1.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    let set = Set::new(SetInstr::Set, "c", 17, 0, val3.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Try to retrieve three keys - get two
    let keys = vec!["a".to_string(), "b".to_string(), "c".to_string()];
    let cmd = Cmd::Get(Get::new(keys));
    let resp = driver.run(cmd);

    let values = resp.get_values().unwrap();
    let val1 = Value::new("a", 15, val1);
    let val3 = Value::new("c", 17, val3);
    assert_eq!(2, values.len());
    assert_eq!(val1, values[0]);
    assert_eq!(val3, values[1]);
}

#[test]
fn test_cmd_stats() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    // Set a key
    let set = Set::new(SetInstr::Set, "x", 0, 0, vec![8, 9], false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it
    let cmd = Cmd::Get(Get::one("x"));
    driver.run(cmd);

    // Run stats
    let cmd = Cmd::Stats;
    let resp = driver.run(cmd);

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

    let stats = resp.get_stats().unwrap();
    assert_eq!(*stats,
               (vec![st_pid,
                                st_uptime,
                                st_time,
                                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]));
}

// this is a slow test that relies on sleeps
#[ignore]
#[test]
fn test_cmd_relative_exptime() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let key_name = "x";
    let blob = vec![1, 2, 3];

    // Set a key with exptime of 1 second
    let set = Set::new(SetInstr::Set, key_name, 0, 1, blob.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it right away - succeeds
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(blob, resp.get_first_value().unwrap().data);

    // sleep 1.5 secs - long enough to expire key
    sleep_secs(1.5);

    // Retrieve the key again - it's gone
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(0, resp.get_values().unwrap().len());
}

// this is a slow test that relies on sleeps
#[ignore]
#[test]
fn test_cmd_absolute_exptime() {
    let cache = Cache::new(100);
    let mut driver = Driver::new(cache);

    let key_name = "x";
    let blob = vec![1, 2, 3];
    let exp = time_now().round() as u32 + 1;

    // Set a key with exptime of 1 second
    let set = Set::new(SetInstr::Set, key_name, 0, exp, blob.clone(), false);
    let cmd = Cmd::Set(set);
    let resp = driver.run(cmd);
    assert_eq!(resp, Resp::Stored);

    // Retrieve it right away - succeeds
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(blob, resp.get_first_value().unwrap().data);

    // sleep 2.5 secs - long enough to expire key
    sleep_secs(2.5);

    // Retrieve the key again - it's gone
    let cmd = Cmd::Get(Get::one(key_name));
    let resp = driver.run(cmd);
    assert_eq!(0, resp.get_values().unwrap().len());
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use bit_vec::BitVec;
use byteorder::{LittleEndian as LE, ReadBytesExt, WriteBytesExt};
use futures_util::future::Either;
use mysql_common::value::serialize_bin_many;

use std::io::Write;

use crate::{
    connection_like::{
        streamless::Streamless, ConnectionLike, ConnectionLikeWrapper, StmtCacheResult,
    },
    consts::{ColumnType, Command},
    error::*,
    io,
    prelude::FromRow,
    queryable::{query_result::QueryResult, BinaryProtocol},
    Column, Params, Row,
    Value::{self, *},
};

/// Inner statement representation.
#[derive(Eq, PartialEq, Clone, Debug)]
pub struct InnerStmt {
    /// Positions and names of named parameters
    pub named_params: Option<Vec<String>>,
    pub params: Option<Vec<Column>>,
    pub columns: Option<Vec<Column>>,
    pub statement_id: u32,
    pub num_columns: u16,
    pub num_params: u16,
    pub warning_count: u16,
}

impl InnerStmt {
    // TODO: Consume payload?
    pub fn new(pld: &[u8], named_params: Option<Vec<String>>) -> Result<InnerStmt> {
        let mut reader = &pld[1..];
        let statement_id = reader.read_u32::<LE>()?;
        let num_columns = reader.read_u16::<LE>()?;
        let num_params = reader.read_u16::<LE>()?;
        let warning_count = reader.read_u16::<LE>()?;
        Ok(InnerStmt {
            named_params,
            statement_id,
            num_columns,
            num_params,
            warning_count,
            params: None,
            columns: None,
        })
    }
}

/// Prepared statement
pub struct Stmt<T> {
    conn_like: Option<Either<T, Streamless<T>>>,
    inner: InnerStmt,
    /// None => In use elsewhere
    /// Some(Cached) => Should not be closed
    /// Some(NotCached(_)) => Should be closed
    cached: Option<StmtCacheResult>,
}

pub fn new<T>(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    Stmt::new(conn_like, inner, cached)
}

impl<T> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    fn new(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T> {
        Stmt {
            conn_like: Some(Either::Left(conn_like)),
            inner,
            cached: Some(cached),
        }
    }

    async fn send_long_data_for_index(
        self,
        params: Vec<Value>,
        index: usize,
    ) -> Result<(Self, Vec<Value>)> {
        let mut this = self;
        for chunk in 0.. {
            let data_cap = crate::consts::MAX_PAYLOAD_LEN - 10;
            let buf = match params[index] {
                Bytes(ref x) => {
                    let statement_id = this.inner.statement_id;
                    let mut chunks = x.chunks(data_cap);
                    match chunks.nth(chunk) {
                        Some(chunk) => {
                            let mut buf = Vec::with_capacity(chunk.len() + 6);
                            buf.write_u32::<LE>(statement_id).unwrap();
                            buf.write_u16::<LE>(index as u16).unwrap();
                            buf.write_all(chunk).unwrap();
                            Some(buf)
                        }
                        _ => None,
                    }
                }
                _ => unreachable!(),
            };
            match buf {
                Some(buf) => {
                    let chunk_len = buf.len() - 6;
                    this = this
                        .write_command_data(Command::COM_STMT_SEND_LONG_DATA, buf)
                        .await?;
                    if chunk_len < data_cap {
                        break;
                    }
                }
                None => break,
            }
        }
        Ok((this, params))
    }

    async fn send_long_data(
        self,
        mut params: Vec<Value>,
        large_bitmap: BitVec<u8>,
    ) -> Result<(Self, Vec<Value>)> {
        let mut bits = large_bitmap.into_iter().enumerate();
        let mut this = self;
        loop {
            match bits.next() {
                Some((index, true)) => {
                    let (this_, params_) = this.send_long_data_for_index(params, index).await?;
                    this = this_;
                    params = params_;
                }
                Some((_, false)) => {}
                None => return Ok((this, params)),
            }
        }
    }

    async fn execute_positional<U>(self, params: U) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        U: ::std::ops::Deref<Target = [Value]>,
        U: IntoIterator<Item = Value>,
        U: Send + 'static,
    {
        if self.inner.num_params as usize != params.len() {
            let error = DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: params.len() as u16,
            }
            .into();
            return Err(error);
        }

        let params_def = self.inner.params.as_ref().unwrap();
        let bin_payload = serialize_bin_many(&*params_def, &*params)?;
        let (row_data, null_bitmap, large_bitmap) = bin_payload;
        let (this, params) = self
            .send_long_data(params.into_iter().collect(), large_bitmap.clone())
            .await?;

        let mut data = Vec::new();
        write_data(
            &mut data,
            this.inner.statement_id,
            row_data,
            params,
            this.inner.params.as_ref().unwrap(),
            null_bitmap,
        );
        let this = this
            .write_command_data(Command::COM_STMT_EXECUTE, data)
            .await?;
        let this = this.read_result_set(None).await?;
        Ok(this)
    }

    async fn execute_named(self, params: Params) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.named_params.is_none() {
            let error = DriverError::NamedParamsForPositionalQuery.into();
            return Err(error);
        }

        let positional_params =
            match params.into_positional(self.inner.named_params.as_ref().unwrap()) {
                Ok(positional_params) => positional_params,
                Err(error) => return Err(error.into()),
            };

        match positional_params {
            Params::Positional(params) => self.execute_positional(params).await,
            _ => unreachable!(),
        }
    }

    async fn execute_empty(self) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.num_params > 0 {
            let error = DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: 0,
            }
            .into();
            return Err(error);
        }

        let mut data = Vec::with_capacity(4 + 1 + 4);
        data.write_u32::<LE>(self.inner.statement_id).unwrap();
        data.write_u8(0u8).unwrap();
        data.write_u32::<LE>(1u32).unwrap();

        let this = self
            .write_command_data(Command::COM_STMT_EXECUTE, data)
            .await?;
        this.read_result_set(None).await
    }

    /// See `Queryable::execute`
    pub async fn execute<P>(self, params: P) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        P: Into<Params>,
    {
        let params = params.into();
        match params {
            Params::Positional(params) => self.execute_positional(params).await,
            Params::Named(_) => self.execute_named(params).await,
            Params::Empty => self.execute_empty().await,
        }
    }

    /// See `Queryable::first`
    pub async fn first<P, R>(self, params: P) -> Result<(Self, Option<R>)>
    where
        P: Into<Params> + 'static,
        R: FromRow,
    {
        let result = self.execute(params).await?;
        let (this, mut rows) = result.collect_and_drop::<Row>().await?;
        if rows.len() > 1 {
            Ok((this, Some(FromRow::from_row(rows.swap_remove(0)))))
        } else {
            Ok((this, rows.pop().map(FromRow::from_row)))
        }
    }

    /// See `Queryable::batch`
    pub async fn batch<I, P>(self, params_iter: I) -> Result<Self>
    where
        I: IntoIterator<Item = P>,
        I::IntoIter: Send + 'static,
        Params: From<P>,
        P: 'static,
    {
        let mut params_iter = params_iter.into_iter().map(Params::from);
        let mut this = self;
        loop {
            match params_iter.next() {
                Some(params) => {
                    this = this.execute(params).await?.drop_result().await?;
                }
                None => break Ok(this),
            }
        }
    }

    /// This will close statement (if it's not in the cache) and resolve to a wrapped queryable.
    pub async fn close(mut self) -> Result<T> {
        let cached = self.cached.take();
        match self.conn_like {
            Some(Either::Left(conn_like)) => {
                if let Some(StmtCacheResult::NotCached(stmt_id)) = cached {
                    conn_like.close_stmt(stmt_id).await
                } else {
                    Ok(conn_like)
                }
            }
            _ => unreachable!(),
        }
    }

    pub(crate) fn unwrap(mut self) -> (T, Option<StmtCacheResult>) {
        match self.conn_like {
            Some(Either::Left(conn_like)) => (conn_like, self.cached.take()),
            _ => unreachable!(),
        }
    }
}

impl<T: ConnectionLike + 'static> ConnectionLikeWrapper for Stmt<T> {
    type ConnLike = T;

    fn take_stream(self) -> (Streamless<Self>, io::Stream)
    where
        Self: Sized,
    {
        let Stmt {
            conn_like,
            inner,
            cached,
        } = self;
        match conn_like {
            Some(Either::Left(conn_like)) => {
                let (streamless, stream) = conn_like.take_stream();
                let this = Stmt {
                    conn_like: Some(Either::Right(streamless)),
                    inner,
                    cached,
                };
                (Streamless::new(this), stream)
            }
            _ => unreachable!(),
        }
    }

    fn return_stream(&mut self, stream: io::Stream) {
        let conn_like = self.conn_like.take().unwrap();
        match conn_like {
            Either::Right(streamless) => {
                self.conn_like = Some(Either::Left(streamless.return_stream(stream)));
            }
            _ => unreachable!(),
        }
    }

    fn conn_like_ref(&self) -> &Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }

    fn conn_like_mut(&mut self) -> &mut Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref mut conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }
}

fn write_data(
    writer: &mut Vec<u8>,
    stmt_id: u32,
    row_data: Vec<u8>,
    params: Vec<Value>,
    params_def: &[Column],
    null_bitmap: BitVec<u8>,
) {
    let capacity = 9 + null_bitmap.storage().len() + 1 + params.len() * 2 + row_data.len();
    writer.reserve(capacity);
    writer.write_u32::<LE>(stmt_id).unwrap();
    writer.write_u8(0u8).unwrap();
    writer.write_u32::<LE>(1u32).unwrap();
    writer.write_all(null_bitmap.storage().as_ref()).unwrap();
    writer.write_u8(1u8).unwrap();
    for i in 0..params.len() {
        let result = match params[i] {
            NULL => writer.write_all(&[params_def[i].column_type() as u8, 0u8]),
            Bytes(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_VAR_STRING as u8, 0u8]),
            Int(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_LONGLONG as u8, 0u8]),
            UInt(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_LONGLONG as u8, 128u8]),
            Float(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_DOUBLE as u8, 0u8]),
            Date(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_DATETIME as u8, 0u8]),
            Time(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_TIME as u8, 0u8]),
        };
        result.unwrap();
    }
    writer.write_all(row_data.as_ref()).unwrap();
}
<commit_msg>Fix payload parsing in InnerStmt::new<commit_after>// Copyright (c) 2017 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use bit_vec::BitVec;
use byteorder::{LittleEndian as LE, ReadBytesExt, WriteBytesExt};
use futures_util::future::Either;
use mysql_common::value::serialize_bin_many;

use std::io::Write;

use crate::{
    connection_like::{
        streamless::Streamless, ConnectionLike, ConnectionLikeWrapper, StmtCacheResult,
    },
    consts::{ColumnType, Command},
    error::*,
    io,
    prelude::FromRow,
    queryable::{query_result::QueryResult, BinaryProtocol},
    Column, Params, Row,
    Value::{self, *},
};

/// Inner statement representation.
#[derive(Eq, PartialEq, Clone, Debug)]
pub struct InnerStmt {
    /// Positions and names of named parameters
    pub named_params: Option<Vec<String>>,
    pub params: Option<Vec<Column>>,
    pub columns: Option<Vec<Column>>,
    pub statement_id: u32,
    pub num_columns: u16,
    pub num_params: u16,
    pub warning_count: u16,
}

impl InnerStmt {
    // TODO: Consume payload?
    pub fn new(payload: &[u8], named_params: Option<Vec<String>>) -> Result<InnerStmt> {
        let mut payload = &payload[..];

        if payload.read_u8()? != 0x00 {
            panic!("Invalid statement packet status");
        }

        let statement_id = payload.read_u32::<LE>()?;
        let num_columns = payload.read_u16::<LE>()?;
        let num_params = payload.read_u16::<LE>()?;
        payload.read_u8()?;
        let warning_count = payload.read_u16::<LE>()?;
        Ok(InnerStmt {
            named_params,
            statement_id,
            num_columns,
            num_params,
            warning_count,
            params: None,
            columns: None,
        })
    }
}

/// Prepared statement
pub struct Stmt<T> {
    conn_like: Option<Either<T, Streamless<T>>>,
    inner: InnerStmt,
    /// None => In use elsewhere
    /// Some(Cached) => Should not be closed
    /// Some(NotCached(_)) => Should be closed
    cached: Option<StmtCacheResult>,
}

pub fn new<T>(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    Stmt::new(conn_like, inner, cached)
}

impl<T> Stmt<T>
where
    T: ConnectionLike + Sized + 'static,
{
    fn new(conn_like: T, inner: InnerStmt, cached: StmtCacheResult) -> Stmt<T> {
        Stmt {
            conn_like: Some(Either::Left(conn_like)),
            inner,
            cached: Some(cached),
        }
    }

    async fn send_long_data_for_index(
        self,
        params: Vec<Value>,
        index: usize,
    ) -> Result<(Self, Vec<Value>)> {
        let mut this = self;
        for chunk in 0.. {
            let data_cap = crate::consts::MAX_PAYLOAD_LEN - 10;
            let buf = match params[index] {
                Bytes(ref x) => {
                    let statement_id = this.inner.statement_id;
                    let mut chunks = x.chunks(data_cap);
                    match chunks.nth(chunk) {
                        Some(chunk) => {
                            let mut buf = Vec::with_capacity(chunk.len() + 6);
                            buf.write_u32::<LE>(statement_id).unwrap();
                            buf.write_u16::<LE>(index as u16).unwrap();
                            buf.write_all(chunk).unwrap();
                            Some(buf)
                        }
                        _ => None,
                    }
                }
                _ => unreachable!(),
            };
            match buf {
                Some(buf) => {
                    let chunk_len = buf.len() - 6;
                    this = this
                        .write_command_data(Command::COM_STMT_SEND_LONG_DATA, buf)
                        .await?;
                    if chunk_len < data_cap {
                        break;
                    }
                }
                None => break,
            }
        }
        Ok((this, params))
    }

    async fn send_long_data(
        self,
        mut params: Vec<Value>,
        large_bitmap: BitVec<u8>,
    ) -> Result<(Self, Vec<Value>)> {
        let mut bits = large_bitmap.into_iter().enumerate();
        let mut this = self;
        loop {
            match bits.next() {
                Some((index, true)) => {
                    let (this_, params_) = this.send_long_data_for_index(params, index).await?;
                    this = this_;
                    params = params_;
                }
                Some((_, false)) => {}
                None => return Ok((this, params)),
            }
        }
    }

    async fn execute_positional<U>(self, params: U) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        U: ::std::ops::Deref<Target = [Value]>,
        U: IntoIterator<Item = Value>,
        U: Send + 'static,
    {
        if self.inner.num_params as usize != params.len() {
            let error = DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: params.len() as u16,
            }
            .into();
            return Err(error);
        }

        let params_def = self.inner.params.as_ref().unwrap();
        let bin_payload = serialize_bin_many(&*params_def, &*params)?;
        let (row_data, null_bitmap, large_bitmap) = bin_payload;
        let (this, params) = self
            .send_long_data(params.into_iter().collect(), large_bitmap.clone())
            .await?;

        let mut data = Vec::new();
        write_data(
            &mut data,
            this.inner.statement_id,
            row_data,
            params,
            this.inner.params.as_ref().unwrap(),
            null_bitmap,
        );
        let this = this
            .write_command_data(Command::COM_STMT_EXECUTE, data)
            .await?;
        let this = this.read_result_set(None).await?;
        Ok(this)
    }

    async fn execute_named(self, params: Params) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.named_params.is_none() {
            let error = DriverError::NamedParamsForPositionalQuery.into();
            return Err(error);
        }

        let positional_params =
            match params.into_positional(self.inner.named_params.as_ref().unwrap()) {
                Ok(positional_params) => positional_params,
                Err(error) => return Err(error.into()),
            };

        match positional_params {
            Params::Positional(params) => self.execute_positional(params).await,
            _ => unreachable!(),
        }
    }

    async fn execute_empty(self) -> Result<QueryResult<Self, BinaryProtocol>> {
        if self.inner.num_params > 0 {
            let error = DriverError::StmtParamsMismatch {
                required: self.inner.num_params,
                supplied: 0,
            }
            .into();
            return Err(error);
        }

        let mut data = Vec::with_capacity(4 + 1 + 4);
        data.write_u32::<LE>(self.inner.statement_id).unwrap();
        data.write_u8(0u8).unwrap();
        data.write_u32::<LE>(1u32).unwrap();

        let this = self
            .write_command_data(Command::COM_STMT_EXECUTE, data)
            .await?;
        this.read_result_set(None).await
    }

    /// See `Queryable::execute`
    pub async fn execute<P>(self, params: P) -> Result<QueryResult<Self, BinaryProtocol>>
    where
        P: Into<Params>,
    {
        let params = params.into();
        match params {
            Params::Positional(params) => self.execute_positional(params).await,
            Params::Named(_) => self.execute_named(params).await,
            Params::Empty => self.execute_empty().await,
        }
    }

    /// See `Queryable::first`
    pub async fn first<P, R>(self, params: P) -> Result<(Self, Option<R>)>
    where
        P: Into<Params> + 'static,
        R: FromRow,
    {
        let result = self.execute(params).await?;
        let (this, mut rows) = result.collect_and_drop::<Row>().await?;
        if rows.len() > 1 {
            Ok((this, Some(FromRow::from_row(rows.swap_remove(0)))))
        } else {
            Ok((this, rows.pop().map(FromRow::from_row)))
        }
    }

    /// See `Queryable::batch`
    pub async fn batch<I, P>(self, params_iter: I) -> Result<Self>
    where
        I: IntoIterator<Item = P>,
        I::IntoIter: Send + 'static,
        Params: From<P>,
        P: 'static,
    {
        let mut params_iter = params_iter.into_iter().map(Params::from);
        let mut this = self;
        loop {
            match params_iter.next() {
                Some(params) => {
                    this = this.execute(params).await?.drop_result().await?;
                }
                None => break Ok(this),
            }
        }
    }

    /// This will close statement (if it's not in the cache) and resolve to a wrapped queryable.
    pub async fn close(mut self) -> Result<T> {
        let cached = self.cached.take();
        match self.conn_like {
            Some(Either::Left(conn_like)) => {
                if let Some(StmtCacheResult::NotCached(stmt_id)) = cached {
                    conn_like.close_stmt(stmt_id).await
                } else {
                    Ok(conn_like)
                }
            }
            _ => unreachable!(),
        }
    }

    pub(crate) fn unwrap(mut self) -> (T, Option<StmtCacheResult>) {
        match self.conn_like {
            Some(Either::Left(conn_like)) => (conn_like, self.cached.take()),
            _ => unreachable!(),
        }
    }
}

impl<T: ConnectionLike + 'static> ConnectionLikeWrapper for Stmt<T> {
    type ConnLike = T;

    fn take_stream(self) -> (Streamless<Self>, io::Stream)
    where
        Self: Sized,
    {
        let Stmt {
            conn_like,
            inner,
            cached,
        } = self;
        match conn_like {
            Some(Either::Left(conn_like)) => {
                let (streamless, stream) = conn_like.take_stream();
                let this = Stmt {
                    conn_like: Some(Either::Right(streamless)),
                    inner,
                    cached,
                };
                (Streamless::new(this), stream)
            }
            _ => unreachable!(),
        }
    }

    fn return_stream(&mut self, stream: io::Stream) {
        let conn_like = self.conn_like.take().unwrap();
        match conn_like {
            Either::Right(streamless) => {
                self.conn_like = Some(Either::Left(streamless.return_stream(stream)));
            }
            _ => unreachable!(),
        }
    }

    fn conn_like_ref(&self) -> &Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }

    fn conn_like_mut(&mut self) -> &mut Self::ConnLike {
        match self.conn_like {
            Some(Either::Left(ref mut conn_like)) => conn_like,
            _ => unreachable!(),
        }
    }
}

fn write_data(
    writer: &mut Vec<u8>,
    stmt_id: u32,
    row_data: Vec<u8>,
    params: Vec<Value>,
    params_def: &[Column],
    null_bitmap: BitVec<u8>,
) {
    let capacity = 9 + null_bitmap.storage().len() + 1 + params.len() * 2 + row_data.len();
    writer.reserve(capacity);
    writer.write_u32::<LE>(stmt_id).unwrap();
    writer.write_u8(0u8).unwrap();
    writer.write_u32::<LE>(1u32).unwrap();
    writer.write_all(null_bitmap.storage().as_ref()).unwrap();
    writer.write_u8(1u8).unwrap();
    for i in 0..params.len() {
        let result = match params[i] {
            NULL => writer.write_all(&[params_def[i].column_type() as u8, 0u8]),
            Bytes(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_VAR_STRING as u8, 0u8]),
            Int(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_LONGLONG as u8, 0u8]),
            UInt(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_LONGLONG as u8, 128u8]),
            Float(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_DOUBLE as u8, 0u8]),
            Date(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_DATETIME as u8, 0u8]),
            Time(..) => writer.write_all(&[ColumnType::MYSQL_TYPE_TIME as u8, 0u8]),
        };
        result.unwrap();
    }
    writer.write_all(row_data.as_ref()).unwrap();
}
<|endoftext|>