<commit_before>use std::io::prelude::*;
use std::{process, env};
use std::path::Path;
use std::error::Error;
use std::fs::{self, OpenOptions};
/* Alias to avoid naming conflict for write_all!() */
use std::fmt::Write as std_write;

use url::Url;
use url_open::UrlOpen;

use serde_json::{from_str, to_string};

use util::*;
use config::Config;
use timesheet::traits::HasHTML;

use timesheet::session::Session;
use timesheet::session::EventType;

#[derive(Serialize, Deserialize, Debug)]
pub struct Timesheet {
    start            : u64,
    end              : u64,
    config           : Config,
    sessions         : Vec<Session>,
}

impl Timesheet {
    /** Initializes the .trk/timesheet.json file which holds
     * the serialized timesheet
     * Returns Some(newTimesheet) if operation succeeded */
    pub fn init(author_name: Option<&str>) -> Option<Timesheet> {
        /* Check if file already exists (no init permitted) */
        if Timesheet::is_init() {
            println!("Timesheet is already initialized!");
            return None;
        }
        /* File does not exist, initialize */
        let git_author_name = git_author();
        let author_name = match author_name {
            Some(name) => name,
            None => {
                match git_author_name {
                    Some(ref git_name) => git_name,
                    None => {
                        println!("Empty name not permitted. \
                                  Please run with 'trk init <name>'");
                        process::exit(0);
                    }
                }
            }
        };
        let mut config = Config::new();
        config.user_name = Some(author_name.to_string());
        let now = get_seconds();
        let sheet = Timesheet {
            start        : now,
            end          : now + 1,
            config       : config,
            sessions     : Vec::<Session>::new(),
        };
        let sheet = if sheet.write_files() {
            git_init_trk();
            Some(sheet)
        } else {
            None
        };
        sheet
    }

    fn is_init() -> bool {
        Path::new("./.trk/timesheet.json").exists() && Timesheet::load_from_file().is_some()
    }

    pub fn new_session(&mut self, timestamp: Option<u64>) -> bool {
        let possible = self.sessions
            .last_mut()
            .map_or(true, |session| {
                if session.is_running() {
                    println!("Last session is still running.");
                }
                !session.is_running()
        });
        if possible {
            match timestamp {
                Some(timestamp) => {
                    let is_valid_ts = match self.sessions.last() {
                        None => timestamp > self.start,
                        Some(last_session) => timestamp > last_session.end,
                    };
                    if is_valid_ts {
                        self.sessions.push(Session::new(Some(timestamp)));
                    } else {
                        println!("That timestamp is invalid.");
                        process::exit(0);
                    }
                }
                None => {
                    self.sessions.push(Session::new(None));
                }
            };
        }
        possible
    }

    pub fn end_session(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.update_end();
                session.finalize(timestamp);
                self.end = session.end + 1;
            }
            None => println!("No session to finalize."),
        }
    }

    pub fn pause(&mut self, timestamp: Option<u64>, note: Option<String>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, note, EventType::Pause);
            }
            None => println!("No session to pause."),
        }
    }

    pub fn resume(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, None, EventType::Resume);
            }
            None => println!("No session to pause."),
        }
    }

    pub fn note(&mut self, timestamp: Option<u64>, note_text: String) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, Some(note_text), EventType::Note);
            }
            None => println!("No session to add note to."),
        }
    }

    pub fn add_commit(&mut self, hash: String) {
        let new_needed = self.sessions
            .last()
            .map_or(true, |session| !session.is_running());
        if new_needed {
            self.new_session(None);
        }
        match self.sessions.last_mut() {
            Some(session) => {
                let message = git_commit_message(&hash).unwrap_or(String::new());
                session.push_event(None, Some(message), EventType::Commit { hash });
            }
            None => println!("No session to add commit to."),
        }
    }

    pub fn add_branch(&mut self, name: String) {
        self.sessions
            .last_mut()
            .map(|session| session.add_branch(name));
    }

    fn write_to_html(&self, ago: Option<u64>) -> bool {
        // TODO: avoid time-of-check-to-time-of-use race risk
        let path = Path::new("./timesheet.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                file.write_all(self.to_html(ago).as_bytes()).unwrap();
                format_file("timesheet.html");
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not report sheet! {}", why.description());
                false
            }
        }
    }

    fn write_last_session_html(&self) -> bool {
        let session = match self.sessions.last() {
            Some(session) => session,
            None => return true,
        };
        let path = Path::new("./session.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        let mut file = match file {
            Ok(file) => file,
            Err(why) => {
                println!("Could not write report to session.html! {}",
                         why.description());
                return false;
            }
        };

        let stylesheets = if self.config.show_commits {
            r#"<link rel="stylesheet" type="text/css" href="style.css">
"#      } else {
            r#"<link rel="stylesheet" type="text/css" href="style.css">
<link rel="stylesheet" type="text/css" href="no_commit.css">
"#
        };

        let html = format!(r#"<!DOCTYPE html>

  {}
  <title>{} for {}</title>

{}

"#,
               stylesheets,
               "Session",
               "Rafael Bachmann",
               session.to_html());
    file.write_all(html.as_bytes()).unwrap();
    format_file("session.html");
    /* Save was successful */
    true
    }

    fn write_to_json(&self) -> bool {
        if !Path::new("./.trk").exists() {
            match fs::create_dir("./.trk") {
                Ok(_) => {}
                _ => {
                    println!("Could not create .trk directory.");
                    process::exit(0);
                }
            }
        }

        let path = Path::new("./.trk/timesheet.json");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                /* Convert the sheet to a JSON string. */
                let serialized =
                    to_string(&self).expect("Could not write serialized time sheet.");
                file.write_all(serialized.as_bytes()).unwrap();
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not open timesheet.json file: {}",
                         why.description());
                false
            }
        }
    }

    pub fn write_files(&self) -> bool {
        /* TODO: avoid time-of-check-to-time-of-use race risk */
        self.write_to_json() && self.write_to_html(None) && self.write_last_session_html()
    }

    /** Return a Some(Timesheet) struct if a timesheet.json file
     * is present and valid in the .trk directory, and None otherwise.
     * */
pub fn load_from_file() -> Option<Timesheet> {
    let mut path = env::current_dir().unwrap();
    loop {
        path.push(".trk");
        if path.exists() {
            env::set_current_dir(&path).is_ok();
            break;
        } else {
            path.pop();
            if !path.pop() {
                return None;
            }
        }
    }

    path.push("timesheet.json");
    let file = OpenOptions::new().read(true).open(&path);
    path.pop();
    env::set_current_dir(path).unwrap();
    match file {
        Ok(mut file) => {
            let mut serialized = String::new();
            match file.read_to_string(&mut serialized) {
                Ok(..) => from_str(&serialized).unwrap_or(None),
                Err(..) => {
                    println!("IO error while reading the timesheet file.");
                    process::exit(0);
                }
            }
        }
        Err(..) => None,
    }
}

    pub fn clear() {
        /* Try to get user name */
        let sheet = Timesheet::load_from_file();
        /* In case there is a sheet, there must also be a name */
        let name: Option<String> = sheet.map(|s| s.config.user_name.unwrap());

        let path = Path::new("./.trk/timesheet.json");
        if path.exists() {
            fs::remove_file(&path).unwrap_or_else(|why| {
                println!("Could not remove sessions file: {}", why.description());
            });
        }
        Timesheet::init(name.as_ref().map(|s| s.as_str()));
    }

    pub fn timesheet_status(&self) -> String {
        let mut status = format!("Sheet running for {}\n",
                sec_to_hms_string(get_seconds() - self.start));
        match self.sessions.len() {
            0 => write!(&mut status, "No sessions yet.\n").unwrap(),
            n => {
                write!(&mut status,
                       "{} session(s) so far.\nLast session:\n{}",
                       n,
                       self.sessions[n - 1].status())
                       .unwrap()
            }
        };
        status
    }

    pub fn last_session_status(&self) -> String {
        let status = self.sessions.last().map(|session| session.status());
        status.unwrap_or("No session yet.".to_string())
    }

    fn open_local_html(&self, filename: String) {
        let file_url = match env::current_dir() {
            Ok(dir) => {
                match dir.join(&filename).to_str() {
                    Some(path) => format!("file://{}", path),
                    None => {
                        println!("Invalid filename: {}.", filename);
                        process::exit(0)
                    }
                }
            }
            Err(why) => {
                println!("Couldn't obtain current directory: {}", why.description());
                process::exit(0)
            }
        };
        match Url::parse(&file_url) {
            Ok(url) => url.open(),
            Err(why) => println!("Couldn't open file: {}", why.description()),
        }
    }

    pub fn report_last_session(&self) {
        self.write_to_html(None);
        self.open_local_html("session.html".to_string());
    }

    pub fn report_sheet(&self, ago: Option<u64>) {
        self.write_to_html(ago);
        self.open_local_html("timesheet.html".to_string());
        /* Leave complete sheet html */
        self.write_to_html(None);
    }

    pub fn show_commits(&mut self, on_off: bool) {
        self.config.show_commits = on_off;
    }

    pub fn set_repo_url(&mut self, repo: String) {
        let repo =
            if repo == "" { None } else { Some(repo) };
        self.config.repository = repo;
    }

    pub fn pause_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.pause_time())
    }

    pub fn work_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.work_time())
    }

    fn to_html(&self, ago: Option<u64>) -> String {
        let timestamp = ago.unwrap_or(self.start);
        let mut sessions_html = String::new();
        for session in &self.sessions {
            if session.start > timestamp {
                sessions_html.push_str(&format!("{}
", session.to_html()));
            }
        }

        let stylesheets = if self.config.show_commits {
            "<link rel=\"stylesheet\" type=\"text/css\" href=\"style.css\">\n".to_string()
        } else {
            r#"<link rel="stylesheet" type="text/css" href="style.css">
<link rel="stylesheet" type="text/css" href="no_commit.css">
"#.to_string()
        };

        let mut html = format!(r#"<!DOCTYPE html>

        {}
        <title>{} for {}</title>

    {}"#,
              stylesheets,
              "Timesheet",
              "Rafael Bachmann",
              sessions_html);

        write!(&mut html,
               r#"

Worked for {}

Paused for {}

"#,
               sec_to_hms_string(self.work_time()),
               sec_to_hms_string(self.pause_time()))
               .unwrap();
        write!(&mut html, "\n").unwrap();
        html
    }
}
<commit_msg>include stylesheets in binary file and write them to .trk dir if necessary<commit_after>use std::io::prelude::*;
use std::{process, env};
use std::path::Path;
use std::error::Error;
use std::fs::{self, OpenOptions};
/* Alias to avoid naming conflict for write_all!() */
use std::fmt::Write as std_write;

use url::Url;
use url_open::UrlOpen;

use serde_json::{from_str, to_string};

use util::*;
use config::Config;
use timesheet::traits::HasHTML;

use timesheet::session::Session;
use timesheet::session::EventType;

#[derive(Serialize, Deserialize, Debug)]
pub struct Timesheet {
    start            : u64,
    end              : u64,
    config           : Config,
    sessions         : Vec<Session>,
}

impl Timesheet {
    /** Initializes the .trk/timesheet.json file which holds
     * the serialized timesheet
     * Returns Some(newTimesheet) if operation succeeded */
    pub fn init(author_name: Option<&str>) -> Option<Timesheet> {
        /* Check if file already exists (no init permitted) */
        if Timesheet::is_init() {
            println!("Timesheet is already initialized!");
            return None;
        }
        /* File does not exist, initialize */
        let git_author_name = git_author();
        let author_name = match author_name {
            Some(name) => name,
            None => {
                match git_author_name {
                    Some(ref git_name) => git_name,
                    None => {
                        println!("Empty name not permitted. \
                                  Please run with 'trk init <name>'");
                        process::exit(0);
                    }
                }
            }
        };
        let mut config = Config::new();
        config.user_name = Some(author_name.to_string());
        let now = get_seconds();
        let sheet = Timesheet {
            start        : now,
            end          : now + 1,
            config       : config,
            sessions     : Vec::<Session>::new(),
        };
        let sheet = if sheet.write_files() {
            git_init_trk();
            Some(sheet)
        } else {
            None
        };
        sheet
    }

    fn is_init() -> bool {
        Path::new("./.trk/timesheet.json").exists() && Timesheet::load_from_file().is_some()
    }

    pub fn new_session(&mut self, timestamp: Option<u64>) -> bool {
        let possible = self.sessions
            .last_mut()
            .map_or(true, |session| {
                if session.is_running() {
                    println!("Last session is still running.");
                }
                !session.is_running()
        });
        if possible {
            match timestamp {
                Some(timestamp) => {
                    let is_valid_ts = match self.sessions.last() {
                        None => timestamp > self.start,
                        Some(last_session) => timestamp > last_session.end,
                    };
                    if is_valid_ts {
                        self.sessions.push(Session::new(Some(timestamp)));
                    } else {
                        println!("That timestamp is invalid.");
                        process::exit(0);
                    }
                }
                None => {
                    self.sessions.push(Session::new(None));
                }
            };
        }
        possible
    }

    pub fn end_session(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.update_end();
                session.finalize(timestamp);
                self.end = session.end + 1;
            }
            None => println!("No session to finalize."),
        }
    }

    pub fn pause(&mut self, timestamp: Option<u64>, note: Option<String>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, note, EventType::Pause);
            }
            None => println!("No session to pause."),
        }
    }

    pub fn resume(&mut self, timestamp: Option<u64>) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, None, EventType::Resume);
            }
            None => println!("No session to pause."),
        }
    }

    pub fn note(&mut self, timestamp: Option<u64>, note_text: String) {
        match self.sessions.last_mut() {
            Some(session) => {
                session.push_event(timestamp, Some(note_text), EventType::Note);
            }
            None => println!("No session to add note to."),
        }
    }

    pub fn add_commit(&mut self, hash: String) {
        let new_needed = self.sessions
            .last()
            .map_or(true, |session| !session.is_running());
        if new_needed {
            self.new_session(None);
        }
        match self.sessions.last_mut() {
            Some(session) => {
                let message = git_commit_message(&hash).unwrap_or(String::new());
                session.push_event(None, Some(message), EventType::Commit { hash });
            }
            None => println!("No session to add commit to."),
        }
    }

    pub fn add_branch(&mut self, name: String) {
        self.sessions
            .last_mut()
            .map(|session| session.add_branch(name));
    }

    fn write_to_html(&self, ago: Option<u64>) -> bool {
        // TODO: avoid time-of-check-to-time-of-use race risk
        let path = Path::new("./timesheet.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                file.write_all(self.to_html(ago).as_bytes()).unwrap();
                format_file("timesheet.html");
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not report sheet! {}", why.description());
                false
            }
        }
    }

    fn write_last_session_html(&self) -> bool {
        let session = match self.sessions.last() {
            Some(session) => session,
            None => return true,
        };
        let path = Path::new("./session.html");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        let mut file = match file {
            Ok(file) => file,
            Err(why) => {
                println!("Could not write report to session.html! {}",
                         why.description());
                return false;
            }
        };

        let stylesheets = if self.config.show_commits {
            r#"<link rel="stylesheet" type="text/css" href=".trk/style.css">
"#      } else {
            r#"<link rel="stylesheet" type="text/css" href=".trk/style.css">
<link rel="stylesheet" type="text/css" href=".trk/no_commit.css">
"#
        };

        let html = format!(r#"<!DOCTYPE html>

  {}
  <title>{} for {}</title>

{}

"#,
               stylesheets,
               "Session",
               "Rafael Bachmann",
               session.to_html());
    file.write_all(html.as_bytes()).unwrap();
    format_file("session.html");
    /* Save was successful */
    true
    }

    fn write_to_json(&self) -> bool {
        if !Path::new("./.trk").exists() {
            match fs::create_dir("./.trk") {
                Ok(_) => {}
                _ => {
                    println!("Could not create .trk directory.");
                    process::exit(0);
                }
            }
        }

        let path = Path::new("./.trk/timesheet.json");
        let file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(&path);

        match file {
            Ok(mut file) => {
                /* Convert the sheet to a JSON string. */
                let serialized =
                    to_string(&self).expect("Could not write serialized time sheet.");
                file.write_all(serialized.as_bytes()).unwrap();
                /* Save was successful */
                true
            }
            Err(why) => {
                println!("Could not open timesheet.json file: {}",
                         why.description());
                false
            }
        }
    }

    fn write_stylesheets(filename: &str, content: &'static str) -> bool {
        // TODO: avoid time-of-check-to-time-of-use race risk
        let mut file_path = env::current_dir().unwrap();
        file_path.push(filename);
        if !file_path.exists() {
            let file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&file_path);
            match file {
                Ok(mut file) => {
                    println!("writing {}", filename);
                    file.write_all(content.as_bytes()).unwrap();
                    /* Save was successful */
                    true
                }
                Err(why) => {
                    println!("Could not report sheet! {}", why.description());
                    false
                }
            }
        } else { true }
    }

    pub fn write_files(&self) -> bool {
        /* TODO: avoid time-of-check-to-time-of-use race risk */
        self.write_to_json() && self.write_to_html(None) && self.write_last_session_html()
    }

    /** Return a Some(Timesheet) struct if a timesheet.json file
     * is present and valid in the .trk directory, and None otherwise.
     * */
pub fn load_from_file() -> Option<Timesheet> {
    let mut path = env::current_dir().unwrap();
    loop {
        path.push(".trk");
        if path.exists() {
            env::set_current_dir(&path).is_ok();
            break;
        } else {
            path.pop();
            if !path.pop() {
                return None;
            }
        }
    }

    path.push("timesheet.json");
    let file = OpenOptions::new().read(true).open(&path);
    let result = match file {
        Ok(mut file) => {
            let mut serialized = String::new();
            match file.read_to_string(&mut serialized) {
                Ok(..) => {
                    let style:           &'static str = include_str!("../../style.css");
                    let no_commit_style: &'static str = include_str!("../../no_commit.css");
                    Timesheet::write_stylesheets("style.css", style);
                    Timesheet::write_stylesheets("no_commit.css", no_commit_style);
                    from_str(&serialized).unwrap_or(None)
                }
                Err(..) => {
                    println!("IO error while reading the timesheet file.");
                    process::exit(0);
                }
            }
        }
        Err(..) => None,
    };
    path.pop();
    env::set_current_dir(path).unwrap();
    result
}

    pub fn clear() {
        /* Try to get user name */
        let sheet = Timesheet::load_from_file();
        /* In case there is a sheet, there must also be a name */
        let name: Option<String> = sheet.map(|s| s.config.user_name.unwrap());

        let path = Path::new("./.trk/timesheet.json");
        if path.exists() {
            fs::remove_file(&path).unwrap_or_else(|why| {
                println!("Could not remove sessions file: {}", why.description());
            });
        }
        Timesheet::init(name.as_ref().map(|s| s.as_str()));
    }

    pub fn timesheet_status(&self) -> String {
        let mut status = format!("Sheet running for {}\n",
                sec_to_hms_string(get_seconds() - self.start));
        match self.sessions.len() {
            0 => write!(&mut status, "No sessions yet.\n").unwrap(),
            n => {
                write!(&mut status,
                       "{} session(s) so far.\nLast session:\n{}",
                       n,
                       self.sessions[n - 1].status())
                       .unwrap()
            }
        };
        status
    }

    pub fn last_session_status(&self) -> String {
        let status = self.sessions.last().map(|session| session.status());
        status.unwrap_or("No session yet.".to_string())
    }

    fn open_local_html(&self, filename: String) {
        let file_url = match env::current_dir() {
            Ok(dir) => {
                match dir.join(&filename).to_str() {
                    Some(path) => format!("file://{}", path),
                    None => {
                        println!("Invalid filename: {}.", filename);
                        process::exit(0)
                    }
                }
            }
            Err(why) => {
                println!("Couldn't obtain current directory: {}", why.description());
                process::exit(0)
            }
        };
        match Url::parse(&file_url) {
            Ok(url) => url.open(),
            Err(why) => println!("Couldn't open file: {}", why.description()),
        }
    }

    pub fn report_last_session(&self) {
        self.write_to_html(None);
        self.open_local_html("session.html".to_string());
    }

    pub fn report_sheet(&self, ago: Option<u64>) {
        self.write_to_html(ago);
        self.open_local_html("timesheet.html".to_string());
        /* Leave complete sheet html */
        self.write_to_html(None);
    }

    pub fn show_commits(&mut self, on_off: bool) {
        self.config.show_commits = on_off;
    }

    pub fn set_repo_url(&mut self, repo: String) {
        let repo =
            if repo == "" { None } else { Some(repo) };
        self.config.repository = repo;
    }

    pub fn pause_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.pause_time())
    }

    pub fn work_time(&self) -> u64 {
        self.sessions.iter().fold(
            0, |total, session| total + session.work_time())
    }

    fn to_html(&self, ago: Option<u64>) -> String {
        let timestamp = ago.unwrap_or(self.start);
        let mut sessions_html = String::new();
        for session in &self.sessions {
            if session.start > timestamp {
                sessions_html.push_str(&format!("{}
", session.to_html()));
            }
        }

        let stylesheets = if self.config.show_commits {
            "<link rel=\"stylesheet\" type=\"text/css\" href=\".trk/style.css\">\n".to_string()
        } else {
            r#"<link rel="stylesheet" type="text/css" href=".trk/style.css">
<link rel="stylesheet" type="text/css" href=".trk/no_commit.css">
"#.to_string()
        };

        let mut html = format!(r#"<!DOCTYPE html>

        {}
        <title>{} for {}</title>

    {}"#,
              stylesheets,
              "Timesheet",
              "Rafael Bachmann",
              sessions_html);

        write!(&mut html,
               r#"

Worked for {}

Paused for {}

"#,
               sec_to_hms_string(self.work_time()),
               sec_to_hms_string(self.pause_time()))
               .unwrap();
        write!(&mut html, "\n").unwrap();
        html
    }
}
<|endoftext|>


<commit_before>use std::io::{Read, Write, Result, Error, Seek, SeekFrom};
use std::io::ErrorKind::Other;
use std::path::{Path, PathBuf};
use std::fs::{File, read_dir, OpenOptions};
use std::borrow::ToOwned;
use std::cmp::min;

use byteorder::{LittleEndian, ByteOrder};
use crc::{crc32, Hasher32};
use nix::sys::stat;

use parser;

use parser::{LvmTextMap, into_textmap, textmap_serialize};

const LABEL_SCAN_SECTORS: usize = 4;
const ID_LEN: usize = 32;
const MDA_MAGIC: &'static [u8] = b"\x20\x4c\x56\x4d\x32\x20\x78\x5b\x35\x41\x25\x72\x30\x4e\x2a\x3e";
const INITIAL_CRC: u32 = 0xf597a6cf;
const CRC_SEED: u32 = 0xedb88320;
const SECTOR_SIZE: usize = 512;
const MDA_HEADER_SIZE: usize = 512;

#[derive(Debug)]
struct LabelHeader {
    id: String,
    sector: u64,
    crc: u32,
    offset: u32,
    label: String,
}

#[derive(Debug)]
struct PvArea {
    offset: u64,
    size: u64,
}

#[derive(Debug)]
struct PvHeader {
    uuid: String,
    size: u64, // in bytes
    ext_version: u32,
    ext_flags: u32,
    data_areas: Vec<PvArea>,
    metadata_areas: Vec<PvArea>,
    bootloader_areas: Vec<PvArea>,
}

fn get_label_header(buf: &[u8]) -> Result<LabelHeader> {

    for x in 0..LABEL_SCAN_SECTORS {
        let sec_buf = &buf[x*SECTOR_SIZE..x*SECTOR_SIZE+SECTOR_SIZE];
        if &sec_buf[..8] == b"LABELONE" {
            let crc = LittleEndian::read_u32(&sec_buf[16..20]);
            crc32_ok(crc, &sec_buf[20..SECTOR_SIZE]);

            let sector = LittleEndian::read_u64(&sec_buf[8..16]);
            if sector != x as u64 {
                return Err(Error::new(Other, "Sector field should equal sector count"));
            }

            return Ok(LabelHeader{
                id: String::from_utf8_lossy(&sec_buf[..8]).into_owned(),
                sector: sector,
                crc: crc,
                offset: LittleEndian::read_u32(&sec_buf[20..24]) + (x*SECTOR_SIZE as usize) as u32,
                label: String::from_utf8_lossy(&sec_buf[24..32]).into_owned(),
            })
        }
    }

    Err(Error::new(Other, "Label not found"))
}

#[derive(Debug)]
struct PvAreaIter<'a> {
    area: &'a[u8],
}

fn iter_pv_area<'a>(buf: &'a[u8]) -> PvAreaIter<'a> {
    PvAreaIter { area: buf }
}

impl<'a> Iterator for PvAreaIter<'a> {
    type Item = PvArea;

    fn next (&mut self) -> Option<PvArea> {
        let off = LittleEndian::read_u64(&self.area[..8]);
        let size = LittleEndian::read_u64(&self.area[8..16]);

        if off == 0 {
            None
        }
        else {
            self.area = &self.area[16..];
            Some(PvArea {
                offset: off,
                size: size,
            })
        }
    }
}

//
// PV HEADER LAYOUT:
// - static header (uuid and size)
// - 0+ data areas (actually max 1, usually 1; size 0 == "rest of blkdev")
// - blank entry
// - 0+ metadata areas (max 1, usually 1)
// - blank entry
// - 8 bytes of pvextension header
// - if version > 0
//   - 0+ bootloader areas (usually 0)
//
fn get_pv_header(buf: &[u8]) -> Result<PvHeader> {

    let mut da_buf = &buf[ID_LEN+8..];

    let da_vec: Vec<_> = iter_pv_area(da_buf).collect();

    // move slice past any actual entries plus blank
    // terminating entry
    da_buf = &da_buf[(da_vec.len()+1)*16..];

    let md_vec: Vec<_> = iter_pv_area(da_buf).collect();

    da_buf = &da_buf[(md_vec.len()+1)*16..];

    let ext_version = LittleEndian::read_u32(&da_buf[..4]);
    let mut ext_flags = 0;
    let mut ba_vec = Vec::new();

    if ext_version != 0 {
        ext_flags = LittleEndian::read_u32(&da_buf[4..8]);

        da_buf = &da_buf[8..];

        ba_vec = iter_pv_area(da_buf).collect();
    }

    Ok(PvHeader{
        uuid: String::from_utf8_lossy(&buf[..ID_LEN]).into_owned(),
        size: LittleEndian::read_u64(&buf[ID_LEN..ID_LEN+8]),
        ext_version: ext_version,
        ext_flags: ext_flags,
        data_areas: da_vec,
        metadata_areas: md_vec,
        bootloader_areas: ba_vec,
    })
}

fn crc32_ok(val: u32, buf: &[u8]) -> bool {
    let mut digest = crc32::Digest::new(CRC_SEED);
    digest.value = INITIAL_CRC;
    digest.write(&buf);
    let crc32 = digest.sum32();

    // TODO: all our crcs are failing, how come?
    if val != crc32 {
        println!("CRC32: input {:x} != calculated {:x}", val, crc32);
    }
    val == crc32
}

fn crc32_calc(buf: &[u8]) -> u32 {
    let mut digest = crc32::Digest::new(CRC_SEED);
    digest.value = INITIAL_CRC;
    digest.write(&buf);
    digest.sum32()
}

#[derive(Debug, PartialEq, Clone)]
struct RawLocn {
    offset: u64,
    size: u64,
    checksum: u32,
    flags: u32,
}

#[derive(Debug)]
struct RawLocnIter<'a> {
    area: &'a[u8],
}

fn iter_raw_locn<'a>(buf: &'a[u8]) -> RawLocnIter<'a> {
    RawLocnIter { area: buf }
}

impl<'a> Iterator for RawLocnIter<'a> {
    type Item = RawLocn;

    fn next (&mut self) -> Option<RawLocn> {
        let off = LittleEndian::read_u64(&self.area[..8]);
        let size = LittleEndian::read_u64(&self.area[8..16]);
        let checksum = LittleEndian::read_u32(&self.area[16..20]);
        let flags = LittleEndian::read_u32(&self.area[20..24]);

        if off == 0 {
            None
        }
        else {
            self.area = &self.area[24..];
            Some(RawLocn {
                offset: off,
                size: size,
                checksum: checksum,
                flags: flags,
            })
        }
    }
}

fn find_pv_in_dev(path: &Path) -> Result<PvHeader> {

    let mut f = try!(File::open(path));

    let mut buf = vec![0; LABEL_SCAN_SECTORS * SECTOR_SIZE];

    try!(f.read(&mut buf));

    let label_header = try!(get_label_header(&buf));
    let pvheader = try!(get_pv_header(&buf[label_header.offset as usize..]));

    return Ok(pvheader);
}

fn round_up_to_sector_size(num: u64) -> u64 {
    let rem = num % SECTOR_SIZE as u64;

    num + SECTOR_SIZE as u64 - rem
}

#[derive(Debug)]
pub struct MDA {
    file: File,
    area: Vec<u8>,
    area_offset: u64,
}

impl MDA {
    pub fn new(path: &Path) -> Result<MDA> {
        let pvheader = try!(find_pv_in_dev(&path));

        let mut f = try!(OpenOptions::new().read(true).write(true).open(path));

        let mda_areas = pvheader.metadata_areas.len();
        if mda_areas != 1 {
            return Err(Error::new(
                Other, format!("Expecting 1 mda, found {}", mda_areas)));
        }

        let md = &pvheader.metadata_areas[0];
        try!(f.seek(SeekFrom::Start(md.offset)));
        let mut buf = vec![0; md.size as usize];
        try!(f.read(&mut buf));

        crc32_ok(LittleEndian::read_u32(&buf[..4]), &buf[4..MDA_HEADER_SIZE]);

        if &buf[4..20] != MDA_MAGIC {
            return Err(Error::new(
                Other, format!("'{}' doesn't match MDA_MAGIC",
                               String::from_utf8_lossy(&buf[4..20]))));
        }

        let ver = LittleEndian::read_u32(&buf[20..24]);
        if ver != 1 {
            return Err(Error::new(Other, "Bad version, expected 1"));
        }

        // TODO: validate these somehow
        //println!("mdah start {}", LittleEndian::read_u64(&buf[24..32]));
        //println!("mdah size {}", LittleEndian::read_u64(&buf[32..40]));

        Ok(MDA {
            file: f,
            area: buf,
            area_offset: md.offset,
        })
    }

    fn get_rlocn0(&self) -> RawLocn {
        let raw_locns: Vec<_> = iter_raw_locn(&self.area[40..]).collect();
        let rlocn_len = raw_locns.len();
        if rlocn_len != 1 {
            panic!("Expecting 1 rlocn, found {}", rlocn_len);
        }

        raw_locns[0].clone()
    }

    fn set_rlocn0(&mut self, rl: &RawLocn) -> Result<()> {

        {
            let mut raw_locn = &mut self.area[40..];

            LittleEndian::write_u64(&mut raw_locn[..8], rl.offset);
            LittleEndian::write_u64(&mut raw_locn[8..16], rl.size);
            LittleEndian::write_u32(&mut raw_locn[16..20], rl.checksum);
            LittleEndian::write_u32(&mut raw_locn[20..24], rl.flags);
        }

        let csum = crc32_calc(&self.area[4..MDA_HEADER_SIZE]);
        LittleEndian::write_u32(&mut self.area[..4], csum);

        try!(self.file.seek(SeekFrom::Start(self.area_offset)));
        try!(self.file.write_all(&mut self.area[..MDA_HEADER_SIZE]));

        Ok(())
    }

    pub fn read_metadata(&self) -> Result<LvmTextMap> {
        let rl = self.get_rlocn0();
        let rl_start = rl.offset as usize;
        let rl_end = rl_start + rl.size as usize;

        if rl_end <= self.area.len() {
            crc32_ok(rl.checksum, &self.area[rl_start..rl_end]);
            parser::into_textmap(&self.area[rl_start..rl_end])
        } else {
            // Split across end/beginning of md area
            let mut text: Vec<u8> = Vec::new();
            let remaining = rl_end - self.area.len();
            text.extend(&self.area[rl_start..rl_end-remaining].to_owned());
            text.extend(
                &self.area[MDA_HEADER_SIZE..MDA_HEADER_SIZE+remaining].to_owned());
            crc32_ok(rl.checksum, &text);
            parser::into_textmap(&text)
        }
    }

    pub fn write_textmap_to_next_rlocn(&mut self, map: &LvmTextMap) -> Result<()> {
        let mut raw_locn = self.get_rlocn0();

        let mut text = textmap_serialize(map);
        // Ends with one null
        text.push(b'\0');

        let start_off =
            round_up_to_sector_size(raw_locn.offset + raw_locn.size)
            % self.area.len() as u64;
        let tail_space = self.area.len() as u64 - start_off;

        assert_eq!(start_off % SECTOR_SIZE as u64, 0);
        assert_eq!(tail_space % SECTOR_SIZE as u64, 0);

        let written = if tail_space != 0 {
            try!(self.file.seek(
                SeekFrom::Start(self.area_offset + start_off)));
            try!(self.file.write_all(
                &text[..min(tail_space as usize, text.len())]));
            min(tail_space as usize, text.len())
        } else {
            0
        };

        if written != text.len() {
            println!("writing2 {} {}", written, text.len());
            try!(self.file.seek(
                SeekFrom::Start(self.area_offset as u64 + MDA_HEADER_SIZE as u64)));
            try!(self.file.write_all(
                &text[written as usize..]));
        }

        println!("rlocn {:#?}", raw_locn);
        raw_locn.offset = start_off;
        raw_locn.size = text.len() as u64;
        raw_locn.checksum = crc32_calc(&text);
        println!("rlocn {:#?}", raw_locn);

        try!(self.set_rlocn0(&raw_locn));

        println!("new rlocn {:#?}", self.get_rlocn0());

        Ok(())
    }
}

pub fn scan_for_pvs(dirs: &[&Path]) -> Result<Vec<PathBuf>> {

    let mut ret_vec = Vec::new();

    for dir in dirs {
        ret_vec.extend(try!(read_dir(dir))
            .into_iter()
            .filter_map(|dir_e| if dir_e.is_ok()
                        { Some(dir_e.unwrap().path()) } else {None} )
            .filter(|path| {
                (stat::stat(path).unwrap().st_mode & 0x6000) == 0x6000 }) // S_IFBLK
            .filter(|path| { find_pv_in_dev(&path).is_ok() })
            .collect::<Vec<_>>());
    }

    Ok(ret_vec)
}
<commit_msg>fix CRCs to use existing crc32-rs APIs<commit_after>use std::io::{Read, Write, Result, Error, Seek, SeekFrom};
use std::io::ErrorKind::Other;
use std::path::{Path, PathBuf};
use std::fs::{File, read_dir, OpenOptions};
use std::borrow::ToOwned;
use std::cmp::min;

use byteorder::{LittleEndian, ByteOrder};
use crc::crc32;
use nix::sys::stat;

use parser;

use parser::{LvmTextMap, into_textmap, textmap_serialize};

const LABEL_SCAN_SECTORS: usize = 4;
const ID_LEN: usize = 32;
const MDA_MAGIC: &'static [u8] = b"\x20\x4c\x56\x4d\x32\x20\x78\x5b\x35\x41\x25\x72\x30\x4e\x2a\x3e";
const INITIAL_CRC: u32 = 0xf597a6cf;
const CRC_SEED: u32 = 0xedb88320;
const SECTOR_SIZE: usize = 512;
const MDA_HEADER_SIZE: usize = 512;

#[derive(Debug)]
struct LabelHeader {
    id: String,
    sector: u64,
    crc: u32,
    offset: u32,
    label: String,
}

#[derive(Debug)]
struct PvArea {
    offset: u64,
    size: u64,
}

#[derive(Debug)]
struct PvHeader {
    uuid: String,
    size: u64, // in bytes
    ext_version: u32,
    ext_flags: u32,
    data_areas: Vec<PvArea>,
    metadata_areas: Vec<PvArea>,
    bootloader_areas: Vec<PvArea>,
}

fn get_label_header(buf: &[u8]) -> Result<LabelHeader> {

    for x in 0..LABEL_SCAN_SECTORS {
        let sec_buf = &buf[x*SECTOR_SIZE..x*SECTOR_SIZE+SECTOR_SIZE];
        if &sec_buf[..8] == b"LABELONE" {
            let crc = LittleEndian::read_u32(&sec_buf[16..20]);
            crc32_ok(crc, &sec_buf[20..SECTOR_SIZE]);

            let sector = LittleEndian::read_u64(&sec_buf[8..16]);
            if sector != x as u64 {
                return Err(Error::new(Other, "Sector field should equal sector count"));
            }

            return Ok(LabelHeader{
                id: String::from_utf8_lossy(&sec_buf[..8]).into_owned(),
                sector: sector,
                crc: crc,
                offset: LittleEndian::read_u32(&sec_buf[20..24]) + (x*SECTOR_SIZE as usize) as u32,
                label: String::from_utf8_lossy(&sec_buf[24..32]).into_owned(),
            })
        }
    }

    Err(Error::new(Other, "Label not found"))
}

#[derive(Debug)]
struct PvAreaIter<'a> {
    area: &'a[u8],
}

fn iter_pv_area<'a>(buf: &'a[u8]) -> PvAreaIter<'a> {
    PvAreaIter { area: buf }
}

impl<'a> Iterator for PvAreaIter<'a> {
    type Item = PvArea;

    fn next (&mut self) -> Option<PvArea> {
        let off = LittleEndian::read_u64(&self.area[..8]);
        let size = LittleEndian::read_u64(&self.area[8..16]);

        if off == 0 {
            None
        }
        else {
            self.area = &self.area[16..];
            Some(PvArea {
                offset: off,
                size: size,
            })
        }
    }
}

//
// PV HEADER LAYOUT:
// - static header (uuid and size)
// - 0+ data areas (actually max 1, usually 1; size 0 == "rest of blkdev")
// - blank entry
// - 0+ metadata areas (max 1, usually 1)
// - blank entry
// - 8 bytes of pvextension header
// - if version > 0
//   - 0+ bootloader areas (usually 0)
//
fn get_pv_header(buf: &[u8]) -> Result<PvHeader> {

    let mut da_buf = &buf[ID_LEN+8..];

    let da_vec: Vec<_> = iter_pv_area(da_buf).collect();

    // move slice past any actual entries plus blank
    // terminating entry
    da_buf = &da_buf[(da_vec.len()+1)*16..];

    let md_vec: Vec<_> = iter_pv_area(da_buf).collect();

    da_buf = &da_buf[(md_vec.len()+1)*16..];

    let ext_version = LittleEndian::read_u32(&da_buf[..4]);
    let mut ext_flags = 0;
    let mut ba_vec = Vec::new();

    if ext_version != 0 {
        ext_flags = LittleEndian::read_u32(&da_buf[4..8]);

        da_buf = &da_buf[8..];

        ba_vec = iter_pv_area(da_buf).collect();
    }

    Ok(PvHeader{
        uuid: String::from_utf8_lossy(&buf[..ID_LEN]).into_owned(),
        size: LittleEndian::read_u64(&buf[ID_LEN..ID_LEN+8]),
        ext_version: ext_version,
        ext_flags: ext_flags,
        data_areas: da_vec,
        metadata_areas: md_vec,
        bootloader_areas: ba_vec,
    })
}

fn crc32_ok(val: u32, buf: &[u8]) -> bool {
    let crc32 = crc32_calc(buf);

    if val != crc32 {
        println!("CRC32: input {:x} != calculated {:x}", val, crc32);
    }
    val == crc32
}

fn crc32_calc(buf: &[u8]) -> u32 {
    let table = crc32::make_table(CRC_SEED);

    // For some reason, we need to negate the initial CRC value
    // and the result, to match what LVM2 is generating.
    !crc32::update(!INITIAL_CRC, &table, buf)
}

#[derive(Debug, PartialEq, Clone)]
struct RawLocn {
    offset: u64,
    size: u64,
    checksum: u32,
    flags: u32,
}

#[derive(Debug)]
struct RawLocnIter<'a> {
    area: &'a[u8],
}

fn iter_raw_locn<'a>(buf: &'a[u8]) -> RawLocnIter<'a> {
    RawLocnIter { area: buf }
}

impl<'a> Iterator for RawLocnIter<'a> {
    type Item = RawLocn;

    fn next (&mut self) -> Option<RawLocn> {
        let off = LittleEndian::read_u64(&self.area[..8]);
        let size = LittleEndian::read_u64(&self.area[8..16]);
        let checksum = LittleEndian::read_u32(&self.area[16..20]);
        let flags = LittleEndian::read_u32(&self.area[20..24]);

        if off == 0 {
            None
        }
        else {
            self.area = &self.area[24..];
            Some(RawLocn {
                offset: off,
                size: size,
                checksum: checksum,
                flags: flags,
            })
        }
    }
}

fn find_pv_in_dev(path: &Path) -> Result<PvHeader> {

    let mut f = try!(File::open(path));

    let mut buf = vec![0; LABEL_SCAN_SECTORS * SECTOR_SIZE];

    try!(f.read(&mut buf));

    let label_header = try!(get_label_header(&buf));
    let pvheader = try!(get_pv_header(&buf[label_header.offset as usize..]));

    return Ok(pvheader);
}

fn round_up_to_sector_size(num: u64) -> u64 {
    let rem = num % SECTOR_SIZE as u64;

    num + SECTOR_SIZE as u64 - rem
}

#[derive(Debug)]
pub struct MDA {
    file: File,
    area: Vec<u8>,
    area_offset: u64,
}

impl MDA {
    pub fn new(path: &Path) -> Result<MDA> {
        let pvheader = try!(find_pv_in_dev(&path));

        let mut f = try!(OpenOptions::new().read(true).write(true).open(path));

        let mda_areas = pvheader.metadata_areas.len();
        if mda_areas != 1 {
            return Err(Error::new(
                Other, format!("Expecting 1 mda, found {}", mda_areas)));
        }

        let md = &pvheader.metadata_areas[0];
        try!(f.seek(SeekFrom::Start(md.offset)));
        let mut buf = vec![0; md.size as usize];
        try!(f.read(&mut buf));

        if !crc32_ok(LittleEndian::read_u32(&buf[..4]), &buf[4..MDA_HEADER_SIZE]) {
            return Err(Error::new(Other, "MDA header checksum failure"));
        }

        if &buf[4..20] != MDA_MAGIC {
            return Err(Error::new(
                Other, format!("'{}' doesn't match MDA_MAGIC",
                               String::from_utf8_lossy(&buf[4..20]))));
        }

        let ver = LittleEndian::read_u32(&buf[20..24]);
        if ver != 1 {
            return Err(Error::new(Other, "Bad version, expected 1"));
        }

        // TODO: validate these somehow
        //println!("mdah start {}", LittleEndian::read_u64(&buf[24..32]));
        //println!("mdah size {}", LittleEndian::read_u64(&buf[32..40]));

        Ok(MDA {
            file: f,
            area: buf,
            area_offset: md.offset,
        })
    }

    fn get_rlocn0(&self) -> RawLocn {
        let raw_locns: Vec<_> = iter_raw_locn(&self.area[40..]).collect();
        let rlocn_len = raw_locns.len();
        if rlocn_len != 1 {
            panic!("Expecting 1 rlocn, found {}", rlocn_len);
        }

        raw_locns[0].clone()
    }

    fn set_rlocn0(&mut self, rl: &RawLocn) -> Result<()> {

        {
            let mut raw_locn = &mut self.area[40..];

            LittleEndian::write_u64(&mut raw_locn[..8], rl.offset);
            LittleEndian::write_u64(&mut raw_locn[8..16], rl.size);
            LittleEndian::write_u32(&mut raw_locn[16..20], rl.checksum);
            LittleEndian::write_u32(&mut raw_locn[20..24], rl.flags);
        }

        let csum = crc32_calc(&self.area[4..MDA_HEADER_SIZE]);
        LittleEndian::write_u32(&mut self.area[..4], csum);

        try!(self.file.seek(SeekFrom::Start(self.area_offset)));
        try!(self.file.write_all(&mut self.area[..MDA_HEADER_SIZE]));

        Ok(())
    }

    pub fn read_metadata(&self) -> Result<LvmTextMap> {
        let rl = self.get_rlocn0();
        let rl_start = rl.offset as usize;
        let rl_end = rl_start + rl.size as usize;

        if rl_end <= self.area.len() {
            crc32_ok(rl.checksum, &self.area[rl_start..rl_end]);
            parser::into_textmap(&self.area[rl_start..rl_end])
        } else {
            // Split across end/beginning of md area
            let mut text: Vec<u8> = Vec::new();
            let remaining = rl_end - self.area.len();
            text.extend(&self.area[rl_start..rl_end-remaining].to_owned());
            text.extend(
                &self.area[MDA_HEADER_SIZE..MDA_HEADER_SIZE+remaining].to_owned());
            crc32_ok(rl.checksum, &text);
            parser::into_textmap(&text)
        }
    }

    pub fn write_textmap_to_next_rlocn(&mut self, map: &LvmTextMap) -> Result<()> {
        let mut raw_locn = self.get_rlocn0();

        let mut text = textmap_serialize(map);
        // Ends with one null
        text.push(b'\0');

        let start_off =
            round_up_to_sector_size(raw_locn.offset + raw_locn.size)
            % self.area.len() as u64;
        let tail_space = self.area.len() as u64 - start_off;

        assert_eq!(start_off % SECTOR_SIZE as u64, 0);
        assert_eq!(tail_space % SECTOR_SIZE as u64, 0);

        let written = if tail_space != 0 {
            try!(self.file.seek(
                SeekFrom::Start(self.area_offset + start_off)));
            try!(self.file.write_all(
                &text[..min(tail_space as usize, text.len())]));
            min(tail_space as usize, text.len())
        } else {
            0
        };

        if written != text.len() {
            println!("writing2 {} {}", written, text.len());
            try!(self.file.seek(
                SeekFrom::Start(self.area_offset as u64 + MDA_HEADER_SIZE as u64)));
            try!(self.file.write_all(
                &text[written as usize..]));
        }

        println!("rlocn {:#?}", raw_locn);
        raw_locn.offset = start_off;
        raw_locn.size = text.len() as u64;
        raw_locn.checksum = crc32_calc(&text);
        println!("rlocn {:#?}", raw_locn);

        try!(self.set_rlocn0(&raw_locn));

        println!("new rlocn {:#?}", self.get_rlocn0());

        Ok(())
    }
}

pub fn scan_for_pvs(dirs: &[&Path]) -> Result<Vec<PathBuf>> {

    let mut ret_vec = Vec::new();

    for dir in dirs {
        ret_vec.extend(try!(read_dir(dir))
            .into_iter()
            .filter_map(|dir_e| if dir_e.is_ok()
                        { Some(dir_e.unwrap().path()) } else {None} )
            .filter(|path| {
                (stat::stat(path).unwrap().st_mode & 0x6000) == 0x6000 }) // S_IFBLK
            .filter(|path| { find_pv_in_dev(&path).is_ok() })
            .collect::<Vec<_>>());
    }

    Ok(ret_vec)
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 <Rick Richardson r@12sidedtech.com>
//
// This software may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details.

use net_stream::NetStream;

use mio::{
    EventLoop,
    EventLoopSender,
    Handler,
    NonBlock,
    IoWriter,
    IoReader,
    IoAcceptor,
    Timeout};

pub use mio::Token;
use mio::net::{SockAddr, Socket};
use mio::net::tcp::{TcpAcceptor, TcpSocket};
use mio::util::Slab;
use mio::event;

use iobuf::{Iobuf, RWIobuf, AROIobuf, Allocator};

use std::io::net::addrinfo::get_host_addresses;
use std::result::Result;
use std::sync::Arc;
use std::sync::mpsc::{Receiver,SyncSender, sync_channel};

use std::time::Duration;

use collections::dlist::DList;

use reactive::Subscriber;
/// The basic sendable buffer which also contains
/// its own addressing. When the buffer is received,
/// Token reflects the socket whence it came,
/// when it is sent, Token is the socket out of which
/// the buffer should be sent
#[derive(Show)]
pub struct StreamBuf (pub AROIobuf, pub Token);

unsafe impl Send for StreamBuf {}

impl Clone for StreamBuf {
    fn clone(&self) -> StreamBuf {
        StreamBuf (self.0.clone(), self.1)
    }
}

pub type TimerCB<'a> = FnMut(&mut Reactor)->bool + 'a;

pub type Reactor = EventLoop<Token, StreamBuf>;

pub type Sender = EventLoopSender<StreamBuf>;

struct Connection {
        sock: TcpSocket,
        outbuf: DList<AROIobuf>,
        interest: event::Interest,
        conn_tx: SyncSender<StreamBuf>
}

impl Connection {
    pub fn new(s: TcpSocket, tx: SyncSender<StreamBuf>) -> Connection {
        Connection {
            sock: s,
            outbuf: DList::new(),
            interest: event::HUP,
            conn_tx: tx
        }
    }
}

pub struct NetEngine<'a> {
    inner: EngineInner<'a>,
    event_loop: Reactor
}

impl<'a> NetEngine<'a> {
    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new(buf_sz, max_conns, max_queue_sz),
                    event_loop: EventLoop::new().unwrap() }
    }

    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    /// * a: Specialized allocator for the buffers to be read
    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new_with_alloc(buf_sz, max_conns, max_queue_sz, a),
                    event_loop: EventLoop::new().unwrap() }
     }

    /// connect to the supplied hostname and port
    /// any data that arrives on the connection will be put into a Buf
    /// and sent down the supplied Sender channel along with the Token of the connection
    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize) -> Result<NetStream<'b>, String> {
        self.inner.connect(hostname, port, &mut self.event_loop)
    }

    /// listen on the supplied ip address and port
    /// any new connections will be accepted and polled for read events
    /// all datagrams that arrive will be put into StreamBufs with their
    /// corresponding token, and added to the default outbound data queue
    /// this can be called multiple times for different ips/ports
    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize) -> Result<Receiver<StreamBuf>, String> {
        self.inner.listen(addr, port, &mut self.event_loop)
    }

    /// fetch the event_loop channel for notifying the event_loop of new outbound data
    pub fn channel(&self) -> EventLoopSender<StreamBuf> {
        self.event_loop.channel()
    }

    pub fn timeout(&mut self, timeout: Duration, callback: Box<TimerCB<'a>>) {
        let tok = self.inner.timeouts.insert((callback, None)).map_err(|_|()).unwrap();
        let handle = self.event_loop.timeout(tok, timeout).unwrap();
        self.inner.timeouts.get_mut(tok).unwrap().1 = Some(handle);
    }

    /// process all incoming and outgoing events in a loop
    pub fn run(mut self) {
        self.event_loop.run(self.inner).map_err(|_| ()).unwrap();
    }

    /// process all incoming and outgoing events in a loop
    pub fn run_once(mut self) {
        self.event_loop.run_once(self.inner).map_err(|_| ()).unwrap();
    }

    /// calculates the 11th digit of pi
    pub fn shutdown(mut self) {
        self.event_loop.shutdown();
    }
}

struct EngineInner<'a> {
    listeners: Slab<(TcpAcceptor, SyncSender<StreamBuf>)>,
    timeouts: Slab<(Box<TimerCB<'a>>, Option<Timeout>)>,
    conns: Slab<Connection>,
    alloc: Option<Arc<Box<Allocator>>>,
    buf_sz: usize,
    queue_sz: usize
}

impl<'a> EngineInner<'a> {

    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> EngineInner<'a> {
        EngineInner {
            listeners: Slab::new_starting_at(Token(0), 128),
            timeouts: Slab::new_starting_at(Token(0), 128),
            conns: Slab::new_starting_at(Token(2), max_conns),
            buf_sz: buf_sz,
            queue_sz: max_queue_sz,
            alloc: None
        }
    }

    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> EngineInner<'a> {
        let ns = EngineInner::new(buf_sz, max_conns, max_queue_sz);
        EngineInner { alloc: Some(a), ..ns }
    }

    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize,
                   event_loop: &mut Reactor) -> Result<NetStream<'b>, String>
    {
        let ip = get_host_addresses(hostname).unwrap()[0]; //TODO manage receiving multiple IPs per hostname, random sample or something
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                let (tx, rx) = sync_channel(self.queue_sz);
                match self.conns.insert(Connection::new(s, tx)) {
                    Ok(tok) => match event_loop.register_opt(&self.conns.get(tok).unwrap().sock, tok, event::READABLE, event::PollOpt::edge()) {
                        Ok(..) => match self.conns.get(tok).unwrap().sock.connect(&SockAddr::InetAddr(ip, port as u16)) {
                            Ok(..) => {
                                debug!("Connected to server for token {:?}", tok);
                                Ok(NetStream::new(tok, rx, event_loop.channel().clone()))
                            }
                            Err(e) => Err(format!("Failed to connect to {:?}:{:?}, error: {:?}", hostname, port, e))
                        },
                        Err(e)      => Err(format!("Failed to register with the event loop, error: {:?}", e))
                    },
                    _ => Err(format!("Failed to insert into connection slab"))
                }
            },
            Err(e) => Err(format!("Failed to create new socket, error:{:?}", e))
        }
    }

    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize,
                  event_loop: &mut Reactor) -> Result<Receiver<StreamBuf>, String>
    {
        let ip = get_host_addresses(addr).unwrap()[0];
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                match s.bind(&SockAddr::InetAddr(ip, port as u16)) {
                Ok(l) => match l.listen(255) {
                    Ok(a) => {
                        let (tx, rx) = sync_channel(self.queue_sz);
                        match self.listeners.insert((a, tx)) {
                            Ok(token) => {
                                event_loop.register_opt(&self.listeners.get_mut(token).unwrap().0,
                                                        token,
                                                        event::READABLE,
                                                        event::PollOpt::edge()).
                                                            map_err(|e| format!("event registration failed: {:?}", e)).
                                                            map(move |_| rx)
                            },
                            Err(_) => Err(format!("failed to insert into listener slab"))
                        }
                    },
                    Err(e) => {Err(format!("Failed to listen to socket {:?}:{:?}, error:{:?}", addr, port, e)) }
                },
                Err(e) => Err(format!("Failed to bind to {:?}:{:?}, error:{:?}", addr, port, e))
            }},
            Err(e) => Err(format!("Failed to create TCP socket, error:{:?}", e))
        }
    }

}

impl<'a> Handler<Token, StreamBuf> for EngineInner<'a> {

    fn readable(&mut self, event_loop: &mut Reactor, token: Token, hint: event::ReadHint) {
        debug!("mio_processor::readable top, token: {:?}", token);
        let mut close = false;
        if self.listeners.contains(token) {
            let (ref mut list, ref tx) = *self.listeners.get_mut(token).unwrap();
            match list.accept() {
                Ok(NonBlock::Ready(sock)) => {
                    //s.set_tcp_nodelay(true); TODO: re-add to mio
                    match self.conns.insert(Connection::new(sock, tx.clone())) {
                        Ok(tok) =>  {
                            event_loop.register_opt(&self.conns.get(tok).unwrap().sock,
                                                    tok, event::READABLE | event::HUP,
                                                    event::PollOpt::edge()).unwrap();
                                      debug!("readable accepted socket for token {:?}", tok); }
                        Err(..)  => error!("Failed to insert into Slab")
                    }},
                _ => error!("Failed to accept socket")
            }
            event_loop.reregister(list, token, event::READABLE, event::PollOpt::edge()).unwrap();
            return;
        }

        match self.conns.get_mut(token) {
            None    => error!("Got a readable event for token {:?},
                               but it is not present in MioHandler connections", token),
            Some(c) => {

                let mut b =
                    if let Some(alloc) = self.alloc.as_ref() {
                        RWIobuf::new_with_allocator(self.buf_sz, alloc.clone())
                    } else {
                        RWIobuf::new(self.buf_sz)
                    };

                let result = unsafe { c.sock.read_slice(b.as_mut_window_slice()) };
                match result {
                    Ok(NonBlock::Ready(n)) => {debug!("read {:?} bytes", n);
                              b.advance(n as u32).unwrap();
                              b.flip_lo();
                              let abuf = b.atomic_read_only().unwrap();
                              c.conn_tx.send( StreamBuf (abuf, token) ).unwrap(); },
                    Ok(NonBlock::WouldBlock) => {
                        debug!("Got Readable event for socket, but failed to write any bytes");
                    },
                    Err(e) => error!("error writing to socket: {:?}", e)
                };

                if hint.contains(event::HUPHINT) {
                    close = true;
                }
                else {
                    c.interest.insert(event::READABLE);
                    event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                }
            }
        }

        if close {
            self.conns.remove(token);
        }
    }

    fn writable(&mut self, event_loop: &mut Reactor, token: Token) {
        debug!("mio_processor::writable, token: {:?}", token);
        match self.conns.get_mut(token) {
            None    => error!("{:?} not present in connections", token),
            Some(c) => { let mut writable = true;
                         let mut index : usize = 0;
                         while writable && c.outbuf.len() > 0 {
                             let (result, sz) =
                             unsafe {
                                 let buf = c.outbuf.front_mut().unwrap();
                                 let sz = buf.len() - index as u32;
                                 let b : &[u8] = &buf.as_window_slice()[index..];
                                 (c.sock.write_slice(b), sz)
                             };
                             match result {
                                 Ok(NonBlock::Ready(n)) =>
                                    {
                                        debug!("Wrote {:?} out of {:?} bytes to socket", n, sz);
                                        if n == sz as usize {
                                            c.outbuf.pop_front(); // we have written the contents of this buffer so lets get rid of it
                                            index = 0;
                                        }
                                        else { // this is unlikely to happen, we didn't write all of the
                                               // buffer to the socket, so lets try again
                                               index += n;
                                        }
                                    },
                                 Ok(NonBlock::WouldBlock) => { // this is also very unlikely, we got a writable message, but failed
                                                                    // to write anything at all.
                                                                    debug!("Got Writable event for socket, but failed to write any bytes");
                                                                    writable = false;
                                },
                                Err(e)              => error!("error writing to socket: {:?}", e)
                             }
                         }
                         if c.outbuf.len() > 0 {
                             c.interest.insert(event::WRITABLE);
                             event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                         }
            },
        }

    }

    fn notify(&mut self, event_loop: &mut Reactor, msg: StreamBuf) {
        debug!("mio_processor::notify top");
        let StreamBuf (buf, tok) = msg;
        match self.conns.get_mut(tok) {
            Some(c) => { c.outbuf.push_back(buf);
                         c.interest.insert(event::WRITABLE);
                         event_loop.reregister(&c.sock, tok, c.interest, event::PollOpt::edge()).unwrap();},
            None => {}
        }
    }

    fn timeout(&mut self, event_loop: &mut Reactor, tok: Token) {
        let (ref mut cb, ref handle) = *self.timeouts.get_mut(tok).unwrap();
        if !(*cb).call_mut((event_loop,)) {
            event_loop.clear_timeout(handle.unwrap());
        }
    }
}

<commit_msg>fixed logging message<commit_after>// Copyright (C) 2015 <Rick Richardson r@12sidedtech.com>
//
// This software may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details.

use net_stream::NetStream;

use mio::{
    EventLoop,
    EventLoopSender,
    Handler,
    NonBlock,
    IoWriter,
    IoReader,
    IoAcceptor,
    Timeout};

pub use mio::Token;
use mio::net::{SockAddr, Socket};
use mio::net::tcp::{TcpAcceptor, TcpSocket};
use mio::util::Slab;
use mio::event;

use iobuf::{Iobuf, RWIobuf, AROIobuf, Allocator};

use std::io::net::addrinfo::get_host_addresses;
use std::result::Result;
use std::sync::Arc;
use std::sync::mpsc::{Receiver,SyncSender, sync_channel};

use std::time::Duration;

use collections::dlist::DList;

use reactive::Subscriber;
/// The basic sendable buffer which also contains
/// its own addressing. When the buffer is received,
/// Token reflects the socket whence it came,
/// when it is sent, Token is the socket out of which
/// the buffer should be sent
#[derive(Show)]
pub struct StreamBuf (pub AROIobuf, pub Token);

unsafe impl Send for StreamBuf {}

impl Clone for StreamBuf {
    fn clone(&self) -> StreamBuf {
        StreamBuf (self.0.clone(), self.1)
    }
}

pub type TimerCB<'a> = FnMut(&mut Reactor)->bool + 'a;

pub type Reactor = EventLoop<Token, StreamBuf>;

pub type Sender = EventLoopSender<StreamBuf>;

struct Connection {
        sock: TcpSocket,
        outbuf: DList<AROIobuf>,
        interest: event::Interest,
        conn_tx: SyncSender<StreamBuf>
}

impl Connection {
    pub fn new(s: TcpSocket, tx: SyncSender<StreamBuf>) -> Connection {
        Connection {
            sock: s,
            outbuf: DList::new(),
            interest: event::HUP,
            conn_tx: tx
        }
    }
}

pub struct NetEngine<'a> {
    inner: EngineInner<'a>,
    event_loop: Reactor
}

impl<'a> NetEngine<'a> {
    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new(buf_sz, max_conns, max_queue_sz),
                    event_loop: EventLoop::new().unwrap() }
    }

    /// * buf_sz: the size of the buffer that will be used to read from and send to sockets
    ///         this should be set to the size of your MTU, likely 1500 bytes
    /// * max_conns: Size for a pre-allocated array of containers for Connection objects
    /// * max_queue_sz: Size of the default sync channel through which buffers are sent
    /// * a: Specialized allocator for the buffers to be read
    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> NetEngine<'a> {
        NetEngine { inner: EngineInner::new_with_alloc(buf_sz, max_conns, max_queue_sz, a),
                    event_loop: EventLoop::new().unwrap() }
     }

    /// connect to the supplied hostname and port
    /// any data that arrives on the connection will be put into a Buf
    /// and sent down the supplied Sender channel along with the Token of the connection
    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize) -> Result<NetStream<'b>, String> {
        self.inner.connect(hostname, port, &mut self.event_loop)
    }

    /// listen on the supplied ip address and port
    /// any new connections will be accepted and polled for read events
    /// all datagrams that arrive will be put into StreamBufs with their
    /// corresponding token, and added to the default outbound data queue
    /// this can be called multiple times for different ips/ports
    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize) -> Result<Receiver<StreamBuf>, String> {
        self.inner.listen(addr, port, &mut self.event_loop)
    }

    /// fetch the event_loop channel for notifying the event_loop of new outbound data
    pub fn channel(&self) -> EventLoopSender<StreamBuf> {
        self.event_loop.channel()
    }

    pub fn timeout(&mut self, timeout: Duration, callback: Box<TimerCB<'a>>) {
        let tok = self.inner.timeouts.insert((callback, None)).map_err(|_|()).unwrap();
        let handle = self.event_loop.timeout(tok, timeout).unwrap();
        self.inner.timeouts.get_mut(tok).unwrap().1 = Some(handle);
    }

    /// process all incoming and outgoing events in a loop
    pub fn run(mut self) {
        self.event_loop.run(self.inner).map_err(|_| ()).unwrap();
    }

    /// process all incoming and outgoing events in a loop
    pub fn run_once(mut self) {
        self.event_loop.run_once(self.inner).map_err(|_| ()).unwrap();
    }

    /// calculates the 11th digit of pi
    pub fn shutdown(mut self) {
        self.event_loop.shutdown();
    }
}

struct EngineInner<'a> {
    listeners: Slab<(TcpAcceptor, SyncSender<StreamBuf>)>,
    timeouts: Slab<(Box<TimerCB<'a>>, Option<Timeout>)>,
    conns: Slab<Connection>,
    alloc: Option<Arc<Box<Allocator>>>,
    buf_sz: usize,
    queue_sz: usize
}

impl<'a> EngineInner<'a> {

    pub fn new(buf_sz: usize, max_conns: usize, max_queue_sz: usize) -> EngineInner<'a> {
        EngineInner {
            listeners: Slab::new_starting_at(Token(0), 128),
            timeouts: Slab::new_starting_at(Token(0), 128),
            conns: Slab::new_starting_at(Token(2), max_conns),
            buf_sz: buf_sz,
            queue_sz: max_queue_sz,
            alloc: None
        }
    }

    pub fn new_with_alloc(buf_sz: usize,
               max_conns: usize,
               max_queue_sz: usize,
               a: Arc<Box<Allocator>>) -> EngineInner<'a> {
        let ns = EngineInner::new(buf_sz, max_conns, max_queue_sz);
        EngineInner { alloc: Some(a), ..ns }
    }

    pub fn connect<'b>(&mut self,
                   hostname: &str,
                   port: usize,
                   event_loop: &mut Reactor) -> Result<NetStream<'b>, String>
    {
        let ip = get_host_addresses(hostname).unwrap()[0]; //TODO manage receiving multiple IPs per hostname, random sample or something
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                let (tx, rx) = sync_channel(self.queue_sz);
                match self.conns.insert(Connection::new(s, tx)) {
                    Ok(tok) => match event_loop.register_opt(&self.conns.get(tok).unwrap().sock, tok, event::READABLE, event::PollOpt::edge()) {
                        Ok(..) => match self.conns.get(tok).unwrap().sock.connect(&SockAddr::InetAddr(ip, port as u16)) {
                            Ok(..) => {
                                debug!("Connected to server for token {:?}", tok);
                                Ok(NetStream::new(tok, rx, event_loop.channel().clone()))
                            }
                            Err(e) => Err(format!("Failed to connect to {:?}:{:?}, error: {:?}", hostname, port, e))
                        },
                        Err(e)      => Err(format!("Failed to register with the event loop, error: {:?}", e))
                    },
                    _ => Err(format!("Failed to insert into connection slab"))
                }
            },
            Err(e) => Err(format!("Failed to create new socket, error:{:?}", e))
        }
    }

    pub fn listen<'b>(&mut self,
                  addr: &'b str,
                  port: usize,
                  event_loop: &mut Reactor) -> Result<Receiver<StreamBuf>, String>
    {
        let ip = get_host_addresses(addr).unwrap()[0];
        match TcpSocket::v4() {
            Ok(s) => {
                //s.set_tcp_nodelay(true); TODO: re-add to mio
                match s.bind(&SockAddr::InetAddr(ip, port as u16)) {
                Ok(l) => match l.listen(255) {
                    Ok(a) => {
                        let (tx, rx) = sync_channel(self.queue_sz);
                        match self.listeners.insert((a, tx)) {
                            Ok(token) => {
                                event_loop.register_opt(&self.listeners.get_mut(token).unwrap().0,
                                                        token,
                                                        event::READABLE,
                                                        event::PollOpt::edge()).
                                                            map_err(|e| format!("event registration failed: {:?}", e)).
                                                            map(move |_| rx)
                            },
                            Err(_) => Err(format!("failed to insert into listener slab"))
                        }
                    },
                    Err(e) => {Err(format!("Failed to listen to socket {:?}:{:?}, error:{:?}", addr, port, e)) }
                },
                Err(e) => Err(format!("Failed to bind to {:?}:{:?}, error:{:?}", addr, port, e))
            }},
            Err(e) => Err(format!("Failed to create TCP socket, error:{:?}", e))
        }
    }

}

impl<'a> Handler<Token, StreamBuf> for EngineInner<'a> {

    fn readable(&mut self, event_loop: &mut Reactor, token: Token, hint: event::ReadHint) {
        debug!("mio_processor::readable top, token: {:?}", token);
        let mut close = false;
        if self.listeners.contains(token) {
            let (ref mut list, ref tx) = *self.listeners.get_mut(token).unwrap();
            match list.accept() {
                Ok(NonBlock::Ready(sock)) => {
                    //s.set_tcp_nodelay(true); TODO: re-add to mio
                    match self.conns.insert(Connection::new(sock, tx.clone())) {
                        Ok(tok) =>  {
                            event_loop.register_opt(&self.conns.get(tok).unwrap().sock,
                                                    tok, event::READABLE | event::HUP,
                                                    event::PollOpt::edge()).unwrap();
                                      debug!("readable accepted socket for token {:?}", tok); }
                        Err(..)  => error!("Failed to insert into Slab")
                    }},
                _ => error!("Failed to accept socket")
            }
            event_loop.reregister(list, token, event::READABLE, event::PollOpt::edge()).unwrap();
            return;
        }

        match self.conns.get_mut(token) {
            None    => error!("Got a readable event for token {:?},
                               but it is not present in MioHandler connections", token),
            Some(c) => {

                let mut b =
                    if let Some(alloc) = self.alloc.as_ref() {
                        RWIobuf::new_with_allocator(self.buf_sz, alloc.clone())
                    } else {
                        RWIobuf::new(self.buf_sz)
                    };

                let result = unsafe { c.sock.read_slice(b.as_mut_window_slice()) };
                match result {
                    Ok(NonBlock::Ready(n)) => {debug!("read {:?} bytes", n);
                              b.advance(n as u32).unwrap();
                              b.flip_lo();
                              let abuf = b.atomic_read_only().unwrap();
                              c.conn_tx.send( StreamBuf (abuf, token) ).unwrap(); },
                    Ok(NonBlock::WouldBlock) => {
                        debug!("Got Readable event for socket, but failed to write any bytes");
                    },
                    Err(e) => error!("error reading from socket: {:?}", e)
                };

                if hint.contains(event::HUPHINT) {
                    close = true;
                }
                else {
                    c.interest.insert(event::READABLE);
                    event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                }
            }
        }

        if close {
            self.conns.remove(token);
        }
    }

    fn writable(&mut self, event_loop: &mut Reactor, token: Token) {
        debug!("mio_processor::writable, token: {:?}", token);
        match self.conns.get_mut(token) {
            None    => error!("{:?} not present in connections", token),
            Some(c) => { let mut writable = true;
                         let mut index : usize = 0;
                         while writable && c.outbuf.len() > 0 {
                             let (result, sz) =
                             unsafe {
                                 let buf = c.outbuf.front_mut().unwrap();
                                 let sz = buf.len() - index as u32;
                                 let b : &[u8] = &buf.as_window_slice()[index..];
                                 (c.sock.write_slice(b), sz)
                             };
                             match result {
                                 Ok(NonBlock::Ready(n)) =>
                                    {
                                        debug!("Wrote {:?} out of {:?} bytes to socket", n, sz);
                                        if n == sz as usize {
                                            c.outbuf.pop_front(); // we have written the contents of this buffer so lets get rid of it
                                            index = 0;
                                        }
                                        else { // this is unlikely to happen, we didn't write all of the
                                               // buffer to the socket, so lets try again
                                               index += n;
                                        }
                                    },
                                 Ok(NonBlock::WouldBlock) => { // this is also very unlikely, we got a writable message, but failed
                                                                    // to write anything at all.
                                                                    debug!("Got Writable event for socket, but failed to write any bytes");
                                                                    writable = false;
                                },
                                Err(e)              => error!("error writing to socket: {:?}", e)
                             }
                         }
                         if c.outbuf.len() > 0 {
                             c.interest.insert(event::WRITABLE);
                             event_loop.reregister(&c.sock, token, c.interest, event::PollOpt::edge()).unwrap();
                         }
            },
        }

    }

    fn notify(&mut self, event_loop: &mut Reactor, msg: StreamBuf) {
        debug!("mio_processor::notify top");
        let StreamBuf (buf, tok) = msg;
        match self.conns.get_mut(tok) {
            Some(c) => { c.outbuf.push_back(buf);
                         c.interest.insert(event::WRITABLE);
                         event_loop.reregister(&c.sock, tok, c.interest, event::PollOpt::edge()).unwrap();},
            None => {}
        }
    }

    fn timeout(&mut self, event_loop: &mut Reactor, tok: Token) {
        let (ref mut cb, ref handle) = *self.timeouts.get_mut(tok).unwrap();
        if !(*cb).call_mut((event_loop,)) {
            event_loop.clear_timeout(handle.unwrap());
        }
    }
}

<|endoftext|>


<commit_before>pub mod builtins;

use std::collections::HashMap;
use std::path::PathBuf;
use std::io;
use std::io::{Read, Write};
use std::collections::hash_map::Entry;

#[derive(PartialEq, Debug, Clone)]
pub struct State {
    cwd: PathBuf,
    environment: HashMap<String, String>,
    aliases: HashMap<String, String>,
    argv: Vec<String>,
    argc: usize,
}

impl State {
    pub fn new(cwd: String) -> State {
        State {
            cwd: PathBuf::from(cwd),
            environment: HashMap::new(),
            aliases: HashMap::new(),
            argv: Vec::new(),
            argc: 0,
        }
    }
}

pub fn run(initial_state: State) {
    let mut builtins = builtins::load();
    let mut s = initial_state.clone();

    println!("Welcome to rsh! {:?}", s);

    loop {

        print!("\n");
        print!("{} -> ", s.cwd.to_str().unwrap());

        // this forces the prompt to print
        io::stdout().flush();

        // read the user input
        let mut input = String::new();
        io::stdin().read_line(&mut input).unwrap();

        s.argv = parse_args(&input);
        s.argc = s.argv.len();

        print!("\n");
        println!("Input: {}\nState: {:?}", input, s);

        let first_arg = s.argv.get(0).unwrap().clone();
        if let Entry::Occupied(f) = builtins.entry(String::from(first_arg)) {
            let bn = f.get();
            s = bn(s.clone());
        }
    }
}

#[derive(PartialEq)]
enum BuildType {
    None,
    Single,
    Double,
}

struct ParseResult {
    result: Vec<String>,
    completed: bool,
    build_string: String,
    build_type: BuildType,
}

fn parse_args(args: &String) -> Vec<String> {
    let mut parse_result: ParseResult = ParseResult {
        result: Vec::new(),
        completed: false,
        build_string: String::from(""),
        build_type: BuildType::None,
    };

    if args.len() == 0 {
        return parse_result.result;
    }

    loop {
        parse_string_into_vec(&args, &mut parse_result);

        if (parse_result.completed) {
            break;
        }
    }

    let result: Vec<String> = parse_result.result
        .into_iter()
        .filter(|s| s.len() > 0)
        .collect();

    result
}

fn parse_string_into_vec(string: &String, parse_result: &mut ParseResult) {
    let &mut ParseResult { ref mut build_string,
                           ref mut build_type,
                           ref mut completed,
                           ref mut result } = parse_result;
    let mut iter = string.chars().peekable();

    while let Some(c) = iter.next() {
        match c {
            '\'' => {
                match *build_type {
                    BuildType::Single => {
                        *build_type = BuildType::None;
                        if iter.peek() == Some(&' ') {
                            result.push(build_string.clone());
                            *build_string = String::from("");
                        }
                    }

                    BuildType::None => {
                        *build_type = BuildType::Single;
                    }

                    _ => {
                        build_string.push(c);
                    }
                }
            }

            '\"' => {
                match *build_type {
                    BuildType::Double => {
                        *build_type = BuildType::None;
                        if iter.peek() == Some(&' ') {
                            result.push(build_string.clone());
                            *build_string = String::from("");
                        }
                    }

                    BuildType::None => {
                        *build_type = BuildType::Double;
                    }

                    _ => {
                        build_string.push(c);
                    }
                }
            }

            ' ' => {
                match *build_type {
                    BuildType::None => {
                        result.push(build_string.clone());
                        *build_string = String::from("");
                    }

                    _ => {
                        build_string.push(c);
                    }
                }
            }

            _ => {
                build_string.push(c);
            }
        }
    }

    if *build_type == BuildType::None {
        result.push(build_string.clone());
        *completed = true;
    }
}

#[test]
fn parse_args_test() {
    // parse empty string
    {
        let expected: Vec<String> = Vec::new();
        let result = parse_args(&String::from(""));
        assert_eq!(result, expected);
    }

    // parse single-word string
    {
        let expected: Vec<String> = vec!["echo".to_string()];
        let result = parse_args(&String::from("echo"));
        assert_eq!(result, expected);
    }

    // parse single-word string inside quotes
    {
        let expected = vec!["echo".to_string()];
        let result = parse_args(&String::from("\"echo\""));
        assert_eq!(result, expected);
    }

    // parse multi-word string with closed quotes section
    {
        let expected = vec!["echo".to_string(), "-n".into(), "Hello Dear World".into()];
        let result = parse_args(&String::from("echo -n \"Hello Dear World\""));
        assert_eq!(result, expected);
    }

    // parse multi-word string with multiple closed quotes sections
    {
        let expected = vec!["echo".to_string(), "Hello".into(), "Dear World".into()];
        let result = parse_args(&String::from("echo \"Hello\" \"Dear World\""));
        assert_eq!(result, expected);
    }

    // parse multi-word string with no spaces around single quotes
    {
        let expected = vec!["echo".to_string(), "helloworld".into()];
        let result = parse_args(&String::from("echo 'hello'world"));
        assert_eq!(result, expected);
    }

    // parse multi-word string with no spaces around double quotes
    {
        let expected = vec!["echo".to_string(), "helloworld".into()];
        let result = parse_args(&String::from("echo \"hello\"world"));
        assert_eq!(result, expected);
    }

    // allow double quotes inside single quotes
    {
        let expected = vec!["\"".to_string()];
        let result = parse_args(&String::from("'\"'"));
        assert_eq!(result, expected);
    }

    // allow single quotes inside double quotes
    {
        let expected = vec!["\'".to_string()];
        let result = parse_args(&String::from("\"\'\""));
        assert_eq!(result, expected);
    }
}
<commit_msg>Request further input if outer quotes are not closed<commit_after>pub mod builtins;

use std::collections::HashMap;
use std::path::PathBuf;
use std::io;
use std::io::{Read, Write};
use std::collections::hash_map::Entry;

#[derive(PartialEq, Debug, Clone)]
pub struct State {
    cwd: PathBuf,
    environment: HashMap<String, String>,
    aliases: HashMap<String, String>,
    argv: Vec<String>,
    argc: usize,
}

impl State {
    pub fn new(cwd: String) -> State {
        State {
            cwd: PathBuf::from(cwd),
            environment: HashMap::new(),
            aliases: HashMap::new(),
            argv: Vec::new(),
            argc: 0,
        }
    }
}

pub fn run(initial_state: State) {
    let mut builtins = builtins::load();
    let mut s = initial_state.clone();

    println!("Welcome to rsh! {:?}", s);

    loop {

        print!("\n");
        print!("{} -> ", s.cwd.to_str().unwrap());

        // this forces the prompt to print
        io::stdout().flush();

        // read the user input
        let mut input = String::new();
        io::stdin().read_line(&mut input).unwrap();

        s.argv = parse_args(&input);
        s.argc = s.argv.len();

        print!("\n");
        println!("Input: {}\nState: {:?}", input, s);

        let first_arg = s.argv.get(0).unwrap().clone();
        if let Entry::Occupied(f) = builtins.entry(String::from(first_arg)) {
            let bn = f.get();
            s = bn(s.clone());
        }
    }
}

#[derive(PartialEq)]
enum BuildType {
    None,
    Single,
    Double,
}

struct ParseResult {
    result: Vec<String>,
    completed: bool,
    build_string: String,
    build_type: BuildType,
}

fn parse_args(args: &String) -> Vec<String> {
    let mut parse_result: ParseResult = ParseResult {
        result: Vec::new(),
        completed: false,
        build_string: String::from(""),
        build_type: BuildType::None,
    };

    if args.len() == 0 {
        return parse_result.result;
    }

    let mut argstr = args.clone();
    loop {
        parse_string_into_vec(&argstr, &mut parse_result);

        if (parse_result.completed) {
            break;
        }

        match parse_result.build_type {
            BuildType::Single => print!("quote> "),
            BuildType::Double => print!("dquote> "),
            _ => {}
        }
        io::stdout().flush();
        argstr = String::from("");
        io::stdin().read_line(&mut argstr).unwrap();
    }

    let result: Vec<String> = parse_result.result
        .into_iter()
        .filter(|s| s.len() > 0)
        .collect();

    result
}

fn parse_string_into_vec(string: &String, parse_result: &mut ParseResult) {
    let &mut ParseResult { ref mut build_string,
                           ref mut build_type,
                           ref mut completed,
                           ref mut result } = parse_result;
    let mut iter = string.chars().peekable();

    while let Some(c) = iter.next() {
        match c {
            '\'' => {
                match *build_type {
                    BuildType::Single => {
                        *build_type = BuildType::None;
                        if iter.peek() == Some(&' ') {
                            result.push(build_string.clone());
                            *build_string = String::from("");
                        }
                    }

                    BuildType::None => {
                        *build_type = BuildType::Single;
                    }

                    _ => {
                        build_string.push(c);
                    }
                }
            }

            '\"' => {
                match *build_type {
                    BuildType::Double => {
                        *build_type = BuildType::None;
                        if iter.peek() == Some(&' ') {
                            result.push(build_string.clone());
                            *build_string = String::from("");
                        }
                    }

                    BuildType::None => {
                        *build_type = BuildType::Double;
                    }

                    _ => {
                        build_string.push(c);
                    }
                }
            }

            ' ' => {
                match *build_type {
                    BuildType::None => {
                        result.push(build_string.clone());
                        *build_string = String::from("");
                    }

                    _ => {
                        build_string.push(c);
                    }
                }
            }

            _ => {
                build_string.push(c);
            }
        }
    }

    if *build_type == BuildType::None {
        result.push(build_string.clone());
        *completed = true;
    }
}

#[test]
fn parse_args_test() {
    // parse empty string
    {
        let expected: Vec<String> = Vec::new();
        let result = parse_args(&String::from(""));
        assert_eq!(result, expected);
    }

    // parse single-word string
    {
        let expected: Vec<String> = vec!["echo".to_string()];
        let result = parse_args(&String::from("echo"));
        assert_eq!(result, expected);
    }

    // parse single-word string inside quotes
    {
        let expected = vec!["echo".to_string()];
        let result = parse_args(&String::from("\"echo\""));
        assert_eq!(result, expected);
    }

    // parse multi-word string with closed quotes section
    {
        let expected = vec!["echo".to_string(), "-n".into(), "Hello Dear World".into()];
        let result = parse_args(&String::from("echo -n \"Hello Dear World\""));
        assert_eq!(result, expected);
    }

    // parse multi-word string with multiple closed quotes sections
    {
        let expected = vec!["echo".to_string(), "Hello".into(), "Dear World".into()];
        let result = parse_args(&String::from("echo \"Hello\" \"Dear World\""));
        assert_eq!(result, expected);
    }

    // parse multi-word string with no spaces around single quotes
    {
        let expected = vec!["echo".to_string(), "helloworld".into()];
        let result = parse_args(&String::from("echo 'hello'world"));
        assert_eq!(result, expected);
    }

    // parse multi-word string with no spaces around double quotes
    {
        let expected = vec!["echo".to_string(), "helloworld".into()];
        let result = parse_args(&String::from("echo \"hello\"world"));
        assert_eq!(result, expected);
    }

    // allow double quotes inside single quotes
    {
        let expected = vec!["\"".to_string()];
        let result = parse_args(&String::from("'\"'"));
        assert_eq!(result, expected);
    }

    // allow single quotes inside double quotes
    {
        let expected = vec!["\'".to_string()];
        let result = parse_args(&String::from("\"\'\""));
        assert_eq!(result, expected);
    }
}
<|endoftext|>