<commit_before>
import os::getcwd;
import os_fs;
import str;

native "rust" mod rustrt {
    fn rust_file_is_dir(str path) -> int;
}

fn path_sep() -> str { ret str::from_char(os_fs::path_sep); }

type path = str;

fn dirname(path p) -> path {
    let int i = str::rindex(p, os_fs::path_sep as u8);
    if (i == -1) {
        i = str::rindex(p, os_fs::alt_path_sep as u8);
        if (i == -1) { ret p; }
    }
    ret str::substr(p, 0u, i as uint);
}

fn basename(path p) -> path {
    let int i = str::rindex(p, os_fs::path_sep as u8);
    if (i == -1) {
        i = str::rindex(p, os_fs::alt_path_sep as u8);
        if (i == -1) { ret p; }
    }
    auto len = str::byte_len(p);
    if (i + 1 as uint >= len) { ret p; }
    ret str::slice(p, i + 1 as uint, len);
}

// FIXME: Need some typestate to avoid bounds check when len(pre) == 0
fn connect(path pre, path post) -> path {
    auto len = str::byte_len(pre);
    ret if (pre.(len - 1u) == os_fs::path_sep as u8) {

            // Trailing '/'?
            pre + post
        } else { pre + path_sep() + post };
}

fn file_is_dir(path p) -> bool { ret rustrt::rust_file_is_dir(p) != 0; }

fn list_dir(path p) -> vec[str] {
    auto pl = str::byte_len(p);
    if (pl == 0u || p.(pl - 1u) as char != os_fs::path_sep) {
        p += path_sep();
    }
    let vec[str] full_paths = [];
    for (str filename in os_fs::list_dir(p)) {
        if (!str::eq(filename, ".")) {
            if (!str::eq(filename, "..")) {
                vec::push[str](full_paths, p + filename);
            }
        }
    }
    ret full_paths;
}

fn path_is_absolute(path p) -> bool {
    ret os_fs::path_is_absolute(p);
}

// FIXME: under Windows, we should prepend the current drive letter to paths
// that start with a slash.
fn make_absolute(path p) -> path {
    if(path_is_absolute(p)) {
        ret p;
    }
    else {
        ret connect(getcwd(), p);
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<commit_msg>Fix logic bug in std::fs::dirname.<commit_after>
import os::getcwd;
import os_fs;
import str;

native "rust" mod rustrt {
    fn rust_file_is_dir(str path) -> int;
}

fn path_sep() -> str { ret str::from_char(os_fs::path_sep); }

type path = str;

fn dirname(path p) -> path {
    let int i = str::rindex(p, os_fs::path_sep as u8);
    if (i == -1) {
        i = str::rindex(p, os_fs::alt_path_sep as u8);
        if (i == -1) { ret "."; }
    }
    ret str::substr(p, 0u, i as uint);
}

fn basename(path p) -> path {
    let int i = str::rindex(p, os_fs::path_sep as u8);
    if (i == -1) {
        i = str::rindex(p, os_fs::alt_path_sep as u8);
        if (i == -1) { ret p; }
    }
    auto len = str::byte_len(p);
    if (i + 1 as uint >= len) { ret p; }
    ret str::slice(p, i + 1 as uint, len);
}

// FIXME: Need some typestate to avoid bounds check when len(pre) == 0
fn connect(path pre, path post) -> path {
    auto len = str::byte_len(pre);
    ret if (pre.(len - 1u) == os_fs::path_sep as u8) {

            // Trailing '/'?
            pre + post
        } else { pre + path_sep() + post };
}

fn file_is_dir(path p) -> bool { ret rustrt::rust_file_is_dir(p) != 0; }

fn list_dir(path p) -> vec[str] {
    auto pl = str::byte_len(p);
    if (pl == 0u || p.(pl - 1u) as char != os_fs::path_sep) {
        p += path_sep();
    }
    let vec[str] full_paths = [];
    for (str filename in os_fs::list_dir(p)) {
        if (!str::eq(filename, ".")) {
            if (!str::eq(filename, "..")) {
                vec::push[str](full_paths, p + filename);
            }
        }
    }
    ret full_paths;
}

fn path_is_absolute(path p) -> bool {
    ret os_fs::path_is_absolute(p);
}

// FIXME: under Windows, we should prepend the current drive letter to paths
// that start with a slash.
fn make_absolute(path p) -> path {
    if(path_is_absolute(p)) {
        ret p;
    }
    else {
        ret connect(getcwd(), p);
    }
}

// Local Variables:
// mode: rust;
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C $RBUILD 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
<|endoftext|>


<commit_before>extern crate chrono;
extern crate walkdir;

use self::chrono::{Duration, UTC};
use self::walkdir::{DirEntry as WalkDirEntry, WalkDir, WalkDirIterator};

use std::borrow::Borrow;
use std::collections::{BTreeSet, VecDeque};
use std::fs::{self, DirEntry};
use std::io::Error as IoError;
use std::os::unix::fs::MetadataExt;
use std::path::{MAIN_SEPARATOR, Path, PathBuf};
use std::result::Result;

macro_rules! path_buf {
    ($($x: expr),+) => ({
        let mut path = PathBuf::new();
        $(path.push($x);)+
        path
    });
}

pub fn trim_current_dir_prefix<S: AsRef<str>>(path_name: S) -> String {
    let pattern = format!(".{}", MAIN_SEPARATOR);

    path_name
        .as_ref()
        .trim_left_matches(&pattern)
        .to_string()
}

fn into_string<D: Borrow<DirEntry>>(entry: D) -> Option<String> {
    entry
        .borrow()
        .file_name()
        .into_string()
        .ok()
}

pub fn la<P: AsRef<Path>>(path: P) -> Result<Vec<String>, IoError> {
    let dirs = try!(fs::read_dir(path));
    let dirs = dirs
        .filter_map(Result::ok)
        .filter_map(into_string)
        .collect::<Vec<String>>();
    Ok(dirs)
}

pub fn is_recently_accessed<P: AsRef<Path>, D: Borrow<Duration>>(p: P, moratorium: D) -> bool {
    let threshold = UTC::now() - *moratorium.borrow();

    let accessed_time = match p.as_ref().metadata() {
        Ok(m)  => m.atime(),
        Err(_) => unreachable!("Wrong to use this function!!"),
    };

    accessed_time > threshold.timestamp()
}

fn is_hidden_name<S: AsRef<str>>(file_name: S) -> bool {
    let file_name = file_name.as_ref();
    file_name.starts_with(".") && file_name.len() > 1 && file_name != ".."
}

fn is_hidden_for_std(entry: &DirEntry) -> bool {
    into_string(entry).map_or(false, is_hidden_name)
}

fn is_hidden_for_walkdir(entry: &WalkDirEntry) -> bool {
    entry
        .file_name()
        .to_str()
        .map_or(false, is_hidden_name)
}

pub fn walk_dir<P: AsRef<Path>>(root: P) -> BTreeSet<String> {
    let walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(|e| !is_hidden_for_walkdir(e))
        .filter_map(Result::ok)
        .filter(|e| !e.file_type().is_dir())
        .collect::<Vec<WalkDirEntry>>();

    walker
        .into_iter()
        .filter_map(|e| e.path().to_str().map(trim_current_dir_prefix))
        .collect::<BTreeSet<String>>()
}

pub fn potentially_empty_dirs<P: AsRef<Path>>(root: P) -> BTreeSet<PathBuf> {
    fn potentially_empty_dirs(mut result: BTreeSet<PathBuf>, mut target_dirs: VecDeque<PathBuf>) -> BTreeSet<PathBuf> {
        match target_dirs.pop_front() {
            None => result,
            Some(mut target_dir) => {
                let read_result = fs::read_dir(&target_dir);
                let ignore      = read_result.is_err();

                let entries = match read_result {
                    Ok(rd) => rd
                        .filter_map(Result::ok)
                        .collect::<Vec<DirEntry>>(),
                    Err(_) => Vec::new(),
                };
                let include_file_or_hidden_dir = entries
                    .iter()
                    .any(|e| e.file_type().ok().map_or(true, |t| t.is_file()) || is_hidden_for_std(e));

                if ignore || include_file_or_hidden_dir {
                    loop {
                        if !(result.remove(&target_dir) && target_dir.pop()) {
                            break;
                        }
                    }
                }

                let dirs = entries
                    .iter()
                    .filter(|e| e.file_type().ok().map_or(false, |t| t.is_dir()) && !is_hidden_for_std(e))
                    .map(|e| e.path())
                    .collect::<BTreeSet<PathBuf>>();
                for dir in dirs.clone().into_iter() {
                    result.insert(dir);
                }

                let mut dirs = dirs.into_iter().collect::<VecDeque<PathBuf>>();
                target_dirs.append(&mut dirs);

                potentially_empty_dirs(result, target_dirs)
            },
        }
    }

    let mut result     : BTreeSet<PathBuf> = BTreeSet::new();
    let mut target_dirs: VecDeque<PathBuf> = VecDeque::new();

    let root = root.as_ref().to_path_buf();
    result.insert(root.clone());
    target_dirs.push_back(root);

    potentially_empty_dirs(result, target_dirs)
}

#[test]
fn is_hidden_name_should_return_true() {
    let file_names = [
        ".a",
        ".ab",
        ".abc",
        "..a",
        "...",
    ];
    for file_name in &file_names {
        assert!(is_hidden_name(file_name));
    }
}
#[test]
fn is_hidden_name_should_return_false() {
    let file_names = [
        "",

        ".",
        "..",
        "a",
        "ab",
        "abc",
    ];
    for file_name in &file_names {
        assert!(!is_hidden_name(file_name));
    }
}

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

    use std::collections::BTreeSet;
    use std::fs::{self, File};
    use std::io::Write;
    use std::path::{Path, PathBuf};

    struct Helper {
        d1: String,
        d2: String,
        d3: String,
        d4: String,
        d5: String,
        d6: String,
        f1: String,
        f2: String,
        f3: String,
        f4: String,
    }
    impl Helper {
        fn new(suffix: &str) -> Self {
            Helper {
                d1: format!("directory1_{}" , suffix),
                d2: format!("directory2_{}" , suffix),
                d3: format!("directory3_{}" , suffix),
                d4: format!("directory4_{}" , suffix),
                d5: format!("directory5_{}" , suffix),
                d6: format!(".directory6_{}", suffix),
                f1: format!("file1_{}"      , suffix),
                f2: format!("file2_{}"      , suffix),
                f3: format!("file3_{}"      , suffix),
                f4: format!(".file4_{}"     , suffix),
            }
        }

        fn path_to_d1(&self) -> PathBuf {
            PathBuf::new().join(&self.d1)
        }
        fn path_to_d2(&self) -> PathBuf {
            self.path_to_d1().join(&self.d2)
        }
        fn path_to_d3(&self) -> PathBuf {
            self.path_to_d2().join(&self.d3)
        }
        fn path_to_d4(&self) -> PathBuf {
            self.path_to_d3().join(&self.d4)
        }
        fn path_to_d5(&self) -> PathBuf {
            self.path_to_d1().join(&self.d5)
        }
        fn path_to_d6(&self) -> PathBuf {
            self.path_to_d1().join(&self.d6)
        }

        fn path_to_f1(&self) -> PathBuf {
            self.path_to_d1().join(&self.f1)
        }
        fn path_to_f2(&self) -> PathBuf {
            self.path_to_d2().join(&self.f2)
        }
        fn path_to_f3(&self) -> PathBuf {
            self.path_to_d3().join(&self.f3)
        }
        fn path_to_f4(&self) -> PathBuf {
            self.path_to_d5().join(&self.f4)
        }

        fn remove_dirs_and_files(&self) {
            fs::remove_dir_all(self.path_to_d1()).ok();
        }

        fn create_dirs_and_files(&self) {
            self.remove_dirs_and_files();

            fs::create_dir_all(self.path_to_d4()).ok();
            fs::create_dir_all(self.path_to_d5()).ok();
            fs::create_dir_all(self.path_to_d6()).ok();

            let mut f = File::create(self.path_to_f1()).unwrap();
            f.write("\n".as_ref()).ok();
            fs::copy(self.path_to_f1(), self.path_to_f2()).ok();
            fs::copy(self.path_to_f1(), self.path_to_f3()).ok();
            fs::copy(self.path_to_f1(), self.path_to_f4()).ok();
        }

        fn to_string_forcely(path: PathBuf) -> String {
            path.to_str().unwrap().to_string()
        }
    }

    #[test]
    fn path_buf_macro_should_create_path_buf() {
        let paths = [
            (Path::new("path").to_path_buf()          , path_buf!["path"              ]),
            (Path::new("path").join("to")             , path_buf!["path", "to"        ]),
            (Path::new("path").join("to").join("file"), path_buf!["path", "to", "file"]),
        ];
        for &(ref correct, ref calculated) in &paths {
            assert_eq!(correct, calculated);
        }
    }

    #[test]
    fn trim_current_dir_prefix_should_return_trimmed() {
        let path_names = [
            (""   , "./"    ),
            (""   , "././"  ),
            (""   , "./././"),
            ("a"  , "./a"   ),
            ("a/" , "./a/"  ),
            ("a/b", "./a/b" ),
        ];
        for &(ref correct, ref input) in &path_names {
            assert_eq!(correct.to_string(), trim_current_dir_prefix(input));
        }
    }
    #[test]
    fn trim_current_dir_prefix_should_return_as_it_is() {
        let path_names = [
            "",
            "/",
            "/a",
            "/a/b",
            "a",
            "a/",
            "a/b",
        ];
        for path_name in &path_names {
            assert_eq!(path_name.to_string(), trim_current_dir_prefix(path_name));
        }
    }

    #[test]
    fn la_should_return_ok() {
        fn to_btree_set(vec: Vec<String>) -> BTreeSet<String> {
            vec.into_iter().collect::<BTreeSet<String>>()
        }

        let helper = Helper::new("la_Ok");
        helper.create_dirs_and_files();

        let empty_btree_set: BTreeSet<String> = BTreeSet::new();

        let correct = to_btree_set(vec![helper.d2.clone(), helper.d5.clone(), helper.d6.clone(), helper.f1.clone()]);
        assert_eq!(correct, to_btree_set(la(helper.path_to_d1()).unwrap()));
        let correct = to_btree_set(vec![helper.d3.clone(), helper.f2.clone()]);
        assert_eq!(correct, to_btree_set(la(helper.path_to_d2()).unwrap()));
        let correct = to_btree_set(vec![helper.d4.clone(), helper.f3.clone()]);
        assert_eq!(correct        , to_btree_set(la(helper.path_to_d3()).unwrap()));
        assert_eq!(empty_btree_set, to_btree_set(la(helper.path_to_d4()).unwrap()));
        let correct = to_btree_set(vec![helper.f4.clone()]);
        assert_eq!(correct        , to_btree_set(la(helper.path_to_d5()).unwrap()));
        assert_eq!(empty_btree_set, to_btree_set(la(helper.path_to_d6()).unwrap()));

        helper.remove_dirs_and_files();
    }

    #[test]
    fn walk_dir_should_return_b_tree_set() {
        let helper = Helper::new("walk_dir_BTreeSet");
        helper.create_dirs_and_files();

        let mut correct = BTreeSet::new();
        correct.insert(Helper::to_string_forcely(helper.path_to_f1()));
        correct.insert(Helper::to_string_forcely(helper.path_to_f2()));
        correct.insert(Helper::to_string_forcely(helper.path_to_f3()));

        assert_eq!(correct, walk_dir(helper.path_to_d1()));

        helper.remove_dirs_and_files();
    }

    #[test]
    fn potentially_empty_dirs_should_return_btree_set() {
        let helper = Helper::new("potentially_empty_dirs_BTreeSet");
        helper.create_dirs_and_files();

        let root = helper.path_to_d1();
        let mut correct = BTreeSet::new();

        correct.insert(helper.path_to_d4());
        assert_eq!(correct, potentially_empty_dirs(&root));

        fs::remove_file(helper.path_to_f3()).ok();
        correct.insert(helper.path_to_d3());
        assert_eq!(correct, potentially_empty_dirs(&root));

        fs::remove_file(helper.path_to_f2()).ok();
        correct.insert(helper.path_to_d2());
        assert_eq!(correct, potentially_empty_dirs(&root));

        fs::remove_file(helper.path_to_f4()).ok();
        correct.insert(helper.path_to_d5());
        assert_eq!(correct, potentially_empty_dirs(&root));

        helper.remove_dirs_and_files();
    }
}
<commit_msg>[Test] Refactor<commit_after>extern crate chrono;
extern crate walkdir;

use self::chrono::{Duration, UTC};
use self::walkdir::{DirEntry as WalkDirEntry, WalkDir, WalkDirIterator};

use std::borrow::Borrow;
use std::collections::{BTreeSet, VecDeque};
use std::fs::{self, DirEntry};
use std::io::Error as IoError;
use std::os::unix::fs::MetadataExt;
use std::path::{MAIN_SEPARATOR, Path, PathBuf};
use std::result::Result;

macro_rules! path_buf {
    ($($x: expr),+) => ({
        let mut path = PathBuf::new();
        $(path.push($x);)+
        path
    });
}

pub fn trim_current_dir_prefix<S: AsRef<str>>(path_name: S) -> String {
    let pattern = format!(".{}", MAIN_SEPARATOR);

    path_name
        .as_ref()
        .trim_left_matches(&pattern)
        .to_string()
}

fn into_string<D: Borrow<DirEntry>>(entry: D) -> Option<String> {
    entry
        .borrow()
        .file_name()
        .into_string()
        .ok()
}

pub fn la<P: AsRef<Path>>(path: P) -> Result<Vec<String>, IoError> {
    let dirs = try!(fs::read_dir(path));
    let dirs = dirs
        .filter_map(Result::ok)
        .filter_map(into_string)
        .collect::<Vec<String>>();
    Ok(dirs)
}

pub fn is_recently_accessed<P: AsRef<Path>, D: Borrow<Duration>>(p: P, moratorium: D) -> bool {
    let threshold = UTC::now() - *moratorium.borrow();

    let accessed_time = match p.as_ref().metadata() {
        Ok(m)  => m.atime(),
        Err(_) => unreachable!("Wrong to use this function!!"),
    };

    accessed_time > threshold.timestamp()
}

fn is_hidden_name<S: AsRef<str>>(file_name: S) -> bool {
    let file_name = file_name.as_ref();
    file_name.starts_with(".") && file_name.len() > 1 && file_name != ".."
}

fn is_hidden_for_std(entry: &DirEntry) -> bool {
    into_string(entry).map_or(false, is_hidden_name)
}

fn is_hidden_for_walkdir(entry: &WalkDirEntry) -> bool {
    entry
        .file_name()
        .to_str()
        .map_or(false, is_hidden_name)
}

pub fn walk_dir<P: AsRef<Path>>(root: P) -> BTreeSet<String> {
    let walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(|e| !is_hidden_for_walkdir(e))
        .filter_map(Result::ok)
        .filter(|e| !e.file_type().is_dir())
        .collect::<Vec<WalkDirEntry>>();

    walker
        .into_iter()
        .filter_map(|e| e.path().to_str().map(trim_current_dir_prefix))
        .collect::<BTreeSet<String>>()
}

pub fn potentially_empty_dirs<P: AsRef<Path>>(root: P) -> BTreeSet<PathBuf> {
    fn potentially_empty_dirs(mut result: BTreeSet<PathBuf>, mut target_dirs: VecDeque<PathBuf>) -> BTreeSet<PathBuf> {
        match target_dirs.pop_front() {
            None => result,
            Some(mut target_dir) => {
                let read_result = fs::read_dir(&target_dir);
                let ignore      = read_result.is_err();

                let entries = match read_result {
                    Ok(rd) => rd
                        .filter_map(Result::ok)
                        .collect::<Vec<DirEntry>>(),
                    Err(_) => Vec::new(),
                };
                let include_file_or_hidden_dir = entries
                    .iter()
                    .any(|e| e.file_type().ok().map_or(true, |t| t.is_file()) || is_hidden_for_std(e));

                if ignore || include_file_or_hidden_dir {
                    loop {
                        if !(result.remove(&target_dir) && target_dir.pop()) {
                            break;
                        }
                    }
                }

                let dirs = entries
                    .iter()
                    .filter(|e| e.file_type().ok().map_or(false, |t| t.is_dir()) && !is_hidden_for_std(e))
                    .map(|e| e.path())
                    .collect::<BTreeSet<PathBuf>>();
                for dir in dirs.clone().into_iter() {
                    result.insert(dir);
                }

                let mut dirs = dirs.into_iter().collect::<VecDeque<PathBuf>>();
                target_dirs.append(&mut dirs);

                potentially_empty_dirs(result, target_dirs)
            },
        }
    }

    let mut result     : BTreeSet<PathBuf> = BTreeSet::new();
    let mut target_dirs: VecDeque<PathBuf> = VecDeque::new();

    let root = root.as_ref().to_path_buf();
    result.insert(root.clone());
    target_dirs.push_back(root);

    potentially_empty_dirs(result, target_dirs)
}

#[test]
fn is_hidden_name_should_return_true() {
    let file_names = [
        ".a",
        ".ab",
        ".abc",
        "..a",
        "...",
    ];
    for file_name in &file_names {
        assert!(is_hidden_name(file_name));
    }
}
#[test]
fn is_hidden_name_should_return_false() {
    let file_names = [
        "",

        ".",
        "..",
        "a",
        "ab",
        "abc",
    ];
    for file_name in &file_names {
        assert!(!is_hidden_name(file_name));
    }
}

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

    use std::collections::BTreeSet;
    use std::fs::{self, File};
    use std::io::Write;
    use std::path::{Path, PathBuf};

    struct Helper {
        d1: String,
        d2: String,
        d3: String,
        d4: String,
        d5: String,
        d6: String,
        f1: String,
        f2: String,
        f3: String,
        f4: String,
    }
    impl Helper {
        fn new(suffix: &str) -> Self {
            Helper {
                d1: format!("directory1_{}" , suffix),
                d2: format!("directory2_{}" , suffix),
                d3: format!("directory3_{}" , suffix),
                d4: format!("directory4_{}" , suffix),
                d5: format!("directory5_{}" , suffix),
                d6: format!(".directory6_{}", suffix),
                f1: format!("file1_{}"      , suffix),
                f2: format!("file2_{}"      , suffix),
                f3: format!("file3_{}"      , suffix),
                f4: format!(".file4_{}"     , suffix),
            }
        }

        fn path_to_d1(&self) -> PathBuf {
            PathBuf::new().join(&self.d1)
        }
        fn path_to_d2(&self) -> PathBuf {
            self.path_to_d1().join(&self.d2)
        }
        fn path_to_d3(&self) -> PathBuf {
            self.path_to_d2().join(&self.d3)
        }
        fn path_to_d4(&self) -> PathBuf {
            self.path_to_d3().join(&self.d4)
        }
        fn path_to_d5(&self) -> PathBuf {
            self.path_to_d1().join(&self.d5)
        }
        fn path_to_d6(&self) -> PathBuf {
            self.path_to_d1().join(&self.d6)
        }

        fn path_to_f1(&self) -> PathBuf {
            self.path_to_d1().join(&self.f1)
        }
        fn path_to_f2(&self) -> PathBuf {
            self.path_to_d2().join(&self.f2)
        }
        fn path_to_f3(&self) -> PathBuf {
            self.path_to_d3().join(&self.f3)
        }
        fn path_to_f4(&self) -> PathBuf {
            self.path_to_d5().join(&self.f4)
        }

        fn remove_dirs_and_files(&self) {
            fs::remove_dir_all(self.path_to_d1()).ok();
        }

        fn create_dirs_and_files(&self) {
            self.remove_dirs_and_files();

            fs::create_dir_all(self.path_to_d4()).ok();
            fs::create_dir_all(self.path_to_d5()).ok();
            fs::create_dir_all(self.path_to_d6()).ok();

            let mut f = File::create(self.path_to_f1()).unwrap();
            f.write("\n".as_ref()).ok();
            fs::copy(self.path_to_f1(), self.path_to_f2()).ok();
            fs::copy(self.path_to_f1(), self.path_to_f3()).ok();
            fs::copy(self.path_to_f1(), self.path_to_f4()).ok();
        }

        fn to_string_forcely(path: PathBuf) -> String {
            path.to_str().unwrap().to_string()
        }
    }

    #[test]
    fn path_buf_macro_should_create_path_buf() {
        let paths = [
            (Path::new("path").to_path_buf()          , path_buf!["path"              ]),
            (Path::new("path").join("to")             , path_buf!["path", "to"        ]),
            (Path::new("path").join("to").join("file"), path_buf!["path", "to", "file"]),
        ];
        for &(ref correct, ref calculated) in &paths {
            assert_eq!(correct, calculated);
        }
    }

    #[test]
    fn trim_current_dir_prefix_should_return_trimmed() {
        let path_names = [
            (""   , "./"    ),
            (""   , "././"  ),
            (""   , "./././"),
            ("a"  , "./a"   ),
            ("a/" , "./a/"  ),
            ("a/b", "./a/b" ),
        ];
        for &(ref correct, ref input) in &path_names {
            assert_eq!(correct.to_string(), trim_current_dir_prefix(input));
        }
    }
    #[test]
    fn trim_current_dir_prefix_should_return_as_it_is() {
        let path_names = [
            "",
            "/",
            "/a",
            "/a/b",
            "a",
            "a/",
            "a/b",
        ];
        for path_name in &path_names {
            assert_eq!(path_name.to_string(), trim_current_dir_prefix(path_name));
        }
    }

    #[test]
    fn la_should_return_ok() {
        fn to_btree_set(vec: Vec<String>) -> BTreeSet<String> {
            vec.into_iter().collect::<BTreeSet<String>>()
        }

        let helper = Helper::new("la_Ok");
        helper.create_dirs_and_files();

        let data_set = vec![
            (vec![helper.d2.clone(), helper.d5.clone(), helper.d6.clone(), helper.f1.clone()], helper.path_to_d1()),
            (vec![helper.d3.clone(), helper.f2.clone()]                                      , helper.path_to_d2()),
            (vec![helper.d4.clone(), helper.f3.clone()]                                      , helper.path_to_d3()),
            (Vec::new()                                                                      , helper.path_to_d4()),
            (vec![helper.f4.clone()]                                                         , helper.path_to_d5()),
            (Vec::new()                                                                      , helper.path_to_d6()),
        ];

        for (correct, target_dir) in data_set.into_iter() {
            assert_eq!(to_btree_set(correct), to_btree_set(la(target_dir).unwrap()));
        }

        helper.remove_dirs_and_files();
    }

    #[test]
    fn walk_dir_should_return_b_tree_set() {
        let helper = Helper::new("walk_dir_BTreeSet");
        helper.create_dirs_and_files();

        let mut correct = BTreeSet::new();
        correct.insert(Helper::to_string_forcely(helper.path_to_f1()));
        correct.insert(Helper::to_string_forcely(helper.path_to_f2()));
        correct.insert(Helper::to_string_forcely(helper.path_to_f3()));

        assert_eq!(correct, walk_dir(helper.path_to_d1()));

        helper.remove_dirs_and_files();
    }

    #[test]
    fn potentially_empty_dirs_should_return_btree_set() {
        let helper = Helper::new("potentially_empty_dirs_BTreeSet");
        helper.create_dirs_and_files();

        let root = helper.path_to_d1();
        let mut correct = BTreeSet::new();

        correct.insert(helper.path_to_d4());
        assert_eq!(correct, potentially_empty_dirs(&root));

        fs::remove_file(helper.path_to_f3()).ok();
        correct.insert(helper.path_to_d3());
        assert_eq!(correct, potentially_empty_dirs(&root));

        fs::remove_file(helper.path_to_f2()).ok();
        correct.insert(helper.path_to_d2());
        assert_eq!(correct, potentially_empty_dirs(&root));

        fs::remove_file(helper.path_to_f4()).ok();
        correct.insert(helper.path_to_d5());
        assert_eq!(correct, potentially_empty_dirs(&root));

        helper.remove_dirs_and_files();
    }
}
<|endoftext|>


<commit_before>import os.libc;

native "rust" mod rustrt {
  fn rust_get_stdin() -> os.libc.FILE;
  fn rust_get_stdout() -> os.libc.FILE;
}

// Reading

// TODO This is all buffered. We might need an unbuffered variant as well

tag seek_style {seek_set; seek_end; seek_cur;}

type reader =
    state obj {
          impure fn read_byte() -> u8;
          impure fn read_bytes(uint len) -> vec[u8];
          impure fn read_char() -> int;
          impure fn unread_char(int i);
          impure fn eof() -> bool;
          impure fn read_line() -> str;
          impure fn read_c_str() -> str;
          impure fn read_le_uint(uint size) -> uint;
          impure fn read_le_int(uint size) -> int;

          impure fn seek(int offset, seek_style whence);
    };

state obj FILE_reader(os.libc.FILE f, bool must_close) {
    impure fn read_byte() -> u8 {
        ret os.libc.fgetc(f) as u8;
    }
    impure fn read_bytes(uint len) -> vec[u8] {
        auto buf = _vec.alloc[u8](len);
        auto read = os.libc.fread(_vec.buf[u8](buf), 1u, len, f);
        _vec.len_set[u8](buf, read);
        ret buf;
    }
    impure fn read_char() -> int {
        ret os.libc.fgetc(f);
    }
    impure fn unread_char(int ch) {
        os.libc.ungetc(ch, f);
    }
    impure fn eof() -> bool {
      auto ch = os.libc.fgetc(f);
      if (ch == -1) {ret true;}
      os.libc.ungetc(ch, f);
      ret false;
    }
    impure fn read_line() -> str {
      auto buf = "";
      while (true) {
        auto ch = os.libc.fgetc(f);
        if (ch == -1) { ret buf; }
        if (ch == 10) { ret buf; }
        buf += _str.unsafe_from_bytes(vec(ch as u8));
      }
      ret buf;
    }
    impure fn read_c_str() -> str {
        auto buf = "";
        while (true) {
            auto ch = os.libc.fgetc(f);
            if (ch < 1) { ret buf; }
            buf += _str.unsafe_from_bytes(vec(ch as u8));
        }
        ret buf;
    }
    // TODO deal with eof?
    impure fn read_le_uint(uint size) -> uint {
        auto val = 0u;
        auto pos = 0u;
        while (size > 0u) {
            val += (os.libc.fgetc(f) as uint) << pos;
            pos += 8u;
            size -= 1u;
        }
        ret val;
    }
    impure fn read_le_int(uint size) -> int {
        auto val = 0u;
        auto pos = 0u;
        while (size > 0u) {
            val += (os.libc.fgetc(f) as uint) << pos;
            pos += 8u;
            size -= 1u;
        }
        ret val as int; // TODO does that work?
    }
    impure fn seek(int offset, seek_style whence) {
        auto wh;
        alt (whence) {
            case (seek_set) {wh = 0;}
            case (seek_cur) {wh = 1;}
            case (seek_end) {wh = 2;}
        }
        check(os.libc.fseek(f, offset, wh) == 0);
    }
    drop {
        if (must_close) {os.libc.fclose(f);}
    }
}

fn stdin() -> reader {
    ret FILE_reader(rustrt.rust_get_stdin(), false);
}

fn file_reader(str path) -> reader {
    auto f = os.libc.fopen(_str.buf(path), _str.buf("r"));
    check (f as uint != 0u);
    ret FILE_reader(f, true);
}

// Writing

// TODO This is all unbuffered. We might need a buffered variant as well

tag fileflag {
    append;
    create;
    truncate;
    none;
}

type buf_writer = state obj {
  fn write(vec[u8] v);
};

state obj fd_buf_writer(int fd, bool must_close) {
    fn write(vec[u8] v) {
        auto len = _vec.len[u8](v);
        auto count = 0u;
        auto vbuf;
        while (count < len) {
            vbuf = _vec.buf_off[u8](v, count);
            auto nout = os.libc.write(fd, vbuf, len);
            if (nout < 0) {
                log "error dumping buffer";
                log sys.rustrt.last_os_error();
                fail;
            }
            count += nout as uint;
        }
    }

    drop {
        if (must_close) {os.libc.close(fd);}
    }
}

fn file_buf_writer(str path, vec[fileflag] flags) -> buf_writer {
    let int fflags =
        os.libc_constants.O_WRONLY() |
        os.libc_constants.O_BINARY();

    for (fileflag f in flags) {
        alt (f) {
            case (append)   { fflags |= os.libc_constants.O_APPEND(); }
            case (create)   { fflags |= os.libc_constants.O_CREAT(); }
            case (truncate) { fflags |= os.libc_constants.O_TRUNC(); }
            case (none) {}
        }
    }

    auto fd = os.libc.open(_str.buf(path),
                           fflags,
                           os.libc_constants.S_IRUSR() |
                           os.libc_constants.S_IWUSR());

    if (fd < 0) {
        log "error opening file for writing";
        log sys.rustrt.last_os_error();
        fail;
    }
    ret fd_buf_writer(fd, true);
}

type writer =
    state obj {
          impure fn write_str(str s);
          impure fn write_int(int n);
          impure fn write_uint(uint n);
          impure fn write_bytes(vec[u8] bytes);
          impure fn write_le_uint(uint n, uint size);
          impure fn write_le_int(int n, uint size);
    };

fn uint_to_le_bytes(uint n, uint size) -> vec[u8] {
    let vec[u8] bytes = vec();
    while (size > 0u) {
        bytes += vec((n & 255u) as u8);
        n >>= 8u;
        size -= 1u;
    }
    ret bytes;
}

state obj new_writer(buf_writer out) {
    impure fn write_str(str s) {
        out.write(_str.bytes(s));
    }
    impure fn write_int(int n) {
        out.write(_str.bytes(_int.to_str(n, 10u)));
    }
    impure fn write_uint(uint n) {
        out.write(_str.bytes(_uint.to_str(n, 10u)));
    }
    impure fn write_bytes(vec[u8] bytes) {
        out.write(bytes);
    }
    impure fn write_le_uint(uint n, uint size) {
        out.write(uint_to_le_bytes(n, size));
    }
    impure fn write_le_int(int n, uint size) {
        out.write(uint_to_le_bytes(n as uint, size));
    }
}

fn file_writer(str path, vec[fileflag] flags) -> writer {
    ret new_writer(file_buf_writer(path, flags));
}

// FIXME it would be great if this could be a const
fn stdout() -> writer {
    ret new_writer(fd_buf_writer(1, false));
}

type str_writer =
    state obj {
          fn get_writer() -> writer;
          fn get_str() -> str;
    };

type byte_buf = @rec(mutable vec[u8] buf);

state obj byte_buf_writer(byte_buf buf) {
    fn write(vec[u8] v) {buf.buf += v;}
}

// TODO awkward! it's not possible to implement a writer with an extra method
fn string_writer() -> str_writer {
    let vec[u8] b = vec();
    let byte_buf buf = @rec(mutable buf = b);
    state obj str_writer_wrap(writer wr, byte_buf buf) {
        fn get_writer() -> writer {ret wr;}
        fn get_str() -> str {ret _str.unsafe_from_bytes(buf.buf);}
    }
    ret str_writer_wrap(new_writer(byte_buf_writer(buf)), buf);
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C .. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
//
<commit_msg>Add slightly nicer failure message to io.rs when it can't open a file.<commit_after>import os.libc;

native "rust" mod rustrt {
  fn rust_get_stdin() -> os.libc.FILE;
  fn rust_get_stdout() -> os.libc.FILE;
}

// Reading

// TODO This is all buffered. We might need an unbuffered variant as well

tag seek_style {seek_set; seek_end; seek_cur;}

type reader =
    state obj {
          impure fn read_byte() -> u8;
          impure fn read_bytes(uint len) -> vec[u8];
          impure fn read_char() -> int;
          impure fn unread_char(int i);
          impure fn eof() -> bool;
          impure fn read_line() -> str;
          impure fn read_c_str() -> str;
          impure fn read_le_uint(uint size) -> uint;
          impure fn read_le_int(uint size) -> int;

          impure fn seek(int offset, seek_style whence);
    };

state obj FILE_reader(os.libc.FILE f, bool must_close) {
    impure fn read_byte() -> u8 {
        ret os.libc.fgetc(f) as u8;
    }
    impure fn read_bytes(uint len) -> vec[u8] {
        auto buf = _vec.alloc[u8](len);
        auto read = os.libc.fread(_vec.buf[u8](buf), 1u, len, f);
        _vec.len_set[u8](buf, read);
        ret buf;
    }
    impure fn read_char() -> int {
        ret os.libc.fgetc(f);
    }
    impure fn unread_char(int ch) {
        os.libc.ungetc(ch, f);
    }
    impure fn eof() -> bool {
      auto ch = os.libc.fgetc(f);
      if (ch == -1) {ret true;}
      os.libc.ungetc(ch, f);
      ret false;
    }
    impure fn read_line() -> str {
      auto buf = "";
      while (true) {
        auto ch = os.libc.fgetc(f);
        if (ch == -1) { ret buf; }
        if (ch == 10) { ret buf; }
        buf += _str.unsafe_from_bytes(vec(ch as u8));
      }
      ret buf;
    }
    impure fn read_c_str() -> str {
        auto buf = "";
        while (true) {
            auto ch = os.libc.fgetc(f);
            if (ch < 1) { ret buf; }
            buf += _str.unsafe_from_bytes(vec(ch as u8));
        }
        ret buf;
    }
    // TODO deal with eof?
    impure fn read_le_uint(uint size) -> uint {
        auto val = 0u;
        auto pos = 0u;
        while (size > 0u) {
            val += (os.libc.fgetc(f) as uint) << pos;
            pos += 8u;
            size -= 1u;
        }
        ret val;
    }
    impure fn read_le_int(uint size) -> int {
        auto val = 0u;
        auto pos = 0u;
        while (size > 0u) {
            val += (os.libc.fgetc(f) as uint) << pos;
            pos += 8u;
            size -= 1u;
        }
        ret val as int; // TODO does that work?
    }
    impure fn seek(int offset, seek_style whence) {
        auto wh;
        alt (whence) {
            case (seek_set) {wh = 0;}
            case (seek_cur) {wh = 1;}
            case (seek_end) {wh = 2;}
        }
        check(os.libc.fseek(f, offset, wh) == 0);
    }
    drop {
        if (must_close) {os.libc.fclose(f);}
    }
}

fn stdin() -> reader {
    ret FILE_reader(rustrt.rust_get_stdin(), false);
}

fn file_reader(str path) -> reader {
    auto f = os.libc.fopen(_str.buf(path), _str.buf("r"));
    if (f as uint == 0u) {
        log "error opening " + path;
        fail;
    }
    ret FILE_reader(f, true);
}

// Writing

// TODO This is all unbuffered. We might need a buffered variant as well

tag fileflag {
    append;
    create;
    truncate;
    none;
}

type buf_writer = state obj {
  fn write(vec[u8] v);
};

state obj fd_buf_writer(int fd, bool must_close) {
    fn write(vec[u8] v) {
        auto len = _vec.len[u8](v);
        auto count = 0u;
        auto vbuf;
        while (count < len) {
            vbuf = _vec.buf_off[u8](v, count);
            auto nout = os.libc.write(fd, vbuf, len);
            if (nout < 0) {
                log "error dumping buffer";
                log sys.rustrt.last_os_error();
                fail;
            }
            count += nout as uint;
        }
    }

    drop {
        if (must_close) {os.libc.close(fd);}
    }
}

fn file_buf_writer(str path, vec[fileflag] flags) -> buf_writer {
    let int fflags =
        os.libc_constants.O_WRONLY() |
        os.libc_constants.O_BINARY();

    for (fileflag f in flags) {
        alt (f) {
            case (append)   { fflags |= os.libc_constants.O_APPEND(); }
            case (create)   { fflags |= os.libc_constants.O_CREAT(); }
            case (truncate) { fflags |= os.libc_constants.O_TRUNC(); }
            case (none) {}
        }
    }

    auto fd = os.libc.open(_str.buf(path),
                           fflags,
                           os.libc_constants.S_IRUSR() |
                           os.libc_constants.S_IWUSR());

    if (fd < 0) {
        log "error opening file for writing";
        log sys.rustrt.last_os_error();
        fail;
    }
    ret fd_buf_writer(fd, true);
}

type writer =
    state obj {
          impure fn write_str(str s);
          impure fn write_int(int n);
          impure fn write_uint(uint n);
          impure fn write_bytes(vec[u8] bytes);
          impure fn write_le_uint(uint n, uint size);
          impure fn write_le_int(int n, uint size);
    };

fn uint_to_le_bytes(uint n, uint size) -> vec[u8] {
    let vec[u8] bytes = vec();
    while (size > 0u) {
        bytes += vec((n & 255u) as u8);
        n >>= 8u;
        size -= 1u;
    }
    ret bytes;
}

state obj new_writer(buf_writer out) {
    impure fn write_str(str s) {
        out.write(_str.bytes(s));
    }
    impure fn write_int(int n) {
        out.write(_str.bytes(_int.to_str(n, 10u)));
    }
    impure fn write_uint(uint n) {
        out.write(_str.bytes(_uint.to_str(n, 10u)));
    }
    impure fn write_bytes(vec[u8] bytes) {
        out.write(bytes);
    }
    impure fn write_le_uint(uint n, uint size) {
        out.write(uint_to_le_bytes(n, size));
    }
    impure fn write_le_int(int n, uint size) {
        out.write(uint_to_le_bytes(n as uint, size));
    }
}

fn file_writer(str path, vec[fileflag] flags) -> writer {
    ret new_writer(file_buf_writer(path, flags));
}

// FIXME it would be great if this could be a const
fn stdout() -> writer {
    ret new_writer(fd_buf_writer(1, false));
}

type str_writer =
    state obj {
          fn get_writer() -> writer;
          fn get_str() -> str;
    };

type byte_buf = @rec(mutable vec[u8] buf);

state obj byte_buf_writer(byte_buf buf) {
    fn write(vec[u8] v) {buf.buf += v;}
}

// TODO awkward! it's not possible to implement a writer with an extra method
fn string_writer() -> str_writer {
    let vec[u8] b = vec();
    let byte_buf buf = @rec(mutable buf = b);
    state obj str_writer_wrap(writer wr, byte_buf buf) {
        fn get_writer() -> writer {ret wr;}
        fn get_str() -> str {ret _str.unsafe_from_bytes(buf.buf);}
    }
    ret str_writer_wrap(new_writer(byte_buf_writer(buf)), buf);
}

//
// Local Variables:
// mode: rust
// fill-column: 78;
// indent-tabs-mode: nil
// c-basic-offset: 4
// buffer-file-coding-system: utf-8-unix
// compile-command: "make -k -C .. 2>&1 | sed -e 's/\\/x\\//x:\\//g'";
// End:
//
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io::Write;

#[cfg(not(feature = "serde_type"))]
use serialize::json::ToJson;
#[cfg(feature = "serde_type")]
use serde::ser::Serialize as ToJson;

use template::{Template};
use render::{Renderable, RenderError, RenderContext};
use helpers::{HelperDef};
use context::{Context};
use helpers;
use support::str::StringWriter;
use TemplateError;

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String>;

fn get_default_escape_fn() -> EscapeFn {
    Box::new(|data| {
        data.replace("&", "&")
            .replace("\"", """)
            .replace("<", "<")
            .replace(">", ">")
    })
}

pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    escape_fn: EscapeFn
}

impl Registry {
    pub fn new() -> Registry {
        let mut r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            escape_fn: get_default_escape_fn(),
        };

        r.register_helper("if", Box::new(helpers::IF_HELPER));
        r.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        r.register_helper("each", Box::new(helpers::EACH_HELPER));
        r.register_helper("with", Box::new(helpers::WITH_HELPER));
        r.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        r.register_helper("raw", Box::new(helpers::RAW_HELPER));
        r.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        r.register_helper("block", Box::new(helpers::BLOCK_HELPER));
        r.register_helper("partial", Box::new(helpers::PARTIAL_HELPER));
        r.register_helper("log", Box::new(helpers::LOG_HELPER));

        r
    }

    pub fn register_template(&mut self, name: &str, mut template: Template) {
        template.name = Some(name.to_owned());
        self.templates.insert(name.to_string(), template);
    }

    pub fn register_template_string(&mut self, name: &str, tpl_str: String) -> Result<(), TemplateError>{
        try!(Template::compile_with_name(tpl_str, name.to_owned())
             .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    pub fn register_helper(&mut self, name: &str, def: Box<HelperDef + 'static>) -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String>(&mut self, escape_fn: F) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = get_default_escape_fn();
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    pub fn render<T>(&self, name: &str, ctx: &T) -> Result<String, RenderError> where T: ToJson {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.renderw(name, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    pub fn renderw(&self, name: &str, context: &Context, writer: &mut Write) -> Result<(), RenderError> {
        let template = self.get_template(&name.to_string());

        if let Some(t) = template {
            let mut render_context = RenderContext::new(writer);
            render_context.root_template = t.name.clone();
            (*t).render(context, self, &mut render_context)
        } else {
            Err(RenderError::new(format!("Template not found: {}", name)))
        }
    }
}

#[cfg(test)]
mod test {
    use template::{Template};
    use registry::{Registry};
    use render::{RenderContext, Renderable, RenderError, Helper};
    use helpers::{HelperDef};
    use context::{Context};
    use support::str::StringWriter;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
            try!(h.template().unwrap().render(c, r, rc));
            Ok(())
        }
    }

    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    fn test_registry_operations() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let t2 = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index2", t2.clone());

        assert_eq!((*r.get_template(&("index".to_string())).unwrap()).to_string(),
                   t.to_string());
        assert_eq!(r.templates.len(), 2);

        r.unregister_template(&("index".to_string()));
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 10+1);
    }

    #[test]
    fn test_renderw() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.renderw("index", &context, &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}")).unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }
}
<commit_msg>(fix) escape function for multi-thread envinronment<commit_after>use std::collections::HashMap;
use std::io::Write;

#[cfg(not(feature = "serde_type"))]
use serialize::json::ToJson;
#[cfg(feature = "serde_type")]
use serde::ser::Serialize as ToJson;

use template::{Template};
use render::{Renderable, RenderError, RenderContext};
use helpers::{HelperDef};
use context::{Context};
use helpers;
use support::str::StringWriter;
use TemplateError;

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

fn get_default_escape_fn() -> EscapeFn {
    Box::new(|data| {
        data.replace("&", "&")
            .replace("\"", """)
            .replace("<", "<")
            .replace(">", ">")
    })
}

pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    escape_fn: EscapeFn
}

impl Registry {
    pub fn new() -> Registry {
        let mut r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            escape_fn: get_default_escape_fn(),
        };

        r.register_helper("if", Box::new(helpers::IF_HELPER));
        r.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        r.register_helper("each", Box::new(helpers::EACH_HELPER));
        r.register_helper("with", Box::new(helpers::WITH_HELPER));
        r.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        r.register_helper("raw", Box::new(helpers::RAW_HELPER));
        r.register_helper(">", Box::new(helpers::INCLUDE_HELPER));
        r.register_helper("block", Box::new(helpers::BLOCK_HELPER));
        r.register_helper("partial", Box::new(helpers::PARTIAL_HELPER));
        r.register_helper("log", Box::new(helpers::LOG_HELPER));

        r
    }

    pub fn register_template(&mut self, name: &str, mut template: Template) {
        template.name = Some(name.to_owned());
        self.templates.insert(name.to_string(), template);
    }

    pub fn register_template_string(&mut self, name: &str, tpl_str: String) -> Result<(), TemplateError>{
        try!(Template::compile_with_name(tpl_str, name.to_owned())
             .and_then(|t| Ok(self.templates.insert(name.to_string(), t))));
        Ok(())
    }

    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    pub fn register_helper(&mut self, name: &str, def: Box<HelperDef + 'static>) -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(&mut self, escape_fn: F) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = get_default_escape_fn();
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    pub fn render<T>(&self, name: &str, ctx: &T) -> Result<String, RenderError> where T: ToJson {
        let mut writer = StringWriter::new();
        let context = Context::wraps(ctx);
        {
            try!(self.renderw(name, &context, &mut writer));
        }
        Ok(writer.to_string())
    }

    pub fn renderw(&self, name: &str, context: &Context, writer: &mut Write) -> Result<(), RenderError> {
        let template = self.get_template(&name.to_string());

        if let Some(t) = template {
            let mut render_context = RenderContext::new(writer);
            render_context.root_template = t.name.clone();
            (*t).render(context, self, &mut render_context)
        } else {
            Err(RenderError::new(format!("Template not found: {}", name)))
        }
    }
}

#[cfg(test)]
mod test {
    use template::{Template};
    use registry::{Registry};
    use render::{RenderContext, Renderable, RenderError, Helper};
    use helpers::{HelperDef};
    use context::{Context};
    use support::str::StringWriter;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(&self, c: &Context, h: &Helper, r: &Registry, rc: &mut RenderContext) -> Result<(), RenderError> {
            try!(h.template().unwrap().render(c, r, rc));
            Ok(())
        }
    }

    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    fn test_registry_operations() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let t2 = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index2", t2.clone());

        assert_eq!((*r.get_template(&("index".to_string())).unwrap()).to_string(),
                   t.to_string());
        assert_eq!(r.templates.len(), 2);

        r.unregister_template(&("index".to_string()));
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 10+1);
    }

    #[test]
    fn test_renderw() {
        let mut r = Registry::new();

        let t = Template::compile("

".to_string()).ok().unwrap();
        r.register_template("index", t.clone());

        let mut sw = StringWriter::new();
        let context = Context::null();

        {
            r.renderw("index", &context, &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());

    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}")).unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }
}
<|endoftext|>


<commit_before>// Copyright(C) 2016 Chris Liebert
extern crate glium;
extern crate image;
extern crate nalgebra;

use std::collections::HashMap;

use common;
use common::{Camera, Mesh, Scene, Shader, Vertex8f32};
use loader::DBLoader;

use glium::backend::glutin_backend::GlutinFacade;
use glium::{Program, Surface, Version};

use self::nalgebra::{Matrix4, PerspectiveMatrix3};

implement_vertex!(Vertex8f32, position, normal, texcoord);

pub struct Renderer {
	pub index_buffer: glium::index::NoIndices,
	pub modelview_matrix_array: [[f32; 4]; 4],
	pub projection_matrix: Matrix4<f32>,
	pub scene: Scene,
	pub textures: HashMap<String, glium::texture::CompressedSrgbTexture2d>,
	pub vertex_buffers: Vec<glium::vertex::VertexBuffer<common::Vertex8f32>>,
}

impl Renderer {
	pub fn new(display: &GlutinFacade, scene: Scene) -> Renderer {

		let mut vertex_buffers: Vec<glium::vertex::VertexBuffer<common::Vertex8f32>> =
			Vec::with_capacity(scene.meshes.len());

		for i in 0..scene.meshes.len() {
			vertex_buffers.push(glium::vertex::VertexBuffer::new(display, &scene.meshes[i].vertices).unwrap());
		}

		let index_buffer: glium::index::NoIndices =
			glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList);

		let mut textures: HashMap<String, glium::texture::CompressedSrgbTexture2d> = HashMap::new();

		for i in 0..scene.images.len() {
			//TODO: determine image format by name extension if nessisary
			//let image = image::load_from_memory_with_format(&scene.images[i].image, image::PNG).unwrap().to_rgba();
			let image = image::load_from_memory(&scene.images[i].image).unwrap().to_rgba();

			let image_dimensions = image.dimensions();
			let image = glium::texture::RawImage2d::from_raw_rgba_reversed(image.into_raw(),
			                                                               image_dimensions);
			let opengl_texture: glium::texture::CompressedSrgbTexture2d =
				glium::texture::CompressedSrgbTexture2d::new(display, image).unwrap();
			textures.insert(scene.images[i].name.clone(), opengl_texture);
		}

		// Get the screen width and height in pixels
		let screen_width: u32;
		let screen_height: u32;
		match display.get_window() {
			Some(window) =>  {
				let pixel_dimensions: (u32,u32) = window.get_inner_size_pixels().unwrap();
				screen_width = pixel_dimensions.0;
				screen_height = pixel_dimensions.1;
			},
			None => {
				panic!("Error retrieving window when querying display size.");
			}
		}

		// Set up camera
		let projection_matrix = PerspectiveMatrix3::new(screen_width as f32 / screen_height as f32,
		                                          45.0,
		                                          0.1,
		                                          1000.0);
		let projection_matrix_array: Matrix4<f32> = projection_matrix.to_matrix();

		let modelview_matrix_array = [[1.0f32, -0.0f32, -0.0f32, 0.0f32],
		                          [0.0f32, 1.0f32, -0.0f32, 0.0f32],
		                          [0.0f32, 0.0f32, 1.0f32, 0.0f32],
		                          [-5.0f32, -3.0f32, -12.0f32, 1.0f32]];
		Renderer {
			index_buffer: index_buffer,
			modelview_matrix_array: modelview_matrix_array,
			projection_matrix: projection_matrix_array,
			scene: scene,
			textures: textures,
			vertex_buffers: vertex_buffers,
		}
	}

	pub fn get_mesh(&self, name: &str) -> Option<&Mesh> {
		for i in 0..self.vertex_buffers.len() as usize {
			if self.scene.meshes[i].name.eq(name) {
				return Some(&self.scene.meshes[i]);
			}
		}
		return None;
	}

	#[allow(unused_assignments)]
	pub fn create_shader_program(&self, shader_name: &str, dbloader: &DBLoader, display: &GlutinFacade) -> glium::program::Program {
		let supported_glsl_version: Version = display.get_supported_glsl_version();
		let api: glium::Api = supported_glsl_version.0;
		let major_version : u8 = supported_glsl_version.1;
		let minor_version : u8 = supported_glsl_version.2;
		let mut glsl_version_string: String = major_version.to_string();
		glsl_version_string.push_str(&minor_version.to_string());
		glsl_version_string.push('0');
		let mut use_gles: bool = false;// Use OpenGLES instead of OpenGL (for mobile devices)

		match api {
				glium::Api::Gl => {
					use_gles = false;
				},
				glium::Api::GlEs => {
					use_gles = true;
					glsl_version_string.push_str(" es");
				},
		};

		let glsl_version_number: u32 = match glsl_version_string.parse() {
			Ok(s) => s,
			Err(e) =>  {
				// This is not likely to happen
				panic!("Unable to parse supported glsl version string: {}", e);
			},
		};

		println!("Using glsl version {}", &glsl_version_string);

		let shader: Shader = dbloader.load_shader(shader_name, &glsl_version_string);

		let program: Program;

		if use_gles {
			program = program!(display,
				glsl_version_number es => {
					vertex: &shader.vertex_source,
					fragment: &shader.fragment_source,
				},
			)
			.unwrap();
		} else {
			program = program!(display,
				glsl_version_number => {
					vertex: &shader.vertex_source,
					fragment: &shader.fragment_source,
				},
			)
			.unwrap();
		}

		// Shader compilation no longer required
		display.release_shader_compiler();
		return program;
	}

	pub fn render(&self, display: &GlutinFacade, program: &glium::program::Program, camera: &Camera) {
		let mut target = display.draw();
		// TODO: generate this texture instead of loading from sqlite
		let default_blank_texture = &self.textures["DEFAULT_BLANK_TEXTURE.png"];
		target.clear_color(0.0, 0.0, 0.0, 1.0);
		for i in 0..self.vertex_buffers.len() as usize {
			let material_index: usize = self.scene.meshes[i].material_index.clone();
			let diffuse = self.scene.materials[material_index].diffuse.clone();
			let diffuse_texname: String =
				self.scene.materials[material_index].diffuse_texname.clone();
			let opengl_texture: &glium::texture::CompressedSrgbTexture2d;
			if diffuse_texname.len() == 0 {
				opengl_texture = &default_blank_texture;
			} else {
				match self.textures.get(&diffuse_texname) {
					Some(t) => {
						opengl_texture = &t;
					},
					None => {
						println!("Unable to load {}, using blank texture instead.", diffuse_texname);
						opengl_texture = &default_blank_texture;
					},
				}
			}
			let uniforms: glium::uniforms::UniformsStorage<_, _> = uniform! {
				projection: *camera.projection_matrix.borrow(),
				modelview: *camera.modelview_matrix.borrow(),
				light_position_worldspace: [2.0, 10.0, 1.0f32],
				diffuse: diffuse,
				diffuse_texture: opengl_texture,
				model: *self.scene.meshes[i].matrix.borrow(),
			};

			target.draw(&self.vertex_buffers[i],
				&self.index_buffer,
				program,
				&uniforms,
				&Default::default())
				.unwrap();
		}
		target.finish().unwrap();
	}
}
<commit_msg>Removed duplicate view matricies that moved to Camera struct<commit_after>// Copyright(C) 2016 Chris Liebert
extern crate glium;
extern crate image;
extern crate nalgebra;

use std::collections::HashMap;

use common;
use common::{Camera, Mesh, Scene, Shader, Vertex8f32};
use loader::DBLoader;

use glium::backend::glutin_backend::GlutinFacade;
use glium::{Program, Surface, Version};

implement_vertex!(Vertex8f32, position, normal, texcoord);

pub struct Renderer {
	pub index_buffer: glium::index::NoIndices,
	pub scene: Scene,
	pub textures: HashMap<String, glium::texture::CompressedSrgbTexture2d>,
	pub vertex_buffers: Vec<glium::vertex::VertexBuffer<common::Vertex8f32>>,
}

impl Renderer {
	pub fn new(display: &GlutinFacade, scene: Scene) -> Renderer {

		let mut vertex_buffers: Vec<glium::vertex::VertexBuffer<common::Vertex8f32>> =
			Vec::with_capacity(scene.meshes.len());

		for i in 0..scene.meshes.len() {
			vertex_buffers.push(glium::vertex::VertexBuffer::new(display, &scene.meshes[i].vertices).unwrap());
		}

		let index_buffer: glium::index::NoIndices =
			glium::index::NoIndices(glium::index::PrimitiveType::TrianglesList);

		let mut textures: HashMap<String, glium::texture::CompressedSrgbTexture2d> = HashMap::new();

		for i in 0..scene.images.len() {
			//TODO: determine image format by name extension if nessisary
			//let image = image::load_from_memory_with_format(&scene.images[i].image, image::PNG).unwrap().to_rgba();
			let image = image::load_from_memory(&scene.images[i].image).unwrap().to_rgba();

			let image_dimensions = image.dimensions();
			let image = glium::texture::RawImage2d::from_raw_rgba_reversed(image.into_raw(),
			                                                               image_dimensions);
			let opengl_texture: glium::texture::CompressedSrgbTexture2d =
				glium::texture::CompressedSrgbTexture2d::new(display, image).unwrap();
			textures.insert(scene.images[i].name.clone(), opengl_texture);
		}

		// Get the screen width and height in pixels
		let screen_width: u32;
		let screen_height: u32;
		match display.get_window() {
			Some(window) =>  {
				let pixel_dimensions: (u32,u32) = window.get_inner_size_pixels().unwrap();
				screen_width = pixel_dimensions.0;
				screen_height = pixel_dimensions.1;
			},
			None => {
				panic!("Error retrieving window when querying display size.");
			}
		}

		Renderer {
			index_buffer: index_buffer,
			scene: scene,
			textures: textures,
			vertex_buffers: vertex_buffers,
		}
	}

	pub fn get_mesh(&self, name: &str) -> Option<&Mesh> {
		for i in 0..self.vertex_buffers.len() as usize {
			if self.scene.meshes[i].name.eq(name) {
				return Some(&self.scene.meshes[i]);
			}
		}
		return None;
	}

	#[allow(unused_assignments)]
	pub fn create_shader_program(&self, shader_name: &str, dbloader: &DBLoader, display: &GlutinFacade) -> glium::program::Program {
		let supported_glsl_version: Version = display.get_supported_glsl_version();
		let api: glium::Api = supported_glsl_version.0;
		let major_version : u8 = supported_glsl_version.1;
		let minor_version : u8 = supported_glsl_version.2;
		let mut glsl_version_string: String = major_version.to_string();
		glsl_version_string.push_str(&minor_version.to_string());
		glsl_version_string.push('0');
		let mut use_gles: bool = false;// Use OpenGLES instead of OpenGL (for mobile devices)

		match api {
				glium::Api::Gl => {
					use_gles = false;
				},
				glium::Api::GlEs => {
					use_gles = true;
					glsl_version_string.push_str(" es");
				},
		};

		let glsl_version_number: u32 = match glsl_version_string.parse() {
			Ok(s) => s,
			Err(e) =>  {
				// This is not likely to happen
				panic!("Unable to parse supported glsl version string: {}", e);
			},
		};

		println!("Using glsl version {}", &glsl_version_string);

		let shader: Shader = dbloader.load_shader(shader_name, &glsl_version_string);

		let program: Program;

		if use_gles {
			program = program!(display,
				glsl_version_number es => {
					vertex: &shader.vertex_source,
					fragment: &shader.fragment_source,
				},
			)
			.unwrap();
		} else {
			program = program!(display,
				glsl_version_number => {
					vertex: &shader.vertex_source,
					fragment: &shader.fragment_source,
				},
			)
			.unwrap();
		}

		// Shader compilation no longer required
		display.release_shader_compiler();
		return program;
	}

	pub fn render(&self, display: &GlutinFacade, program: &glium::program::Program, camera: &Camera) {
		let mut target = display.draw();
		// TODO: generate this texture instead of loading from sqlite
		let default_blank_texture = &self.textures["DEFAULT_BLANK_TEXTURE.png"];
		target.clear_color(0.0, 0.0, 0.0, 1.0);
		for i in 0..self.vertex_buffers.len() as usize {
			let material_index: usize = self.scene.meshes[i].material_index.clone();
			let diffuse = self.scene.materials[material_index].diffuse.clone();
			let diffuse_texname: String =
				self.scene.materials[material_index].diffuse_texname.clone();
			let opengl_texture: &glium::texture::CompressedSrgbTexture2d;
			if diffuse_texname.len() == 0 {
				opengl_texture = &default_blank_texture;
			} else {
				match self.textures.get(&diffuse_texname) {
					Some(t) => {
						opengl_texture = &t;
					},
					None => {
						println!("Unable to load {}, using blank texture instead.", diffuse_texname);
						opengl_texture = &default_blank_texture;
					},
				}
			}
			let uniforms: glium::uniforms::UniformsStorage<_, _> = uniform! {
				projection: *camera.projection_matrix.borrow(),
				modelview: *camera.modelview_matrix.borrow(),
				light_position_worldspace: [2.0, 10.0, 1.0f32],
				diffuse: diffuse,
				diffuse_texture: opengl_texture,
				model: *self.scene.meshes[i].matrix.borrow(),
			};

			target.draw(&self.vertex_buffers[i],
				&self.index_buffer,
				program,
				&uniforms,
				&Default::default())
				.unwrap();
		}
		target.finish().unwrap();
	}
}
<|endoftext|>


<commit_before>use failure::{Error, Fail};
use strum_macros::Display;
use log::{info};
use bitflags::{bitflags};

use super::arch::{Arch, RVA, VA};
use super::loaders::pe::PELoader;
use super::loaders::sc::ShellcodeLoader;
use super::analysis::{Analyzer};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given buffer uses a bitness incompatible with the architecture")]
    MismatchedBitness,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw, // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

bitflags! {
    pub struct Permissions: u8 {
        const R = 0b00000001;
        const W = 0b00000010;
        const X = 0b00000100;
        const RW = Self::R.bits | Self::W.bits;
        const RX =  Self::R.bits | Self::X.bits;
        const WX =  Self::W.bits | Self::X.bits;
        const RWX =  Self::R.bits | Self::W.bits | Self::X.bits;
    }
}

#[derive(Debug)]
pub struct Section {
    pub addr: RVA,
    pub buf: Vec<u8>,
    pub perms: Permissions,
    pub name: String,
}

impl Section {
    pub fn contains(self: &Section, rva: RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        let max = self.addr + self.buf.len();
        if rva >= max {
            return false;
        }

        true
    }

    pub fn is_executable(&self) -> bool {
        self.perms.intersects(Permissions::X)
    }
}

pub struct LoadedModule {
    pub base_address: VA,
    pub sections: Vec<Section>,
}

pub trait Loader {
    /// Fetch the number of bits for a pointer in this architecture.
    fn get_arch(&self) -> Arch;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!(
            "{}/{}/{}",
            self.get_plat(),
            self.get_arch(),
            self.get_file_format()
        );
    }

    /// Returns True if this Loader knows how to load the given bytes.
    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given bytes into a Module and suggest the appropriate Analyzers.
    ///
    /// While the loader is parsing a file, it should determine what
    ///  the most appropriate analyzers are, e.g. a PE loader may inspect the headers
    ///  to determine if there is Control Flow Guard metadata that can be analyzed.
    fn load(&self, buf: &[u8]) -> Result<(LoadedModule, Vec<Box<dyn Analyzer>>), Error>;
}

pub fn default_loaders() -> Vec<Box<dyn Loader>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader>> = vec![];
    // the order here matters!
    // the default `load` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.

    loaders.push(Box::new(PELoader::new(Arch::X32)));
    loaders.push(Box::new(PELoader::new(Arch::X64)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X32)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X64)));

    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// match taste(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/x32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste(buf: &[u8]) -> impl Iterator<Item = Box<dyn Loader>> {
    default_loaders()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample using the first matching loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// load(b"\xEB\xFE")
///   .map(|(loader, module, analyzers)| {
///     assert_eq!(loader.get_name(),       "Windows/x32/Raw");
///     assert_eq!(module.base_address,     0x0);
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load(buf: &[u8]) -> Result<(Box<dyn Loader>, LoadedModule, Vec<Box<dyn Analyzer>>), Error> {
    match taste(buf).nth(0) {
        Some(loader) => {
            info!("auto-detected loader: {}", loader.get_name());
            loader.load(buf).map(|(module, analyzers)| (loader, module, analyzers))
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<commit_msg>loader: tests passing again<commit_after>use failure::{Error, Fail};
use strum_macros::Display;
use log::{info};
use bitflags::{bitflags};

use super::arch::{Arch, RVA, VA};
use super::loaders::pe::PELoader;
use super::loaders::sc::ShellcodeLoader;
use super::analysis::{Analyzer};

#[derive(Debug, Fail)]
pub enum LoaderError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given buffer uses a bitness incompatible with the architecture")]
    MismatchedBitness,
}

#[derive(Display, Clone, Copy)]
pub enum FileFormat {
    Raw, // shellcode
    PE,
}

#[derive(Display, Clone, Copy)]
pub enum Platform {
    Windows,
}

bitflags! {
    pub struct Permissions: u8 {
        const R = 0b00000001;
        const W = 0b00000010;
        const X = 0b00000100;
        const RW = Self::R.bits | Self::W.bits;
        const RX =  Self::R.bits | Self::X.bits;
        const WX =  Self::W.bits | Self::X.bits;
        const RWX =  Self::R.bits | Self::W.bits | Self::X.bits;
    }
}

#[derive(Debug)]
pub struct Section {
    pub addr: RVA,
    pub buf: Vec<u8>,
    pub perms: Permissions,
    pub name: String,
}

impl Section {
    pub fn contains(self: &Section, rva: RVA) -> bool {
        if rva < self.addr {
            return false;
        }

        let max = self.addr + self.buf.len();
        if rva >= max {
            return false;
        }

        true
    }

    pub fn is_executable(&self) -> bool {
        self.perms.intersects(Permissions::X)
    }
}

pub struct LoadedModule {
    pub base_address: VA,
    pub sections: Vec<Section>,
}

pub trait Loader {
    /// Fetch the number of bits for a pointer in this architecture.
    fn get_arch(&self) -> Arch;
    fn get_plat(&self) -> Platform;
    fn get_file_format(&self) -> FileFormat;

    fn get_name(&self) -> String {
        return format!(
            "{}/{}/{}",
            self.get_plat(),
            self.get_arch(),
            self.get_file_format()
        );
    }

    /// Returns True if this Loader knows how to load the given bytes.
    fn taste(&self, buf: &[u8]) -> bool;

    /// Load the given bytes into a Module and suggest the appropriate Analyzers.
    ///
    /// While the loader is parsing a file, it should determine what
    ///  the most appropriate analyzers are, e.g. a PE loader may inspect the headers
    ///  to determine if there is Control Flow Guard metadata that can be analyzed.
    fn load(&self, buf: &[u8]) -> Result<(LoadedModule, Vec<Box<dyn Analyzer>>), Error>;
}

pub fn default_loaders() -> Vec<Box<dyn Loader>> {
    // we might like these to come from a lazy_static global,
    //  however, then these have to be Sync.
    // I'm not sure if that's a good idea yet.
    let mut loaders: Vec<Box<dyn Loader>> = vec![];
    // the order here matters!
    // the default `load` routine will pick the first matching loader,
    //  so the earlier entries here have higher precedence.

    loaders.push(Box::new(PELoader::new(Arch::X32)));
    loaders.push(Box::new(PELoader::new(Arch::X64)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X32)));
    loaders.push(Box::new(ShellcodeLoader::new(Platform::Windows, Arch::X64)));

    loaders
}

/// Find the loaders that support loading the given sample.
///
/// The result is an iterator so that the caller can use the first
///  matching result without waiting for all Loaders to taste the bytes.
///
/// Loaders are tasted in the order defined in `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// match taste(b"\xEB\xFE").nth(0) {
///   Some(loader) => assert_eq!(loader.get_name(), "Windows/x32/Raw"),
///   None => panic!("no matching loaders"),
/// };
/// ```
pub fn taste(buf: &[u8]) -> impl Iterator<Item = Box<dyn Loader>> {
    default_loaders()
        .into_iter()
        .filter(move |loader| loader.taste(buf))
}

/// Load the given sample using the first matching loader from `default_loaders`.
///
/// Example:
///
/// ```
/// use lancelot::arch::*;
/// use lancelot::loader::*;
///
/// load(b"\xEB\xFE")
///   .map(|(loader, module, analyzers)| {
///     assert_eq!(loader.get_name(),       "Windows/x32/Raw");
///     assert_eq!(module.base_address,     VA(0x0));
///     assert_eq!(module.sections[0].name, "raw");
///   })
///   .map_err(|e| panic!(e));
/// ```
pub fn load(buf: &[u8]) -> Result<(Box<dyn Loader>, LoadedModule, Vec<Box<dyn Analyzer>>), Error> {
    match taste(buf).nth(0) {
        Some(loader) => {
            info!("auto-detected loader: {}", loader.get_name());
            loader.load(buf).map(|(module, analyzers)| (loader, module, analyzers))
        },
        None => Err(LoaderError::NotSupported.into()),
    }
}
<|endoftext|>


<commit_before>use futures::future::{ok, Future};

use std::rc::Rc;

use super::{box_up_err, peer_err, BoxedNewPeerFuture, Peer};
use super::{ConstructParams, L2rUser, PeerConstructor, Specifier, Options};

extern crate native_tls;
extern crate readwrite;
extern crate tokio_tls;

use self::native_tls::{TlsConnector,TlsAcceptor,Pkcs12};
use self::tokio_tls::{TlsConnectorExt,TlsAcceptorExt};

use ::std::ffi::{OsStr,OsString};

pub fn interpret_pkcs12(x: &OsStr) -> ::std::result::Result<Vec<u8>,OsString> {
    match (||{
        use std::io::Read;
        let mut f = ::std::fs::File::open(x)?;
        let mut v = Vec::with_capacity(2048);
        f.read_to_end(&mut v)?;
        Ok(v)
    })() {
        Err(e) => {
            let e : Box<::std::error::Error> = e;
            let o : OsString = format!("{}",e).into();
            Err(o)
        },
        Ok(x) => Ok(x),
    }
}

#[derive(Debug)]
pub struct TlsConnect<T: Specifier>(pub T);
impl<T: Specifier> Specifier for TlsConnect<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p, l2r| ssl_connect(p, l2r, cp.program_options.tls_domain.clone()))
    }
    specifier_boilerplate!(noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = TlsConnectClass,
    target = TlsConnect,
    prefixes = ["ssl-c","ssl:","tls:","ssl-connect:","tls-connect:","c-ssl:","connect-ssl:","c-tls:","connect-tls:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
Overlay to add TLS encryption atop of existing connection [A]

Example: manually connect to a secure websocket

    websocat -t - ws-c:tls-c:tcp:174.129.224.73:1080 --ws-c-uri ws://echo.websocket.org --tls-domain echo.websocket.org

For a user-friendly solution, see --socks5 command-line option
"#
);

#[derive(Debug)]
pub struct TlsAccept<T: Specifier>(pub T);
impl<T: Specifier> Specifier for TlsAccept<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p, l2r| ssl_accept(p, l2r, cp.program_options.clone()))
    }
    specifier_boilerplate!(noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = TlsAcceptClass,
    target = TlsAccept,
    prefixes = ["ssl-a:","tls-a:","ssl-accept:","tls-accept:","a-ssl:","accept-ssl:","accept-tls:","accept-tls:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
Accept an TLS connection using arbitrary backing stream. [A]

Example: The same as in TlsListenClass's example, but with manual acceptor

    websocat -E -b --pkcs12-der=q.pkcs12 tls-a:tcp-l:127.0.0.1:1234 mirror:
"#
);

specifier_alias!(
    name = TlsListenClass,
    prefixes = ["ssl-listen:","ssl-l:","tls-l:","tls-listen:","l-ssl:","listen-ssl:","listen-tls:","listen-tls:"],
    alias = "tls-accept:tcp-l:",
    help = r#"
Listen for SSL conections on a TCP port

Example: Non-websocket SSL echo server

    websocat -E -b --pkcs12-der=q.pkcs12 ssl-listen:127.0.0.1:1234 mirror:
    socat - ssl:127.0.0.1:1234,verify=0
"#
);

specifier_alias!(
    name = WssListenClass,
    prefixes = ["wss-l:","l-wss:","wss-listen:","wss-listen:"],
    alias = "ws-u:tls-accept:tcp-l:",
    help = r#"
Listen for secure WebSocket conections on a TCP port

Example: wss:// echo server + client for testing

    websocat -E -t --pkcs12-der=q.pkcs12 wss-listen:127.0.0.1:1234 mirror:
    websocat --ws-c-uri=wss://localhost/ -t - ws-c:cmd:'socat - ssl:127.0.0.1:1234,verify=0'

See [moreexamples.md](./moreexamples.md) for info about generation of `q.pkcs12`.
"#
);

use tokio_io::AsyncRead;

pub fn ssl_connect(inner_peer: Peer, _l2r: L2rUser, dom: Option<String>) -> BoxedNewPeerFuture {
    let squashed_peer = readwrite::ReadWriteAsync::new(inner_peer.0, inner_peer.1);

    fn gettlsc() -> native_tls::Result<TlsConnector> {
        Ok(TlsConnector::builder()?.build()?)
    }

    let tls = match gettlsc() {
        Ok(x) => x,
        Err(e) => return peer_err(e),
    };

    info!("Connecting to TLS");
    if let Some(dom) = dom {
        Box::new(
            tls.connect_async(dom.as_str(), squashed_peer)
                .map_err(box_up_err)
                .and_then(move |tls_stream| {
                    info!("Connected to TLS");
                    let (r, w) = tls_stream.split();
                    ok(Peer::new(r, w))
                }),
        )
    } else {
        Box::new(tls.danger_connect_async_without_providing_domain_for_certificate_verification_and_server_name_indication(squashed_peer).map_err(box_up_err).and_then(move |tls_stream| {
            warn!("Connected to TLS without proper verification of certificate. Use --tls-domain option.");
            let (r,w) = tls_stream.split();
            ok(Peer::new(r,w))
        }))
    }
}

pub fn ssl_accept(inner_peer: Peer, _l2r: L2rUser, progopt:Rc<Options>) -> BoxedNewPeerFuture {
    let squashed_peer = readwrite::ReadWriteAsync::new(inner_peer.0, inner_peer.1);

    fn gettlsa(cert:&[u8]) -> native_tls::Result<TlsAcceptor> {
        let pkcs12  = Pkcs12::from_der(&cert[..], "")?;
        Ok(TlsAcceptor::builder(pkcs12)?.build()?)
    }

    let der = progopt.pkcs12_der.as_ref()
        .expect("lint should have cought the missing pkcs12_der option");
    let tls = match gettlsa(der) {
        Ok(x) => x,
        Err(e) => return peer_err(e),
    };

    info!("Accepting a TLS connection");
    Box::new(
            tls.accept_async(squashed_peer)
                .map_err(box_up_err)
                .and_then(move |tls_stream| {
                    info!("Connected to TLS");
                    let (r, w) = tls_stream.split();
                    ok(Peer::new(r, w))
                }),
    )
}
<commit_msg>Change SSL-related primary prefix names<commit_after>use futures::future::{ok, Future};

use std::rc::Rc;

use super::{box_up_err, peer_err, BoxedNewPeerFuture, Peer};
use super::{ConstructParams, L2rUser, PeerConstructor, Specifier, Options};

extern crate native_tls;
extern crate readwrite;
extern crate tokio_tls;

use self::native_tls::{TlsConnector,TlsAcceptor,Pkcs12};
use self::tokio_tls::{TlsConnectorExt,TlsAcceptorExt};

use ::std::ffi::{OsStr,OsString};

pub fn interpret_pkcs12(x: &OsStr) -> ::std::result::Result<Vec<u8>,OsString> {
    match (||{
        use std::io::Read;
        let mut f = ::std::fs::File::open(x)?;
        let mut v = Vec::with_capacity(2048);
        f.read_to_end(&mut v)?;
        Ok(v)
    })() {
        Err(e) => {
            let e : Box<::std::error::Error> = e;
            let o : OsString = format!("{}",e).into();
            Err(o)
        },
        Ok(x) => Ok(x),
    }
}

#[derive(Debug)]
pub struct TlsConnect<T: Specifier>(pub T);
impl<T: Specifier> Specifier for TlsConnect<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p, l2r| ssl_connect(p, l2r, cp.program_options.tls_domain.clone()))
    }
    specifier_boilerplate!(noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = TlsConnectClass,
    target = TlsConnect,
    prefixes = ["ssl-connect:","ssl-c","ssl:","tls:","tls-connect:","c-ssl:","connect-ssl:","c-tls:","connect-tls:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
Overlay to add TLS encryption atop of existing connection [A]

Example: manually connect to a secure websocket

    websocat -t - ws-c:tls-c:tcp:174.129.224.73:1080 --ws-c-uri ws://echo.websocket.org --tls-domain echo.websocket.org

For a user-friendly solution, see --socks5 command-line option
"#
);

#[derive(Debug)]
pub struct TlsAccept<T: Specifier>(pub T);
impl<T: Specifier> Specifier for TlsAccept<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p, l2r| ssl_accept(p, l2r, cp.program_options.clone()))
    }
    specifier_boilerplate!(noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = TlsAcceptClass,
    target = TlsAccept,
    prefixes = ["ssl-accept:","ssl-a:","tls-a:","tls-accept:","a-ssl:","accept-ssl:","accept-tls:","accept-tls:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
Accept an TLS connection using arbitrary backing stream. [A]

Example: The same as in TlsListenClass's example, but with manual acceptor

    websocat -E -b --pkcs12-der=q.pkcs12 tls-a:tcp-l:127.0.0.1:1234 mirror:
"#
);

specifier_alias!(
    name = TlsListenClass,
    prefixes = ["ssl-listen:","ssl-l:","tls-l:","tls-listen:","l-ssl:","listen-ssl:","listen-tls:","listen-tls:"],
    alias = "tls-accept:tcp-l:",
    help = r#"
Listen for SSL conections on a TCP port

Example: Non-websocket SSL echo server

    websocat -E -b --pkcs12-der=q.pkcs12 ssl-listen:127.0.0.1:1234 mirror:
    socat - ssl:127.0.0.1:1234,verify=0
"#
);

specifier_alias!(
    name = WssListenClass,
    prefixes = ["wss-listen:", "wss-l:","l-wss:","wss-listen:"],
    alias = "ws-u:tls-accept:tcp-l:",
    help = r#"
Listen for secure WebSocket conections on a TCP port

Example: wss:// echo server + client for testing

    websocat -E -t --pkcs12-der=q.pkcs12 wss-listen:127.0.0.1:1234 mirror:
    websocat --ws-c-uri=wss://localhost/ -t - ws-c:cmd:'socat - ssl:127.0.0.1:1234,verify=0'

See [moreexamples.md](./moreexamples.md) for info about generation of `q.pkcs12`.
"#
);

use tokio_io::AsyncRead;

pub fn ssl_connect(inner_peer: Peer, _l2r: L2rUser, dom: Option<String>) -> BoxedNewPeerFuture {
    let squashed_peer = readwrite::ReadWriteAsync::new(inner_peer.0, inner_peer.1);

    fn gettlsc() -> native_tls::Result<TlsConnector> {
        Ok(TlsConnector::builder()?.build()?)
    }

    let tls = match gettlsc() {
        Ok(x) => x,
        Err(e) => return peer_err(e),
    };

    info!("Connecting to TLS");
    if let Some(dom) = dom {
        Box::new(
            tls.connect_async(dom.as_str(), squashed_peer)
                .map_err(box_up_err)
                .and_then(move |tls_stream| {
                    info!("Connected to TLS");
                    let (r, w) = tls_stream.split();
                    ok(Peer::new(r, w))
                }),
        )
    } else {
        Box::new(tls.danger_connect_async_without_providing_domain_for_certificate_verification_and_server_name_indication(squashed_peer).map_err(box_up_err).and_then(move |tls_stream| {
            warn!("Connected to TLS without proper verification of certificate. Use --tls-domain option.");
            let (r,w) = tls_stream.split();
            ok(Peer::new(r,w))
        }))
    }
}

pub fn ssl_accept(inner_peer: Peer, _l2r: L2rUser, progopt:Rc<Options>) -> BoxedNewPeerFuture {
    let squashed_peer = readwrite::ReadWriteAsync::new(inner_peer.0, inner_peer.1);

    fn gettlsa(cert:&[u8]) -> native_tls::Result<TlsAcceptor> {
        let pkcs12  = Pkcs12::from_der(&cert[..], "")?;
        Ok(TlsAcceptor::builder(pkcs12)?.build()?)
    }

    let der = progopt.pkcs12_der.as_ref()
        .expect("lint should have cought the missing pkcs12_der option");
    let tls = match gettlsa(der) {
        Ok(x) => x,
        Err(e) => return peer_err(e),
    };

    info!("Accepting a TLS connection");
    Box::new(
            tls.accept_async(squashed_peer)
                .map_err(box_up_err)
                .and_then(move |tls_stream| {
                    info!("Connected to TLS");
                    let (r, w) = tls_stream.split();
                    ok(Peer::new(r, w))
                }),
    )
}
<|endoftext|>