<commit_before>
use std::ffi;
use std::ffi::AsOsStr;
use std::convert::AsRef;
use std::path::Path;

pub mod record;
pub mod header;

use htslib;
use bcf::header::{HeaderView, SampleSubset};

pub use bcf::header::Header;
pub use bcf::record::Record;

pub struct Reader {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
}

impl Reader {
   pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let htsfile = bcf_open(path, b"r");
        let header = unsafe { htslib::vcf::bcf_hdr_read(htsfile) };
        Reader { inner: htsfile, header: HeaderView::new(header) }
    }

    pub fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::vcf::bcf_read(self.inner, self.header.inner, record.inner) } {
            0  => {
                record.header = self.header.inner;
                Ok(())
            },
            -1 => Err(ReadError::NoMoreRecord),
            _  => Err(ReadError::Invalid),
        }
    }

    pub fn records(&self) -> Records {
        Records { reader: self }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Writer {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
    subset: Option<SampleSubset>,
}

impl Writer {
    pub fn new<P: AsRef<Path>>(path: &P, header: &Header, uncompressed: bool, vcf: bool) -> Self {
        let mode: &[u8] = match (uncompressed, vcf) {
            (true, true)   => b"w",
            (false, true)  => b"wz",
            (true, false)  => b"wu",
            (false, false) => b"wb",
        };

        let htsfile = bcf_open(path, mode);
        unsafe { htslib::vcf::bcf_hdr_write(htsfile, header.inner) };
        Writer {
            inner: htsfile,
            header: HeaderView::new(unsafe { htslib::vcf::bcf_hdr_dup(header.inner) }),
            subset: header.subset.clone()
        }
    }

    /// Translate record to header of this writer.
    pub fn translate(&mut self, record: &mut record::Record) {
        match self.subset {
            Some(ref mut subset) => unsafe {
                htslib::vcf::bcf_subset(self.header.inner, record.inner, subset.len() as i32, subset.as_mut_ptr());
            },
            None         => ()
        }
        unsafe {
            htslib::vcf::bcf_translate(self.header.inner, record.header, record.inner);
        }
        record.header = self.header.inner;
    }

    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::vcf::bcf_write(self.inner, self.header.inner, record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for Writer {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Records<'a> {
    reader: &'a Reader,
}

impl<'a> Iterator for Records<'a> {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.reader.read(&mut record) {
            Err(ReadError::NoMoreRecord) => None,
            Err(e)                       => Some(Err(e)),
            Ok(())                       => Some(Ok(record)),
        }
    }
}

/// Wrapper for opening a BCF file.
fn bcf_open<P: AsRef<Path>>(path: &P, mode: &[u8]) -> *mut htslib::vcf::htsFile {
    unsafe {
        htslib::vcf::hts_open(
            path.as_ref().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

pub enum ReadError {
    Invalid,
    NoMoreRecord,
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use std::path::Path;

    fn _test_read<P: AsRef<Path>>(path: &P) {
        let bcf = Reader::new(path);
        assert_eq!(bcf.header.samples(), [b"NA12878.subsample-0.25-0"]);

        for (i, rec) in bcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.sample_count(), 1);

            assert_eq!(record.rid().expect("Error reading rid."), 0);
            assert_eq!(record.pos(), 10021 + i as u32);
            assert_eq!(record.qual(), 0f32);
            assert_eq!(record.info(b"MQ0F").float().ok().expect("Error reading info."), [1.0]);
            if i == 59 {
                assert_eq!(record.info(b"SGB").float().ok().expect("Error reading info."), [-0.379885]);
            }
            // the artificial "not observed" allele is present in each record.
            assert_eq!(record.alleles().iter().last().unwrap(), b"<X>");

            let mut fmt = record.format(b"PL");
            let pl = fmt.integer().ok().expect("Error reading format.");
            assert_eq!(pl.len(), 1);
            if i == 59 {
                assert_eq!(pl[0].len(), 6);
            }
            else {
                assert_eq!(pl[0].len(), 3);
            }
        }
    }

    #[test]
    fn test_read() {
        _test_read(&"test.bcf");
    }

    #[test]
    fn test_write() {
        let bcf = Reader::new(&"test.bcf");
        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bcfpath = tmp.path().join("test.bcf");
        println!("{:?}", bcfpath);
        {
            let header = Header::with_template(&bcf.header);
            let mut writer = Writer::new(&bcfpath, &header, false, false);
            for rec in bcf.records() {
                let mut record = rec.ok().expect("Error reading record.");
                writer.translate(&mut record);
                record.trim_alleles().ok().expect("Error trimming alleles.");
                writer.write(&record).ok().expect("Error writing record");
            }
        }
        {
            _test_read(&bcfpath);
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<commit_msg>Separate translate and subset.<commit_after>
use std::ffi;
use std::ffi::AsOsStr;
use std::convert::AsRef;
use std::path::Path;

pub mod record;
pub mod header;

use htslib;
use bcf::header::{HeaderView, SampleSubset};

pub use bcf::header::Header;
pub use bcf::record::Record;

pub struct Reader {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
}

impl Reader {
   pub fn new<P: AsRef<Path>>(path: &P) -> Self {
        let htsfile = bcf_open(path, b"r");
        let header = unsafe { htslib::vcf::bcf_hdr_read(htsfile) };
        Reader { inner: htsfile, header: HeaderView::new(header) }
    }

    pub fn read(&self, record: &mut record::Record) -> Result<(), ReadError> {
        match unsafe { htslib::vcf::bcf_read(self.inner, self.header.inner, record.inner) } {
            0  => {
                record.header = self.header.inner;
                Ok(())
            },
            -1 => Err(ReadError::NoMoreRecord),
            _  => Err(ReadError::Invalid),
        }
    }

    pub fn records(&self) -> Records {
        Records { reader: self }
    }
}

impl Drop for Reader {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Writer {
    inner: *mut htslib::vcf::htsFile,
    pub header: HeaderView,
    subset: Option<SampleSubset>,
}

impl Writer {
    pub fn new<P: AsRef<Path>>(path: &P, header: &Header, uncompressed: bool, vcf: bool) -> Self {
        let mode: &[u8] = match (uncompressed, vcf) {
            (true, true)   => b"w",
            (false, true)  => b"wz",
            (true, false)  => b"wu",
            (false, false) => b"wb",
        };

        let htsfile = bcf_open(path, mode);
        unsafe { htslib::vcf::bcf_hdr_write(htsfile, header.inner) };
        Writer {
            inner: htsfile,
            header: HeaderView::new(unsafe { htslib::vcf::bcf_hdr_dup(header.inner) }),
            subset: header.subset.clone()
        }
    }

    /// Translate record to header of this writer.
    pub fn translate(&mut self, record: &mut record::Record) {
        unsafe {
            htslib::vcf::bcf_translate(self.header.inner, record.header, record.inner);
        }
        record.header = self.header.inner;
    }

    /// Subset samples of record to match header of this writer.
    pub fn subset(&mut self, record: &mut record::Record) {
        match self.subset {
            Some(ref mut subset) => unsafe {
                htslib::vcf::bcf_subset(self.header.inner, record.inner, subset.len() as i32, subset.as_mut_ptr());
            },
            None         => ()
        }
    }

    pub fn write(&mut self, record: &record::Record) -> Result<(), ()> {
        if unsafe { htslib::vcf::bcf_write(self.inner, self.header.inner, record.inner) } == -1 {
            Err(())
        }
        else {
            Ok(())
        }
    }
}

impl Drop for Writer {
    fn drop(&mut self) {
        unsafe {
            htslib::vcf::bcf_hdr_destroy(self.header.inner);
            htslib::vcf::hts_close(self.inner);
        }
    }
}

pub struct Records<'a> {
    reader: &'a Reader,
}

impl<'a> Iterator for Records<'a> {
    type Item = Result<record::Record, ReadError>;

    fn next(&mut self) -> Option<Result<record::Record, ReadError>> {
        let mut record = record::Record::new();
        match self.reader.read(&mut record) {
            Err(ReadError::NoMoreRecord) => None,
            Err(e)                       => Some(Err(e)),
            Ok(())                       => Some(Ok(record)),
        }
    }
}

/// Wrapper for opening a BCF file.
fn bcf_open<P: AsRef<Path>>(path: &P, mode: &[u8]) -> *mut htslib::vcf::htsFile {
    unsafe {
        htslib::vcf::hts_open(
            path.as_ref().as_os_str().to_cstring().unwrap().as_ptr(),
            ffi::CString::new(mode).unwrap().as_ptr()
        )
    }
}

pub enum ReadError {
    Invalid,
    NoMoreRecord,
}

#[cfg(test)]
mod tests {
    extern crate tempdir;
    use super::*;
    use std::path::Path;

    fn _test_read<P: AsRef<Path>>(path: &P) {
        let bcf = Reader::new(path);
        assert_eq!(bcf.header.samples(), [b"NA12878.subsample-0.25-0"]);

        for (i, rec) in bcf.records().enumerate() {
            let mut record = rec.ok().expect("Error reading record.");
            assert_eq!(record.sample_count(), 1);

            assert_eq!(record.rid().expect("Error reading rid."), 0);
            assert_eq!(record.pos(), 10021 + i as u32);
            assert_eq!(record.qual(), 0f32);
            assert_eq!(record.info(b"MQ0F").float().ok().expect("Error reading info."), [1.0]);
            if i == 59 {
                assert_eq!(record.info(b"SGB").float().ok().expect("Error reading info."), [-0.379885]);
            }
            // the artificial "not observed" allele is present in each record.
            assert_eq!(record.alleles().iter().last().unwrap(), b"<X>");

            let mut fmt = record.format(b"PL");
            let pl = fmt.integer().ok().expect("Error reading format.");
            assert_eq!(pl.len(), 1);
            if i == 59 {
                assert_eq!(pl[0].len(), 6);
            }
            else {
                assert_eq!(pl[0].len(), 3);
            }
        }
    }

    #[test]
    fn test_read() {
        _test_read(&"test.bcf");
    }

    #[test]
    fn test_write() {
        let bcf = Reader::new(&"test.bcf");
        let tmp = tempdir::TempDir::new("rust-htslib").ok().expect("Cannot create temp dir");
        let bcfpath = tmp.path().join("test.bcf");
        println!("{:?}", bcfpath);
        {
            let header = Header::with_template(&bcf.header);
            let mut writer = Writer::new(&bcfpath, &header, false, false);
            for rec in bcf.records() {
                let mut record = rec.ok().expect("Error reading record.");
                writer.translate(&mut record);
                writer.subset(&mut record);
                record.trim_alleles().ok().expect("Error trimming alleles.");
                writer.write(&record).ok().expect("Error writing record");
            }
        }
        {
            _test_read(&bcfpath);
        }
        tmp.close().ok().expect("Failed to delete temp dir");
    }
}
<|endoftext|>


<commit_before>#![feature(slicing_syntax)]

extern crate getopts;

extern crate coreutils;

use std::os;

fn main() {

    let args = os::args();
    let program = args[0].clone();

    let opts = [
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => {
            println!("{}", f.to_string());
            println!("-------------------------------");
            println!("Usage: {}", getopts::usage(program.as_slice(), opts));
            fail!()
        },
    };

    if matches.opt_present("help") {
        println!("Usage: {}", getopts::usage(program.as_slice(), opts));
        return;
    }

    if matches.opt_present("version") {
        println!("{}: {}", program, coreutils::version());
        return;
    }

    let output = match args.len() {
        1 => "y".to_string(),
        _ => args[1..].connect(" "),
    };

    loop {
        println!("{}", output);
    }
}
<commit_msg>Handle broken pipes and other cleanup in yes<commit_after>extern crate getopts;

extern crate coreutils;

use std::os;
use std::io;

fn main() {

    let args = os::args();
    let program = args[0].clone();

    let opts = [
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m },
        Err(f) => {
            println!("{}", f.to_string());
            println!("-------------------------------");
            println!("Usage: {}", getopts::usage(program.as_slice(), opts));
            return;
        },
    };

    if matches.opt_present("help") {
        println!("Usage: {}", getopts::usage(program.as_slice(), opts));
        return;
    }

    if matches.opt_present("version") {
        println!("{}: {}", program, coreutils::version());
        return;
    }

    let output = match args.len() {
        1 => "y".to_string(),
        _ => args.tail().connect(" "),
    };

    let mut stdout = io::stdio::stdout();

    loop {
        match writeln!(stdout, "{}", output).err() {
            Some(io::IoError { kind: io::BrokenPipe, .. }) => break,
            _ => (),
        }
    }
}
<|endoftext|>


<commit_before>use diesel;
use diesel::prelude::*;

use data::types::*;
use helpers::db::DB;

use models::resource::{Resource, UpdatedResource};

use schema::resources::dsl as resources_dsl;

/*
Generic function which accepts an Enum as type identifier.
*/
pub fn check_resources(costs: &Option<Vec<(ResourceTypes, i64)>>,
                       resources: Vec<Resource>,
                       db: &DB)
                       -> bool {
    match costs.as_ref() {
        // There are no costs for this module/research
        None => return true,
        // There are some costs
        Some(costs) => {
            for &(ref resource_type, amount) in costs {
                // Try to get the correct entry from existing resources.
                let existing = resources
                    .iter()
                    .filter(|x| x.name == resource_type.to_string())
                    .next();
                match existing {
                    // There is no resource for this resource_type,
                    // thereby it's not enough.
                    None => return false,
                    // There is a resource for this resource type
                    // We need to check if we got enough of it.
                    Some(existing_resource) => {
                        if existing_resource.amount < amount {
                            return false;
                        }
                    }
                }
            }
            update_resources(resources, costs, true, db);
            return true;
        }
    }
}

pub fn update_resources(resources: Vec<Resource>,
                        costs: &Vec<(ResourceTypes, i64)>,
                        subtract: bool,
                        db: &DB) {

    for &(ref resource_type, amount) in costs.iter() {
        // Try to get the correct entry from existing resources.
        let resource = resources
            .iter()
            .filter(|x| x.name == resource_type.to_string())
            .next();
        match resource {
            // There is no resource for this resource_type,
            // thereby it's not enough.
            // There is a resource for this resource type
            Some(resource) => {
                update_resource(resource, amount, subtract, db);
            }
            None => (),
        }
    }
}

pub fn update_resource(resource: &Resource, amount: i64, subtract: bool, db: &DB) {
    let mut new_amount: i64;
    if subtract {
        new_amount = resource.amount - amount;
    } else {
        new_amount = resource.amount + amount;
        if new_amount > resource.max_amount {
            new_amount = resource.max_amount;
        }
    }
    let updated_resource = UpdatedResource {
        amount: Some(new_amount),
        max_amount: None,
    };

    diesel::update(resources_dsl::resources.filter(resources_dsl::name.eq(resource.name.clone())))
        .set(&updated_resource)
        .get_result::<Resource>(&**db)
        .expect("Failed to update resource.");
}
<commit_msg>More documentation, fix bug in update_resource function<commit_after>use diesel;
use diesel::prelude::*;

use data::types::*;
use helpers::db::DB;

use models::resource::{Resource, UpdatedResource};

use schema::resources::dsl as resources_dsl;

/// This function checks if there are enough resources for a given set
/// of costs from a research or model entry.
/// It also subtracts and updates the resources the specified value from the database.
///
/// The first parameter is a vector of resources which represent the costs.
/// The second parameter is a vector of all `Resource` database models from a pod or a base.
pub fn check_resources(costs: &Option<Vec<(ResourceTypes, i64)>>,
                       resources: Vec<Resource>,
                       db: &DB)
                       -> bool {
    match costs.as_ref() {
        // There are no costs for this module/research
        None => return true,
        // There are some costs
        Some(costs) => {
            for &(ref resource_type, amount) in costs {
                // Try to get the correct entry from existing resources.
                let existing = resources
                    .iter()
                    .filter(|x| x.name == resource_type.to_string())
                    .next();
                match existing {
                    // There is no resource for this resource_type,
                    // thereby it's not enough.
                    None => return false,
                    // There is a resource for this resource type
                    // We need to check if we got enough of it.
                    Some(existing_resource) => {
                        if existing_resource.amount < amount {
                            return false;
                        }
                    }
                }
            }
            update_resources(resources, costs, true, db);
            return true;
        }
    }
}

pub fn update_resources(resources: Vec<Resource>,
                        costs: &Vec<(ResourceTypes, i64)>,
                        subtract: bool,
                        db: &DB) {

    for &(ref resource_type, amount) in costs.iter() {
        // Try to get the correct entry from existing resources.
        let resource = resources
            .iter()
            .filter(|x| x.name == resource_type.to_string())
            .next();
        match resource {
            // There is no resource for this resource_type,
            // thereby it's not enough.
            // There is a resource for this resource type
            Some(resource) => {
                update_resource(resource, amount, subtract, db);
            }
            None => (),
        }
    }
}

pub fn update_resource(resource: &Resource, amount: i64, subtract: bool, db: &DB) {
    let mut new_amount: i64;
    if subtract {
        new_amount = resource.amount - amount;
    } else {
        new_amount = resource.amount + amount;
        if new_amount > resource.max_amount {
            new_amount = resource.max_amount;
        }
    }

    let updated_resource = UpdatedResource {
        amount: Some(new_amount),
        max_amount: None,
    };

    diesel::update(resources_dsl::resources.filter(resources_dsl::id.eq(resource.id)))
        .set(&updated_resource)
        .get_result::<Resource>(&**db)
        .expect("Failed to update resource.");
}
<|endoftext|>


<commit_before>use constants::*;
use std::ffi::OsString;
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};

pub fn initialize() {
    println!("Initialize ...");

    let dir_name    = ".kic";
    let path_to_dir = Path::new(dir_name);
    if path_to_dir.exists() && path_to_dir.is_dir() {
        println!("  OK: \"{}\" directory has already exist.", dir_name);
    } else {
        match fs::create_dir(path_to_dir) {
            Ok(_)    => println!("  OK: Created \"{}\" directory.", dir_name),
            Err(why) => panic!("{:?}", why),
        }
    }

    let config_file    = "config";
    let path_to_config = path_to_dir.clone().join(config_file);
    create_setting_file(path_to_config, DEFAULT_CONFIG);

    let ignore_file    = "ignore";
    let path_to_ignore = path_to_dir.clone().join(ignore_file);
    let ignore_contents = walk_dir(".").iter()
        .fold(String::new(), |c, f| c + &f + "\n");
    create_setting_file(path_to_ignore, ignore_contents);
}

fn create_setting_file<S: AsRef<str>>(path_to_file: PathBuf, contents: S) {
    let file_name = path_to_file.file_name().and_then(|f| f.to_str()).unwrap();
    if path_to_file.exists() && path_to_file.is_file() {
        println!("  OK: \"{}\" file has already exist.", file_name);
    } else {
        match File::create(&path_to_file) {
            Ok(mut fp) => {
                println!("  OK: Created \"{}\" file.", file_name);
                match fp.write(contents.as_ref().as_bytes()) {
                    Ok(_)    => {},
                    Err(why) => panic!("{:?}", why),
                }
            },
            Err(why) => panic!("{:?}", why),
        }
    }
}

fn walk_dir<P: AsRef<Path>>(path: P) -> Vec<String> {
    fn walk_dir<P: AsRef<Path>>(path: P) -> Vec<OsString> {
        let dirs = match fs::read_dir(path) {
            Ok(ds)   => ds,
            Err(why) => panic!("{:?}", why),
        };
        let dirs = dirs.filter(|d| d.is_ok())
            .flat_map(|d| {
                let path = d.unwrap().path();
                if is_hidden(&path) {
                    Vec::new()
                } else if path.is_file() {
                    vec![path.into_os_string()]
                } else {
                    walk_dir(&path)
                }
            })
            .collect::<Vec<OsString>>();
        dirs
    }

    walk_dir(path).iter()
        .map(|f| f.clone().into_string())
        .filter(|f| f.is_ok())
        .map(|f| f.unwrap())
        .collect::<Vec<String>>()
}

fn is_hidden(path: &PathBuf) -> bool {
    let file_name = path.file_name().and_then(|f| f.to_str());
    file_name.map_or(false, |f| f.starts_with("."))
}

pub fn set_params() {
    println!("set parameters");
}

pub fn sweep() {
    println!("sweep");
}

pub fn burn() {
    println!("burn");
}

pub fn register_with_cron() {
    println!("register");
}

pub fn unregister_cron() {
    println!("unregister");
}

pub fn destroy() {
    println!("destroy");
}
<commit_msg>Refactor walk_dir()<commit_after>use constants::*;
use std::ffi::OsString;
use std::fs;
use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};

pub fn initialize() {
    println!("Initialize ...");

    let dir_name    = ".kic";
    let path_to_dir = Path::new(dir_name);
    if path_to_dir.exists() && path_to_dir.is_dir() {
        println!("  OK: \"{}\" directory has already exist.", dir_name);
    } else {
        match fs::create_dir(path_to_dir) {
            Ok(_)    => println!("  OK: Created \"{}\" directory.", dir_name),
            Err(why) => panic!("{:?}", why),
        }
    }

    let config_file    = "config";
    let path_to_config = path_to_dir.clone().join(config_file);
    create_setting_file(path_to_config, DEFAULT_CONFIG);

    let ignore_file    = "ignore";
    let path_to_ignore = path_to_dir.clone().join(ignore_file);
    let ignore_contents = walk_dir(".").iter()
        .fold(String::new(), |c, f| c + &f + "\n");
    create_setting_file(path_to_ignore, ignore_contents);
}

fn create_setting_file<S: AsRef<str>>(path_to_file: PathBuf, contents: S) {
    let file_name = path_to_file.file_name().and_then(|f| f.to_str()).unwrap();
    if path_to_file.exists() && path_to_file.is_file() {
        println!("  OK: \"{}\" file has already exist.", file_name);
    } else {
        match File::create(&path_to_file) {
            Ok(mut fp) => {
                println!("  OK: Created \"{}\" file.", file_name);
                match fp.write(contents.as_ref().as_bytes()) {
                    Ok(_)    => {},
                    Err(why) => panic!("{:?}", why),
                }
            },
            Err(why) => panic!("{:?}", why),
        }
    }
}

fn walk_dir<P: AsRef<Path>>(path: P) -> Vec<String> {
    fn walk_dir(path: PathBuf) -> Vec<OsString> {
        let paths = match fs::read_dir(path) {
            Ok(rd)   => rd,
            Err(why) => panic!("{:?}", why),
        };
        paths
            .filter(|de| de.is_ok())
            .flat_map(|de| {
                let path = de.unwrap().path();
                if is_hidden(&path) {
                    Vec::new()
                } else if path.is_file() {
                    vec![path.into_os_string()]
                } else {
                    walk_dir(path)
                }
            })
            .collect::<Vec<OsString>>()
    }

    walk_dir(path.as_ref().to_path_buf()).iter()
        .map(|f| f.clone().into_string())
        .filter(|f| f.is_ok())
        .map(|f| f.unwrap())
        .collect::<Vec<String>>()
}

fn is_hidden(path: &PathBuf) -> bool {
    let file_name = path.file_name().and_then(|f| f.to_str());
    file_name.map_or(false, |f| f.starts_with("."))
}

pub fn set_params() {
    println!("set parameters");
}

pub fn sweep() {
    println!("sweep");
}

pub fn burn() {
    println!("burn");
}

pub fn register_with_cron() {
    println!("register");
}

pub fn unregister_cron() {
    println!("unregister");
}

pub fn destroy() {
    println!("destroy");
}
<|endoftext|>


<commit_before>use ::JobSpec;
use std::os::unix::process::CommandExt;
use std::process::Command;
use std;
use ::nom::IResult;

fn token_char(ch: char) -> bool {
    match ch {
        ch if ch as u8 <= 32 || 127 <= ch as u8 => false,
        '"' | '\'' | '>' | '<' | '|' | ';' | '{' | '}' | ' ' => false,
        _ => true,
    }
}

named!(bare_token<&str, &str>, take_while1_s!(token_char));
named!(quoted_token<&str, &str>, delimited!(tag_s!("\""),
                                            take_until_s!("\""),
                                            tag_s!("\"")));
named!(place_holder<&str, &str>, tag_s!("{}"));
named!(command_token<&str, &str>,
       alt!(bare_token | quoted_token | place_holder));

named!(command< &str, Vec<&str> >,
       terminated!(ws!(many1!(command_token)), eof!()));

pub fn parse_cmd(text: &str) -> Vec<&str> {
    match command(text) {
        IResult::Done(_, result) => result,
        IResult::Error(error) => panic!("Error {:?}", error),
        IResult::Incomplete(needed) => panic!("Needed {:?}", needed)
    }
}

/// Single Command
pub struct ShellCommand {
    command: Command,
    setpgid: bool
}

impl ShellCommand {
    pub fn new(format: &str, args: &[&str]) -> ShellCommand {
        let vec = parse_cmd(format);
        let mut command = Command::new(vec[0]);
        command.args(&vec.as_slice()[1..]);
        ShellCommand {
            command: command,
            setpgid: false
        }
    }
}

impl JobSpec for ShellCommand {
    fn exec(mut self) -> ! {
        self.command.exec();
        std::process::exit(1);
    }

    fn setpgid(mut self) -> Self {
        self.setpgid = true;
        self
    }

    fn getpgid(&self) -> bool {
        return self.setpgid;
    }
}

#[macro_export]
macro_rules! cmd {
    ($format:expr) => ($crate::command::ShellCommand::new($format, &[]));
    ($format:expr, $($arg:expr),+) =>
        ($crate::command::ShellCommand::new($format, &[$($arg),+]));
}

#[test]
fn test_parse_cmd() {
    let tokens = parse_cmd(r#"cmd 1 2
                              3 "
  4""#);
    assert_eq!("cmd", tokens[0]);
    assert_eq!("1", tokens[1]);
    assert_eq!("2", tokens[2]);
    assert_eq!("3", tokens[3]);
    assert_eq!("\n  4", tokens[4]);
}
<commit_msg>Remove unneeded condition<commit_after>use ::JobSpec;
use std::os::unix::process::CommandExt;
use std::process::Command;
use std;
use ::nom::IResult;

fn token_char(ch: char) -> bool {
    match ch {
        ch if ch as u8 <= 32 || 127 <= ch as u8 => false,
        '"' | '\'' | '>' | '<' | '|' | ';' | '{' | '}' => false,
        _ => true,
    }
}

named!(bare_token<&str, &str>, take_while1_s!(token_char));
named!(quoted_token<&str, &str>, delimited!(tag_s!("\""),
                                            take_until_s!("\""),
                                            tag_s!("\"")));
named!(place_holder<&str, &str>, tag_s!("{}"));
named!(command_token<&str, &str>,
       alt!(bare_token | quoted_token | place_holder));

named!(command< &str, Vec<&str> >,
       terminated!(ws!(many1!(command_token)), eof!()));

pub fn parse_cmd(text: &str) -> Vec<&str> {
    match command(text) {
        IResult::Done(_, result) => result,
        IResult::Error(error) => panic!("Error {:?}", error),
        IResult::Incomplete(needed) => panic!("Needed {:?}", needed)
    }
}

/// Single Command
pub struct ShellCommand {
    command: Command,
    setpgid: bool
}

impl ShellCommand {
    pub fn new(format: &str, args: &[&str]) -> ShellCommand {
        let vec = parse_cmd(format);
        let mut command = Command::new(vec[0]);
        command.args(&vec.as_slice()[1..]);
        ShellCommand {
            command: command,
            setpgid: false
        }
    }
}

impl JobSpec for ShellCommand {
    fn exec(mut self) -> ! {
        self.command.exec();
        std::process::exit(1);
    }

    fn setpgid(mut self) -> Self {
        self.setpgid = true;
        self
    }

    fn getpgid(&self) -> bool {
        return self.setpgid;
    }
}

#[macro_export]
macro_rules! cmd {
    ($format:expr) => ($crate::command::ShellCommand::new($format, &[]));
    ($format:expr, $($arg:expr),+) =>
        ($crate::command::ShellCommand::new($format, &[$($arg),+]));
}

#[test]
fn test_parse_cmd() {
    let tokens = parse_cmd(r#"cmd 1 2
                              3 "
  4""#);
    assert_eq!("cmd", tokens[0]);
    assert_eq!("1", tokens[1]);
    assert_eq!("2", tokens[2]);
    assert_eq!("3", tokens[3]);
    assert_eq!("\n  4", tokens[4]);
}
<|endoftext|>


<commit_before>//! Connection helper.

extern crate tokio_dns;
extern crate tokio_core;

use std::io::Result as IoResult;

use self::tokio_core::net::TcpStream;
use self::tokio_core::reactor::Remote;
use self::tokio_dns::tcp_connect;

use futures::{future, Future};
use tungstenite::Error;
use tungstenite::client::url_mode;
use tungstenite::handshake::client::Response;

use stream::NoDelay;
use super::{WebSocketStream, Request, client_async};

impl NoDelay for TcpStream {
    fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
        TcpStream::set_nodelay(self, nodelay)
    }
}

#[cfg(feature="tls")]
mod encryption {
    extern crate native_tls;
    extern crate tokio_tls;

    use super::tokio_core::net::TcpStream;

    use self::native_tls::TlsConnector;
    use self::tokio_tls::{TlsConnectorExt, TlsStream};

    use std::io::{Read, Write, Result as IoResult};

    use futures::{future, Future};

    use tungstenite::Error;
    use tungstenite::stream::Mode;

    use stream::NoDelay;

    pub use stream::Stream as StreamSwitcher;
    pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;

    impl<T: Read + Write + NoDelay> NoDelay for TlsStream<T> {
        fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
            self.get_mut().get_mut().set_nodelay(nodelay)
        }
    }

    pub fn wrap_stream(socket: TcpStream, domain: String, mode: Mode)
        -> Box<Future<Item=AutoStream, Error=Error>>
    {
        match mode {
            Mode::Plain => Box::new(future::ok(StreamSwitcher::Plain(socket))),
            Mode::Tls => {
                Box::new(future::result(TlsConnector::builder())
                            .and_then(move |builder| future::result(builder.build()))
                            .and_then(move |connector| connector.connect_async(&domain, socket))
                            .map(|s| StreamSwitcher::Tls(s))
                            .map_err(|e| Error::Tls(e)))
            }
        }
    }
}

#[cfg(not(feature="tls"))]
mod encryption {
    use super::tokio_core::net::TcpStream;

    use futures::{future, Future};

    use tungstenite::Error;
    use tungstenite::stream::Mode;

    pub type AutoStream = TcpStream;

    pub fn wrap_stream(socket: TcpStream, _domain: String, mode: Mode)
        -> Box<Future<Item=AutoStream, Error=Error>>
    {
        match mode {
            Mode::Plain => Box::new(future::ok(socket)),
            Mode::Tls => Box::new(future::err(Error::Url("TLS support not compiled in.".into()))),
        }
    }
}

use self::encryption::{AutoStream, wrap_stream};

/// Connect to a given URL.
pub fn connect_async<R>(request: R, handle: Remote)
    -> Box<Future<Item=(WebSocketStream<AutoStream>, Response), Error=Error>>
where
    R: Into<Request<'static>>
{
    let request: Request = request.into();

    // Make sure we check domain and mode first. URL must be valid.
    let mode = match url_mode(&request.url) {
        Ok(m) => m,
        Err(e) => return Box::new(future::err(e.into())),
    };
    let domain = match request.url.host_str() {
        Some(d) => d.to_string(),
        None => return Box::new(future::err(Error::Url("No host name in the URL".into()))),
    };
    let port = request.url.port_or_known_default().expect("Bug: port unknown");

    Box::new(tcp_connect((domain.as_str(), port), handle).map_err(|e| e.into())
                .and_then(move |socket| wrap_stream(socket, domain, mode))
                .and_then(|mut stream| {
                    NoDelay::set_nodelay(&mut stream, true)
                        .map(move |()| stream)
                        .map_err(|e| e.into())
                })
                .and_then(move |stream| client_async(request, stream)))
}
<commit_msg>connect: parametrize AutoStream with Stream<commit_after>//! Connection helper.

extern crate tokio_dns;
extern crate tokio_core;

use std::io::Result as IoResult;

use self::tokio_core::net::TcpStream;
use self::tokio_core::reactor::Remote;
use self::tokio_dns::tcp_connect;

use futures::{future, Future};
use tungstenite::Error;
use tungstenite::client::url_mode;
use tungstenite::handshake::client::Response;

use stream::NoDelay;
use super::{WebSocketStream, Request, client_async};

impl NoDelay for TcpStream {
    fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
        TcpStream::set_nodelay(self, nodelay)
    }
}

#[cfg(feature="tls")]
mod encryption {
    extern crate native_tls;
    extern crate tokio_tls;

    use self::native_tls::TlsConnector;
    use self::tokio_tls::{TlsConnectorExt, TlsStream};

    use std::io::{Read, Write, Result as IoResult};

    use futures::{future, Future};
    use tokio_io::{AsyncRead, AsyncWrite};

    use tungstenite::Error;
    use tungstenite::stream::Mode;

    use stream::NoDelay;

    pub use stream::Stream as StreamSwitcher;
    pub type AutoStream<S> = StreamSwitcher<S, TlsStream<S>>;

    impl<T: Read + Write + NoDelay> NoDelay for TlsStream<T> {
        fn set_nodelay(&mut self, nodelay: bool) -> IoResult<()> {
            self.get_mut().get_mut().set_nodelay(nodelay)
        }
    }

    pub fn wrap_stream<S>(socket: S, domain: String, mode: Mode)
        -> Box<Future<Item=AutoStream<S>, Error=Error>>
    where
        S: 'static + AsyncRead + AsyncWrite,
    {
        match mode {
            Mode::Plain => Box::new(future::ok(StreamSwitcher::Plain(socket))),
            Mode::Tls => {
                Box::new(future::result(TlsConnector::builder())
                            .and_then(move |builder| future::result(builder.build()))
                            .and_then(move |connector| connector.connect_async(&domain, socket))
                            .map(|s| StreamSwitcher::Tls(s))
                            .map_err(|e| Error::Tls(e)))
            }
        }
    }
}

#[cfg(not(feature="tls"))]
mod encryption {
    use futures::{future, Future};
    use tokio_io::{AsyncRead, AsyncWrite};

    use tungstenite::Error;
    use tungstenite::stream::Mode;

    pub type AutoStream<S> = S;

    pub fn wrap_stream<S>(socket: S, _domain: String, mode: Mode)
        -> Box<Future<Item=AutoStream<S>, Error=Error>>
    where
        S: 'static + AsyncRead + AsyncWrite,
    {
        match mode {
            Mode::Plain => Box::new(future::ok(socket)),
            Mode::Tls => Box::new(future::err(Error::Url("TLS support not compiled in.".into()))),
        }
    }
}

use self::encryption::{AutoStream, wrap_stream};

/// Connect to a given URL.
pub fn connect_async<R>(request: R, handle: Remote)
    -> Box<Future<Item=(WebSocketStream<AutoStream<TcpStream>>, Response), Error=Error>>
where
    R: Into<Request<'static>>
{
    let request: Request = request.into();

    // Make sure we check domain and mode first. URL must be valid.
    let mode = match url_mode(&request.url) {
        Ok(m) => m,
        Err(e) => return Box::new(future::err(e.into())),
    };
    let domain = match request.url.host_str() {
        Some(d) => d.to_string(),
        None => return Box::new(future::err(Error::Url("No host name in the URL".into()))),
    };
    let port = request.url.port_or_known_default().expect("Bug: port unknown");

    Box::new(tcp_connect((domain.as_str(), port), handle).map_err(|e| e.into())
                .and_then(move |socket| wrap_stream(socket, domain, mode))
                .and_then(|mut stream| {
                    NoDelay::set_nodelay(&mut stream, true)
                        .map(move |()| stream)
                        .map_err(|e| e.into())
                })
                .and_then(move |stream| client_async(request, stream)))
}
<|endoftext|>


<commit_before>//! Shadowsocks Server Context

use std::{
    io,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use bloomfilter::Bloom;
use spin::Mutex;
use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;

#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
use crate::{
    config::{Config, ConfigType, ServerConfig},
    relay::{dns_resolver::resolve, socks5::Address},
};

// Entries for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_SERVER: usize = 1_000_000;

// Entries for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_CLIENT: usize = 10_000;

// Error rate for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_SERVER: f64 = 1e-6;

// Error rate for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_CLIENT: f64 = 1e-15;

// A bloom filter borrowed from shadowsocks-libev's `ppbloom`
//
// It contains 2 bloom filters and each one holds 1/2 entries.
// Use them as a ring buffer.
struct PingPongBloom {
    blooms: [Bloom<[u8]>; 2],
    bloom_count: [usize; 2],
    item_count: usize,
    current: usize,
}

impl PingPongBloom {
    fn new(ty: ConfigType) -> PingPongBloom {
        let (mut item_count, fp_p) = if ty.is_local() {
            (BF_NUM_ENTRIES_FOR_CLIENT, BF_ERROR_RATE_FOR_CLIENT)
        } else {
            (BF_NUM_ENTRIES_FOR_SERVER, BF_ERROR_RATE_FOR_SERVER)
        };

        item_count /= 2;

        PingPongBloom {
            blooms: [
                Bloom::new_for_fp_rate(item_count, fp_p),
                Bloom::new_for_fp_rate(item_count, fp_p),
            ],
            bloom_count: [0, 0],
            item_count,
            current: 0,
        }
    }

    // Check if data in `buf` exist.
    //
    // Set into the current bloom filter if not exist.
    //
    // Return `true` if data exist in bloom filter.
    fn check_and_set(&mut self, buf: &[u8]) -> bool {
        for bloom in &self.blooms {
            if bloom.check(buf) {
                return true;
            }
        }

        if self.bloom_count[self.current] >= self.item_count {
            // Current bloom filter is full,
            // Create a new one and use that one as current.

            self.current = (self.current + 1) % 2;

            self.bloom_count[self.current] = 0;
            self.blooms[self.current].clear();
        }

        // Cannot be optimized by `check_and_set`
        // Because we have to check every filters in `blooms` before `set`
        self.blooms[self.current].set(buf);
        self.bloom_count[self.current] += 1;

        false
    }
}

/// Server's global running status
///
/// Shared between UDP and TCP servers
pub struct ServerState {
    #[cfg(feature = "trust-dns")]
    dns_resolver: Option<TokioAsyncResolver>,
}

impl ServerState {
    #[cfg(feature = "trust-dns")]
    pub async fn new_shared(config: &Config, rt: Handle) -> SharedServerState {
        let state = ServerState {
            dns_resolver: match create_resolver(config.get_dns_config(), rt).await {
                Ok(resolver) => Some(resolver),
                Err(..) => None,
            },
        };

        Arc::new(state)
    }

    #[cfg(not(feature = "trust-dns"))]
    pub async fn new_shared(_config: &Config, _rt: Handle) -> SharedServerState {
        Arc::new(ServerState { dns_resolver: None })
    }

    /// Get the global shared resolver
    #[cfg(feature = "trust-dns")]
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.dns_resolver.as_ref()
    }
}

/// `ServerState` wrapped in `Arc`
pub type SharedServerState = Arc<ServerState>;

/// Shared basic configuration for the whole server
pub struct Context {
    config: Config,
    server_state: SharedServerState,
    server_running: AtomicBool,
    nonce_ppbloom: Mutex<PingPongBloom>,
}

/// Unique context thw whole server
pub type SharedContext = Arc<Context>;

impl Context {
    /// Create a non-shared Context
    fn new(config: Config, server_state: SharedServerState) -> Context {
        let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type));

        Context {
            config,
            server_state,
            server_running: AtomicBool::new(true),
            nonce_ppbloom,
        }
    }

    /// Create a shared Context, wrapped in `Arc`
    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
        SharedContext::new(Context::new(config, server_state))
    }

    /// Config for TCP server
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Clone the internal ServerState
    pub fn clone_server_state(&self) -> SharedServerState {
        self.server_state.clone()
    }

    /// Mutable Config for TCP server
    ///
    /// NOTE: Only for launching plugins
    pub fn config_mut(&mut self) -> &mut Config {
        &mut self.config
    }

    /// Get ServerConfig by index
    pub fn server_config(&self, idx: usize) -> &ServerConfig {
        &self.config.server[idx]
    }

    /// Get mutable ServerConfig by index
    pub fn server_config_mut(&mut self, idx: usize) -> &mut ServerConfig {
        &mut self.config.server[idx]
    }

    #[cfg(feature = "trust-dns")]
    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.server_state.dns_resolver()
    }

    /// Perform a DNS resolution
    pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
        resolve(self, host, port).await
    }

    /// Check if the server is still in running state
    pub fn server_running(&self) -> bool {
        self.server_running.load(Ordering::Acquire)
    }

    /// Stops the server, kills all detached running tasks
    pub fn set_server_stopped(&self) {
        self.server_running.store(false, Ordering::Release)
    }

    /// Check if nonce exist or not
    ///
    /// If not, set into the current bloom filter
    pub fn check_nonce_and_set(&self, nonce: &[u8]) -> bool {
        // Plain cipher doesn't have a nonce
        // Always treated as non-duplicated
        if nonce.is_empty() {
            return false;
        }

        let mut ppbloom = self.nonce_ppbloom.lock();
        ppbloom.check_and_set(nonce)
    }

    /// Check client ACL (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_client_blocked(addr),
        }
    }

    /// Check outbound address ACL (for server)
    pub fn check_outbound_blocked(&self, addr: &Address) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_outbound_blocked(addr),
        }
    }

    /// Check resolved outbound address ACL (for server)
    pub fn check_resolved_outbound_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_resolved_outbound_blocked(addr),
        }
    }

    /// Check target address ACL (for client)
    pub async fn check_target_bypassed(&self, target: &Address) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => false,
            Some(ref a) => a.check_target_bypassed(self, target).await,
        }
    }
}
<commit_msg>Fix the build without trust-dns<commit_after>//! Shadowsocks Server Context

use std::{
    io,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use bloomfilter::Bloom;
use spin::Mutex;
use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;

#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
use crate::{
    config::{Config, ConfigType, ServerConfig},
    relay::{dns_resolver::resolve, socks5::Address},
};

// Entries for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_SERVER: usize = 1_000_000;

// Entries for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_CLIENT: usize = 10_000;

// Error rate for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_SERVER: f64 = 1e-6;

// Error rate for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_CLIENT: f64 = 1e-15;

// A bloom filter borrowed from shadowsocks-libev's `ppbloom`
//
// It contains 2 bloom filters and each one holds 1/2 entries.
// Use them as a ring buffer.
struct PingPongBloom {
    blooms: [Bloom<[u8]>; 2],
    bloom_count: [usize; 2],
    item_count: usize,
    current: usize,
}

impl PingPongBloom {
    fn new(ty: ConfigType) -> PingPongBloom {
        let (mut item_count, fp_p) = if ty.is_local() {
            (BF_NUM_ENTRIES_FOR_CLIENT, BF_ERROR_RATE_FOR_CLIENT)
        } else {
            (BF_NUM_ENTRIES_FOR_SERVER, BF_ERROR_RATE_FOR_SERVER)
        };

        item_count /= 2;

        PingPongBloom {
            blooms: [
                Bloom::new_for_fp_rate(item_count, fp_p),
                Bloom::new_for_fp_rate(item_count, fp_p),
            ],
            bloom_count: [0, 0],
            item_count,
            current: 0,
        }
    }

    // Check if data in `buf` exist.
    //
    // Set into the current bloom filter if not exist.
    //
    // Return `true` if data exist in bloom filter.
    fn check_and_set(&mut self, buf: &[u8]) -> bool {
        for bloom in &self.blooms {
            if bloom.check(buf) {
                return true;
            }
        }

        if self.bloom_count[self.current] >= self.item_count {
            // Current bloom filter is full,
            // Create a new one and use that one as current.

            self.current = (self.current + 1) % 2;

            self.bloom_count[self.current] = 0;
            self.blooms[self.current].clear();
        }

        // Cannot be optimized by `check_and_set`
        // Because we have to check every filters in `blooms` before `set`
        self.blooms[self.current].set(buf);
        self.bloom_count[self.current] += 1;

        false
    }
}

/// Server's global running status
///
/// Shared between UDP and TCP servers
pub struct ServerState {
    #[cfg(feature = "trust-dns")]
    dns_resolver: Option<TokioAsyncResolver>,
}

impl ServerState {
    #[cfg(feature = "trust-dns")]
    pub async fn new_shared(config: &Config, rt: Handle) -> SharedServerState {
        let state = ServerState {
            dns_resolver: match create_resolver(config.get_dns_config(), rt).await {
                Ok(resolver) => Some(resolver),
                Err(..) => None,
            },
        };

        Arc::new(state)
    }

    #[cfg(not(feature = "trust-dns"))]
    pub async fn new_shared(_config: &Config, _rt: Handle) -> SharedServerState {
        Arc::new(ServerState { })
    }

    /// Get the global shared resolver
    #[cfg(feature = "trust-dns")]
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.dns_resolver.as_ref()
    }
}

/// `ServerState` wrapped in `Arc`
pub type SharedServerState = Arc<ServerState>;

/// Shared basic configuration for the whole server
pub struct Context {
    config: Config,
    server_state: SharedServerState,
    server_running: AtomicBool,
    nonce_ppbloom: Mutex<PingPongBloom>,
}

/// Unique context thw whole server
pub type SharedContext = Arc<Context>;

impl Context {
    /// Create a non-shared Context
    fn new(config: Config, server_state: SharedServerState) -> Context {
        let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type));

        Context {
            config,
            server_state,
            server_running: AtomicBool::new(true),
            nonce_ppbloom,
        }
    }

    /// Create a shared Context, wrapped in `Arc`
    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
        SharedContext::new(Context::new(config, server_state))
    }

    /// Config for TCP server
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Clone the internal ServerState
    pub fn clone_server_state(&self) -> SharedServerState {
        self.server_state.clone()
    }

    /// Mutable Config for TCP server
    ///
    /// NOTE: Only for launching plugins
    pub fn config_mut(&mut self) -> &mut Config {
        &mut self.config
    }

    /// Get ServerConfig by index
    pub fn server_config(&self, idx: usize) -> &ServerConfig {
        &self.config.server[idx]
    }

    /// Get mutable ServerConfig by index
    pub fn server_config_mut(&mut self, idx: usize) -> &mut ServerConfig {
        &mut self.config.server[idx]
    }

    #[cfg(feature = "trust-dns")]
    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.server_state.dns_resolver()
    }

    /// Perform a DNS resolution
    pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
        resolve(self, host, port).await
    }

    /// Check if the server is still in running state
    pub fn server_running(&self) -> bool {
        self.server_running.load(Ordering::Acquire)
    }

    /// Stops the server, kills all detached running tasks
    pub fn set_server_stopped(&self) {
        self.server_running.store(false, Ordering::Release)
    }

    /// Check if nonce exist or not
    ///
    /// If not, set into the current bloom filter
    pub fn check_nonce_and_set(&self, nonce: &[u8]) -> bool {
        // Plain cipher doesn't have a nonce
        // Always treated as non-duplicated
        if nonce.is_empty() {
            return false;
        }

        let mut ppbloom = self.nonce_ppbloom.lock();
        ppbloom.check_and_set(nonce)
    }

    /// Check client ACL (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_client_blocked(addr),
        }
    }

    /// Check outbound address ACL (for server)
    pub fn check_outbound_blocked(&self, addr: &Address) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_outbound_blocked(addr),
        }
    }

    /// Check resolved outbound address ACL (for server)
    pub fn check_resolved_outbound_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_resolved_outbound_blocked(addr),
        }
    }

    /// Check target address ACL (for client)
    pub async fn check_target_bypassed(&self, target: &Address) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => false,
            Some(ref a) => a.check_target_bypassed(self, target).await,
        }
    }
}
<|endoftext|>


<commit_before>use super::super::bus;
use super::cp0::CP0;
use super::instruction::Instruction;

use std::fmt;

const NUM_GPREG: usize = 32;
const NUM_FPREG: usize = 32;

const PIF_ROM_START: u64 = 0xffff_ffff_bfc0_0000;

pub struct Cpu {
    reg_gprs: [u64; NUM_GPREG],
    reg_fprs: [f64; NUM_FPREG],

    reg_pc: u64,

    reg_hi: u64,
    reg_lo: u64,

    reg_llbit: bool,

    reg_fcr0: u32,
    reg_fcr31: u32,

    cp0: CP0,

    bus: bus::Bus,
}

impl fmt::Debug for Cpu {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        const REGS_PER_LINE: usize = 2;
        const REG_NAMES: [&'static str; NUM_GPREG] =
            ["r0", "at", "v0", "v1", "a0", "a1", "a2", "a3", "t0", "t1", "t2", "t3", "t4", "t5",
             "t6", "t7", "s0", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "t8", "t9", "k0", "k1",
             "gp", "sp", "s8", "ra"];

        try!(write!(f, "\nCPU General Purpose Registers:"));
        for reg_num in 0..NUM_GPREG {
            if (reg_num % REGS_PER_LINE) == 0 {
                try!(writeln!(f, ""));
            }
            try!(write!(f,
                        "{reg_name}/gpr{num:02}: {value:#018X} ",
                        num = reg_num,
                        reg_name = REG_NAMES[reg_num],
                        value = self.reg_gprs[reg_num]));
        }

        try!(write!(f, "\n\nCPU Floating Point Registers:"));
        for reg_num in 0..NUM_FPREG {
            if (reg_num % REGS_PER_LINE) == 0 {
                try!(writeln!(f, ""));
            }
            try!(write!(f,
                "fpr{num:02}: {value:21} ",
                num = reg_num,
                value = self.reg_fprs[reg_num],));
        }

        try!(writeln!(f, "\n\nCPU Special Registers:"));
        try!(writeln!(f,
                      "\
            reg_pc: {:#018X}\nreg_hi: {:#018X}\nreg_lo: \
                       {:#018X}\nreg_llbit: {}\nreg_fcr0:  {:#018X}\nreg_fcr31: {:#018X}\n",
                      self.reg_pc,
                      self.reg_hi,
                      self.reg_lo,
                      self.reg_llbit,
                      self.reg_fcr0,
                      self.reg_fcr31));

        try!(writeln!(f, "{:#?}", self.cp0));
        writeln!(f, "{:#?}", self.bus)
    }
}

impl Cpu {
    pub fn new(bus: bus::Bus) -> Cpu {
        Cpu {
            reg_gprs: [0; NUM_GPREG],
            reg_fprs: [0.0; NUM_FPREG],

            reg_pc: 0,

            reg_hi: 0,
            reg_lo: 0,

            reg_llbit: false,

            reg_fcr0: 0,
            reg_fcr31: 0,

            cp0: CP0::default(),

            bus: bus,
        }
    }

    pub fn power_on_reset(&mut self) {
        self.cp0.power_on_reset();

        self.reg_pc = PIF_ROM_START;
    }

    pub fn run(&mut self) {
        loop {
            self.run_instruction();
        }
    }

    pub fn run_instruction(&mut self) {
        let instruction = self.read_instruction(self.reg_pc);

        let opcode = instruction.get_opcode();
        let rs = instruction.get_bits(21, 5);
        let rt = instruction.get_bits(16, 5);
        let imval = instruction.get_immediate();

        match opcode {
            0b010000 => {
                // MTC0
                let rd = instruction.get_bits(11, 5);
                let data = self.read_gpr(rt as usize);
                self.cp0.write_reg(rd, data);
            }
            0b001101 => {
                // ORI
                let res = self.read_gpr(rs as usize) | (imval as u64);
                self.write_gpr(rt as usize, res);
            }
            0b001111 => {
                // LUI
                // assume 32 bit mode
                self.write_gpr(rt as usize, (((imval as u32) << 16) as i32) as u64);
            }
            0b100011 => {
                // LW
                let base = self.read_gpr(rs as usize);
                let vaddr = base.wrapping_add((instruction.get_immediate() as i16) as u64);
                if vaddr & 0b11 != 0 {
                    panic!("Address error exception");
                }
                panic!("Word {:#x}, base {:#x}, vaddr: {:#x}, imval: {:#x}",
                       instruction,
                       base,
                       vaddr,
                       imval);
                let word = self.read_word(vaddr);
                let mem = (word as i32) as u64;
                self.write_gpr(rt as usize, mem);

            }
            _ => {
                panic!("Unrecognised instruction {:#x}", instruction);
            }
        }

        self.reg_pc += 4;
    }

    fn read_word(&self, addr: u64) -> u32 {
        let paddr = vaddr_to_paddr(addr);
        self.bus.read_word(paddr as u32)
    }

    fn read_instruction(&self, addr: u64) -> Instruction {
        Instruction::new(self.read_word(addr))
    }

    fn write_gpr(&mut self, index: usize, value: u64) {
        if index != 0 {
            self.reg_gprs[index] = value;
        }
    }

    fn read_gpr(&self, index: usize) -> u64 {
        match index {
            0 => 0,
            _ => self.reg_gprs[index],
        }
    }
}

fn vaddr_to_paddr(vaddr: u64) -> u64 {
    let addr_bit_values = (vaddr >> 29) & 0b111;
    if addr_bit_values == 0b101 {
        // kseg1
        vaddr - 0xffff_ffff_a000_0000
    } else {
        panic!("Unrecognised virtual address {:#x}", vaddr);
    }
}
<commit_msg>Moved cpu debug to bottom of file<commit_after>use super::super::bus;
use super::cp0::CP0;
use super::instruction::Instruction;

use std::fmt;

const NUM_GPREG: usize = 32;
const NUM_FPREG: usize = 32;

const PIF_ROM_START: u64 = 0xffff_ffff_bfc0_0000;

pub struct Cpu {
    reg_gprs: [u64; NUM_GPREG],
    reg_fprs: [f64; NUM_FPREG],

    reg_pc: u64,

    reg_hi: u64,
    reg_lo: u64,

    reg_llbit: bool,

    reg_fcr0: u32,
    reg_fcr31: u32,

    cp0: CP0,

    bus: bus::Bus,
}

impl Cpu {
    pub fn new(bus: bus::Bus) -> Cpu {
        Cpu {
            reg_gprs: [0; NUM_GPREG],
            reg_fprs: [0.0; NUM_FPREG],

            reg_pc: 0,

            reg_hi: 0,
            reg_lo: 0,

            reg_llbit: false,

            reg_fcr0: 0,
            reg_fcr31: 0,

            cp0: CP0::default(),

            bus: bus,
        }
    }

    pub fn power_on_reset(&mut self) {
        self.cp0.power_on_reset();

        self.reg_pc = PIF_ROM_START;
    }

    pub fn run(&mut self) {
        loop {
            self.run_instruction();
        }
    }

    pub fn run_instruction(&mut self) {
        let instruction = self.read_instruction(self.reg_pc);

        let opcode = instruction.get_opcode();
        let rs = instruction.get_bits(21, 5);
        let rt = instruction.get_bits(16, 5);
        let imval = instruction.get_immediate();

        match opcode {
            0b010000 => {
                // MTC0
                let rd = instruction.get_bits(11, 5);
                let data = self.read_gpr(rt as usize);
                self.cp0.write_reg(rd, data);
            }
            0b001101 => {
                // ORI
                let res = self.read_gpr(rs as usize) | (imval as u64);
                self.write_gpr(rt as usize, res);
            }
            0b001111 => {
                // LUI
                // assume 32 bit mode
                self.write_gpr(rt as usize, (((imval as u32) << 16) as i32) as u64);
            }
            0b100011 => {
                // LW
                let base = self.read_gpr(rs as usize);
                let vaddr = base.wrapping_add((instruction.get_immediate() as i16) as u64);
                if vaddr & 0b11 != 0 {
                    panic!("Address error exception");
                }
                panic!("Word {:#x}, base {:#x}, vaddr: {:#x}, imval: {:#x}",
                       instruction,
                       base,
                       vaddr,
                       imval);
                let word = self.read_word(vaddr);
                let mem = (word as i32) as u64;
                self.write_gpr(rt as usize, mem);

            }
            _ => {
                panic!("Unrecognised instruction {:#x}", instruction);
            }
        }

        self.reg_pc += 4;
    }

    fn read_word(&self, addr: u64) -> u32 {
        let paddr = vaddr_to_paddr(addr);
        self.bus.read_word(paddr as u32)
    }

    fn read_instruction(&self, addr: u64) -> Instruction {
        Instruction::new(self.read_word(addr))
    }

    fn write_gpr(&mut self, index: usize, value: u64) {
        if index != 0 {
            self.reg_gprs[index] = value;
        }
    }

    fn read_gpr(&self, index: usize) -> u64 {
        match index {
            0 => 0,
            _ => self.reg_gprs[index],
        }
    }
}

impl fmt::Debug for Cpu {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        const REGS_PER_LINE: usize = 2;
        const REG_NAMES: [&'static str; NUM_GPREG] =
            ["r0", "at", "v0", "v1", "a0", "a1", "a2", "a3", "t0", "t1", "t2", "t3", "t4", "t5",
             "t6", "t7", "s0", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "t8", "t9", "k0", "k1",
             "gp", "sp", "s8", "ra"];

        try!(write!(f, "\nCPU General Purpose Registers:"));
        for reg_num in 0..NUM_GPREG {
            if (reg_num % REGS_PER_LINE) == 0 {
                try!(writeln!(f, ""));
            }
            try!(write!(f,
                        "{reg_name}/gpr{num:02}: {value:#018X} ",
                        num = reg_num,
                        reg_name = REG_NAMES[reg_num],
                        value = self.reg_gprs[reg_num]));
        }

        try!(write!(f, "\n\nCPU Floating Point Registers:"));
        for reg_num in 0..NUM_FPREG {
            if (reg_num % REGS_PER_LINE) == 0 {
                try!(writeln!(f, ""));
            }
            try!(write!(f,
                "fpr{num:02}: {value:21} ",
                num = reg_num,
                value = self.reg_fprs[reg_num],));
        }

        try!(writeln!(f, "\n\nCPU Special Registers:"));
        try!(writeln!(f,
                      "\
            reg_pc: {:#018X}\nreg_hi: {:#018X}\nreg_lo: \
                       {:#018X}\nreg_llbit: {}\nreg_fcr0:  {:#010X}\nreg_fcr31: {:#010X}\n",
                      self.reg_pc,
                      self.reg_hi,
                      self.reg_lo,
                      self.reg_llbit,
                      self.reg_fcr0,
                      self.reg_fcr31));

        try!(writeln!(f, "{:#?}", self.cp0));
        writeln!(f, "{:#?}", self.bus)
    }
}

fn vaddr_to_paddr(vaddr: u64) -> u64 {
    let addr_bit_values = (vaddr >> 29) & 0b111;
    if addr_bit_values == 0b101 {
        // kseg1
        vaddr - 0xffff_ffff_a000_0000
    } else {
        panic!("Unrecognised virtual address {:#x}", vaddr);
    }
}
<|endoftext|>


<commit_before>use std::{collections::VecDeque, io::SeekFrom};

use log::{debug, error, trace};
use nom::{self, Err, IResult, Needed, Offset};

use av_data::{packet::Packet, params::*, rational::Rational64, timeinfo::TimeInfo};
use av_format::{
    buffer::Buffered,
    common::GlobalInfo,
    demuxer::{Demuxer, Descr, Descriptor, Event},
    error::*,
    stream::Stream,
};

use crate::{
    ebml::{self, custom_error, ebml_header, EBMLHeader},
    elements::{
        segment, segment_element, simple_block, Cluster, Info, SeekHead, SegmentElement,
        TrackEntry, Tracks,
    },
};

#[derive(Debug, Clone, Default)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
    pub params: Option<DemuxerParams>,
}

#[derive(Debug, Clone, Default)]
pub struct DemuxerParams {
    pub track_numbers: Option<Vec<u64>>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
            params: None,
        }
    }

    pub fn with_params(params: DemuxerParams) -> MkvDemuxer {
        MkvDemuxer {
            params: Some(params),
            ..Default::default()
        }
    }

    pub fn parse_until_tracks<'a>(
        &mut self,
        original_input: &'a [u8],
    ) -> IResult<&'a [u8], (), ebml::Error<'a>> {
        let (i1, header) = ebml_header(original_input)?;

        self.header = Some(header);

        let (mut input, _) = segment(i1)?;

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            let (i3, element) = segment_element(input)?;

            match element {
                SegmentElement::SeekHead(s) => {
                    trace!("got seek head: {:#?}", s);
                    self.seek_head = if self.seek_head.is_none() {
                        Some(s)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Info(i) => {
                    trace!("got info: {:#?}", i);
                    self.info = if self.info.is_none() {
                        Some(i)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Tracks(t) => {
                    trace!("got tracks: {:#?}", t);
                    self.tracks = if self.tracks.is_none() {
                        // Only keep tracks we're interested in
                        if let Some(params) = &self.params {
                            if let Some(track_numbers) = &params.track_numbers {
                                let mut t = t;
                                t.tracks = t
                                    .tracks
                                    .into_iter()
                                    .filter(|tr| track_numbers.contains(&tr.track_number))
                                    .collect();
                                Some(t)
                            } else {
                                None
                            }
                        } else {
                            None
                        }
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                el => {
                    debug!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<dyn Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self
                    .info
                    .as_ref()
                    .and_then(|info| info.duration)
                    .map(|d| d as u64);
                if let Some(ref mut t) = self.tracks {
                    for tr in t.tracks.iter_mut() {
                        info.add_stream(track_to_stream(self.info.as_ref().unwrap(), tr));
                        tr.stream_index = info.streams.last().unwrap().index;
                    }
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => buf.data().len() + usize::from(size),
                    _ => 1024,
                };
                Err(Error::MoreDataNeeded(sz))
            }
            e => {
                error!("{:?}", e);
                Err(Error::InvalidData)
            }
        }
    }

    fn read_event(&mut self, buf: &Box<dyn Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            debug!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets(self.tracks.as_ref().unwrap());
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        _el => {}
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                }
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(usize::from(size)))
                }
                e => {
                    error!("{:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1" => Some("av1".to_owned()),
        "V_VP8" => Some("vp8".to_owned()),
        "V_VP9" => Some("vp9".to_owned()),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    use av_data::audiosample::*;
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let map = if audio.channel_positions.is_none() {
            Some(ChannelMap::default_map(audio.channels as usize))
        } else {
            unimplemented!("Convert matroska map to rust-av map")
        };
        let a = AudioInfo {
            rate: rate as usize,
            map,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = t
        .track_timecode_scale
        .map(|ts| {
            if ts != 0. {
                (ts * info.timecode_scale as f64) as i64
            } else {
                info.timecode_scale as i64
            }
        })
        .unwrap_or(info.timecode_scale as i64);

    Stream {
        id: t.track_uid as isize,
        index: t.stream_index,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None,
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self, tracks: &Tracks) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                debug!("parsing simple block: {:?}", block);
                if let Some(index) = tracks.lookup(block.track_number) {
                    let packet = Packet {
                        data: i.into(),
                        t: TimeInfo {
                            pts: Some(i64::from(block.timecode)),
                            dts: None,
                            duration: None,
                            timebase: None,
                            user_private: None,
                        },
                        pos: None,
                        stream_index: index as isize,
                        is_key: block.keyframe,
                        is_corrupted: false,
                    };

                    v.push(Event::NewPacket(packet));
                }
            } else {
                error!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<dyn Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe(&self) -> &Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &dyn Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    },
};

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
    use std::io::Cursor;

    use log::info;
    use nom::Offset;

    use av_format::{buffer::*, demuxer::Context};

    use super::*;

    const webm: &[u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        info!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                info!("offset: {} bytes", webm.offset(i));
            }
            e => {
                info!("could not parse: {:?}", e);
            }
        }

        info!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => info!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => info!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(
            Box::new(MkvDemuxer::new()),
            Box::new(AccReader::new(Cursor::new(webm))),
        );
        info!("read headers: {:?}", context.read_headers().unwrap());
        info!("streams: {:?}", context.info.streams);
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
    }
}
<commit_msg>Simplify probe<commit_after>use std::{collections::VecDeque, io::SeekFrom};

use log::{debug, error, trace};
use nom::{self, Err, IResult, Needed, Offset};

use av_data::{packet::Packet, params::*, rational::Rational64, timeinfo::TimeInfo};
use av_format::{
    buffer::Buffered,
    common::GlobalInfo,
    demuxer::{Demuxer, Descr, Descriptor, Event},
    error::*,
    stream::Stream,
};

use crate::{
    ebml::{self, custom_error, ebml_header, EBMLHeader},
    elements::{
        segment, segment_element, simple_block, Cluster, Info, SeekHead, SegmentElement,
        TrackEntry, Tracks,
    },
};

#[derive(Debug, Clone, Default)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
    pub params: Option<DemuxerParams>,
}

#[derive(Debug, Clone, Default)]
pub struct DemuxerParams {
    pub track_numbers: Option<Vec<u64>>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
            params: None,
        }
    }

    pub fn with_params(params: DemuxerParams) -> MkvDemuxer {
        MkvDemuxer {
            params: Some(params),
            ..Default::default()
        }
    }

    pub fn parse_until_tracks<'a>(
        &mut self,
        original_input: &'a [u8],
    ) -> IResult<&'a [u8], (), ebml::Error<'a>> {
        let (i1, header) = ebml_header(original_input)?;

        self.header = Some(header);

        let (mut input, _) = segment(i1)?;

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            let (i3, element) = segment_element(input)?;

            match element {
                SegmentElement::SeekHead(s) => {
                    trace!("got seek head: {:#?}", s);
                    self.seek_head = if self.seek_head.is_none() {
                        Some(s)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Info(i) => {
                    trace!("got info: {:#?}", i);
                    self.info = if self.info.is_none() {
                        Some(i)
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                SegmentElement::Tracks(t) => {
                    trace!("got tracks: {:#?}", t);
                    self.tracks = if self.tracks.is_none() {
                        // Only keep tracks we're interested in
                        if let Some(params) = &self.params {
                            if let Some(track_numbers) = &params.track_numbers {
                                let mut t = t;
                                t.tracks = t
                                    .tracks
                                    .into_iter()
                                    .filter(|tr| track_numbers.contains(&tr.track_number))
                                    .collect();
                                Some(t)
                            } else {
                                None
                            }
                        } else {
                            None
                        }
                    } else {
                        return Err(Err::Error(custom_error(input, 1)));
                    };
                }
                el => {
                    debug!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn read_headers(&mut self, buf: &Box<dyn Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self
                    .info
                    .as_ref()
                    .and_then(|info| info.duration)
                    .map(|d| d as u64);
                if let Some(ref mut t) = self.tracks {
                    for tr in t.tracks.iter_mut() {
                        info.add_stream(track_to_stream(self.info.as_ref().unwrap(), tr));
                        tr.stream_index = info.streams.last().unwrap().index;
                    }
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
            Err(Err::Incomplete(needed)) => {
                let sz = match needed {
                    Needed::Size(size) => buf.data().len() + usize::from(size),
                    _ => 1024,
                };
                Err(Error::MoreDataNeeded(sz))
            }
            e => {
                error!("{:?}", e);
                Err(Error::InvalidData)
            }
        }
    }

    fn read_event(&mut self, buf: &Box<dyn Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            debug!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets(self.tracks.as_ref().unwrap());
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        _el => {}
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                }
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(usize::from(size)))
                }
                e => {
                    error!("{:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1" => Some("av1".to_owned()),
        "V_VP8" => Some("vp8".to_owned()),
        "V_VP9" => Some("vp9".to_owned()),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    use av_data::audiosample::*;
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let map = if audio.channel_positions.is_none() {
            Some(ChannelMap::default_map(audio.channels as usize))
        } else {
            unimplemented!("Convert matroska map to rust-av map")
        };
        let a = AudioInfo {
            rate: rate as usize,
            map,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = t
        .track_timecode_scale
        .map(|ts| {
            if ts != 0. {
                (ts * info.timecode_scale as f64) as i64
            } else {
                info.timecode_scale as i64
            }
        })
        .unwrap_or(info.timecode_scale as i64);

    Stream {
        id: t.track_uid as isize,
        index: t.stream_index,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None,
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self, tracks: &Tracks) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                debug!("parsing simple block: {:?}", block);
                if let Some(index) = tracks.lookup(block.track_number) {
                    let packet = Packet {
                        data: i.into(),
                        t: TimeInfo {
                            pts: Some(i64::from(block.timecode)),
                            dts: None,
                            duration: None,
                            timebase: None,
                            user_private: None,
                        },
                        pos: None,
                        stream_index: index as isize,
                        is_key: block.keyframe,
                        is_corrupted: false,
                    };

                    v.push(Event::NewPacket(packet));
                }
            } else {
                error!("error parsing simple block");
            }
        }

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<dyn Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe(&self) -> &Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        ebml_header(&data[..100]).map(|_| 100).unwrap_or(0)
    }
}

pub const MKV_DESC: &dyn Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    },
};

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
    use std::io::Cursor;

    use log::info;
    use nom::Offset;

    use av_format::{buffer::*, demuxer::Context};

    use super::*;

    const webm: &[u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        info!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                info!("offset: {} bytes", webm.offset(i));
            }
            e => {
                info!("could not parse: {:?}", e);
            }
        }

        info!("demuxer: {:#?}", demuxer);
    }

    #[test]
    fn parse_headers_incremental_buffer() {
        let mut demuxer = MkvDemuxer::new();

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => info!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => info!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

    #[test]
    fn context() {
        let mut context = Context::new(
            Box::new(MkvDemuxer::new()),
            Box::new(AccReader::new(Cursor::new(webm))),
        );
        info!("read headers: {:?}", context.read_headers().unwrap());
        info!("streams: {:?}", context.info.streams);
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
        info!("event: {:?}", context.read_event().unwrap());
    }
}
<|endoftext|>