<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;
use std::ffi::OsString;
use std::fs;
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()
}

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

pub fn is_empty_dir<P: AsRef<Path>>(path: P) -> bool {
    match fs::read_dir(path) {
        Ok(rd) => rd.count() == 0,
        Err(_) => false,
    }
}

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(file_name: &str) -> bool {
    file_name.starts_with(".") && file_name.len() > 1 && file_name != ".."
}

fn is_hidden(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(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 dirs_ordered_by_descending_depth<P: AsRef<Path>>(root: P) -> Vec<PathBuf> {
    let mut walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(|e| !is_hidden(e))
        .filter_entry(|e| e.file_type().is_dir())
        .filter_map(Result::ok)
        .collect::<Vec<WalkDirEntry>>();

    walker.sort_by(|l, r| {
        let l_depth = &l.depth();
        let r_depth = &r.depth();
        r_depth.cmp(l_depth)
    });

    walker
        .into_iter()
        .map(|e| e.path().to_path_buf())
        .collect::<Vec<PathBuf>>()
}

#[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,
        f1: String,
        f2: String,
    }
    impl Helper {
        fn new(suffix: &str) -> Self {
            Helper {
                d1: format!("directory1_{}", suffix),
                d2: format!("directory2_{}", suffix),
                d3: format!("directory3_{}", suffix),
                f1: format!("file1_{}"     , suffix),
                f2: format!("file2_{}"     , 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_f1(&self) -> PathBuf {
            self.path_to_d1().join(&self.f1)
        }
        fn path_to_f2(&self) -> PathBuf {
            self.path_to_d2().join(&self.f2)
        }

        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_d3()).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();
        }

        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 ls_should_return_ok() {
        let helper = Helper::new("ls_Ok");
        helper.create_dirs_and_files();

        let empty_vec: Vec<String> = Vec::new();

        assert_eq!(vec![helper.d2.clone(), helper.f1.clone()], ls(helper.path_to_d1()).unwrap());
        assert_eq!(vec![helper.d3.clone(), helper.f2.clone()], ls(helper.path_to_d2()).unwrap());
        assert_eq!(empty_vec                                 , ls(helper.path_to_d3()).unwrap());

        helper.remove_dirs_and_files();
    }

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

        assert!(is_empty_dir(helper.path_to_d3()));

        helper.remove_dirs_and_files();
    }
    #[test]
    fn is_empty_dir_should_return_false() {
        let helper = Helper::new("is_empty_dir_false");
        helper.create_dirs_and_files();

        assert!(!is_empty_dir(helper.path_to_d1()));
        assert!(!is_empty_dir(helper.path_to_d2()));

        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()));

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

        helper.remove_dirs_and_files();
    }

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

        let correct = vec![helper.path_to_d3(), helper.path_to_d2(), helper.path_to_d1()];

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

        helper.remove_dirs_and_files();
    }
}
<commit_msg>[Feature] Implement potential_empty_dirs() roughly<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::ffi::OsString;
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()
}

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

pub fn is_empty_dir<P: AsRef<Path>>(path: P) -> bool {
    match fs::read_dir(path) {
        Ok(rd) => rd.count() == 0,
        Err(_) => false,
    }
}

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(file_name: &str) -> bool {
    file_name.starts_with(".") && file_name.len() > 1 && file_name != ".."
}

fn is_hidden(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(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 potential_empty_dirs<P: AsRef<Path>>(root: P) -> Result<BTreeSet<PathBuf>, IoError> {
    fn potential_empty_dirs(mut result: BTreeSet<PathBuf>, mut target_dirs: VecDeque<PathBuf>) -> Result<BTreeSet<PathBuf>, IoError> {
        match target_dirs.pop_front() {
            None => Ok(result),
            Some(mut target_dir) => {
                let entries = try!(fs::read_dir(&target_dir))
                    .filter_map(Result::ok)
                    .collect::<Vec<DirEntry>>();

                let include_file = entries
                    .iter()
                    .any(|entry| entry
                        .file_type()
                        .map(|t| t.is_file())
                        .unwrap_or(true)
                    );
                if include_file {
                    loop {
                        result.remove(&target_dir);
                        if !target_dir.pop(){
                            break;
                        }
                    }
                }

                let dirs = entries
                    .into_iter()
                    .filter_map(|entry| {
                        let is_dir = entry
                            .file_type()
                            .map(|t| t.is_dir())
                            .unwrap_or(false);
                        if is_dir {
                            Some(entry.path())
                        } else {
                            None
                        }
                    })
                    .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);

                potential_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);

    potential_empty_dirs(result, target_dirs)
}

pub fn dirs_ordered_by_descending_depth<P: AsRef<Path>>(root: P) -> Vec<PathBuf> {
    let mut walker = WalkDir::new(root)
        .into_iter()
        .filter_entry(|e| !is_hidden(e))
        .filter_entry(|e| e.file_type().is_dir())
        .filter_map(Result::ok)
        .collect::<Vec<WalkDirEntry>>();

    walker.sort_by(|l, r| {
        let l_depth = &l.depth();
        let r_depth = &r.depth();
        r_depth.cmp(l_depth)
    });

    walker
        .into_iter()
        .map(|e| e.path().to_path_buf())
        .collect::<Vec<PathBuf>>()
}

#[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,
        f1: String,
        f2: String,
    }
    impl Helper {
        fn new(suffix: &str) -> Self {
            Helper {
                d1: format!("directory1_{}", suffix),
                d2: format!("directory2_{}", suffix),
                d3: format!("directory3_{}", suffix),
                f1: format!("file1_{}"     , suffix),
                f2: format!("file2_{}"     , 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_f1(&self) -> PathBuf {
            self.path_to_d1().join(&self.f1)
        }
        fn path_to_f2(&self) -> PathBuf {
            self.path_to_d2().join(&self.f2)
        }

        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_d3()).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();
        }

        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 ls_should_return_ok() {
        let helper = Helper::new("ls_Ok");
        helper.create_dirs_and_files();

        let empty_vec: Vec<String> = Vec::new();

        assert_eq!(vec![helper.d2.clone(), helper.f1.clone()], ls(helper.path_to_d1()).unwrap());
        assert_eq!(vec![helper.d3.clone(), helper.f2.clone()], ls(helper.path_to_d2()).unwrap());
        assert_eq!(empty_vec                                 , ls(helper.path_to_d3()).unwrap());

        helper.remove_dirs_and_files();
    }

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

        assert!(is_empty_dir(helper.path_to_d3()));

        helper.remove_dirs_and_files();
    }
    #[test]
    fn is_empty_dir_should_return_false() {
        let helper = Helper::new("is_empty_dir_false");
        helper.create_dirs_and_files();

        assert!(!is_empty_dir(helper.path_to_d1()));
        assert!(!is_empty_dir(helper.path_to_d2()));

        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()));

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

        helper.remove_dirs_and_files();
    }

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

        let correct = vec![helper.path_to_d3(), helper.path_to_d2(), helper.path_to_d1()];

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

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


<commit_before>// Copyright 2014-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.
/// The standard logging macro.
///
/// This macro will generically log with the specified `LogLevel` and `format!`
/// based argument list.
///
/// The `log_level` cfg can be used to statically disable logging at various
/// levels.
#[macro_export]
macro_rules! log {
    (target: $target:expr, $lvl:expr, $($arg:tt)+) => ({
        static _LOC: $crate::LogLocation = $crate::LogLocation {
            __line: line!(),
            __file: file!(),
            __module_path: module_path!(),
        };
        let lvl = $lvl;
        if !cfg!(log_level = "off") &&
                (lvl <= $crate::LogLevel::Error || !cfg!(log_level = "error")) &&
                (lvl <= $crate::LogLevel::Warn || !cfg!(log_level = "warn")) &&
                (lvl <= $crate::LogLevel::Debug || !cfg!(log_level = "debug")) &&
                (lvl <= $crate::LogLevel::Info || !cfg!(log_level = "info")) &&
                lvl <= $crate::max_log_level() {
            $crate::__log(lvl, $target, &_LOC, format_args!($($arg)+))
        }
    });
    ($lvl:expr, $($arg:tt)+) => (log!(target: module_path!(), $lvl, $($arg)+))
}

/// Logs a message at the error level.
///
/// Logging at this level is disabled if the `log_level = "off"` cfg is
/// present.
#[macro_export]
macro_rules! error {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Error, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Error, $($arg)*);
    )
}

/// Logs a message at the warn level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"` or `log_level = "error"`.
#[macro_export]
macro_rules! warn {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Warn, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Warn, $($arg)*);
    )
}

/// Logs a message at the info level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"`, `log_level = "error"`, or
/// `log_level = "warn"`.
#[macro_export]
macro_rules! info {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Info, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Info, $($arg)*);
    )
}

/// Logs a message at the debug level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"`, `log_level = "error"`, `log_level = "warn"`,
/// or `log_level = "info"`.
#[macro_export]
macro_rules! debug {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Debug, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Debug, $($arg)*);
    )
}

/// Logs a message at the trace level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"`, `log_level = "error"`, `log_level = "warn"`,
/// `log_level = "info"`, or `log_level = "debug"`.
#[macro_export]
macro_rules! trace {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Trace, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Trace, $($arg)*);
    )
}

/// Determines if a message logged at the specified level in that module will
/// be logged.
///
/// This can be used to avoid expensive computation of log message arguments if
/// the message would be ignored anyway.
///
/// # Examples
///
/// ```rust
/// # #[macro_use]
/// # extern crate log;
/// use log::LogLevel::Debug;
///
/// # fn foo() {
/// if log_enabled!(Debug) {
///     debug!("expensive debug data: {}", expensive_call());
/// }
/// # }
/// # fn expensive_call() -> u32 { 0 }
/// # fn main() {}
/// ```
#[macro_export]
macro_rules! log_enabled {
    (target: $target:expr, $lvl:expr) => ({
        let lvl = $lvl;
        !cfg!(log_level = "off") &&
            (lvl <= $crate::LogLevel::Error || !cfg!(log_level = "error")) &&
            (lvl <= $crate::LogLevel::Warn || !cfg!(log_level = "warn")) &&
            (lvl <= $crate::LogLevel::Debug || !cfg!(log_level = "debug")) &&
            (lvl <= $crate::LogLevel::Info || !cfg!(log_level = "info")) &&
            lvl <= $crate::max_log_level() &&
            $crate::__enabled(lvl, $target)
    });
    ($lvl:expr) => (log_enabled!(target: module_path!(), $lvl))
}
<commit_msg>[docs] Use a better example for log_enabled!<commit_after>// Copyright 2014-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.
/// The standard logging macro.
///
/// This macro will generically log with the specified `LogLevel` and `format!`
/// based argument list.
///
/// The `log_level` cfg can be used to statically disable logging at various
/// levels.
#[macro_export]
macro_rules! log {
    (target: $target:expr, $lvl:expr, $($arg:tt)+) => ({
        static _LOC: $crate::LogLocation = $crate::LogLocation {
            __line: line!(),
            __file: file!(),
            __module_path: module_path!(),
        };
        let lvl = $lvl;
        if !cfg!(log_level = "off") &&
                (lvl <= $crate::LogLevel::Error || !cfg!(log_level = "error")) &&
                (lvl <= $crate::LogLevel::Warn || !cfg!(log_level = "warn")) &&
                (lvl <= $crate::LogLevel::Debug || !cfg!(log_level = "debug")) &&
                (lvl <= $crate::LogLevel::Info || !cfg!(log_level = "info")) &&
                lvl <= $crate::max_log_level() {
            $crate::__log(lvl, $target, &_LOC, format_args!($($arg)+))
        }
    });
    ($lvl:expr, $($arg:tt)+) => (log!(target: module_path!(), $lvl, $($arg)+))
}

/// Logs a message at the error level.
///
/// Logging at this level is disabled if the `log_level = "off"` cfg is
/// present.
#[macro_export]
macro_rules! error {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Error, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Error, $($arg)*);
    )
}

/// Logs a message at the warn level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"` or `log_level = "error"`.
#[macro_export]
macro_rules! warn {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Warn, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Warn, $($arg)*);
    )
}

/// Logs a message at the info level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"`, `log_level = "error"`, or
/// `log_level = "warn"`.
#[macro_export]
macro_rules! info {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Info, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Info, $($arg)*);
    )
}

/// Logs a message at the debug level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"`, `log_level = "error"`, `log_level = "warn"`,
/// or `log_level = "info"`.
#[macro_export]
macro_rules! debug {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Debug, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Debug, $($arg)*);
    )
}

/// Logs a message at the trace level.
///
/// Logging at this level is disabled if any of the following cfgs are present:
/// `log_level = "off"`, `log_level = "error"`, `log_level = "warn"`,
/// `log_level = "info"`, or `log_level = "debug"`.
#[macro_export]
macro_rules! trace {
    (target: $target:expr, $($arg:tt)*) => (
        log!(target: $target, $crate::LogLevel::Trace, $($arg)*);
    );
    ($($arg:tt)*) => (
        log!($crate::LogLevel::Trace, $($arg)*);
    )
}

/// Determines if a message logged at the specified level in that module will
/// be logged.
///
/// This can be used to avoid expensive computation of log message arguments if
/// the message would be ignored anyway.
///
/// # Examples
///
/// ```rust
/// # #[macro_use]
/// # extern crate log;
/// use log::LogLevel::Debug;
///
/// # fn foo() {
/// if log_enabled!(Debug) {
///     let data = expensive_call();
///     debug!("expensive debug data: {} {}", data.x, data.y);
/// }
/// # }
/// # struct Data { x: u32, y: u32 }
/// # fn expensive_call() -> Data { Data { x: 0, y: 0 } }
/// # fn main() {}
/// ```
#[macro_export]
macro_rules! log_enabled {
    (target: $target:expr, $lvl:expr) => ({
        let lvl = $lvl;
        !cfg!(log_level = "off") &&
            (lvl <= $crate::LogLevel::Error || !cfg!(log_level = "error")) &&
            (lvl <= $crate::LogLevel::Warn || !cfg!(log_level = "warn")) &&
            (lvl <= $crate::LogLevel::Debug || !cfg!(log_level = "debug")) &&
            (lvl <= $crate::LogLevel::Info || !cfg!(log_level = "info")) &&
            lvl <= $crate::max_log_level() &&
            $crate::__enabled(lvl, $target)
    });
    ($lvl:expr) => (log_enabled!(target: module_path!(), $lvl))
}
<|endoftext|>


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

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

/// Connect events to sending a message.
///
/// Rule #1:
/// Send `$msg` to `$other_component` when the GTK+ `$event` is emitted on `$widget`.
///
/// Rule #2:
/// Optionally send `$msg.0` when the GTK+ `$event` is emitted on `$widget`.
/// Return `$msg.1` in the GTK+ callback.
/// This variant gives more control to the caller since it expects a `$msg` returning `(Option<MSG>,
/// ReturnValue)` where the `ReturnValue` is the value to return in the GTK+ callback.
/// Option<MSG> can be None if no message needs to be emitted.
///
/// Rule #3:
/// Send `$msg` when the GTK+ `$event` is emitted on `$widget`.
///
/// Rule #4:
/// Send `$msg` to `$widget` when the `$message` is received on `$stream`.
#[macro_export]
macro_rules! connect {
    // Connect to a GTK+ widget event, sending a message to another widget.
    ($widget:expr, $event:ident($($args:pat),*), $other_component:expr, $msg:expr) => {
        connect_stream!($widget, $event($($args),*), $other_component.stream(), $msg);
    };

    // Connect to a GTK+ widget event.
    // This variant gives more control to the caller since it expects a `$msg` returning (Option<MSG>,
    // ReturnValue) where the ReturnValue is the value to return in the GTK+ callback.
    // Option<MSG> can be None if no message needs to be emitted.
    ($relm:expr, $widget:expr, $event:ident($($args:pat),*), return $msg:expr) => {{
        connect_stream!(return $relm.stream(), $widget, $event($($args),*), $msg);
    }};

    // Connect to a GTK+ widget event.
    ($relm:expr, $widget:expr, $event:ident($($args:pat),*), $msg:expr) => {{
        let stream = $relm.stream().clone();
        let _ = $widget.$event(move |$($args),*| {
            let msg: Option<_> = $crate::IntoOption::into_option($msg);
            if let Some(msg) = msg {
                stream.emit(msg);
            }
        });
    }};

    // Connect to a message reception.
    // TODO: create another macro rule accepting multiple patterns.
    ($src_component:ident @ $message:pat, $dst_component:expr, $msg:expr) => {
        let stream = $src_component.stream().clone();
        connect_stream!(stream@$message, $dst_component.stream(), $msg);
    };
}

/// Connect events to sending a message.
/// Similar to `connect!` but wants a stream instead of a component.
///
/// Rule #1:
/// Send `$msg` to `$other_stream` when the GTK+ `$event` is emitted on `$widget`.
///
/// Rule #2:
/// Send `$msg` to `$widget` when the `$message` is received on `$stream`.
#[macro_export]
macro_rules! connect_stream {
    // Connect to a GTK+ widget event.
    // This variant gives more control to the caller since it expects a `$msg` returning (Option<MSG>,
    // ReturnValue) where the ReturnValue is the value to return in the GTK+ callback.
    // Option<MSG> can be None if no message needs to be emitted.
    (return $stream:expr, $widget:expr, $event:ident($($args:pat),*), $msg:expr) => {{
        let stream = $stream.clone();
        let _ = $widget.$event(move |$($args),*| {
            let (msg, return_value) = $crate::IntoPair::into_pair($msg);
            let msg: Option<_> = $crate::IntoOption::into_option(msg);
            if let Some(msg) = msg {
                stream.emit(msg);
            }
            return_value
        });
    }};

    // Connect to a GTK+ widget event, sending a message to another widget.
    ($widget:expr, $event:ident($($args:pat),*), $other_stream:expr, $msg:expr) => {
        let stream = $other_stream.clone();
        let _ = $widget.$event(move |$($args),*| {
            let msg: Option<_> = $crate::IntoOption::into_option($msg);
            if let Some(msg) = msg {
                stream.emit(msg);
            }
        });
    };

    // Connect to a message reception.
    // TODO: create another macro rule accepting multiple patterns.
    ($src_stream:ident @ $message:pat, $dst_stream:expr, $msg:expr) => {
        let stream = $dst_stream.clone();
        $src_stream.observe(move |msg| {
            #[allow(unreachable_patterns)]
            match msg {
                &$message =>  {
                    let msg: Option<_> = $crate::IntoOption::into_option($msg);
                    if let Some(msg) = msg {
                        stream.emit(msg);
                    }
                },
                _ => (),
            }
        });
    };
}
<|endoftext|>


<commit_before>/// The Matrix type represents a 2-dimensional Matrix.
pub struct Matrix {
    cols: usize,
    contents: Vec<f64>
}

impl Matrix {
    /// Creates a new Matrix with the given width and height
    pub fn new(width: usize, height: usize) -> Matrix {
        Matrix {
            contents: vec![0.0; width * height],
            cols: width,
        }
    }

    /// Returns a reference to the specified coordinates of the Matrix
    pub fn get(&self, col: usize, row: usize) -> &f64 {
        debug_assert!(col * row < self.contents.len());
        self.contents.get(row * self.cols + col).unwrap()
    }

    /// Sets the coordinates of the Matrix to the specified value
    pub fn set(&mut self, col: usize, row: usize, val: f64) {
        debug_assert!(col * row < self.contents.len());
        self.contents[row * self.cols + col] = val;
    }
}

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

    #[test]
    fn test_get() {
        let mat = Matrix::new(10, 5);
        assert_eq!(mat.get(0, 0), &0.0);
        assert_eq!(mat.get(9, 4), &0.0);
    }

    #[test]
    fn test_set() {
        let mut mat = Matrix::new(10, 5);
        mat.set(9, 4, 1.0);
        assert_eq!(mat.get(0, 0), &0.0);
        assert_eq!(mat.get(9, 4), &1.0);
    }
}
<commit_msg>Optimize matrix initialization<commit_after>/// The Matrix type represents a 2-dimensional Matrix.
pub struct Matrix {
    cols: usize,
    contents: Vec<f64>
}

impl Matrix {
    /// Creates a new Matrix with the given width and height
    pub fn new(width: usize, height: usize) -> Matrix {
        Matrix {
            contents: (0..width * height).map(|_i| 0.0).collect(),
            cols: width
        }
    }

    /// Returns a reference to the specified coordinates of the Matrix
    pub fn get(&self, col: usize, row: usize) -> &f64 {
        debug_assert!(col * row < self.contents.len());
        self.contents.get(row * self.cols + col).unwrap()
    }

    /// Sets the coordinates of the Matrix to the specified value
    pub fn set(&mut self, col: usize, row: usize, val: f64) {
        debug_assert!(col * row < self.contents.len());
        self.contents[row * self.cols + col] = val;
    }
}

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

    #[test]
    fn test_get() {
        let mat = Matrix::new(10, 5);
        assert_eq!(mat.get(0, 0), &0.0);
        assert_eq!(mat.get(9, 4), &0.0);
    }

    #[test]
    fn test_set() {
        let mut mat = Matrix::new(10, 5);
        mat.set(9, 4, 1.0);
        assert_eq!(mat.get(0, 0), &0.0);
        assert_eq!(mat.get(9, 4), &1.0);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Nathan Zadoks <nathan@nathan7.eu>
// See the LICENSE file included in this distribution.
//! This module provides stacks using the operating system's
//! anonymous memory mapping facility, usually `mmap`.
//! The stacks it provides come with a guard page, which is not included
//! in the stack limit.

extern crate std;
use core::prelude::*;
use self::std::io::Error as IoError;
use stack;
mod sys;

pub struct StackSource;

unsafe impl Send for StackSource {}
unsafe impl Sync for StackSource {}

#[allow(raw_pointer_derive)]
#[derive(Debug)]
pub struct Stack {
  ptr: *mut u8,
  len: usize
}

unsafe impl Send for Stack {}

impl stack::StackSource for StackSource {
  type Output = Stack;
  type Error = IoError;

  fn get_stack(size: usize) -> Result<Stack, IoError> {
    let page_size = sys::page_size();

    // round the page size up,
    // using the fact that it is a power of two
    let len = (size + page_size - 1) & !(page_size - 1);

    let stack = unsafe {
      let ptr = try!(match sys::map_stack(size) {
        None => Err(IoError::last_os_error()),
        Some(ptr) => Ok(ptr)
      });

      Stack { ptr: ptr as *mut u8, len: len }
    };

    try!(unsafe {
      if sys::protect_stack(stack.ptr) { Ok(()) }
      else { Err(IoError::last_os_error()) }
    });

    Ok(stack)
  }
}

impl stack::Stack for Stack {
  fn top(&mut self) -> *mut u8 {
    unsafe {
      self.ptr.offset(self.len as isize)
    }
  }

  fn limit(&self) -> *const u8 {
    unsafe {
      self.ptr.offset(sys::page_size() as isize)
    }
  }
}

impl Drop for Stack {
  fn drop(&mut self) {
    unsafe {
      if !sys::unmap_stack(self.ptr, self.len) {
        panic!("munmap for stack {:p} of size {} failed: {}",
               self.ptr, self.len, IoError::last_os_error())
      }
    }
  }
}
<commit_msg>derive Copy and Clone for os::StackSource<commit_after>// Copyright (c) 2015, Nathan Zadoks <nathan@nathan7.eu>
// See the LICENSE file included in this distribution.
//! This module provides stacks using the operating system's
//! anonymous memory mapping facility, usually `mmap`.
//! The stacks it provides come with a guard page, which is not included
//! in the stack limit.

extern crate std;
use core::prelude::*;
use self::std::io::Error as IoError;
use stack;
mod sys;

#[derive(Copy, Clone)]
pub struct StackSource;

unsafe impl Send for StackSource {}
unsafe impl Sync for StackSource {}

#[allow(raw_pointer_derive)]
#[derive(Debug)]
pub struct Stack {
  ptr: *mut u8,
  len: usize
}

unsafe impl Send for Stack {}

impl stack::StackSource for StackSource {
  type Output = Stack;
  type Error = IoError;

  fn get_stack(size: usize) -> Result<Stack, IoError> {
    let page_size = sys::page_size();

    // round the page size up,
    // using the fact that it is a power of two
    let len = (size + page_size - 1) & !(page_size - 1);

    let stack = unsafe {
      let ptr = try!(match sys::map_stack(size) {
        None => Err(IoError::last_os_error()),
        Some(ptr) => Ok(ptr)
      });

      Stack { ptr: ptr as *mut u8, len: len }
    };

    try!(unsafe {
      if sys::protect_stack(stack.ptr) { Ok(()) }
      else { Err(IoError::last_os_error()) }
    });

    Ok(stack)
  }
}

impl stack::Stack for Stack {
  fn top(&mut self) -> *mut u8 {
    unsafe {
      self.ptr.offset(self.len as isize)
    }
  }

  fn limit(&self) -> *const u8 {
    unsafe {
      self.ptr.offset(sys::page_size() as isize)
    }
  }
}

impl Drop for Stack {
  fn drop(&mut self) {
    unsafe {
      if !sys::unmap_stack(self.ptr, self.len) {
        panic!("munmap for stack {:p} of size {} failed: {}",
               self.ptr, self.len, IoError::last_os_error())
      }
    }
  }
}
<|endoftext|>


<commit_before>use crate::{
    script::{EvalErr, EvalErrType, InitErr},
    serializer::*,
};
use std::io::{self, Cursor};

#[derive(Copy, Clone, Debug)]
pub struct Config {
    pub skip_reward: bool,
}

impl Config {
    pub const fn strict() -> Self {
        Self { skip_reward: false }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum BlockErr {
    InvalidBlockHeight,
    InvalidMerkleRoot,
    InvalidSignature,
    InvalidPrevHash,
    InvalidHash,
    Tx(TxErr),
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum TxErr {
    ScriptTooLarge,
    ScriptEval(EvalErr),
    ScriptHashMismatch,
    ScriptRetFalse,
    Arithmetic,
    InsufficientBalance,
    InvalidFeeAmount,
    TooManySignatures,
    TxProhibited,
    TxExpired,
    TxDupe,
}

impl TxErr {
    pub fn serialize(self, buf: &mut Vec<u8>) {
        match self {
            TxErr::ScriptTooLarge => buf.push(0),
            TxErr::ScriptEval(err) => {
                buf.push(1);
                buf.push_u32(err.pos);
                buf.push(err.err as u8);
            }
            TxErr::ScriptHashMismatch => buf.push(2),
            TxErr::ScriptRetFalse => buf.push(3),
            TxErr::Arithmetic => buf.push(4),
            TxErr::InsufficientBalance => buf.push(5),
            TxErr::InvalidFeeAmount => buf.push(6),
            TxErr::TooManySignatures => buf.push(7),
            TxErr::TxProhibited => buf.push(8),
            TxErr::TxExpired => buf.push(9),
            TxErr::TxDupe => buf.push(10),
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        Ok(match tag {
            0 => TxErr::ScriptTooLarge,
            1 => {
                let pos = cursor.take_u32()?;
                let kind = match cursor.take_u8()? {
                    t if t == EvalErrType::UnexpectedEOF as u8 => EvalErrType::UnexpectedEOF,
                    t if t == EvalErrType::UnknownOp as u8 => EvalErrType::UnknownOp,
                    t if t == EvalErrType::InvalidItemOnStack as u8 => {
                        EvalErrType::InvalidItemOnStack
                    }
                    t if t == EvalErrType::StackOverflow as u8 => EvalErrType::StackOverflow,
                    t if t == EvalErrType::StackUnderflow as u8 => EvalErrType::StackUnderflow,
                    _ => {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "failed to deserialize EvalErrType",
                        ))
                    }
                };
                TxErr::ScriptEval(EvalErr::new(pos, kind))
            }
            2 => TxErr::ScriptHashMismatch,
            3 => TxErr::ScriptRetFalse,
            4 => TxErr::Arithmetic,
            5 => TxErr::InsufficientBalance,
            6 => TxErr::InvalidFeeAmount,
            7 => TxErr::TooManySignatures,
            8 => TxErr::TxProhibited,
            9 => TxErr::TxExpired,
            10 => TxErr::TxDupe,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to deserialize TxErr",
                ))
            }
        })
    }
}

impl From<InitErr> for TxErr {
    fn from(err: InitErr) -> Self {
        match err {
            InitErr::ScriptTooLarge => TxErr::ScriptTooLarge,
            InitErr::TooManySignatures => TxErr::TooManySignatures,
        }
    }
}
<commit_msg>Rename ScriptTooLarge to TxTooLarge<commit_after>use crate::{
    script::{EvalErr, EvalErrType, InitErr},
    serializer::*,
};
use std::io::{self, Cursor};

#[derive(Copy, Clone, Debug)]
pub struct Config {
    pub skip_reward: bool,
}

impl Config {
    pub const fn strict() -> Self {
        Self { skip_reward: false }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum BlockErr {
    InvalidBlockHeight,
    InvalidMerkleRoot,
    InvalidSignature,
    InvalidPrevHash,
    InvalidHash,
    Tx(TxErr),
}

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum TxErr {
    ScriptEval(EvalErr),
    ScriptHashMismatch,
    ScriptRetFalse,
    Arithmetic,
    InsufficientBalance,
    InvalidFeeAmount,
    TooManySignatures,
    TxTooLarge,
    TxProhibited,
    TxExpired,
    TxDupe,
}

impl TxErr {
    pub fn serialize(self, buf: &mut Vec<u8>) {
        match self {
            TxErr::ScriptEval(err) => {
                buf.push(0);
                buf.push_u32(err.pos);
                buf.push(err.err as u8);
            }
            TxErr::ScriptHashMismatch => buf.push(1),
            TxErr::ScriptRetFalse => buf.push(2),
            TxErr::Arithmetic => buf.push(3),
            TxErr::InsufficientBalance => buf.push(4),
            TxErr::InvalidFeeAmount => buf.push(5),
            TxErr::TooManySignatures => buf.push(6),
            TxErr::TxTooLarge => buf.push(7),
            TxErr::TxProhibited => buf.push(8),
            TxErr::TxExpired => buf.push(9),
            TxErr::TxDupe => buf.push(10),
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        Ok(match tag {
            0 => {
                let pos = cursor.take_u32()?;
                let kind = match cursor.take_u8()? {
                    t if t == EvalErrType::UnexpectedEOF as u8 => EvalErrType::UnexpectedEOF,
                    t if t == EvalErrType::UnknownOp as u8 => EvalErrType::UnknownOp,
                    t if t == EvalErrType::InvalidItemOnStack as u8 => {
                        EvalErrType::InvalidItemOnStack
                    }
                    t if t == EvalErrType::StackOverflow as u8 => EvalErrType::StackOverflow,
                    t if t == EvalErrType::StackUnderflow as u8 => EvalErrType::StackUnderflow,
                    _ => {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "failed to deserialize EvalErrType",
                        ))
                    }
                };
                TxErr::ScriptEval(EvalErr::new(pos, kind))
            }
            1 => TxErr::ScriptHashMismatch,
            2 => TxErr::ScriptRetFalse,
            3 => TxErr::Arithmetic,
            4 => TxErr::InsufficientBalance,
            5 => TxErr::InvalidFeeAmount,
            6 => TxErr::TooManySignatures,
            7 => TxErr::TxTooLarge,
            8 => TxErr::TxProhibited,
            9 => TxErr::TxExpired,
            10 => TxErr::TxDupe,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to deserialize TxErr",
                ))
            }
        })
    }
}

impl From<InitErr> for TxErr {
    fn from(err: InitErr) -> Self {
        match err {
            InitErr::ScriptTooLarge => TxErr::TxTooLarge,
            InitErr::TooManySignatures => TxErr::TooManySignatures,
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    collections::{BTreeMap, HashMap},
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::{net::TcpSocketOpts, relay::socks5::Address};
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    storage::RingBuffer,
    time::{Duration, Instant},
    wire::{IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    sync::{mpsc, Notify},
    task::JoinHandle,
    time,
};

use crate::local::{
    context::ServiceContext,
    loadbalancing::PingBalancer,
    net::AutoProxyClientStream,
    utils::{establish_tcp_tunnel, to_ipv4_mapped},
};

use super::virt_device::VirtTunDevice;

// NOTE: Default value is taken from Linux
// recv: /proc/sys/net/ipv4/tcp_rmem 87380 bytes
// send: /proc/sys/net/ipv4/tcp_wmem 16384 bytes
const DEFAULT_TCP_SEND_BUFFER_SIZE: u32 = 16384;
const DEFAULT_TCP_RECV_BUFFER_SIZE: u32 = 87380;

struct TcpSocketControl {
    send_buffer: RingBuffer<'static, u8>,
    send_waker: Option<Waker>,
    recv_buffer: RingBuffer<'static, u8>,
    recv_waker: Option<Waker>,
    is_closed: bool,
}

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
    sockets: HashMap<SocketHandle, Arc<ParkingMutex<TcpSocketControl>>>,
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    control: Arc<ParkingMutex<TcpSocketControl>>,
    manager_notify: Arc<Notify>,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut control = self.control.lock();
        control.is_closed = true;
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager, tcp_opts: &TcpSocketOpts) -> TcpConnection {
        let send_buffer_size = tcp_opts.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
        let recv_buffer_size = tcp_opts.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

        let (control, manager_notify) = {
            let mut manager = manager.lock();
            let socket_handle = manager.iface.add_socket(socket);

            let control = Arc::new(ParkingMutex::new(TcpSocketControl {
                send_buffer: RingBuffer::new(vec![0u8; send_buffer_size as usize]),
                send_waker: None,
                recv_buffer: RingBuffer::new(vec![0u8; recv_buffer_size as usize]),
                recv_waker: None,
                is_closed: false,
            }));

            manager.sockets.insert(socket_handle.clone(), control.clone());
            (control, manager.manager_notify.clone())
        };

        TcpConnection {
            control,
            manager_notify,
        }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        // If socket is already closed, just return EOF directly.
        if control.is_closed {
            return Ok(()).into();
        }

        // Read from buffer

        if control.recv_buffer.is_empty() {
            // Nothing could be read. Wait for notify.
            if let Some(old_waker) = control.recv_waker.replace(cx.waker().clone()) {
                if !old_waker.will_wake(cx.waker()) {
                    old_waker.wake();
                }
            }

            return Poll::Pending;
        }

        let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
        let n = control.recv_buffer.dequeue_slice(recv_buf);
        buf.advance(n);

        self.manager_notify.notify_one();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut control = self.control.lock();
        if control.is_closed {
            return Err(io::ErrorKind::BrokenPipe.into()).into();
        }

        // Write to buffer

        if control.send_buffer.is_full() {
            if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
                if !old_waker.will_wake(cx.waker()) {
                    old_waker.wake();
                }
            }

            return Poll::Pending;
        }

        let n = control.send_buffer.enqueue_slice(buf);

        self.manager_notify.notify_one();
        Ok(n).into()
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Ok(()).into()
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        if control.is_closed {
            return Ok(()).into();
        }

        control.is_closed = true;
        if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
            if !old_waker.will_wake(cx.waker()) {
                old_waker.wake();
            }
        }

        Poll::Pending
    }
}

pub struct TcpTun {
    context: Arc<ServiceContext>,
    manager: SharedTcpSocketManager,
    manager_handle: JoinHandle<()>,
    manager_notify: Arc<Notify>,
    balancer: PingBalancer,
    iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
    iface_tx: mpsc::Sender<Vec<u8>>,
}

impl Drop for TcpTun {
    fn drop(&mut self) {
        self.manager_handle.abort();
    }
}

impl TcpTun {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> TcpTun {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        capabilities.max_transmission_unit = mtu as usize;

        let (virt, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);

        let iface_builder = InterfaceBuilder::new(virt, vec![]);
        let iface_ipaddrs = [
            IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0),
            IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0),
        ];
        let mut iface_routes = Routes::new(BTreeMap::new());
        iface_routes
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("IPv4 route");
        iface_routes
            .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
            .expect("IPv6 route");
        let iface = iface_builder
            .any_ip(true)
            .ip_addrs(iface_ipaddrs)
            .routes(iface_routes)
            .finalize();

        let manager_notify = Arc::new(Notify::new());
        let manager = Arc::new(ParkingMutex::new(TcpSocketManager {
            iface,
            manager_notify: manager_notify.clone(),
            sockets: HashMap::new(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let TcpSocketManager {
                            ref mut iface,
                            ref mut sockets,
                            ..
                        } = *(manager.lock());

                        let before_poll = Instant::now();
                        let updated_sockets = match iface.poll(before_poll) {
                            Ok(u) => u,
                            Err(err) => {
                                error!("VirtDevice::poll error: {}", err);
                                false
                            }
                        };

                        let after_poll = Instant::now();

                        if updated_sockets {
                            trace!("VirtDevice::poll costed {}", after_poll - before_poll);
                        }

                        // Check all the sockets' status
                        let mut sockets_to_remove = Vec::new();

                        for (socket_handle, control) in sockets.iter() {
                            let socket_handle = socket_handle.clone();
                            let socket = iface.get_socket::<TcpSocket>(socket_handle);
                            let mut control = control.lock();

                            #[inline]
                            fn close_socket_control(control: &mut TcpSocketControl) {
                                control.is_closed = true;
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            if !socket.is_open() || socket.state() == TcpState::Closed {
                                sockets_to_remove.push(socket_handle);
                                close_socket_control(&mut *control);
                                continue;
                            }

                            if control.is_closed {
                                // Close the socket.
                                socket.close();
                                // sockets_to_remove.push(socket_handle);
                                // close_socket_control(&mut *control);
                                continue;
                            }

                            // Check if readable
                            let mut has_received = false;
                            while socket.can_recv() && !control.recv_buffer.is_full() {
                                let result = socket.recv(|buffer| {
                                    let n = control.recv_buffer.enqueue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        has_received = true;
                                    }
                                    Err(err) => {
                                        error!("socket recv error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                        break;
                                    }
                                }
                            }

                            if has_received {
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            // Check if writable
                            let mut has_sent = false;
                            while socket.can_send() && !control.send_buffer.is_empty() {
                                let result = socket.send(|buffer| {
                                    let n = control.send_buffer.dequeue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        has_sent = true;
                                    }
                                    Err(err) => {
                                        error!("socket send error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                        break;
                                    }
                                }
                            }

                            if has_sent {
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                            }
                        }

                        for socket_handle in sockets_to_remove {
                            sockets.remove(&socket_handle);
                            iface.remove_socket(socket_handle);
                        }

                        let next_duration = iface.poll_delay(Instant::now()).unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::select! {
                        _ = time::sleep(StdDuration::from(next_duration)) => {}
                        _ = manager_notify.notified() => {}
                    }
                }
            })
        };

        TcpTun {
            context,
            manager,
            manager_handle,
            manager_notify,
            balancer,
            iface_rx,
            iface_tx,
        }
    }

    pub async fn handle_packet(
        &mut self,
        src_addr: SocketAddr,
        dst_addr: SocketAddr,
        tcp_packet: &TcpPacket<&[u8]>,
    ) -> io::Result<()> {
        // TCP first handshake packet, create a new Connection
        if tcp_packet.syn() && !tcp_packet.ack() {
            let accept_opts = self.context.accept_opts();

            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

            let mut socket = TcpSocket::new(
                TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
                TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
            );
            socket.set_keep_alive(accept_opts.tcp.keepalive.map(From::from));

            if let Err(err) = socket.listen(dst_addr) {
                return Err(io::Error::new(ErrorKind::Other, err));
            }

            trace!("created TCP connection for {} <-> {}", src_addr, dst_addr);

            let connection = TcpConnection::new(socket, self.manager.clone(), &accept_opts.tcp);

            // establish a tunnel
            let context = self.context.clone();
            let balancer = self.balancer.clone();
            tokio::spawn(async move {
                if let Err(err) = handle_redir_client(context, balancer, connection, src_addr, dst_addr).await {
                    error!("TCP tunnel failure, {} <-> {}, error: {}", src_addr, dst_addr, err);
                }
            });
        }

        Ok(())
    }

    pub async fn drive_interface_state(&mut self, frame: &[u8]) {
        if let Err(..) = self.iface_tx.send(frame.to_vec()).await {
            panic!("interface send channel closed unexpectly");
        }

        // Wake up and poll the interface.
        self.manager_notify.notify_one();
    }

    pub async fn recv_packet(&mut self) -> Vec<u8> {
        match self.iface_rx.recv().await {
            Some(v) => v,
            None => unreachable!("channel closed unexpectedly"),
        }
    }
}

/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir<'a>(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    mut stream: TcpConnection,
    peer_addr: SocketAddr,
    addr: &Address,
) -> io::Result<()> {
    let server = balancer.best_tcp_server();
    let svr_cfg = server.server_config();

    let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;

    establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}

async fn handle_redir_client(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    s: TcpConnection,
    peer_addr: SocketAddr,
    mut daddr: SocketAddr,
) -> io::Result<()> {
    // Get forward address from socket
    //
    // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
    if let SocketAddr::V6(ref a) = daddr {
        if let Some(v4) = to_ipv4_mapped(a.ip()) {
            daddr = SocketAddr::new(IpAddr::from(v4), a.port());
        }
    }
    let target_addr = Address::from(daddr);
    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
<commit_msg>TCP socket timeout in 7200s following Linux's setting<commit_after>use std::{
    collections::{BTreeMap, HashMap},
    io::{self, ErrorKind},
    mem,
    net::{IpAddr, SocketAddr},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
    time::Duration as StdDuration,
};

use log::{error, trace};
use parking_lot::Mutex as ParkingMutex;
use shadowsocks::{net::TcpSocketOpts, relay::socks5::Address};
use smoltcp::{
    iface::{Interface, InterfaceBuilder, Routes, SocketHandle},
    phy::{DeviceCapabilities, Medium},
    socket::{TcpSocket, TcpSocketBuffer, TcpState},
    storage::RingBuffer,
    time::{Duration, Instant},
    wire::{IpAddress, IpCidr, Ipv4Address, Ipv6Address, TcpPacket},
};
use tokio::{
    io::{AsyncRead, AsyncWrite, ReadBuf},
    sync::{mpsc, Notify},
    task::JoinHandle,
    time,
};

use crate::local::{
    context::ServiceContext,
    loadbalancing::PingBalancer,
    net::AutoProxyClientStream,
    utils::{establish_tcp_tunnel, to_ipv4_mapped},
};

use super::virt_device::VirtTunDevice;

// NOTE: Default value is taken from Linux
// recv: /proc/sys/net/ipv4/tcp_rmem 87380 bytes
// send: /proc/sys/net/ipv4/tcp_wmem 16384 bytes
const DEFAULT_TCP_SEND_BUFFER_SIZE: u32 = 16384;
const DEFAULT_TCP_RECV_BUFFER_SIZE: u32 = 87380;

struct TcpSocketControl {
    send_buffer: RingBuffer<'static, u8>,
    send_waker: Option<Waker>,
    recv_buffer: RingBuffer<'static, u8>,
    recv_waker: Option<Waker>,
    is_closed: bool,
}

struct TcpSocketManager {
    iface: Interface<'static, VirtTunDevice>,
    manager_notify: Arc<Notify>,
    sockets: HashMap<SocketHandle, Arc<ParkingMutex<TcpSocketControl>>>,
}

type SharedTcpSocketManager = Arc<ParkingMutex<TcpSocketManager>>;

struct TcpConnection {
    control: Arc<ParkingMutex<TcpSocketControl>>,
    manager_notify: Arc<Notify>,
}

impl Drop for TcpConnection {
    fn drop(&mut self) {
        let mut control = self.control.lock();
        control.is_closed = true;
    }
}

impl TcpConnection {
    fn new(socket: TcpSocket<'static>, manager: SharedTcpSocketManager, tcp_opts: &TcpSocketOpts) -> TcpConnection {
        let send_buffer_size = tcp_opts.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
        let recv_buffer_size = tcp_opts.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

        let (control, manager_notify) = {
            let mut manager = manager.lock();
            let socket_handle = manager.iface.add_socket(socket);

            let control = Arc::new(ParkingMutex::new(TcpSocketControl {
                send_buffer: RingBuffer::new(vec![0u8; send_buffer_size as usize]),
                send_waker: None,
                recv_buffer: RingBuffer::new(vec![0u8; recv_buffer_size as usize]),
                recv_waker: None,
                is_closed: false,
            }));

            manager.sockets.insert(socket_handle.clone(), control.clone());
            (control, manager.manager_notify.clone())
        };

        TcpConnection {
            control,
            manager_notify,
        }
    }
}

impl AsyncRead for TcpConnection {
    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        // If socket is already closed, just return EOF directly.
        if control.is_closed {
            return Ok(()).into();
        }

        // Read from buffer

        if control.recv_buffer.is_empty() {
            // Nothing could be read. Wait for notify.
            if let Some(old_waker) = control.recv_waker.replace(cx.waker().clone()) {
                if !old_waker.will_wake(cx.waker()) {
                    old_waker.wake();
                }
            }

            return Poll::Pending;
        }

        let recv_buf = unsafe { mem::transmute::<_, &mut [u8]>(buf.unfilled_mut()) };
        let n = control.recv_buffer.dequeue_slice(recv_buf);
        buf.advance(n);

        self.manager_notify.notify_one();
        Ok(()).into()
    }
}

impl AsyncWrite for TcpConnection {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
        let mut control = self.control.lock();
        if control.is_closed {
            return Err(io::ErrorKind::BrokenPipe.into()).into();
        }

        // Write to buffer

        if control.send_buffer.is_full() {
            if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
                if !old_waker.will_wake(cx.waker()) {
                    old_waker.wake();
                }
            }

            return Poll::Pending;
        }

        let n = control.send_buffer.enqueue_slice(buf);

        self.manager_notify.notify_one();
        Ok(n).into()
    }

    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Ok(()).into()
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        let mut control = self.control.lock();

        if control.is_closed {
            return Ok(()).into();
        }

        control.is_closed = true;
        if let Some(old_waker) = control.send_waker.replace(cx.waker().clone()) {
            if !old_waker.will_wake(cx.waker()) {
                old_waker.wake();
            }
        }

        Poll::Pending
    }
}

pub struct TcpTun {
    context: Arc<ServiceContext>,
    manager: SharedTcpSocketManager,
    manager_handle: JoinHandle<()>,
    manager_notify: Arc<Notify>,
    balancer: PingBalancer,
    iface_rx: mpsc::UnboundedReceiver<Vec<u8>>,
    iface_tx: mpsc::Sender<Vec<u8>>,
}

impl Drop for TcpTun {
    fn drop(&mut self) {
        self.manager_handle.abort();
    }
}

impl TcpTun {
    pub fn new(context: Arc<ServiceContext>, balancer: PingBalancer, mtu: u32) -> TcpTun {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        capabilities.max_transmission_unit = mtu as usize;

        let (virt, iface_rx, iface_tx) = VirtTunDevice::new(capabilities);

        let iface_builder = InterfaceBuilder::new(virt, vec![]);
        let iface_ipaddrs = [
            IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0),
            IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0),
        ];
        let mut iface_routes = Routes::new(BTreeMap::new());
        iface_routes
            .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
            .expect("IPv4 route");
        iface_routes
            .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
            .expect("IPv6 route");
        let iface = iface_builder
            .any_ip(true)
            .ip_addrs(iface_ipaddrs)
            .routes(iface_routes)
            .finalize();

        let manager_notify = Arc::new(Notify::new());
        let manager = Arc::new(ParkingMutex::new(TcpSocketManager {
            iface,
            manager_notify: manager_notify.clone(),
            sockets: HashMap::new(),
        }));

        let manager_handle = {
            let manager = manager.clone();
            let manager_notify = manager_notify.clone();
            tokio::spawn(async move {
                loop {
                    let next_duration = {
                        let TcpSocketManager {
                            ref mut iface,
                            ref mut sockets,
                            ..
                        } = *(manager.lock());

                        let before_poll = Instant::now();
                        let updated_sockets = match iface.poll(before_poll) {
                            Ok(u) => u,
                            Err(err) => {
                                error!("VirtDevice::poll error: {}", err);
                                false
                            }
                        };

                        let after_poll = Instant::now();

                        if updated_sockets {
                            trace!("VirtDevice::poll costed {}", after_poll - before_poll);
                        }

                        // Check all the sockets' status
                        let mut sockets_to_remove = Vec::new();

                        for (socket_handle, control) in sockets.iter() {
                            let socket_handle = socket_handle.clone();
                            let socket = iface.get_socket::<TcpSocket>(socket_handle);
                            let mut control = control.lock();

                            #[inline]
                            fn close_socket_control(control: &mut TcpSocketControl) {
                                control.is_closed = true;
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            if !socket.is_open() || socket.state() == TcpState::Closed {
                                sockets_to_remove.push(socket_handle);
                                close_socket_control(&mut *control);
                                continue;
                            }

                            if control.is_closed {
                                // Close the socket.
                                socket.close();
                                // sockets_to_remove.push(socket_handle);
                                // close_socket_control(&mut *control);
                                continue;
                            }

                            // Check if readable
                            let mut has_received = false;
                            while socket.can_recv() && !control.recv_buffer.is_full() {
                                let result = socket.recv(|buffer| {
                                    let n = control.recv_buffer.enqueue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        has_received = true;
                                    }
                                    Err(err) => {
                                        error!("socket recv error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                        break;
                                    }
                                }
                            }

                            if has_received {
                                if let Some(waker) = control.recv_waker.take() {
                                    waker.wake();
                                }
                            }

                            // Check if writable
                            let mut has_sent = false;
                            while socket.can_send() && !control.send_buffer.is_empty() {
                                let result = socket.send(|buffer| {
                                    let n = control.send_buffer.dequeue_slice(buffer);
                                    (n, ())
                                });

                                match result {
                                    Ok(..) => {
                                        has_sent = true;
                                    }
                                    Err(err) => {
                                        error!("socket send error: {}", err);
                                        sockets_to_remove.push(socket_handle);
                                        close_socket_control(&mut *control);
                                        break;
                                    }
                                }
                            }

                            if has_sent {
                                if let Some(waker) = control.send_waker.take() {
                                    waker.wake();
                                }
                            }
                        }

                        for socket_handle in sockets_to_remove {
                            sockets.remove(&socket_handle);
                            iface.remove_socket(socket_handle);
                        }

                        let next_duration = iface.poll_delay(Instant::now()).unwrap_or(Duration::from_millis(50));

                        next_duration
                    };

                    tokio::select! {
                        _ = time::sleep(StdDuration::from(next_duration)) => {}
                        _ = manager_notify.notified() => {}
                    }
                }
            })
        };

        TcpTun {
            context,
            manager,
            manager_handle,
            manager_notify,
            balancer,
            iface_rx,
            iface_tx,
        }
    }

    pub async fn handle_packet(
        &mut self,
        src_addr: SocketAddr,
        dst_addr: SocketAddr,
        tcp_packet: &TcpPacket<&[u8]>,
    ) -> io::Result<()> {
        // TCP first handshake packet, create a new Connection
        if tcp_packet.syn() && !tcp_packet.ack() {
            let accept_opts = self.context.accept_opts();

            let send_buffer_size = accept_opts.tcp.send_buffer_size.unwrap_or(DEFAULT_TCP_SEND_BUFFER_SIZE);
            let recv_buffer_size = accept_opts.tcp.recv_buffer_size.unwrap_or(DEFAULT_TCP_RECV_BUFFER_SIZE);

            let mut socket = TcpSocket::new(
                TcpSocketBuffer::new(vec![0u8; recv_buffer_size as usize]),
                TcpSocketBuffer::new(vec![0u8; send_buffer_size as usize]),
            );
            socket.set_keep_alive(accept_opts.tcp.keepalive.map(From::from));
            // FIXME: This should follows system's setting. 7200 is Linux's default.
            socket.set_timeout(Some(Duration::from_secs(7200)));

            if let Err(err) = socket.listen(dst_addr) {
                return Err(io::Error::new(ErrorKind::Other, err));
            }

            trace!("created TCP connection for {} <-> {}", src_addr, dst_addr);

            let connection = TcpConnection::new(socket, self.manager.clone(), &accept_opts.tcp);

            // establish a tunnel
            let context = self.context.clone();
            let balancer = self.balancer.clone();
            tokio::spawn(async move {
                if let Err(err) = handle_redir_client(context, balancer, connection, src_addr, dst_addr).await {
                    error!("TCP tunnel failure, {} <-> {}, error: {}", src_addr, dst_addr, err);
                }
            });
        }

        Ok(())
    }

    pub async fn drive_interface_state(&mut self, frame: &[u8]) {
        if let Err(..) = self.iface_tx.send(frame.to_vec()).await {
            panic!("interface send channel closed unexpectly");
        }

        // Wake up and poll the interface.
        self.manager_notify.notify_one();
    }

    pub async fn recv_packet(&mut self) -> Vec<u8> {
        match self.iface_rx.recv().await {
            Some(v) => v,
            None => unreachable!("channel closed unexpectedly"),
        }
    }
}

/// Established Client Transparent Proxy
///
/// This method must be called after handshaking with client (for example, socks5 handshaking)
async fn establish_client_tcp_redir<'a>(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    mut stream: TcpConnection,
    peer_addr: SocketAddr,
    addr: &Address,
) -> io::Result<()> {
    let server = balancer.best_tcp_server();
    let svr_cfg = server.server_config();

    let mut remote = AutoProxyClientStream::connect(context, &server, addr).await?;

    establish_tcp_tunnel(svr_cfg, &mut stream, &mut remote, peer_addr, addr).await
}

async fn handle_redir_client(
    context: Arc<ServiceContext>,
    balancer: PingBalancer,
    s: TcpConnection,
    peer_addr: SocketAddr,
    mut daddr: SocketAddr,
) -> io::Result<()> {
    // Get forward address from socket
    //
    // Try to convert IPv4 mapped IPv6 address for dual-stack mode.
    if let SocketAddr::V6(ref a) = daddr {
        if let Some(v4) = to_ipv4_mapped(a.ip()) {
            daddr = SocketAddr::new(IpAddr::from(v4), a.port());
        }
    }
    let target_addr = Address::from(daddr);
    establish_client_tcp_redir(context, balancer, s, peer_addr, &target_addr).await
}
<|endoftext|>