<commit_before>use std::cmp::min;
use std::cmp::Ordering;
use std::env;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, BufRead, Write};
use std::io::Error as IoError;
use std::str::FromStr;
use std::str::from_utf8;
use std::path::Path;
use std::process::{Command, Stdio};

use libc::funcs::posix88::unistd::{geteuid, getegid};
use libc::{pid_t, uid_t, gid_t};

use config::Range;
use config::Settings;
use self::Uidmap::*;

#[derive(Clone)]
pub enum Uidmap {
    Singleton(uid_t, gid_t),
    Ranges(Vec<(uid_t, uid_t, uid_t)>, Vec<(gid_t, gid_t, gid_t)>),
}

fn read_uid_map(username: &str) -> Result<Vec<Range>,String> {
    let file = File::open(&Path::new("/etc/subuid"));
    let mut res = Vec::new();
    let mut reader = BufReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subuid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = FromStr::from_str(parts[1]);
        let count = FromStr::from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_err() || count.is_err() {
            return Err(format!("/etc/subuid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(username) {
            let start: uid_t = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

fn read_gid_map(username: &str) -> Result<Vec<Range>,String> {
    let file = File::open(&Path::new("/etc/subgid"));
    let mut res = Vec::new();
    let mut reader = BufReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subgid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = FromStr::from_str(parts[1]);
        let count = FromStr::from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_err() || count.is_err() {
            return Err(format!("/etc/subgid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(username) {
            let start: gid_t = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

pub fn match_ranges(req: &Vec<Range>, allowed: &Vec<Range>, own_id: uid_t)
    -> Result<Vec<(uid_t, uid_t, uid_t)>, ()>
{
    let mut res = vec!((0, own_id, 1));
    let mut reqiter = req.iter();
    let mut reqval = *reqiter.next().unwrap();
    let mut allowiter = allowed.iter();
    let mut allowval = *allowiter.next().unwrap();
    loop {
        if reqval.start == 0 {
            reqval = reqval.shift(1);
        }
        if allowval.start == 0 {
            allowval = allowval.shift(1);
        }
        let clen = min(reqval.len(), allowval.len());
        if clen > 0 {
            res.push((reqval.start, allowval.start, clen));
        }
        reqval = reqval.shift(clen);
        allowval = allowval.shift(clen);
        if reqval.len() == 0 {
            reqval = match reqiter.next() {
                Some(val) => *val,
                None => break,
            };
        }
        if allowval.len() == 0 {
            allowval = match allowiter.next() {
                Some(val) => *val,
                None => return Err(()),
            };
        }
    }
    return Ok(res);
}

pub fn get_max_uidmap() -> Result<Uidmap, String>
{
    let mut cmd = Command::new("id");
    cmd.arg("--user").arg("--name");
    if let Ok(path) = env::var("HOST_PATH") {
        cmd.env("PATH", path);
    }
    cmd.stdin(Stdio::null()).stderr(Stdio::inherit());
    let username = try!(cmd.output()
        .map_err(|e| format!("Error running `id --user --name`: {}", e))
        .and_then(|out| if out.status == ExitStatus(0) { Ok(out.output) } else
            { Err(format!("Error running `id --user --name`")) })
        .and_then(|val| from_utf8(val.as_slice()).map(|x| x.trim().to_string())
                   .map_err(|e| format!("Can't decode username: {}", e))));
    let uid_map = read_uid_map(username.as_slice()).ok();
    let gid_map = read_gid_map(username.as_slice()).ok();

    let uid = unsafe { geteuid() };
    let gid = unsafe { getegid() };
    if let (Some(uid_map), Some(gid_map)) = (uid_map, gid_map) {
        if uid_map.len() == 0 && gid_map.len() == 0 && uid == 0 {
            let uid_rng = try!(read_uid_ranges("/proc/self/uid_map", true));
            let gid_rng = try!(read_uid_ranges("/proc/self/gid_map", true));
            return Ok(Ranges(
                uid_rng.into_iter()
                    .map(|r| (r.start, r.start, r.end-r.start+1)).collect(),
                gid_rng.into_iter()
                    .map(|r| (r.start, r.start, r.end-r.start+1)).collect()));
        }
        let mut uids = vec!((0, uid, 1));
        for &rng in uid_map.iter() {
            let mut rng = rng;
            if uid >= rng.start && uid <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(uid == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            uids.push((rng.start, rng.start, rng.len()));
        }

        let mut gids = vec!((0, gid, 1));
        for &rng in gid_map.iter() {
            let mut rng = rng;
            if gid >= rng.start && gid <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(gid == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            gids.push((rng.start, rng.start, rng.len()));
        }

        return Ok(Ranges(uids, gids));
    } else {
        warn!(concat!("Your system doesn't have /etc/subuid and /etc/subgid.",
            " Presumably your system is too old. Some features may not work"));
        return Ok(Singleton(uid, gid));
    }
}

pub fn apply_uidmap(pid: pid_t, map: &Uidmap) -> Result<(), IoError> {
    match map {
        &Singleton(uid, gid) => {
            let uid_map = format!("0 {} 1", uid);
            debug!("Writing uid_map: {}", uid_map);
            match File::open_mode(&Path::new("/proc")
                              .join(pid.to_string())
                              .join("uid_map"), Open, Write)
                    .write_str(uid_map.as_slice()) {
                Ok(()) => {}
                Err(e) => return Err(IoError {
                    kind: e.kind,
                    desc: "Error writing uid mapping",
                    detail: e.detail,
                    }),
            }
            let gid_map = format!("0 {} 1", gid);
            debug!("Writing gid_map: {}", gid_map);
            match File::open_mode(&Path::new("/proc")
                              .join(pid.to_string())
                              .join("gid_map"), Open, Write)
                    .write_str(uid_map.as_slice()) {
                Ok(()) => {}
                Err(e) => return Err(IoError {
                    kind: e.kind,
                    desc: "Error writing gid mapping",
                    detail: e.detail
                    }),
            }
        }
        &Ranges(ref uids, ref gids) => {
            let myuid = unsafe { geteuid() };
            if myuid > 0 {
                let mut cmd = Command::new("newuidmap");
                cmd.stdin(Stdio::null()).stdout(Stdio::inherit()).stderr(Stdio::inherit());
                cmd.arg(pid.to_string());
                for &(req, allowed, count) in uids.iter() {
                    cmd
                        .arg(req.to_string())
                        .arg(allowed.to_string())
                        .arg(count.to_string());
                }
                info!("Uid map command: {:?}", cmd);
                match cmd.status() {
                    Ok(ExitStatus(0)) => {},
                    Ok(ExitStatus(x)) => {
                        return Err(IoError {
                            kind: OtherIoError,
                            desc: "Error writing uid mapping",
                            detail: Some(format!(
                                "newuidmap exited with status {}", x)),
                            });
                    }
                    Ok(ExitSignal(x)) => {
                        return Err(IoError {
                            kind: OtherIoError,
                            desc: "Error writing uid mapping",
                            detail: Some(format!(
                                "newuidmap exited with signal {}", x)),
                            });
                    }
                    Err(e) => return Err(IoError {
                        kind: e.kind,
                        desc: "Error writing uid mapping",
                        detail: e.detail
                        }),
                }

                let mut cmd = Command::new("newgidmap");
                cmd.stdin(Stdio::null()).stdout(Stdio::inherit()).stderr(Stdio::inherit());
                cmd.arg(pid.to_string());
                for &(req, allowed, count) in gids.iter() {
                    cmd
                        .arg(req.to_string())
                        .arg(allowed.to_string())
                        .arg(count.to_string());
                }
                info!("Gid map command: {:?}", cmd);
                match cmd.status() {
                    Ok(ExitStatus(0)) => {},
                    Ok(ExitStatus(x)) => {
                        return Err(IoError {
                            kind: OtherIoError,
                            desc: "Error writing gid mapping",
                            detail: Some(format!(
                                "newgidmap exited with status {}", x)),
                            });
                    }
                    Ok(ExitSignal(x)) => {
                        return Err(IoError {
                            kind: OtherIoError,
                            desc: "Error writing gid mapping",
                            detail: Some(format!(
                                "newgidmap exited with signal {}", x)),
                            });
                    }
                    Err(e) => return Err(IoError {
                        kind: e.kind,
                        desc: "Error writing gid mapping",
                        detail: e.detail
                        }),
                }
            } else {
                let mut membuf = MemWriter::new();
                for &(ins, outs, cnt) in uids.iter() {
                    try!(writeln!(&mut membuf, "{} {} {}", ins, outs, cnt));
                }
                let mut file = try!(File::create(&Path::new(
                    format!("/proc/{}/uid_map", pid))));
                let value = membuf.into_inner();
                debug!("Writing uid map ```{}```",
                    String::from_utf8_lossy(value.as_slice()));
                try!(file.write(value.as_slice()));

                let mut membuf = MemWriter::new();
                for &(ins, outs, cnt) in gids.iter() {
                    try!(writeln!(&mut membuf, "{} {} {}", ins, outs, cnt));
                }
                let mut file = try!(File::create(&Path::new(
                    format!("/proc/{}/gid_map", pid))));
                let value = membuf.into_inner();
                debug!("Writing gid map ```{}```",
                    String::from_utf8_lossy(value.as_slice()));
                try!(file.write(value.as_slice()));
            }
        }
    }
    return Ok(());
}

fn read_uid_ranges(path: &str, read_inside: bool) -> Result<Vec<Range>, String>
{
    let mut file = BufReader::new(try!(File::open(&Path::new(path))
        .map_err(|e| format!("Error reading uid/gid map: {}", e))));
    let mut result = vec!();
    for line in file.lines() {
        let line = try!(line
            .map_err(|e| format!("Error reading uid/gid map: {}", e)));
        let mut words = line.as_slice().words();
        let inside = try!(words.next().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("uid/gid map format error")));
        let outside = try!(words.next().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("uid/gid map format error")));
        let count = try!(words.next().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("uid/gid map format error")));
        if read_inside {
            result.push(Range { start: inside, end: inside+count-1 });
        } else {
            result.push(Range { start: outside, end: outside+count-1 });
        }
    }
    return Ok(result);
}

pub fn map_users(settings: &Settings, uids: &Vec<Range>, gids: &Vec<Range>)
    -> Result<Uidmap, String>
{
    let default_uids = vec!(Range { start: 0, end: 0 });
    let default_gids = vec!(Range { start: 0, end: 0 });
    let uids = if uids.len() > 0 { uids } else { &default_uids };
    let gids = if gids.len() > 0 { gids } else { &default_gids };
    if settings.uid_map.is_none() {
        let ranges = try!(read_uid_ranges("/proc/self/uid_map", true));
        let uid_map = try!(match_ranges(uids, &ranges, 0)
            .map_err(|()| {
                return format!("Number of allowed subuids is too small. \
                    Required {}, allowed {}. You either need to increase \
                    allowed numbers in /etc/subuid (preferred) or decrease \
                    needed ranges in vagga.yaml by adding `uids` key \
                    to container config", uids, ranges);
            }));
        let ranges = try!(read_uid_ranges("/proc/self/gid_map", true));
        let gid_map = try!(match_ranges(gids, &ranges, 0)
            .map_err(|()| {
                return format!("Number of allowed subgids is too small. \
                    Required {}, allowed {}. You either need to increase \
                    allowed numbers in /etc/subgid (preferred) or decrease \
                    needed ranges in vagga.yaml by adding `gids` key \
                    to container config", gids, ranges);
            }));
        return Ok(Ranges(uid_map, gid_map));
    } else {
        let &(ref uids, ref gids) = settings.uid_map.as_ref().unwrap();
        return Ok(Ranges(uids.clone(), gids.clone()));
    }
}
<commit_msg>WIP: Fix uidmap.rs errors<commit_after>use std::cmp::min;
use std::cmp::Ordering;
use std::env;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, BufRead, Write};
use std::io::Error as IoError;
use std::io::ErrorKind::{Interrupted, Other};
use std::str::FromStr;
use std::str::from_utf8;
use std::path::Path;
use std::process::{Command, Stdio};

use libc::funcs::posix88::unistd::{geteuid, getegid};
use libc::{pid_t, uid_t, gid_t};

use config::Range;
use config::Settings;
use self::Uidmap::*;

#[derive(Clone)]
pub enum Uidmap {
    Singleton(uid_t, gid_t),
    Ranges(Vec<(uid_t, uid_t, uid_t)>, Vec<(gid_t, gid_t, gid_t)>),
}

fn read_uid_map(username: &str) -> Result<Vec<Range>,String> {
    let file = File::open(&Path::new("/etc/subuid"));
    let mut res = Vec::new();
    let mut reader = BufReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subuid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = FromStr::from_str(parts[1]);
        let count = FromStr::from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_err() || count.is_err() {
            return Err(format!("/etc/subuid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(username) {
            let start: uid_t = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

fn read_gid_map(username: &str) -> Result<Vec<Range>,String> {
    let file = File::open(&Path::new("/etc/subgid"));
    let mut res = Vec::new();
    let mut reader = BufReader::new(file);
    for (num, line) in reader.lines().enumerate() {
        let line = try!(line.map_err(
            |e| format!("Error reading /etc/subgid: {}", e)));
        let parts: Vec<&str> = line.as_slice().split(':').collect();
        let start = FromStr::from_str(parts[1]);
        let count = FromStr::from_str(parts[2].trim_right());
        if parts.len() != 3 || start.is_err() || count.is_err() {
            return Err(format!("/etc/subgid:{}: Bad syntax", num+1));
        }
        if parts[0].eq(username) {
            let start: gid_t = start.unwrap();
            let end = start + count.unwrap() - 1;
            res.push(Range::new(start, end));
        }
    }
    return Ok(res);
}

pub fn match_ranges(req: &Vec<Range>, allowed: &Vec<Range>, own_id: uid_t)
    -> Result<Vec<(uid_t, uid_t, uid_t)>, ()>
{
    let mut res = vec!((0, own_id, 1));
    let mut reqiter = req.iter();
    let mut reqval = *reqiter.next().unwrap();
    let mut allowiter = allowed.iter();
    let mut allowval = *allowiter.next().unwrap();
    loop {
        if reqval.start == 0 {
            reqval = reqval.shift(1);
        }
        if allowval.start == 0 {
            allowval = allowval.shift(1);
        }
        let clen = min(reqval.len(), allowval.len());
        if clen > 0 {
            res.push((reqval.start, allowval.start, clen));
        }
        reqval = reqval.shift(clen);
        allowval = allowval.shift(clen);
        if reqval.len() == 0 {
            reqval = match reqiter.next() {
                Some(val) => *val,
                None => break,
            };
        }
        if allowval.len() == 0 {
            allowval = match allowiter.next() {
                Some(val) => *val,
                None => return Err(()),
            };
        }
    }
    return Ok(res);
}

pub fn get_max_uidmap() -> Result<Uidmap, String>
{
    let mut cmd = Command::new("id");
    cmd.arg("--user").arg("--name");
    if let Ok(path) = env::var("HOST_PATH") {
        cmd.env("PATH", path);
    }
    cmd.stdin(Stdio::null()).stderr(Stdio::inherit());
    let username = try!(cmd.output()
        .map_err(|e| format!("Error running `id --user --name`: {}", e))
        .and_then(|out| if out.status.status() == 0 { Ok(out.output) } else
            { Err(format!("Error running `id --user --name`")) })
        .and_then(|val| from_utf8(val.as_slice()).map(|x| x.trim().to_string())
                   .map_err(|e| format!("Can't decode username: {}", e))));
    let uid_map = read_uid_map(username.as_slice()).ok();
    let gid_map = read_gid_map(username.as_slice()).ok();

    let uid = unsafe { geteuid() };
    let gid = unsafe { getegid() };
    if let (Some(uid_map), Some(gid_map)) = (uid_map, gid_map) {
        if uid_map.len() == 0 && gid_map.len() == 0 && uid == 0 {
            let uid_rng = try!(read_uid_ranges("/proc/self/uid_map", true));
            let gid_rng = try!(read_uid_ranges("/proc/self/gid_map", true));
            return Ok(Ranges(
                uid_rng.into_iter()
                    .map(|r| (r.start, r.start, r.end-r.start+1)).collect(),
                gid_rng.into_iter()
                    .map(|r| (r.start, r.start, r.end-r.start+1)).collect()));
        }
        let mut uids = vec!((0, uid, 1));
        for &rng in uid_map.iter() {
            let mut rng = rng;
            if uid >= rng.start && uid <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(uid == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            uids.push((rng.start, rng.start, rng.len()));
        }

        let mut gids = vec!((0, gid, 1));
        for &rng in gid_map.iter() {
            let mut rng = rng;
            if gid >= rng.start && gid <= rng.end {
                // TODO(tailhook) implement better heuristic
                assert!(gid == rng.start);
                rng = rng.shift(1);
                if rng.len() == 0 { continue; }
            }
            gids.push((rng.start, rng.start, rng.len()));
        }

        return Ok(Ranges(uids, gids));
    } else {
        warn!(concat!("Your system doesn't have /etc/subuid and /etc/subgid.",
            " Presumably your system is too old. Some features may not work"));
        return Ok(Singleton(uid, gid));
    }
}

pub fn apply_uidmap(pid: pid_t, map: &Uidmap) -> Result<(), IoError> {
    match map {
        &Singleton(uid, gid) => {
            let uid_map = format!("0 {} 1", uid);
            debug!("Writing uid_map: {}", uid_map);
            let uid_map_file = Path::new("/proc")
                              .join(pid.to_string())
                              .join("uid_map");
            try!(OpenOptions().write(true).open(&uid_map_file)
                .write_all(&uid_map));
            let gid_map = format!("0 {} 1", gid);
            debug!("Writing gid_map: {}", gid_map);
            let gid_map_file = Path::new("/proc")
                              .join(pid.to_string())
                              .join("gid_map");
            try!(OpenOptions().write().open(&gid_map_file)
                .write_all(uid_map.as_slice()));
        }
        &Ranges(ref uids, ref gids) => {
            let myuid = unsafe { geteuid() };
            if myuid > 0 {
                let mut cmd = Command::new("newuidmap");
                cmd.stdin(Stdio::null());
                cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
                cmd.arg(pid.to_string());
                for &(req, allowed, count) in uids.iter() {
                    cmd
                        .arg(req.to_string())
                        .arg(allowed.to_string())
                        .arg(count.to_string());
                }
                info!("Uid map command: {:?}", cmd);
                match cmd.status() {
                    Ok(s) if s.success() => {},
                    Ok(s) if s.signal().is_some() => {
                        return Err(IoError::new(Interrupted,
                            format!("Error writing uid mapping: \
                                newuidmap was killed on signal {}",
                                s.signal().unwrap())));
                    }
                    Ok(s) => {
                        return Err(IoError::new(Other,
                            format!("Error writing uid mapping: \
                                newuidmap was exit with status {}",
                                s.status())));
                    }
                    Err(e) => return Err(e),
                }

                let mut cmd = Command::new("newgidmap");
                cmd.stdin(Stdio::null());
                cmd.stdout(Stdio::inherit()).stderr(Stdio::inherit());
                cmd.arg(pid.to_string());
                for &(req, allowed, count) in gids.iter() {
                    cmd
                        .arg(req.to_string())
                        .arg(allowed.to_string())
                        .arg(count.to_string());
                }
                info!("Gid map command: {:?}", cmd);
                match cmd.status() {
                    Ok(s) if s.success() => {},
                    Ok(s) if s.signal().is_some() => {
                        return Err(IoError::new(Interrupted,
                            format!("Error writing gid mapping: \
                                newuidmap was killed on signal {}",
                                s.signal().unwrap())));
                    }
                    Ok(s) => {
                        return Err(IoError::new(Other,
                            format!("Error writing gid mapping: \
                                newuidmap was exit with status {}",
                                s.status())));
                    }
                    Err(e) => return Err(e),
                }
            } else {
                let mut membuf = Vec::with_capacity(100);
                for &(ins, outs, cnt) in uids.iter() {
                    try!(writeln!(&mut membuf, "{} {} {}", ins, outs, cnt));
                }
                let mut file = try!(File::create(&Path::new(
                    format!("/proc/{}/uid_map", pid))));
                let value = membuf.into_inner();
                debug!("Writing uid map ```{}```",
                    String::from_utf8_lossy(value.as_slice()));
                try!(file.write(value.as_slice()));

                let mut membuf = Vec::with_capacity(100);
                for &(ins, outs, cnt) in gids.iter() {
                    try!(writeln!(&mut membuf, "{} {} {}", ins, outs, cnt));
                }
                let mut file = try!(File::create(&Path::new(
                    format!("/proc/{}/gid_map", pid))));
                let value = membuf.into_inner();
                debug!("Writing gid map ```{}```",
                    String::from_utf8_lossy(value.as_slice()));
                try!(file.write(value.as_slice()));
            }
        }
    }
    return Ok(());
}

fn read_uid_ranges(path: &str, read_inside: bool) -> Result<Vec<Range>, String>
{
    let mut file = BufReader::new(try!(File::open(&Path::new(path))
        .map_err(|e| format!("Error reading uid/gid map: {}", e))));
    let mut result = vec!();
    for line in file.lines() {
        let line = try!(line
            .map_err(|e| format!("Error reading uid/gid map: {}", e)));
        let mut words = line.as_slice().words();
        let inside = try!(words.next().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("uid/gid map format error")));
        let outside = try!(words.next().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("uid/gid map format error")));
        let count = try!(words.next().and_then(|x| FromStr::from_str(x).ok())
            .ok_or(format!("uid/gid map format error")));
        if read_inside {
            result.push(Range { start: inside, end: inside+count-1 });
        } else {
            result.push(Range { start: outside, end: outside+count-1 });
        }
    }
    return Ok(result);
}

pub fn map_users(settings: &Settings, uids: &Vec<Range>, gids: &Vec<Range>)
    -> Result<Uidmap, String>
{
    let default_uids = vec!(Range { start: 0, end: 0 });
    let default_gids = vec!(Range { start: 0, end: 0 });
    let uids = if uids.len() > 0 { uids } else { &default_uids };
    let gids = if gids.len() > 0 { gids } else { &default_gids };
    if settings.uid_map.is_none() {
        let ranges = try!(read_uid_ranges("/proc/self/uid_map", true));
        let uid_map = try!(match_ranges(uids, &ranges, 0)
            .map_err(|()| {
                return format!("Number of allowed subuids is too small. \
                    Required {}, allowed {}. You either need to increase \
                    allowed numbers in /etc/subuid (preferred) or decrease \
                    needed ranges in vagga.yaml by adding `uids` key \
                    to container config", uids, ranges);
            }));
        let ranges = try!(read_uid_ranges("/proc/self/gid_map", true));
        let gid_map = try!(match_ranges(gids, &ranges, 0)
            .map_err(|()| {
                return format!("Number of allowed subgids is too small. \
                    Required {}, allowed {}. You either need to increase \
                    allowed numbers in /etc/subgid (preferred) or decrease \
                    needed ranges in vagga.yaml by adding `gids` key \
                    to container config", gids, ranges);
            }));
        return Ok(Ranges(uid_map, gid_map));
    } else {
        let &(ref uids, ref gids) = settings.uid_map.as_ref().unwrap();
        return Ok(Ranges(uids.clone(), gids.clone()));
    }
}
<|endoftext|>


<commit_before>#![feature(slice_patterns)]
use std::env;
use std::fs::File;
use std::io::Read;
use std::collections::HashMap;

#[derive(Clone)]
enum Expr {
  Int(i64),
  Str(String)
}

// XXX this should probably be split up. The VM has the instructions, a "Thread" has locals, stack, and arguments.
struct VM<'a> {
  instructions: Vec<&'a str>,
  labels: HashMap<&'a str, usize>,
}

fn parse_labels<'a>(instructions: &Vec<&'a str>) -> HashMap<&'a str, usize> {
  let mut map = HashMap::new();
  for (i, instr) in instructions.iter().enumerate() {
    if instr.starts_with("$label") {
      match instr.split(' ').nth(1) {
        Some(label) => map.insert(label, i),
        None => panic!("VM error: missing label"),
      };
    }
  };
  map
}

impl<'a> VM<'a> {
  pub fn new(instructions: Vec<&'a str>) -> VM {
    VM {
      labels: parse_labels(&instructions),
      instructions: instructions,
    }
  }

  pub fn run(&mut self) {
    let ip = self.labels.get("main").unwrap_or(&0).clone();
    self.run_call(ip, &mut Vec::new());
  }

  // TODO locals
  pub fn run_call(&mut self, start_ip: usize, stack: &mut Vec<Expr>) -> Option<Expr> {
    let mut ip = start_ip;
    let len = self.instructions.len();
    let locals: &mut HashMap<&'a str, Expr> = &mut HashMap::new();

    let mut carry = false; // comparison carry

    let mut arguments = Vec::new();

    while ip < len {
      let instr = self.instructions[ip];
      ip += 1;
      if instr.starts_with("#") || instr.trim() == "" {
        continue;
      }
      match instr.split(" ").collect::<Vec<_>>().as_slice() {
        &["$label", _] => (),
        &["push", "int", n] => stack.push(Expr::Int(n.parse::<i64>().unwrap())),
        &["push", "str", ref n..] => stack.push(Expr::Str(String::from(n.join(" ")))),
        &["push", "arg"] => arguments.push(stack.pop().expect("Nothing on the stack to push to the arguments")),

        &["dup"] => {
          let expr = stack.pop().expect("Nothing on the stack to pop");
          stack.push(expr.clone());
          stack.push(expr);
        },

        // how = "always" | "carry"
        &["jump", how, n] => {
          let new_ip = if n.starts_with("$") {
            self.labels.get(n.trim_matches('$')).expect("VM error: no such label").clone()
          } else {
            n.parse::<usize>().unwrap()
          };

          if new_ip < len {
            // only check carry if we are conditionally jumping
            if how == "always" || (how == "carry" && carry) {
              ip = new_ip;
            }
          } else {
            panic!("VM error: trying to jump past end of bytecode ({} > {})", new_ip, len);
          }
        },

        &["carry", "set", val] => carry = val == "true",
        &["carry", "invert"] => carry = !carry,

        &["call", "primitive", "say"] => if let Some(Expr::Str(arg)) = stack.pop() {
          println!("{}", arg);
        } else {
          panic!("VM error: incorrect arguments to `say`");
        },

        &["call", "primitive", "typeof"] => match stack.pop() {
            Some(Expr::Str(_)) => stack.push(Expr::Str(String::from("str"))),
            Some(Expr::Int(_)) => stack.push(Expr::Str(String::from("int"))),
            None => panic!("VM error: no argument supplied to `typeof`"),
        },

        &["call", "primitive", "convert"] => if let
            (Some(Expr::Str(from)), Some(Expr::Str(to)))
            =
            (stack.pop(), stack.pop()) {

            match (from.as_ref(), to.as_ref()) {
                ("int", "str") => if let Some(Expr::Int(arg)) = stack.pop() {
                  stack.push(Expr::Str(arg.to_string()));
                } else {
                  panic!("VM error: cannot convert int to string");
                },
                _ => panic!("VM error: cannot convert from {} to {}", from, to),
            }
        }
        else {
            panic!("VM error: bad arguments to primitive call `convert`");
        },

        // TODO resolve name in scope, check type, apply
        &["call", name] => {
          let new_ip = self.labels.get(name).expect("VM error: no such label").clone();
          let ret = self.run_call(new_ip, &mut arguments);
          arguments = Vec::new();
          match ret {
            Some(val) => stack.push(val),
            None => (),
          };
        },

        &["ret"] => {
          return None;
        },

        &["ret", "val"] => {
          if stack.len() != 1 {
            panic!("VM error: trying to return a value but the stack size is not 1 ({})", stack.len());
          }
          return stack.pop();
        },

        // NOTE: it's stack[*-1] OP stack[*-2]
        // which means if have stack=[1, 2]
        // you'll have 2 OP 1
        /// XXX this should JUST be a call!
        &["math", op] => if let
            (Some(Expr::Int(arg1)), Some(Expr::Int(arg2)))
            =
            (stack.pop(), stack.pop()) {
          let value = match op {
            "+" => arg1 + arg2,
            "-" => arg1 - arg2,
            //"/" => TODO. this should have very specific rules...
            "*" => arg1 * arg2,
            _ => panic!("VM error: unrecognized `op` operator: {}", op),
          };
          stack.push(Expr::Int(value));
        } else {
          panic!("VM error: not enough arguments to `math`");
        },

        // same note about "2 OP 1"
        // XXX same note about how this should be a `call`
        &["cmp", op] => carry = match (stack.pop(), stack.pop()) {
            (Some(Expr::Int(arg1)), Some(Expr::Int(arg2))) => match op {
                "<" => arg1 < arg2,
                ">" => arg1 > arg2,
                "<=" => arg1 <= arg2,
                ">=" => arg1 >= arg2,
                "=" => arg1 == arg2,
                _ => panic!("VM error: unrecognized `cmp` operator for int: {}", op),
            },
            (Some(Expr::Str(arg1)), Some(Expr::Str(arg2))) => match op {
                "=" => arg1 == arg2,
                _ => panic!("VM error: unrecognized `cmp` operator for str: {}", op),
            },
            _ => panic!("VM error: not enough arguments to `cmp`"),
        },

        // TODO this *shouldn't* be an identifier
        //      the backend should transform SSAF to tons of load/store
        //      (probably)
        //      (see backend#1)
        &["local", "load", name] => match locals.get(name) {
          Some(expr) => stack.push(expr.clone()),
          None       => panic!("VM error: unknown local variable {}", name),
        },

        &["local", "store", name] => if let Some(arg) = stack.pop() {
          // XXX check
          let _ = locals.insert(name, arg);
        } else {
          panic!("VM error: not enough arguments for `local store`");
        },

        &[instr, ..] => panic!("VM error: no such instruction {}", instr),
        &[..] => panic!("VM error: Unrecognized instruction!"),
      }
    }

    None
  }
}

fn main() {
	if let Some(path) = env::args().nth(1) {
		println!("Loading file {}", path);

    match File::open(&path) {
      Ok(mut file) => {
        let mut content = String::new();
        if let Ok(_) = file.read_to_string(&mut content) {
          // remove trailing newlines, split by line
          let instructions = content.trim_matches('\n').split('\n').collect();
          let mut vm = VM::new(instructions);
          vm.run();
        }
      },
      Err(err) => {
        panic!("Can't read file {} (error: {})", path, err);
      }
    }
	} else {
    panic!("Please give a file argument");
  }

  //     let path = Path::new("chry.fa");
  // let mut file = BufferedReader::new(File::open(&path));
  // for line in file.lines() {
  //  print!("{}", line.unwrap());
  //      }
}
<commit_msg>Allow loading asm from stdin<commit_after>#![feature(slice_patterns)]
use std::env;
use std::fs::File;
use std::io::Read;
use std::collections::HashMap;

#[derive(Clone)]
enum Expr {
  Int(i64),
  Str(String)
}

// XXX this should probably be split up. The VM has the instructions, a "Thread" has locals, stack, and arguments.
struct VM<'a> {
  instructions: Vec<&'a str>,
  labels: HashMap<&'a str, usize>,
}

fn parse_labels<'a>(instructions: &Vec<&'a str>) -> HashMap<&'a str, usize> {
  let mut map = HashMap::new();
  for (i, instr) in instructions.iter().enumerate() {
    if instr.starts_with("$label") {
      match instr.split(' ').nth(1) {
        Some(label) => map.insert(label, i),
        None => panic!("VM error: missing label"),
      };
    }
  };
  map
}

impl<'a> VM<'a> {
  pub fn new(instructions: Vec<&'a str>) -> VM {
    VM {
      labels: parse_labels(&instructions),
      instructions: instructions,
    }
  }

  pub fn run(&mut self) {
    let ip = self.labels.get("main").unwrap_or(&0).clone();
    self.run_call(ip, &mut Vec::new());
  }

  // TODO locals
  pub fn run_call(&mut self, start_ip: usize, stack: &mut Vec<Expr>) -> Option<Expr> {
    let mut ip = start_ip;
    let len = self.instructions.len();
    let locals: &mut HashMap<&'a str, Expr> = &mut HashMap::new();

    let mut carry = false; // comparison carry

    let mut arguments = Vec::new();

    while ip < len {
      let instr = self.instructions[ip];
      ip += 1;
      if instr.starts_with("#") || instr.trim() == "" {
        continue;
      }
      match instr.split(" ").collect::<Vec<_>>().as_slice() {
        &["$label", _] => (),
        &["push", "int", n] => stack.push(Expr::Int(n.parse::<i64>().unwrap())),
        &["push", "str", ref n..] => stack.push(Expr::Str(String::from(n.join(" ")))),
        &["push", "arg"] => arguments.push(stack.pop().expect("Nothing on the stack to push to the arguments")),

        &["dup"] => {
          let expr = stack.pop().expect("Nothing on the stack to pop");
          stack.push(expr.clone());
          stack.push(expr);
        },

        // how = "always" | "carry"
        &["jump", how, n] => {
          let new_ip = if n.starts_with("$") {
            self.labels.get(n.trim_matches('$')).expect("VM error: no such label").clone()
          } else {
            n.parse::<usize>().unwrap()
          };

          if new_ip < len {
            // only check carry if we are conditionally jumping
            if how == "always" || (how == "carry" && carry) {
              ip = new_ip;
            }
          } else {
            panic!("VM error: trying to jump past end of bytecode ({} > {})", new_ip, len);
          }
        },

        &["carry", "set", val] => carry = val == "true",
        &["carry", "invert"] => carry = !carry,

        &["call", "primitive", "say"] => if let Some(Expr::Str(arg)) = stack.pop() {
          println!("{}", arg);
        } else {
          panic!("VM error: incorrect arguments to `say`");
        },

        &["call", "primitive", "typeof"] => match stack.pop() {
            Some(Expr::Str(_)) => stack.push(Expr::Str(String::from("str"))),
            Some(Expr::Int(_)) => stack.push(Expr::Str(String::from("int"))),
            None => panic!("VM error: no argument supplied to `typeof`"),
        },

        &["call", "primitive", "convert"] => if let
            (Some(Expr::Str(from)), Some(Expr::Str(to)))
            =
            (stack.pop(), stack.pop()) {

            match (from.as_ref(), to.as_ref()) {
                ("int", "str") => if let Some(Expr::Int(arg)) = stack.pop() {
                  stack.push(Expr::Str(arg.to_string()));
                } else {
                  panic!("VM error: cannot convert int to string");
                },
                _ => panic!("VM error: cannot convert from {} to {}", from, to),
            }
        }
        else {
            panic!("VM error: bad arguments to primitive call `convert`");
        },

        // TODO resolve name in scope, check type, apply
        &["call", name] => {
          let new_ip = self.labels.get(name).expect("VM error: no such label").clone();
          let ret = self.run_call(new_ip, &mut arguments);
          arguments = Vec::new();
          match ret {
            Some(val) => stack.push(val),
            None => (),
          };
        },

        &["ret"] => {
          return None;
        },

        &["ret", "val"] => {
          if stack.len() != 1 {
            panic!("VM error: trying to return a value but the stack size is not 1 ({})", stack.len());
          }
          return stack.pop();
        },

        // NOTE: it's stack[*-1] OP stack[*-2]
        // which means if have stack=[1, 2]
        // you'll have 2 OP 1
        /// XXX this should JUST be a call!
        &["math", op] => if let
            (Some(Expr::Int(arg1)), Some(Expr::Int(arg2)))
            =
            (stack.pop(), stack.pop()) {
          let value = match op {
            "+" => arg1 + arg2,
            "-" => arg1 - arg2,
            //"/" => TODO. this should have very specific rules...
            "*" => arg1 * arg2,
            _ => panic!("VM error: unrecognized `op` operator: {}", op),
          };
          stack.push(Expr::Int(value));
        } else {
          panic!("VM error: not enough arguments to `math`");
        },

        // same note about "2 OP 1"
        // XXX same note about how this should be a `call`
        &["cmp", op] => carry = match (stack.pop(), stack.pop()) {
            (Some(Expr::Int(arg1)), Some(Expr::Int(arg2))) => match op {
                "<" => arg1 < arg2,
                ">" => arg1 > arg2,
                "<=" => arg1 <= arg2,
                ">=" => arg1 >= arg2,
                "=" => arg1 == arg2,
                _ => panic!("VM error: unrecognized `cmp` operator for int: {}", op),
            },
            (Some(Expr::Str(arg1)), Some(Expr::Str(arg2))) => match op {
                "=" => arg1 == arg2,
                _ => panic!("VM error: unrecognized `cmp` operator for str: {}", op),
            },
            _ => panic!("VM error: not enough arguments to `cmp`"),
        },

        // TODO this *shouldn't* be an identifier
        //      the backend should transform SSAF to tons of load/store
        //      (probably)
        //      (see backend#1)
        &["local", "load", name] => match locals.get(name) {
          Some(expr) => stack.push(expr.clone()),
          None       => panic!("VM error: unknown local variable {}", name),
        },

        &["local", "store", name] => if let Some(arg) = stack.pop() {
          // XXX check
          let _ = locals.insert(name, arg);
        } else {
          panic!("VM error: not enough arguments for `local store`");
        },

        &[instr, ..] => panic!("VM error: no such instruction {}", instr),
        &[..] => panic!("VM error: Unrecognized instruction!"),
      }
    }

    None
  }
}

fn main() {
    let mut content = String::new();
    let result;

	if let Some(path) = env::args().nth(1) {
		println!("Loading file {}", path);

        match File::open(&path) {
          Ok(mut file) => result = file.read_to_string(&mut content),
          Err(err) => {
            panic!("Can't read file {} (error: {})", path, err);
          }
        }
    } else {
        result = std::io::stdin().read_to_string(&mut content);
   }

   if let Ok(_) = result {
     // remove trailing newlines, split by line
     let instructions = content.trim_matches('\n').split('\n').collect();
     let mut vm = VM::new(instructions);
     vm.run();
   }

  //     let path = Path::new("chry.fa");
  // let mut file = BufferedReader::new(File::open(&path));
  // for line in file.lines() {
  //  print!("{}", line.unwrap());
  //      }
}
<|endoftext|>


<commit_before>use std::result;
use std::fmt;
use stringutils::status_to_string;

/// Error type
///
/// `cfitsio` passes errors through integer status codes. This struct wraps this and its associated
/// error message.
#[derive(Debug, PartialEq, Eq)]
pub struct FitsError {
    pub status: i32,
    pub message: String,
}

/// Display implementation for `FitsError`
///
/// This enables the error to be printed in a user-facing way
impl fmt::Display for FitsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let full_message = format!("Fits status {}: {}", self.status, self.message);
        full_message.fmt(f)
    }
}

/// Error implementation for `FitsError`
///
/// This enables the fits error type to be treated as a Box<Error>
impl ::std::error::Error for FitsError {
    fn description(&self) -> &str {
        "fits error"
    }
}

/// Macro for returning a FITS error type
macro_rules! fits_try {
    ($status: ident, $e: expr) => {
        match $status {
            0 => Ok($e),
            _ => {
                Err(FitsError {
                    status: $status,
                    message: stringutils::status_to_string($status).unwrap(),
                })
            }
        }
    }
}

/// FITS specific result type
///
/// This is a shortcut for a result with `FitsError` as the error type
pub type FitsResult<T> = result::Result<T, FitsError>;

pub fn status_to_error(status: i32) -> FitsResult<()> {
    Err(FitsError {
        status: status,
        message: status_to_string(status).unwrap(),
    })
}

#[cfg(test)]
mod test {
    use std::error::Error;
    use super::FitsError;

    #[test]
    fn fits_error_implements_error() {
        fn foo() -> ::std::result::Result<(), Box<Error>> {
            let err = FitsError {
                status: 100,
                message: "test".into(),
            };
            Err(err.into())
        }

        if let Err(e) = foo() {
            assert_eq!(format!("{}", e), "Fits status 100: test".to_string());
            assert_eq!(e.description(), "fits error");
        }
    }
}
<commit_msg>Add From<&str> for FitsError<commit_after>use std::result;
use std::fmt;
use stringutils::status_to_string;

/// Error type
///
/// `cfitsio` passes errors through integer status codes. This struct wraps this and its associated
/// error message.
#[derive(Debug, PartialEq, Eq)]
pub struct FitsError {
    pub status: i32,
    pub message: String,
}

/// Display implementation for `FitsError`
///
/// This enables the error to be printed in a user-facing way
impl fmt::Display for FitsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let full_message = format!("Fits status {}: {}", self.status, self.message);
        full_message.fmt(f)
    }
}

/// Error implementation for `FitsError`
///
/// This enables the fits error type to be treated as a Box<Error>
impl ::std::error::Error for FitsError {
    fn description(&self) -> &str {
        "fits error"
    }
}

impl<'a> From<&'a str> for FitsError {
    fn from(s: &'a str) -> Self {
        FitsError {
            status: 600,
            message: s.to_string(),
        }
    }
}

/// Macro for returning a FITS error type
macro_rules! fits_try {
    ($status: ident, $e: expr) => {
        match $status {
            0 => Ok($e),
            _ => {
                Err(FitsError {
                    status: $status,
                    message: stringutils::status_to_string($status).unwrap(),
                })
            }
        }
    }
}

/// FITS specific result type
///
/// This is a shortcut for a result with `FitsError` as the error type
pub type FitsResult<T> = result::Result<T, FitsError>;

pub fn status_to_error(status: i32) -> FitsResult<()> {
    Err(FitsError {
        status: status,
        message: status_to_string(status).unwrap(),
    })
}

#[cfg(test)]
mod test {
    use std::error::Error;
    use super::FitsError;

    #[test]
    fn fits_error_implements_error() {
        fn foo() -> ::std::result::Result<(), Box<Error>> {
            let err = FitsError {
                status: 100,
                message: "test".into(),
            };
            Err(err.into())
        }

        if let Err(e) = foo() {
            assert_eq!(format!("{}", e), "Fits status 100: test".to_string());
            assert_eq!(e.description(), "fits error");
        }
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Execution plan for reading Parquet files

use std::any::Any;
use std::fs::File;
use std::rc::Rc;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::{fmt, thread};

use super::{RecordBatchStream, SendableRecordBatchStream};
use crate::error::{DataFusionError, Result};
use crate::physical_plan::ExecutionPlan;
use crate::physical_plan::{common, Partitioning};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::RecordBatch;
use parquet::file::reader::SerializedFileReader;

use crossbeam::channel::{bounded, Receiver, RecvError, Sender};
use fmt::Debug;
use parquet::arrow::{ArrowReader, ParquetFileArrowReader};

use async_trait::async_trait;
use futures::stream::Stream;

/// Execution plan for scanning a Parquet file
#[derive(Debug, Clone)]
pub struct ParquetExec {
    /// Path to directory containing partitioned Parquet files with the same schema
    filenames: Vec<String>,
    /// Schema after projection is applied
    schema: SchemaRef,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
}

impl ParquetExec {
    /// Create a new Parquet reader execution plan
    pub fn try_new(
        path: &str,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Result<Self> {
        let mut filenames: Vec<String> = vec![];
        common::build_file_list(path, &mut filenames, ".parquet")?;
        if filenames.is_empty() {
            Err(DataFusionError::Plan("No files found".to_string()))
        } else {
            let file = File::open(&filenames[0])?;
            let file_reader = Rc::new(SerializedFileReader::new(file)?);
            let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
            let schema = arrow_reader.get_schema()?;

            let projection = match projection {
                Some(p) => p,
                None => (0..schema.fields().len()).collect(),
            };

            let projected_schema = Schema::new(
                projection
                    .iter()
                    .map(|i| schema.field(*i).clone())
                    .collect(),
            );

            Ok(Self {
                filenames,
                schema: Arc::new(projected_schema),
                projection,
                batch_size,
            })
        }
    }
}

#[async_trait]
impl ExecutionPlan for ParquetExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        // this is a leaf node and has no children
        vec![]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        Partitioning::UnknownPartitioning(self.filenames.len())
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if children.is_empty() {
            Ok(Arc::new(self.clone()))
        } else {
            Err(DataFusionError::Internal(format!(
                "Children cannot be replaced in {:?}",
                self
            )))
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        // because the parquet implementation is not thread-safe, it is necessary to execute
        // on a thread and communicate with channels
        let (response_tx, response_rx): (
            Sender<Option<ArrowResult<RecordBatch>>>,
            Receiver<Option<ArrowResult<RecordBatch>>>,
        ) = bounded(2);

        let filename = self.filenames[partition].clone();
        let projection = self.projection.clone();
        let batch_size = self.batch_size;

        thread::spawn(move || {
            if let Err(e) = read_file(&filename, projection, batch_size, response_tx) {
                println!("Parquet reader thread terminated due to error: {:?}", e);
            }
        });

        Ok(Box::pin(ParquetStream {
            schema: self.schema.clone(),
            response_rx,
        }))
    }
}

fn send_result(
    response_tx: &Sender<Option<ArrowResult<RecordBatch>>>,
    result: Option<ArrowResult<RecordBatch>>,
) -> Result<()> {
    response_tx
        .send(result)
        .map_err(|e| DataFusionError::Execution(e.to_string()))?;
    Ok(())
}

fn read_file(
    filename: &str,
    projection: Vec<usize>,
    batch_size: usize,
    response_tx: Sender<Option<ArrowResult<RecordBatch>>>,
) -> Result<()> {
    let file = File::open(&filename)?;
    let file_reader = Rc::new(SerializedFileReader::new(file)?);
    let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
    let mut batch_reader =
        arrow_reader.get_record_reader_by_columns(projection.clone(), batch_size)?;
    loop {
        match batch_reader.next() {
            Some(Ok(batch)) => send_result(&response_tx, Some(Ok(batch)))?,
            None => {
                // finished reading file
                send_result(&response_tx, None)?;
                break;
            }
            Some(Err(e)) => {
                let err_msg =
                    format!("Error reading batch from {}: {}", filename, e.to_string());
                // send error to operator
                send_result(
                    &response_tx,
                    Some(Err(ArrowError::ParquetError(err_msg.clone()))),
                )?;
                // terminate thread with error
                return Err(DataFusionError::Execution(err_msg));
            }
        }
    }
    Ok(())
}

struct ParquetStream {
    schema: SchemaRef,
    response_rx: Receiver<Option<ArrowResult<RecordBatch>>>,
}

impl Stream for ParquetStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        _: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        match self.response_rx.recv() {
            Ok(batch) => Poll::Ready(batch),
            // RecvError means receiver has exited and closed the channel
            Err(RecvError) => Poll::Ready(None),
        }
    }
}

impl RecordBatchStream for ParquetStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::StreamExt;
    use std::env;

    #[tokio::test]
    async fn test() -> Result<()> {
        let testdata =
            env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
        let filename = format!("{}/alltypes_plain.parquet", testdata);
        let parquet_exec = ParquetExec::try_new(&filename, Some(vec![0, 1, 2]), 1024)?;
        assert_eq!(parquet_exec.output_partitioning().partition_count(), 1);

        let mut results = parquet_exec.execute(0).await?;
        let batch = results.next().await.unwrap()?;

        assert_eq!(8, batch.num_rows());
        assert_eq!(3, batch.num_columns());

        let schema = batch.schema();
        let field_names: Vec<&str> =
            schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);

        let batch = results.next().await;
        assert!(batch.is_none());

        let batch = results.next().await;
        assert!(batch.is_none());

        let batch = results.next().await;
        assert!(batch.is_none());

        Ok(())
    }
}
<commit_msg>ARROW-10454: [Rust] [Datafusion] support creating ParquetExec from filelist and schema<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Execution plan for reading Parquet files

use std::any::Any;
use std::fs::File;
use std::rc::Rc;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::{fmt, thread};

use super::{RecordBatchStream, SendableRecordBatchStream};
use crate::error::{DataFusionError, Result};
use crate::physical_plan::ExecutionPlan;
use crate::physical_plan::{common, Partitioning};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::RecordBatch;
use parquet::file::reader::SerializedFileReader;

use crossbeam::channel::{bounded, Receiver, RecvError, Sender};
use fmt::Debug;
use parquet::arrow::{ArrowReader, ParquetFileArrowReader};

use async_trait::async_trait;
use futures::stream::Stream;

/// Execution plan for scanning a Parquet file
#[derive(Debug, Clone)]
pub struct ParquetExec {
    /// Path to directory containing partitioned Parquet files with the same schema
    filenames: Vec<String>,
    /// Schema after projection is applied
    schema: SchemaRef,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
}

impl ParquetExec {
    /// Create a new Parquet reader execution plan
    pub fn try_new(
        path: &str,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Result<Self> {
        let mut filenames: Vec<String> = vec![];
        common::build_file_list(path, &mut filenames, ".parquet")?;
        if filenames.is_empty() {
            Err(DataFusionError::Plan("No files found".to_string()))
        } else {
            let file = File::open(&filenames[0])?;
            let file_reader = Rc::new(SerializedFileReader::new(file)?);
            let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
            let schema = arrow_reader.get_schema()?;

            Ok(Self::new(filenames, schema, projection, batch_size))
        }
    }

    /// Create a new Parquet reader execution plan with provided files and schema
    pub fn new(
        filenames: Vec<String>,
        schema: Schema,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Self {
        let projection = match projection {
            Some(p) => p,
            None => (0..schema.fields().len()).collect(),
        };

        let projected_schema = Schema::new(
            projection
                .iter()
                .map(|i| schema.field(*i).clone())
                .collect(),
        );

        Self {
            filenames,
            schema: Arc::new(projected_schema),
            projection,
            batch_size,
        }
    }
}

#[async_trait]
impl ExecutionPlan for ParquetExec {
    /// Return a reference to Any that can be used for downcasting
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        // this is a leaf node and has no children
        vec![]
    }

    /// Get the output partitioning of this plan
    fn output_partitioning(&self) -> Partitioning {
        Partitioning::UnknownPartitioning(self.filenames.len())
    }

    fn with_new_children(
        &self,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if children.is_empty() {
            Ok(Arc::new(self.clone()))
        } else {
            Err(DataFusionError::Internal(format!(
                "Children cannot be replaced in {:?}",
                self
            )))
        }
    }

    async fn execute(&self, partition: usize) -> Result<SendableRecordBatchStream> {
        // because the parquet implementation is not thread-safe, it is necessary to execute
        // on a thread and communicate with channels
        let (response_tx, response_rx): (
            Sender<Option<ArrowResult<RecordBatch>>>,
            Receiver<Option<ArrowResult<RecordBatch>>>,
        ) = bounded(2);

        let filename = self.filenames[partition].clone();
        let projection = self.projection.clone();
        let batch_size = self.batch_size;

        thread::spawn(move || {
            if let Err(e) = read_file(&filename, projection, batch_size, response_tx) {
                println!("Parquet reader thread terminated due to error: {:?}", e);
            }
        });

        Ok(Box::pin(ParquetStream {
            schema: self.schema.clone(),
            response_rx,
        }))
    }
}

fn send_result(
    response_tx: &Sender<Option<ArrowResult<RecordBatch>>>,
    result: Option<ArrowResult<RecordBatch>>,
) -> Result<()> {
    response_tx
        .send(result)
        .map_err(|e| DataFusionError::Execution(e.to_string()))?;
    Ok(())
}

fn read_file(
    filename: &str,
    projection: Vec<usize>,
    batch_size: usize,
    response_tx: Sender<Option<ArrowResult<RecordBatch>>>,
) -> Result<()> {
    let file = File::open(&filename)?;
    let file_reader = Rc::new(SerializedFileReader::new(file)?);
    let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
    let mut batch_reader =
        arrow_reader.get_record_reader_by_columns(projection.clone(), batch_size)?;
    loop {
        match batch_reader.next() {
            Some(Ok(batch)) => send_result(&response_tx, Some(Ok(batch)))?,
            None => {
                // finished reading file
                send_result(&response_tx, None)?;
                break;
            }
            Some(Err(e)) => {
                let err_msg =
                    format!("Error reading batch from {}: {}", filename, e.to_string());
                // send error to operator
                send_result(
                    &response_tx,
                    Some(Err(ArrowError::ParquetError(err_msg.clone()))),
                )?;
                // terminate thread with error
                return Err(DataFusionError::Execution(err_msg));
            }
        }
    }
    Ok(())
}

struct ParquetStream {
    schema: SchemaRef,
    response_rx: Receiver<Option<ArrowResult<RecordBatch>>>,
}

impl Stream for ParquetStream {
    type Item = ArrowResult<RecordBatch>;

    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        _: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        match self.response_rx.recv() {
            Ok(batch) => Poll::Ready(batch),
            // RecvError means receiver has exited and closed the channel
            Err(RecvError) => Poll::Ready(None),
        }
    }
}

impl RecordBatchStream for ParquetStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::StreamExt;
    use std::env;

    #[tokio::test]
    async fn test() -> Result<()> {
        let testdata =
            env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
        let filename = format!("{}/alltypes_plain.parquet", testdata);
        let parquet_exec = ParquetExec::try_new(&filename, Some(vec![0, 1, 2]), 1024)?;
        assert_eq!(parquet_exec.output_partitioning().partition_count(), 1);

        let mut results = parquet_exec.execute(0).await?;
        let batch = results.next().await.unwrap()?;

        assert_eq!(8, batch.num_rows());
        assert_eq!(3, batch.num_columns());

        let schema = batch.schema();
        let field_names: Vec<&str> =
            schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);

        let batch = results.next().await;
        assert!(batch.is_none());

        let batch = results.next().await;
        assert!(batch.is_none());

        let batch = results.next().await;
        assert!(batch.is_none());

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


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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. This file may not be copied, modified, or distributed
// except according to those terms.

#![deny(warnings)]

extern crate build_helper;
extern crate cc;

use std::env;
use std::path::PathBuf;
use std::process::Command;
use build_helper::{run, native_lib_boilerplate};

fn main() {
    // FIXME: This is a hack to support building targets that don't
    // support jemalloc alongside hosts that do. The jemalloc build is
    // controlled by a feature of the std crate, and if that feature
    // changes between targets, it invalidates the fingerprint of
    // std's build script (this is a cargo bug); so we must ensure
    // that the feature set used by std is the same across all
    // targets, which means we have to build the alloc_jemalloc crate
    // for targets like emscripten, even if we don't use it.
    let target = env::var("TARGET").expect("TARGET was not set");
    let host = env::var("HOST").expect("HOST was not set");
    if target.contains("bitrig") || target.contains("emscripten") || target.contains("fuchsia") ||
       target.contains("msvc") || target.contains("openbsd") || target.contains("redox") ||
       target.contains("rumprun") || target.contains("wasm32") {
        println!("cargo:rustc-cfg=dummy_jemalloc");
        return;
    }

    // CloudABI ships with a copy of jemalloc that has been patched to
    // work well with sandboxing. Don't attempt to build our own copy,
    // as it won't build.
    if target.contains("cloudabi") {
        return;
    }

    if target.contains("android") {
        println!("cargo:rustc-link-lib=gcc");
    } else if !target.contains("windows") && !target.contains("musl") {
        println!("cargo:rustc-link-lib=pthread");
    }

    if let Some(jemalloc) = env::var_os("JEMALLOC_OVERRIDE") {
        let jemalloc = PathBuf::from(jemalloc);
        println!("cargo:rustc-link-search=native={}",
                 jemalloc.parent().unwrap().display());
        let stem = jemalloc.file_stem().unwrap().to_str().unwrap();
        let name = jemalloc.file_name().unwrap().to_str().unwrap();
        let kind = if name.ends_with(".a") {
            "static"
        } else {
            "dylib"
        };
        println!("cargo:rustc-link-lib={}={}", kind, &stem[3..]);
        return;
    }

    let link_name = if target.contains("windows") { "jemalloc" } else { "jemalloc_pic" };
    let native = match native_lib_boilerplate("jemalloc", "jemalloc", link_name, "lib") {
        Ok(native) => native,
        _ => return,
    };

    let mut cmd = Command::new("sh");
    cmd.arg(native.src_dir.join("configure")
                          .to_str()
                          .unwrap()
                          .replace("C:\\", "/c/")
                          .replace("\\", "/"))
       .current_dir(&native.out_dir)
       // jemalloc generates Makefile deps using GCC's "-MM" flag. This means
       // that GCC will run the preprocessor, and only the preprocessor, over
       // jemalloc's source files. If we don't specify CPPFLAGS, then at least
       // on ARM that step fails with a "Missing implementation for 32-bit
       // atomic operations" error. This is because no "-march" flag will be
       // passed to GCC, and then GCC won't define the
       // "__GCC_HAVE_SYNC_COMPARE_AND_SWAP_4" macro that jemalloc needs to
       // select an atomic operation implementation.
       .env("CPPFLAGS", env::var_os("CFLAGS").unwrap_or_default());

    if target.contains("ios") {
        cmd.arg("--disable-tls");
    } else if target.contains("android") {
        // We force android to have prefixed symbols because apparently
        // replacement of the libc allocator doesn't quite work. When this was
        // tested (unprefixed symbols), it was found that the `realpath`
        // function in libc would allocate with libc malloc (not jemalloc
        // malloc), and then the standard library would free with jemalloc free,
        // causing a segfault.
        //
        // If the test suite passes, however, without symbol prefixes then we
        // should be good to go!
        cmd.arg("--with-jemalloc-prefix=je_");
        cmd.arg("--disable-tls");
    } else if target.contains("dragonfly") || target.contains("musl") {
        cmd.arg("--with-jemalloc-prefix=je_");
    }

    // FIXME: building with jemalloc assertions is currently broken.
    // See <https://github.com/rust-lang/rust/issues/44152>.
    //if cfg!(feature = "debug") {
    //    cmd.arg("--enable-debug");
    //}

    cmd.arg(format!("--host={}", build_helper::gnu_target(&target)));
    cmd.arg(format!("--build={}", build_helper::gnu_target(&host)));

    // for some reason, jemalloc configure doesn't detect this value
    // automatically for this target
    if target == "sparc64-unknown-linux-gnu" {
        cmd.arg("--with-lg-quantum=4");
    }

    run(&mut cmd);

    let mut make = Command::new(build_helper::make(&host));
    make.current_dir(&native.out_dir)
        .arg("build_lib_static");

    // These are intended for mingw32-make which we don't use
    if cfg!(windows) {
        make.env_remove("MAKEFLAGS").env_remove("MFLAGS");
    }

    // mingw make seems... buggy? unclear...
    if !host.contains("windows") {
        make.arg("-j")
            .arg(env::var("NUM_JOBS").expect("NUM_JOBS was not set"));
    }

    run(&mut make);

    // The pthread_atfork symbols is used by jemalloc on android but the really
    // old android we're building on doesn't have them defined, so just make
    // sure the symbols are available.
    if target.contains("androideabi") {
        println!("cargo:rerun-if-changed=pthread_atfork_dummy.c");
        cc::Build::new()
            .flag("-fvisibility=hidden")
            .file("pthread_atfork_dummy.c")
            .compile("pthread_atfork_dummy");
    }
}
<commit_msg>enable jemalloc assertions when configured to do so<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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. This file may not be copied, modified, or distributed
// except according to those terms.

#![deny(warnings)]

extern crate build_helper;
extern crate cc;

use std::env;
use std::path::PathBuf;
use std::process::Command;
use build_helper::{run, native_lib_boilerplate};

fn main() {
    // FIXME: This is a hack to support building targets that don't
    // support jemalloc alongside hosts that do. The jemalloc build is
    // controlled by a feature of the std crate, and if that feature
    // changes between targets, it invalidates the fingerprint of
    // std's build script (this is a cargo bug); so we must ensure
    // that the feature set used by std is the same across all
    // targets, which means we have to build the alloc_jemalloc crate
    // for targets like emscripten, even if we don't use it.
    let target = env::var("TARGET").expect("TARGET was not set");
    let host = env::var("HOST").expect("HOST was not set");
    if target.contains("bitrig") || target.contains("emscripten") || target.contains("fuchsia") ||
       target.contains("msvc") || target.contains("openbsd") || target.contains("redox") ||
       target.contains("rumprun") || target.contains("wasm32") {
        println!("cargo:rustc-cfg=dummy_jemalloc");
        return;
    }

    // CloudABI ships with a copy of jemalloc that has been patched to
    // work well with sandboxing. Don't attempt to build our own copy,
    // as it won't build.
    if target.contains("cloudabi") {
        return;
    }

    if target.contains("android") {
        println!("cargo:rustc-link-lib=gcc");
    } else if !target.contains("windows") && !target.contains("musl") {
        println!("cargo:rustc-link-lib=pthread");
    }

    if let Some(jemalloc) = env::var_os("JEMALLOC_OVERRIDE") {
        let jemalloc = PathBuf::from(jemalloc);
        println!("cargo:rustc-link-search=native={}",
                 jemalloc.parent().unwrap().display());
        let stem = jemalloc.file_stem().unwrap().to_str().unwrap();
        let name = jemalloc.file_name().unwrap().to_str().unwrap();
        let kind = if name.ends_with(".a") {
            "static"
        } else {
            "dylib"
        };
        println!("cargo:rustc-link-lib={}={}", kind, &stem[3..]);
        return;
    }

    let link_name = if target.contains("windows") { "jemalloc" } else { "jemalloc_pic" };
    let native = match native_lib_boilerplate("jemalloc", "jemalloc", link_name, "lib") {
        Ok(native) => native,
        _ => return,
    };

    let mut cmd = Command::new("sh");
    cmd.arg(native.src_dir.join("configure")
                          .to_str()
                          .unwrap()
                          .replace("C:\\", "/c/")
                          .replace("\\", "/"))
       .current_dir(&native.out_dir)
       // jemalloc generates Makefile deps using GCC's "-MM" flag. This means
       // that GCC will run the preprocessor, and only the preprocessor, over
       // jemalloc's source files. If we don't specify CPPFLAGS, then at least
       // on ARM that step fails with a "Missing implementation for 32-bit
       // atomic operations" error. This is because no "-march" flag will be
       // passed to GCC, and then GCC won't define the
       // "__GCC_HAVE_SYNC_COMPARE_AND_SWAP_4" macro that jemalloc needs to
       // select an atomic operation implementation.
       .env("CPPFLAGS", env::var_os("CFLAGS").unwrap_or_default());

    if target.contains("ios") {
        cmd.arg("--disable-tls");
    } else if target.contains("android") {
        // We force android to have prefixed symbols because apparently
        // replacement of the libc allocator doesn't quite work. When this was
        // tested (unprefixed symbols), it was found that the `realpath`
        // function in libc would allocate with libc malloc (not jemalloc
        // malloc), and then the standard library would free with jemalloc free,
        // causing a segfault.
        //
        // If the test suite passes, however, without symbol prefixes then we
        // should be good to go!
        cmd.arg("--with-jemalloc-prefix=je_");
        cmd.arg("--disable-tls");
    } else if target.contains("dragonfly") || target.contains("musl") {
        cmd.arg("--with-jemalloc-prefix=je_");
    }

    if cfg!(feature = "debug") {
        // Enable jemalloc assertions.
        cmd.arg("--enable-debug");
    }

    cmd.arg(format!("--host={}", build_helper::gnu_target(&target)));
    cmd.arg(format!("--build={}", build_helper::gnu_target(&host)));

    // for some reason, jemalloc configure doesn't detect this value
    // automatically for this target
    if target == "sparc64-unknown-linux-gnu" {
        cmd.arg("--with-lg-quantum=4");
    }

    run(&mut cmd);

    let mut make = Command::new(build_helper::make(&host));
    make.current_dir(&native.out_dir)
        .arg("build_lib_static");

    // These are intended for mingw32-make which we don't use
    if cfg!(windows) {
        make.env_remove("MAKEFLAGS").env_remove("MFLAGS");
    }

    // mingw make seems... buggy? unclear...
    if !host.contains("windows") {
        make.arg("-j")
            .arg(env::var("NUM_JOBS").expect("NUM_JOBS was not set"));
    }

    run(&mut make);

    // The pthread_atfork symbols is used by jemalloc on android but the really
    // old android we're building on doesn't have them defined, so just make
    // sure the symbols are available.
    if target.contains("androideabi") {
        println!("cargo:rerun-if-changed=pthread_atfork_dummy.c");
        cc::Build::new()
            .flag("-fvisibility=hidden")
            .file("pthread_atfork_dummy.c")
            .compile("pthread_atfork_dummy");
    }
}
<|endoftext|>