<commit_before>use std::{fmt, mem};
use std::sync::{Arc, Mutex};

use itertools::{Itertools, MultiPeek};

use unicode::{self, UString};

use lang::{Context, Filter};
use lang::context::PrecedenceGroup;
use util::Labeled;

#[derive(Debug)]
pub enum ParseError {
    InvalidToken(char),
    MismatchedParens(Token, Tf),
    NotAllowed(Filter),
    NotFullyParsed(Vec<Tf>),
    UnbalancedParen(Token)
}

#[derive(Debug)]
pub enum Token {
    /// An unrecognized character
    Invalid(char),
    /// An opening parenthesis `(`
    OpenParen,
    /// A closing parenthesis `)`
    CloseParen,
    /// The sequential execution operator `;;`, and all following code
    AndThen(Code),
    /// A sequence of one or more whitespace characters
    Whitespace
}

//#[derive(Debug)] // https://github.com/bluss/rust-itertools/issues/32
enum CodeVariant {
    Empty,
    UString { s: UString, peek_index: usize },
    UStringIter(MultiPeek<unicode::IntoIter>),
    Mutation
}

impl fmt::Debug for CodeVariant { // https://github.com/bluss/rust-itertools/issues/32
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            CodeVariant::Empty => try!(write!(w, "CodeVariant::Empty")),
            CodeVariant::UString { ref s, ref peek_index } => try!(write!(w, "CodeVariant::UString {{ s: {:?}, peek_index: {:?} }}", s, peek_index)),
            CodeVariant::UStringIter(_) => try!(write!(w, "CodeVariant::UStringIter(/* ... */)")),
            CodeVariant::Mutation => try!(write!(w, "CodeVariant::Mutation"))
        }
        Ok(())
    }
}

#[derive(Debug)]
pub struct Code(Mutex<CodeVariant>);

impl Code {
    fn peek(&mut self) -> Option<char> {
        let mut lock = self.0.lock().unwrap();
        match *&mut *lock {
            CodeVariant::Empty => None,
            CodeVariant::UString { ref s, ref mut peek_index } => {
                if *peek_index < s.len() {
                    *peek_index += 1;
                    Some(s[*peek_index - 1])
                } else {
                    None
                }
            }
            CodeVariant::UStringIter(ref mut it) => {
                it.peek().map(|&c| c)
            }
            CodeVariant::Mutation => panic!("code mutex has been emptied")
        }
    }
}

impl Default for Code {
    fn default() -> Code {
        Code(Mutex::new(CodeVariant::Empty))
    }
}

impl<T: Into<UString>> From<T> for Code {
    fn from(code_string: T) -> Code {
        Code(Mutex::new(CodeVariant::UString { s: code_string.into(), peek_index: 0 }))
    }
}

impl Clone for Code {
    fn clone(&self) -> Code {
        let mut lock = self.0.lock().unwrap();
        match *&mut *lock {
            CodeVariant::Empty => Code(Mutex::new(CodeVariant::Empty)),
            CodeVariant::UString { ref s, .. } => Code(Mutex::new(CodeVariant::UString { s: s.clone(), peek_index: 0 })),
            ref mut code_variant @ CodeVariant::UStringIter(_) => {
                if let CodeVariant::UStringIter(it) = mem::replace(code_variant, CodeVariant::Mutation) {
                    let s = it.collect::<UString>();
                    *code_variant = CodeVariant::UString { s: s.clone(), peek_index: 0 };
                    Code(Mutex::new(CodeVariant::UString { s: s, peek_index: 0 }))
                } else {
                    unreachable!()
                }
            }
            CodeVariant::Mutation => panic!("code mutex has been emptied")
        }
    }
}

impl Iterator for Code {
    type Item = char;

    fn next(&mut self) -> Option<char> {
        let mut lock = self.0.lock().unwrap();
        match *&mut *lock {
            CodeVariant::Empty => None,
            ref mut code_variant @ CodeVariant::UString { .. } => {
                if let CodeVariant::UString { s, .. } = mem::replace(code_variant, CodeVariant::Mutation) {
                    let mut iter = s.into_iter().multipeek();
                    let result = iter.next();
                    *code_variant = CodeVariant::UStringIter(iter);
                    result
                } else {
                    unreachable!()
                }
            }
            CodeVariant::UStringIter(ref mut iter) => iter.next(),
            CodeVariant::Mutation => panic!("code mutex has been emptied")
        }
    }
}

struct Tokens {
    code: Code,
    // context: Context
}

impl Tokens {
    fn new<T: Into<Code>>(code: T, _: Context) -> Tokens {
        Tokens {
            code: code.into(),
            // context: context
        }
    }
}

impl Iterator for Tokens {
    type Item = Token;

    fn next(&mut self) -> Option<Token> {
        use self::Token::*;

        match self.code.next() {
            Some(' ') => Some(Whitespace),
            Some('(') => Some(OpenParen),
            Some(')') => Some(CloseParen),
            Some(';') => {
                if self.code.peek() == Some(';') {
                    self.code.next(); // discard the second semicolon
                    Some(AndThen(mem::replace(&mut self.code, Code::default())))
                } else {
                    Some(Invalid(';'))
                }
            }
            Some(c) => Some(Invalid(c)),
            None => None
        }
    }
}

/// A token or filter, used by the in-place parsing algorithm.
#[derive(Debug)]
pub enum Tf {
    Token(Token),
    Filter(Filter)
}

/// Convert a sequence of tokens into an executable filter.
pub fn parse<T: Into<Code>>(code: T, context: Context) -> Result<Filter, ParseError> {
    parse_inner(Tokens::new(code, context.clone()).map(Tf::Token), context)
}

fn parse_inner<I: IntoIterator<Item = Tf>>(tf_iter: I, context: Context) -> Result<Filter, ParseError> {
    let mut tf = tf_iter.into_iter().collect::<Vec<_>>(); // the list of tokens and filters on which the in-place parsing algorithm operates
    // error if any invalid token is found
    if let Some(pos) = tf.iter().position(|i| if let Tf::Token(Token::Invalid(_)) = *i { true } else { false }) {
        if let Tf::Token(Token::Invalid(c)) = tf[pos] {
            return Err(ParseError::InvalidToken(c));
        } else {
            unreachable!();
        }
    }
    // define the macro used for testing if filters are allowed
    macro_rules! try_filter {
        ($f:expr) => {
            match $f {
                f => {
                    if (context.filter_allowed)(&f) {
                        f
                    } else {
                        return Err(ParseError::NotAllowed(f));
                    }
                }
            }
        }
    }
    // remove leading and trailing whitespace as it is semantically irrelevant
    while let Some(&Tf::Token(Token::Whitespace)) = tf.first() { tf.remove(0); }
    while let Some(&Tf::Token(Token::Whitespace)) = tf.last() { tf.pop(); }
    // return an empty filter if the token list is empty
    if tf.len() == 0 { return Ok(try_filter!(Filter::Empty)); }
    // parse operators in decreasing precedence
    for (_, precedence_group) in context.operators.clone().into_iter().rev() { // iterate from highest to lowest precedence
        match precedence_group {
            PrecedenceGroup::AndThen => {
                let mut found = None; // flag any AndThen tokens and remember their contents
                for idx in (0..tf.len()).rev() { // iterate right-to-left for in-place manipulation
                    if let Some(remaining_code) = mem::replace(&mut found, None) {
                        if let Tf::Token(Token::Whitespace) = tf[idx] {
                            // ignore whitespace between `;;` and its left operand
                            tf.remove(idx);
                            continue;
                        }
                        tf[idx] = Tf::Filter(try_filter!(Filter::AndThen {
                            lhs: Box::new(if let Tf::Filter(ref lhs) = tf[idx] { lhs.clone() } else { try_filter!(Filter::Empty) }),
                            remaining_code: remaining_code
                        }));
                    } else {
                        match tf.remove(idx) {
                            Tf::Token(Token::AndThen(remaining_code)) => {
                                found = Some(remaining_code); // found an AndThen (`;;`), will be merged into a syntax tree with the element to its left
                            }
                            tf_item => {
                                tf.insert(idx, tf_item);
                            }
                        }
                    }
                }
                if let Some(remaining_code) = found {
                    // the code begins with an `;;`
                    tf.insert(0, Tf::Filter(try_filter!(Filter::AndThen {
                        lhs: Box::new(try_filter!(Filter::Empty)),
                        remaining_code: remaining_code
                    })));
                }
            }
            PrecedenceGroup::Circumfix => {
                let mut paren_balance = 0; // how many closing parens have not been matched by opening parens
                let mut paren_start = None; // the index of the outermost closing paren
                for idx in (0..tf.len()).rev() { // iterate right-to-left for in-place manipulation
                    match tf[idx] {
                        Tf::Token(Token::CloseParen) => {
                            if paren_balance == 0 {
                                paren_start = Some(idx);
                            }
                            paren_balance += 1;
                        }
                        Tf::Token(Token::OpenParen) => {
                            paren_balance -= 1;
                            if paren_balance < 0 {
                                return Err(ParseError::UnbalancedParen(Token::OpenParen));
                            } else if paren_balance == 0 {
                                if let Some(paren_start) = paren_start {
                                    if let Tf::Token(Token::CloseParen) = tf[paren_start] {
                                        tf.remove(paren_start);
                                        //let inner = tf.drain(idx + 1..paren_start).collect::<Vec<_>>(); //TODO use this when stabilized
                                        let mut inner = vec![];
                                        for _ in idx + 1..paren_start {
                                            inner.push(tf.remove(idx + 1));
                                        }
                                        tf[idx] = Tf::Filter(try_filter!(Filter::Custom {
                                            attributes: vec![try!(parse_inner(inner, context.clone()))],
                                            run: Box::new(Labeled::new("<filter group (α)>", Arc::new(|attrs, input, output| {
                                                assert_eq!(attrs.len(), 1);
                                                attrs[0].run(input, output)
                                            })))
                                        }));
                                    } else {
                                        return Err(ParseError::MismatchedParens(Token::OpenParen, tf.remove(paren_start)));
                                    }
                                } else {
                                    unreachable!();
                                }
                                paren_start = None;
                            }
                        }
                        _ => { continue; }
                    }
                }
                if paren_balance > 0 {
                    if let Some(paren_start) = paren_start {
                        if let Tf::Token(Token::CloseParen) = tf[paren_start] {
                            return Err(ParseError::UnbalancedParen(Token::CloseParen));
                        } else {
                            unreachable!();
                        }
                    } else {
                        unreachable!();
                    }
                }
            }
        }
    }
    if tf.len() == 1 {
        match tf.pop() {
            Some(Tf::Filter(result)) => Ok(result),
            Some(Tf::Token(token)) => Err(ParseError::NotFullyParsed(vec![Tf::Token(token)])),
            None => unreachable!()
        }
    } else {
        Err(ParseError::NotFullyParsed(tf))
    }
}
<commit_msg>Allow tabs, newlines, and carriage returns as whitespace<commit_after>use std::{fmt, mem};
use std::sync::{Arc, Mutex};

use itertools::{Itertools, MultiPeek};

use unicode::{self, UString};

use lang::{Context, Filter};
use lang::context::PrecedenceGroup;
use util::Labeled;

#[derive(Debug)]
pub enum ParseError {
    InvalidToken(char),
    MismatchedParens(Token, Tf),
    NotAllowed(Filter),
    NotFullyParsed(Vec<Tf>),
    UnbalancedParen(Token)
}

#[derive(Debug)]
pub enum Token {
    /// An unrecognized character
    Invalid(char),
    /// An opening parenthesis `(`
    OpenParen,
    /// A closing parenthesis `)`
    CloseParen,
    /// The sequential execution operator `;;`, and all following code
    AndThen(Code),
    /// A sequence of one or more whitespace characters
    Whitespace
}

//#[derive(Debug)] // https://github.com/bluss/rust-itertools/issues/32
enum CodeVariant {
    Empty,
    UString { s: UString, peek_index: usize },
    UStringIter(MultiPeek<unicode::IntoIter>),
    Mutation
}

impl fmt::Debug for CodeVariant { // https://github.com/bluss/rust-itertools/issues/32
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            CodeVariant::Empty => try!(write!(w, "CodeVariant::Empty")),
            CodeVariant::UString { ref s, ref peek_index } => try!(write!(w, "CodeVariant::UString {{ s: {:?}, peek_index: {:?} }}", s, peek_index)),
            CodeVariant::UStringIter(_) => try!(write!(w, "CodeVariant::UStringIter(/* ... */)")),
            CodeVariant::Mutation => try!(write!(w, "CodeVariant::Mutation"))
        }
        Ok(())
    }
}

#[derive(Debug)]
pub struct Code(Mutex<CodeVariant>);

impl Code {
    fn peek(&mut self) -> Option<char> {
        let mut lock = self.0.lock().unwrap();
        match *&mut *lock {
            CodeVariant::Empty => None,
            CodeVariant::UString { ref s, ref mut peek_index } => {
                if *peek_index < s.len() {
                    *peek_index += 1;
                    Some(s[*peek_index - 1])
                } else {
                    None
                }
            }
            CodeVariant::UStringIter(ref mut it) => {
                it.peek().map(|&c| c)
            }
            CodeVariant::Mutation => panic!("code mutex has been emptied")
        }
    }
}

impl Default for Code {
    fn default() -> Code {
        Code(Mutex::new(CodeVariant::Empty))
    }
}

impl<T: Into<UString>> From<T> for Code {
    fn from(code_string: T) -> Code {
        Code(Mutex::new(CodeVariant::UString { s: code_string.into(), peek_index: 0 }))
    }
}

impl Clone for Code {
    fn clone(&self) -> Code {
        let mut lock = self.0.lock().unwrap();
        match *&mut *lock {
            CodeVariant::Empty => Code(Mutex::new(CodeVariant::Empty)),
            CodeVariant::UString { ref s, .. } => Code(Mutex::new(CodeVariant::UString { s: s.clone(), peek_index: 0 })),
            ref mut code_variant @ CodeVariant::UStringIter(_) => {
                if let CodeVariant::UStringIter(it) = mem::replace(code_variant, CodeVariant::Mutation) {
                    let s = it.collect::<UString>();
                    *code_variant = CodeVariant::UString { s: s.clone(), peek_index: 0 };
                    Code(Mutex::new(CodeVariant::UString { s: s, peek_index: 0 }))
                } else {
                    unreachable!()
                }
            }
            CodeVariant::Mutation => panic!("code mutex has been emptied")
        }
    }
}

impl Iterator for Code {
    type Item = char;

    fn next(&mut self) -> Option<char> {
        let mut lock = self.0.lock().unwrap();
        match *&mut *lock {
            CodeVariant::Empty => None,
            ref mut code_variant @ CodeVariant::UString { .. } => {
                if let CodeVariant::UString { s, .. } = mem::replace(code_variant, CodeVariant::Mutation) {
                    let mut iter = s.into_iter().multipeek();
                    let result = iter.next();
                    *code_variant = CodeVariant::UStringIter(iter);
                    result
                } else {
                    unreachable!()
                }
            }
            CodeVariant::UStringIter(ref mut iter) => iter.next(),
            CodeVariant::Mutation => panic!("code mutex has been emptied")
        }
    }
}

struct Tokens {
    code: Code,
    // context: Context
}

impl Tokens {
    fn new<T: Into<Code>>(code: T, _: Context) -> Tokens {
        Tokens {
            code: code.into(),
            // context: context
        }
    }
}

impl Iterator for Tokens {
    type Item = Token;

    fn next(&mut self) -> Option<Token> {
        use self::Token::*;

        match self.code.next() {
            Some('\t') |
            Some('\n') |
            Some('\r') |
            Some(' ') => Some(Whitespace),
            Some('(') => Some(OpenParen),
            Some(')') => Some(CloseParen),
            Some(';') => {
                if self.code.peek() == Some(';') {
                    self.code.next(); // discard the second semicolon
                    Some(AndThen(mem::replace(&mut self.code, Code::default())))
                } else {
                    Some(Invalid(';'))
                }
            }
            Some(c) => Some(Invalid(c)),
            None => None
        }
    }
}

/// A token or filter, used by the in-place parsing algorithm.
#[derive(Debug)]
pub enum Tf {
    Token(Token),
    Filter(Filter)
}

/// Convert a sequence of tokens into an executable filter.
pub fn parse<T: Into<Code>>(code: T, context: Context) -> Result<Filter, ParseError> {
    parse_inner(Tokens::new(code, context.clone()).map(Tf::Token), context)
}

fn parse_inner<I: IntoIterator<Item = Tf>>(tf_iter: I, context: Context) -> Result<Filter, ParseError> {
    let mut tf = tf_iter.into_iter().collect::<Vec<_>>(); // the list of tokens and filters on which the in-place parsing algorithm operates
    // error if any invalid token is found
    if let Some(pos) = tf.iter().position(|i| if let Tf::Token(Token::Invalid(_)) = *i { true } else { false }) {
        if let Tf::Token(Token::Invalid(c)) = tf[pos] {
            return Err(ParseError::InvalidToken(c));
        } else {
            unreachable!();
        }
    }
    // define the macro used for testing if filters are allowed
    macro_rules! try_filter {
        ($f:expr) => {
            match $f {
                f => {
                    if (context.filter_allowed)(&f) {
                        f
                    } else {
                        return Err(ParseError::NotAllowed(f));
                    }
                }
            }
        }
    }
    // remove leading and trailing whitespace as it is semantically irrelevant
    while let Some(&Tf::Token(Token::Whitespace)) = tf.first() { tf.remove(0); }
    while let Some(&Tf::Token(Token::Whitespace)) = tf.last() { tf.pop(); }
    // return an empty filter if the token list is empty
    if tf.len() == 0 { return Ok(try_filter!(Filter::Empty)); }
    // parse operators in decreasing precedence
    for (_, precedence_group) in context.operators.clone().into_iter().rev() { // iterate from highest to lowest precedence
        match precedence_group {
            PrecedenceGroup::AndThen => {
                let mut found = None; // flag any AndThen tokens and remember their contents
                for idx in (0..tf.len()).rev() { // iterate right-to-left for in-place manipulation
                    if let Some(remaining_code) = mem::replace(&mut found, None) {
                        if let Tf::Token(Token::Whitespace) = tf[idx] {
                            // ignore whitespace between `;;` and its left operand
                            tf.remove(idx);
                            continue;
                        }
                        tf[idx] = Tf::Filter(try_filter!(Filter::AndThen {
                            lhs: Box::new(if let Tf::Filter(ref lhs) = tf[idx] { lhs.clone() } else { try_filter!(Filter::Empty) }),
                            remaining_code: remaining_code
                        }));
                    } else {
                        match tf.remove(idx) {
                            Tf::Token(Token::AndThen(remaining_code)) => {
                                found = Some(remaining_code); // found an AndThen (`;;`), will be merged into a syntax tree with the element to its left
                            }
                            tf_item => {
                                tf.insert(idx, tf_item);
                            }
                        }
                    }
                }
                if let Some(remaining_code) = found {
                    // the code begins with an `;;`
                    tf.insert(0, Tf::Filter(try_filter!(Filter::AndThen {
                        lhs: Box::new(try_filter!(Filter::Empty)),
                        remaining_code: remaining_code
                    })));
                }
            }
            PrecedenceGroup::Circumfix => {
                let mut paren_balance = 0; // how many closing parens have not been matched by opening parens
                let mut paren_start = None; // the index of the outermost closing paren
                for idx in (0..tf.len()).rev() { // iterate right-to-left for in-place manipulation
                    match tf[idx] {
                        Tf::Token(Token::CloseParen) => {
                            if paren_balance == 0 {
                                paren_start = Some(idx);
                            }
                            paren_balance += 1;
                        }
                        Tf::Token(Token::OpenParen) => {
                            paren_balance -= 1;
                            if paren_balance < 0 {
                                return Err(ParseError::UnbalancedParen(Token::OpenParen));
                            } else if paren_balance == 0 {
                                if let Some(paren_start) = paren_start {
                                    if let Tf::Token(Token::CloseParen) = tf[paren_start] {
                                        tf.remove(paren_start);
                                        //let inner = tf.drain(idx + 1..paren_start).collect::<Vec<_>>(); //TODO use this when stabilized
                                        let mut inner = vec![];
                                        for _ in idx + 1..paren_start {
                                            inner.push(tf.remove(idx + 1));
                                        }
                                        tf[idx] = Tf::Filter(try_filter!(Filter::Custom {
                                            attributes: vec![try!(parse_inner(inner, context.clone()))],
                                            run: Box::new(Labeled::new("<filter group (α)>", Arc::new(|attrs, input, output| {
                                                assert_eq!(attrs.len(), 1);
                                                attrs[0].run(input, output)
                                            })))
                                        }));
                                    } else {
                                        return Err(ParseError::MismatchedParens(Token::OpenParen, tf.remove(paren_start)));
                                    }
                                } else {
                                    unreachable!();
                                }
                                paren_start = None;
                            }
                        }
                        _ => { continue; }
                    }
                }
                if paren_balance > 0 {
                    if let Some(paren_start) = paren_start {
                        if let Tf::Token(Token::CloseParen) = tf[paren_start] {
                            return Err(ParseError::UnbalancedParen(Token::CloseParen));
                        } else {
                            unreachable!();
                        }
                    } else {
                        unreachable!();
                    }
                }
            }
        }
    }
    if tf.len() == 1 {
        match tf.pop() {
            Some(Tf::Filter(result)) => Ok(result),
            Some(Tf::Token(token)) => Err(ParseError::NotFullyParsed(vec![Tf::Token(token)])),
            None => unreachable!()
        }
    } else {
        Err(ParseError::NotFullyParsed(tf))
    }
}
<|endoftext|>


<commit_before>//! Main module to handle the layout.
//! This is where the i3-specific code is.

use std::sync::{Mutex, MutexGuard, TryLockError};
use std::ptr;

use super::container::{Container, Handle, ContainerType};
use super::node::{Node};
use super::super::rustwlc::handle::{WlcView, WlcOutput};

pub type TreeErr = TryLockError<MutexGuard<'static, Tree>>;
pub type TreeResult = Result<MutexGuard<'static, Tree>, TreeErr>;

const ERR_BAD_TREE: &'static str = "Layout tree was in an invalid configuration";

lazy_static! {
    static ref TREE: Mutex<Tree> = {
        Mutex::new(Tree{
            root: Node::new(Container::new_root()),
            active_container: ptr::null(),
        })
    };
}

pub struct Tree {
    root: Node,
    active_container: *const Node,
}

unsafe impl Send for Tree {}

impl Tree {
    /// Switch to the workspace with the give name
    pub fn switch_workspace(&mut self, name: &str ) {
        trace!("Switching to workspace {}", name);
        if let Some(old_workspace) = self.get_active_workspace() {
            old_workspace.set_visibility(false);
        }
        let current_workspace: *const Node;
        {
            if let Some(_) = self.get_workspace_by_name(name) {
                trace!("Found workspace {}", name);
            } else {
                trace!("Adding workspace {}", name);
                self.add_workspace(name.to_string());
            }
            let new_current_workspace = self.get_workspace_by_name_mut(name).expect(ERR_BAD_TREE);
            new_current_workspace.set_visibility(true);
            // Set the first view to be focused, so the screen refreshes itself
            if new_current_workspace.get_children()[0].get_children().len() > 0 {
                trace!("Focusing view");
                match new_current_workspace.get_children_mut()[0].get_children_mut()[0]
                    .get_val().get_handle().expect(ERR_BAD_TREE) {
                    Handle::View(view) => view.focus(),
                    _ => {},
                }
            } else {
                WlcView::root().focus();
            }
            // Update the tree's pointer to the currently focused container
            current_workspace = &new_current_workspace.get_children()[0] as *const Node;
        }
        self.active_container = current_workspace;
    }

    /// Returns the currently viewed container.
    /// If multiple views are selected, the parent container they share is returned
    pub fn get_active_container(&self) -> Option<&Node> {
        if self.active_container.is_null() {
            None
        } else {
            unsafe {
                Some(&*self.active_container)
            }
        }
    }

    /// Get the monitor (output) that the active container is located on
    pub fn get_active_output(&self) -> Option<&mut Node> {
        if let Some(node) = self.get_active_container() {
            node.get_ancestor_of_type(ContainerType::Output)
        } else {
            None
        }
    }

    /// Get the workspace that the active container is located on
    pub fn get_active_workspace(&self) -> Option<&mut Node> {
        if let Some(container) = self.get_active_container() {
            if let Some(workspace) = container.get_ancestor_of_type(ContainerType::Workspace) {
                return Some(workspace)
            }
        }
        return None
    }

    /// Find the workspace node that has the given name
    pub fn get_workspace_by_name(&self, name: &str) -> Option<&Node> {
        for child in self.root.get_children()[0].get_children() {
            if child.get_val().get_name().expect(ERR_BAD_TREE) != name {
                continue
            }
            return Some(child);
        }
        return None
    }

    /// Find the workspace node that has the given name, with a mutable reference
    pub fn get_workspace_by_name_mut(&mut self, name: &str) -> Option<&mut Node> {
        for child in self.root.get_children_mut()[0].get_children_mut() {
            if child.get_val().get_name().expect(ERR_BAD_TREE) != name {
                continue
            }
            return Some(child);
        }
        return None
    }

    /// Make a new output container with the given WlcOutput.
    /// This is done when a new monitor is added
    pub fn add_output(&mut self, wlc_output: WlcOutput) {
        match self.root.new_child(Container::new_output(wlc_output)) {
            Ok(_) => {},
            Err(e) => error!("Could not add output: {:?}", e),
        }
    }

    /// Make a new workspace container with the given name.
    pub fn add_workspace(&mut self, name: String) {
        let workspace = Container::new_workspace(name.to_string());
        let mut index = 0;
        if let Some(output) = self.get_active_output() {
            for (cur_index, child) in self.root.get_children().iter().enumerate() {
                if child == output {
                    index = cur_index;
                    break;
                }
            }
        }
        match self.root.get_children_mut()[index].new_child(workspace) {
            Ok(workspace) => {
                trace!("Added workspace {:?}", workspace);
                if let Err(e) = workspace.new_child(Container::new_container()) {
                    error!("Could not add container to workspace: {:?}",e );
                }
            },
            Err(e) => error!("Could not add workspace: {:?}", e),
        }
    }

    /// Make a new view container with the given WlcView, and adds it to
    /// the active workspace.
    pub fn add_view(&mut self, wlc_view: WlcView) {
        let mut maybe_new_view: *const Node = ptr::null();
        if let Some(current_workspace) = self.get_active_workspace() {
            trace!("Adding view {:?} to {:?}", wlc_view, current_workspace);
            if current_workspace.get_children().len() == 0 {
                match current_workspace.new_child(Container::new_container()) {
                    Ok(_) => {},
                    Err(e) => {
                        error!("Could not add workspace: {:?}", e);
                        return;
                    }
                }
            }
            let container = &mut current_workspace.get_children_mut()[0];
            match container.new_child(Container::new_view(wlc_view)) {
                Ok(view_node) => maybe_new_view = view_node as *const Node,
                Err(e) => {
                    error!("Could not add view to current workspace: {:?}", e);
                    return;
                }
            }
        };
        if ! maybe_new_view.is_null() {
            self.active_container = maybe_new_view as *const Node;
        }
    }

    /// Remove the view container with the given view
    pub fn remove_view(&self, wlc_view: &WlcView) {
        if let Some(view) = self.root.find_view_by_handle(&wlc_view) {
            let parent = view.get_parent().expect(ERR_BAD_TREE);
            parent.remove_child(view);
        }
    }
}

pub fn try_lock_tree() -> TreeResult {
    trace!("Locking the tree!");
    TREE.try_lock()
}
<commit_msg>Removed unnecessary iteration to find the output<commit_after>//! Main module to handle the layout.
//! This is where the i3-specific code is.

use std::sync::{Mutex, MutexGuard, TryLockError};
use std::ptr;

use super::container::{Container, Handle, ContainerType};
use super::node::{Node};
use super::super::rustwlc::handle::{WlcView, WlcOutput};

pub type TreeErr = TryLockError<MutexGuard<'static, Tree>>;
pub type TreeResult = Result<MutexGuard<'static, Tree>, TreeErr>;

const ERR_BAD_TREE: &'static str = "Layout tree was in an invalid configuration";

lazy_static! {
    static ref TREE: Mutex<Tree> = {
        Mutex::new(Tree{
            root: Node::new(Container::new_root()),
            active_container: ptr::null(),
        })
    };
}

pub struct Tree {
    root: Node,
    active_container: *const Node,
}

unsafe impl Send for Tree {}

impl Tree {
    /// Switch to the workspace with the give name
    pub fn switch_workspace(&mut self, name: &str ) {
        trace!("Switching to workspace {}", name);
        if let Some(old_workspace) = self.get_active_workspace() {
            old_workspace.set_visibility(false);
        }
        let current_workspace: *const Node;
        {
            if let Some(_) = self.get_workspace_by_name(name) {
                trace!("Found workspace {}", name);
            } else {
                trace!("Adding workspace {}", name);
                self.add_workspace(name.to_string());
            }
            let new_current_workspace = self.get_workspace_by_name_mut(name).expect(ERR_BAD_TREE);
            new_current_workspace.set_visibility(true);
            // Set the first view to be focused, so the screen refreshes itself
            if new_current_workspace.get_children()[0].get_children().len() > 0 {
                trace!("Focusing view");
                match new_current_workspace.get_children_mut()[0].get_children_mut()[0]
                    .get_val().get_handle().expect(ERR_BAD_TREE) {
                    Handle::View(view) => view.focus(),
                    _ => {},
                }
            } else {
                WlcView::root().focus();
            }
            // Update the tree's pointer to the currently focused container
            current_workspace = &new_current_workspace.get_children()[0] as *const Node;
        }
        self.active_container = current_workspace;
    }

    /// Returns the currently viewed container.
    /// If multiple views are selected, the parent container they share is returned
    pub fn get_active_container(&self) -> Option<&Node> {
        if self.active_container.is_null() {
            None
        } else {
            unsafe {
                Some(&*self.active_container)
            }
        }
    }

    /// Get the monitor (output) that the active container is located on
    pub fn get_active_output(&self) -> Option<&mut Node> {
        if let Some(node) = self.get_active_container() {
            node.get_ancestor_of_type(ContainerType::Output)
        } else {
            None
        }
    }

    /// Get the workspace that the active container is located on
    pub fn get_active_workspace(&self) -> Option<&mut Node> {
        if let Some(container) = self.get_active_container() {
            if let Some(workspace) = container.get_ancestor_of_type(ContainerType::Workspace) {
                return Some(workspace)
            }
        }
        return None
    }

    /// Find the workspace node that has the given name
    pub fn get_workspace_by_name(&self, name: &str) -> Option<&Node> {
        for child in self.root.get_children()[0].get_children() {
            if child.get_val().get_name().expect(ERR_BAD_TREE) != name {
                continue
            }
            return Some(child);
        }
        return None
    }

    /// Find the workspace node that has the given name, with a mutable reference
    pub fn get_workspace_by_name_mut(&mut self, name: &str) -> Option<&mut Node> {
        for child in self.root.get_children_mut()[0].get_children_mut() {
            if child.get_val().get_name().expect(ERR_BAD_TREE) != name {
                continue
            }
            return Some(child);
        }
        return None
    }

    /// Make a new output container with the given WlcOutput.
    /// This is done when a new monitor is added
    pub fn add_output(&mut self, wlc_output: WlcOutput) {
        match self.root.new_child(Container::new_output(wlc_output)) {
            Ok(_) => {},
            Err(e) => error!("Could not add output: {:?}", e),
        }
    }

    /// Make a new workspace container with the given name.
    pub fn add_workspace(&mut self, name: String) {
        let workspace = Container::new_workspace(name.to_string());
        if let Some(output) = self.get_active_output() {
            match output.new_child(workspace) {
                Ok(workspace) => {
                    trace!("Added workspace {:?}", workspace);
                    if let Err(e) = workspace.new_child(Container::new_container()) {
                        error!("Could not add container to workspace: {:?}",e );
                    }
                },
                Err(e) => error!("Could not add workspace: {:?}", e),
            }
        }
    }

    /// Make a new view container with the given WlcView, and adds it to
    /// the active workspace.
    pub fn add_view(&mut self, wlc_view: WlcView) {
        let mut maybe_new_view: *const Node = ptr::null();
        if let Some(current_workspace) = self.get_active_workspace() {
            trace!("Adding view {:?} to {:?}", wlc_view, current_workspace);
            if current_workspace.get_children().len() == 0 {
                match current_workspace.new_child(Container::new_container()) {
                    Ok(_) => {},
                    Err(e) => {
                        error!("Could not add workspace: {:?}", e);
                        return;
                    }
                }
            }
            let container = &mut current_workspace.get_children_mut()[0];
            match container.new_child(Container::new_view(wlc_view)) {
                Ok(view_node) => maybe_new_view = view_node as *const Node,
                Err(e) => {
                    error!("Could not add view to current workspace: {:?}", e);
                    return;
                }
            }
        };
        if ! maybe_new_view.is_null() {
            self.active_container = maybe_new_view as *const Node;
        }
    }

    /// Remove the view container with the given view
    pub fn remove_view(&self, wlc_view: &WlcView) {
        if let Some(view) = self.root.find_view_by_handle(&wlc_view) {
            let parent = view.get_parent().expect(ERR_BAD_TREE);
            parent.remove_child(view);
        }
    }
}

pub fn try_lock_tree() -> TreeResult {
    trace!("Locking the tree!");
    TREE.try_lock()
}
<|endoftext|>


<commit_before>//! GPT-header object and helper functions.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use crc::{crc32, Hasher32};
use std::fmt;
use std::fs::{File, OpenOptions};
use std::io::{Cursor, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::path::Path;
use uuid;
use log::*;

use crate::disk;
use crate::partition;

/// Header describing a GPT disk.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Header {
    /// GPT header magic signature, hardcoded to "EFI PART".
    pub signature: String, // Offset  0. "EFI PART", 45h 46h 49h 20h 50h 41h 52h 54h
    /// 00 00 01 00
    pub revision: u32, // Offset  8
    /// little endian
    pub header_size_le: u32, // Offset 12
    /// CRC32 of the header with crc32 section zeroed
    pub crc32: u32, // Offset 16
    /// must be 0
    pub reserved: u32, // Offset 20
    /// For main header, 1
    pub current_lba: u64, // Offset 24
    /// LBA for backup header
    pub backup_lba: u64, // Offset 32
    /// First usable LBA for partitions (primary table last LBA + 1)
    pub first_usable: u64, // Offset 40
    /// Last usable LBA (secondary partition table first LBA - 1)
    pub last_usable: u64, // Offset 48
    /// UUID of the disk
    pub disk_guid: uuid::Uuid, // Offset 56
    /// Starting LBA of partition entries
    pub part_start: u64, // Offset 72
    /// Number of partition entries
    pub num_parts: u32, // Offset 80
    /// Size of a partition entry, usually 128
    pub part_size: u32, // Offset 84
    /// CRC32 of the partition table
    pub crc32_parts: u32, // Offset 88
}

impl Header {
    pub(crate) fn compute_new(
        primary: bool,
        pp: &[partition::Partition],
        guid: uuid::Uuid,
        backup_offset: u64,
    ) -> Result<Self> {
        let (cur, bak) = if primary {
            (1, backup_offset)
        } else {
            (backup_offset, 1)
        };
        let first = 34u64;
        let last = backup_offset
            .checked_sub(first)
            .ok_or_else(|| Error::new(ErrorKind::Other, "header underflow - last usable"))?;

        let hdr = Header {
            signature: "EFI PART".to_string(),
            revision: 65536,
            header_size_le: 92,
            crc32: 0,
            reserved: 0,
            current_lba: cur,
            backup_lba: bak,
            first_usable: first,
            last_usable: last,
            disk_guid: guid,
            part_start: 2,
            num_parts: pp.len() as u32,
            part_size: 128,
            crc32_parts: 0,
        };

        Ok(hdr)
    }

    /// Write the primary header.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    pub fn write_primary(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        // This is the primary header. It must start before the backup one.
        if self.current_lba >= self.backup_lba {
            return Err(Error::new(
                ErrorKind::Other,
                "primary header does not start before backup one",
            ));
        }
        self.file_write_header(file, self.current_lba, lb_size)
    }

    /// Write the backup header.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    pub fn write_backup(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        // This is the backup header. It must start after the primary one.
        if self.current_lba <= self.backup_lba {
            return Err(Error::new(
                ErrorKind::Other,
                "backup header does not start after primary one",
            ));
        }
        self.file_write_header(file, self.current_lba, lb_size)
    }

    /// Write an header to an arbitrary LBA.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    fn file_write_header(
        &self,
        file: &mut File,
        lba: u64,
        lb_size: disk::LogicalBlockSize,
    ) -> Result<usize> {
        // Build up byte array in memory
        let parts_checksum = partentry_checksum(file, self, lb_size)?;
        let bytes = self.as_bytes(None, Some(parts_checksum))?;

        // Calculate the CRC32 from the byte array
        let checksum = calculate_crc32(&bytes)?;

        // Write it to disk in 1 shot
        let start = lba.checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - offset"))?;
        let _ = file.seek(SeekFrom::Start(start))?;
        let len = file.write(&self.as_bytes(Some(checksum), Some(parts_checksum))?)?;

        Ok(len)
    }

    fn as_bytes(&self, checksum: Option<u32>, parts_checksum: Option<u32>) -> Result<Vec<u8>> {
        let mut buff: Vec<u8> = Vec::new();

        buff.write_all(self.signature.as_bytes())?;
        buff.write_u32::<LittleEndian>(self.revision)?;
        buff.write_u32::<LittleEndian>(self.header_size_le)?;
        match checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        buff.write_u32::<LittleEndian>(0)?;
        buff.write_u64::<LittleEndian>(self.current_lba)?;
        buff.write_u64::<LittleEndian>(self.backup_lba)?;
        buff.write_u64::<LittleEndian>(self.first_usable)?;
        buff.write_u64::<LittleEndian>(self.last_usable)?;
        buff.write_all(self.disk_guid.as_bytes())?;
        buff.write_u64::<LittleEndian>(self.part_start)?;
        buff.write_u32::<LittleEndian>(self.num_parts)?;
        buff.write_u32::<LittleEndian>(self.part_size)?;
        match parts_checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        Ok(buff)
    }
}

/// Parses a uuid with first 3 portions in little endian.
pub fn parse_uuid(rdr: &mut Cursor<&[u8]>) -> Result<uuid::Uuid> {
    let d1: u32 = rdr.read_u32::<LittleEndian>()?;
    let d2: u16 = rdr.read_u16::<LittleEndian>()?;
    let d3: u16 = rdr.read_u16::<LittleEndian>()?;
    let uuid = uuid::Uuid::from_fields(
        d1,
        d2,
        d3,
        &rdr.get_ref()[rdr.position() as usize..rdr.position() as usize + 8],
    );
    rdr.seek(SeekFrom::Current(8))?;

    match uuid {
        Ok(uuid) => Ok(uuid),
        Err(_) => Err(Error::new(ErrorKind::Other, "Invalid Disk UUID?")),
    }
}

impl fmt::Display for Header {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Disk:\t\t{}\nCRC32:\t\t{}\nTable CRC:\t{}",
            self.disk_guid, self.crc32, self.crc32_parts
        )
    }
}

/// Read a GPT header from a given path.
///
/// ## Example
///
/// ```rust,no_run
/// use gpt::header::read_header;
///
/// let lb_size = gpt::disk::DEFAULT_SECTOR_SIZE;
/// let diskpath = std::path::Path::new("/dev/sdz");
///
/// let h = read_header(diskpath, lb_size).unwrap();
/// ```
pub fn read_header(path: &Path, sector_size: disk::LogicalBlockSize) -> Result<Header> {
    let mut file = File::open(path)?;
    read_primary_header(&mut file, sector_size)
}

pub(crate) fn read_primary_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let offset: u64 = sector_size.into();
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn read_backup_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let h2sect = find_backup_lba(file, sector_size)?;
    let offset = h2sect
        .checked_mul(sector_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "backup header overflow - offset"))?;
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn file_read_header(file: &mut File, offset: u64) -> Result<Header> {
    let _ = file.seek(SeekFrom::Start(offset));
    let mut hdr: [u8; 92] = [0; 92];

    let _ = file.read_exact(&mut hdr);
    let mut reader = Cursor::new(&hdr[..]);

    let sigstr = String::from_utf8_lossy(
        &reader.get_ref()[reader.position() as usize..reader.position() as usize + 8],
    );
    reader.seek(SeekFrom::Current(8))?;

    if sigstr != "EFI PART" {
        return Err(Error::new(ErrorKind::Other, "invalid GPT signature"));
    };

    let h = Header {
        signature: sigstr.to_string(),
        revision: reader.read_u32::<LittleEndian>()?,
        header_size_le: reader.read_u32::<LittleEndian>()?,
        crc32: reader.read_u32::<LittleEndian>()?,
        reserved: reader.read_u32::<LittleEndian>()?,
        current_lba: reader.read_u64::<LittleEndian>()?,
        backup_lba: reader.read_u64::<LittleEndian>()?,
        first_usable: reader.read_u64::<LittleEndian>()?,
        last_usable: reader.read_u64::<LittleEndian>()?,
        disk_guid: parse_uuid(&mut reader)?,
        part_start: reader.read_u64::<LittleEndian>()?,
        num_parts: reader.read_u32::<LittleEndian>()?,
        part_size: reader.read_u32::<LittleEndian>()?,
        crc32_parts: reader.read_u32::<LittleEndian>()?,
    };

    let mut hdr_crc = hdr;
    for crc_byte in hdr_crc.iter_mut().skip(16).take(4) {
        *crc_byte = 0;
    }
    let c = crc32::checksum_ieee(&hdr_crc);
    trace!("header CRC32: {:#x} - computed CRC32: {:#x}", h.crc32, c);
    if crc32::checksum_ieee(&hdr_crc) == h.crc32 {
        Ok(h)
    } else {
        Err(Error::new(ErrorKind::Other, "invalid CRC32 checksum"))
    }
}

pub(crate) fn find_backup_lba(f: &mut File, sector_size: disk::LogicalBlockSize) -> Result<u64> {
    trace!("querying file size to find backup header location");
    let lb_size: u64 = sector_size.into();
    let m = f.metadata()?;
    if m.len() <= lb_size {
        return Err(Error::new(
            ErrorKind::Other,
            "disk image too small for backup header",
        ));
    }
    let bak_offset = m.len().saturating_sub(lb_size);
    let bak_lba = bak_offset / lb_size;
    trace!(
        "backup header: LBA={}, bytes offset={}",
        bak_lba,
        bak_offset
    );

    Ok(bak_lba)
}

fn calculate_crc32(b: &[u8]) -> Result<u32> {
    let mut digest = crc32::Digest::new(crc32::IEEE);
    trace!("Writing buffer to digest calculator");
    digest.write(b);

    Ok(digest.sum32())
}

pub(crate) fn partentry_checksum(
    file: &mut File,
    hdr: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<u32> {
    // Seek to start of partition table.
    let start = hdr.part_start
        .checked_mul(lb_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - partition table start"))?;
    let _ = file.seek(SeekFrom::Start(start))?;

    // Read partition table.
    let pt_len = u64::from(hdr.num_parts)
        .checked_mul(hdr.part_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partition table - size"))?;
    let mut buf = vec![0; pt_len as usize];
    file.read_exact(&mut buf)?;

    // Compute CRC32 over all table bits.
    let mut digest = crc32::Digest::new(crc32::IEEE);
    digest.write(&buf);
    Ok(digest.sum32())
}

/// A helper function to create a new header and write it to disk.
/// If the uuid isn't given a random one will be generated.  Use
/// this in conjunction with Partition::write()
// TODO: Move this to Header::new() and Header::write to write it
// that will match the Partition::write() API
pub fn write_header(
    p: &Path,
    uuid: Option<uuid::Uuid>,
    sector_size: disk::LogicalBlockSize,
) -> Result<uuid::Uuid> {
    debug!("opening {} for writing", p.display());
    let mut file = OpenOptions::new().write(true).read(true).open(p)?;
    let bak = find_backup_lba(&mut file, sector_size)?;
    let guid = match uuid {
        Some(u) => u,
        None => {
            let u = uuid::Uuid::new_v4();
            debug!("Generated random uuid: {}", u);
            u
        }
    };

    let hdr = Header::compute_new(true, &[], guid, bak)?;
    debug!("new header: {:#?}", hdr);
    hdr.write_primary(&mut file, sector_size)?;

    Ok(guid)
}
<commit_msg>header: Determine file length using seek<commit_after>//! GPT-header object and helper functions.

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use crc::{crc32, Hasher32};
use std::fmt;
use std::fs::{File, OpenOptions};
use std::io::{Cursor, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::path::Path;
use uuid;
use log::*;

use crate::disk;
use crate::partition;

/// Header describing a GPT disk.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Header {
    /// GPT header magic signature, hardcoded to "EFI PART".
    pub signature: String, // Offset  0. "EFI PART", 45h 46h 49h 20h 50h 41h 52h 54h
    /// 00 00 01 00
    pub revision: u32, // Offset  8
    /// little endian
    pub header_size_le: u32, // Offset 12
    /// CRC32 of the header with crc32 section zeroed
    pub crc32: u32, // Offset 16
    /// must be 0
    pub reserved: u32, // Offset 20
    /// For main header, 1
    pub current_lba: u64, // Offset 24
    /// LBA for backup header
    pub backup_lba: u64, // Offset 32
    /// First usable LBA for partitions (primary table last LBA + 1)
    pub first_usable: u64, // Offset 40
    /// Last usable LBA (secondary partition table first LBA - 1)
    pub last_usable: u64, // Offset 48
    /// UUID of the disk
    pub disk_guid: uuid::Uuid, // Offset 56
    /// Starting LBA of partition entries
    pub part_start: u64, // Offset 72
    /// Number of partition entries
    pub num_parts: u32, // Offset 80
    /// Size of a partition entry, usually 128
    pub part_size: u32, // Offset 84
    /// CRC32 of the partition table
    pub crc32_parts: u32, // Offset 88
}

impl Header {
    pub(crate) fn compute_new(
        primary: bool,
        pp: &[partition::Partition],
        guid: uuid::Uuid,
        backup_offset: u64,
    ) -> Result<Self> {
        let (cur, bak) = if primary {
            (1, backup_offset)
        } else {
            (backup_offset, 1)
        };
        let first = 34u64;
        let last = backup_offset
            .checked_sub(first)
            .ok_or_else(|| Error::new(ErrorKind::Other, "header underflow - last usable"))?;

        let hdr = Header {
            signature: "EFI PART".to_string(),
            revision: 65536,
            header_size_le: 92,
            crc32: 0,
            reserved: 0,
            current_lba: cur,
            backup_lba: bak,
            first_usable: first,
            last_usable: last,
            disk_guid: guid,
            part_start: 2,
            num_parts: pp.len() as u32,
            part_size: 128,
            crc32_parts: 0,
        };

        Ok(hdr)
    }

    /// Write the primary header.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    pub fn write_primary(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        // This is the primary header. It must start before the backup one.
        if self.current_lba >= self.backup_lba {
            return Err(Error::new(
                ErrorKind::Other,
                "primary header does not start before backup one",
            ));
        }
        self.file_write_header(file, self.current_lba, lb_size)
    }

    /// Write the backup header.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    pub fn write_backup(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        // This is the backup header. It must start after the primary one.
        if self.current_lba <= self.backup_lba {
            return Err(Error::new(
                ErrorKind::Other,
                "backup header does not start after primary one",
            ));
        }
        self.file_write_header(file, self.current_lba, lb_size)
    }

    /// Write an header to an arbitrary LBA.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    fn file_write_header(
        &self,
        file: &mut File,
        lba: u64,
        lb_size: disk::LogicalBlockSize,
    ) -> Result<usize> {
        // Build up byte array in memory
        let parts_checksum = partentry_checksum(file, self, lb_size)?;
        let bytes = self.as_bytes(None, Some(parts_checksum))?;

        // Calculate the CRC32 from the byte array
        let checksum = calculate_crc32(&bytes)?;

        // Write it to disk in 1 shot
        let start = lba.checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - offset"))?;
        let _ = file.seek(SeekFrom::Start(start))?;
        let len = file.write(&self.as_bytes(Some(checksum), Some(parts_checksum))?)?;

        Ok(len)
    }

    fn as_bytes(&self, checksum: Option<u32>, parts_checksum: Option<u32>) -> Result<Vec<u8>> {
        let mut buff: Vec<u8> = Vec::new();

        buff.write_all(self.signature.as_bytes())?;
        buff.write_u32::<LittleEndian>(self.revision)?;
        buff.write_u32::<LittleEndian>(self.header_size_le)?;
        match checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        buff.write_u32::<LittleEndian>(0)?;
        buff.write_u64::<LittleEndian>(self.current_lba)?;
        buff.write_u64::<LittleEndian>(self.backup_lba)?;
        buff.write_u64::<LittleEndian>(self.first_usable)?;
        buff.write_u64::<LittleEndian>(self.last_usable)?;
        buff.write_all(self.disk_guid.as_bytes())?;
        buff.write_u64::<LittleEndian>(self.part_start)?;
        buff.write_u32::<LittleEndian>(self.num_parts)?;
        buff.write_u32::<LittleEndian>(self.part_size)?;
        match parts_checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        Ok(buff)
    }
}

/// Parses a uuid with first 3 portions in little endian.
pub fn parse_uuid(rdr: &mut Cursor<&[u8]>) -> Result<uuid::Uuid> {
    let d1: u32 = rdr.read_u32::<LittleEndian>()?;
    let d2: u16 = rdr.read_u16::<LittleEndian>()?;
    let d3: u16 = rdr.read_u16::<LittleEndian>()?;
    let uuid = uuid::Uuid::from_fields(
        d1,
        d2,
        d3,
        &rdr.get_ref()[rdr.position() as usize..rdr.position() as usize + 8],
    );
    rdr.seek(SeekFrom::Current(8))?;

    match uuid {
        Ok(uuid) => Ok(uuid),
        Err(_) => Err(Error::new(ErrorKind::Other, "Invalid Disk UUID?")),
    }
}

impl fmt::Display for Header {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Disk:\t\t{}\nCRC32:\t\t{}\nTable CRC:\t{}",
            self.disk_guid, self.crc32, self.crc32_parts
        )
    }
}

/// Read a GPT header from a given path.
///
/// ## Example
///
/// ```rust,no_run
/// use gpt::header::read_header;
///
/// let lb_size = gpt::disk::DEFAULT_SECTOR_SIZE;
/// let diskpath = std::path::Path::new("/dev/sdz");
///
/// let h = read_header(diskpath, lb_size).unwrap();
/// ```
pub fn read_header(path: &Path, sector_size: disk::LogicalBlockSize) -> Result<Header> {
    let mut file = File::open(path)?;
    read_primary_header(&mut file, sector_size)
}

pub(crate) fn read_primary_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let offset: u64 = sector_size.into();
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn read_backup_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let h2sect = find_backup_lba(file, sector_size)?;
    let offset = h2sect
        .checked_mul(sector_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "backup header overflow - offset"))?;
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn file_read_header(file: &mut File, offset: u64) -> Result<Header> {
    let _ = file.seek(SeekFrom::Start(offset));
    let mut hdr: [u8; 92] = [0; 92];

    let _ = file.read_exact(&mut hdr);
    let mut reader = Cursor::new(&hdr[..]);

    let sigstr = String::from_utf8_lossy(
        &reader.get_ref()[reader.position() as usize..reader.position() as usize + 8],
    );
    reader.seek(SeekFrom::Current(8))?;

    if sigstr != "EFI PART" {
        return Err(Error::new(ErrorKind::Other, "invalid GPT signature"));
    };

    let h = Header {
        signature: sigstr.to_string(),
        revision: reader.read_u32::<LittleEndian>()?,
        header_size_le: reader.read_u32::<LittleEndian>()?,
        crc32: reader.read_u32::<LittleEndian>()?,
        reserved: reader.read_u32::<LittleEndian>()?,
        current_lba: reader.read_u64::<LittleEndian>()?,
        backup_lba: reader.read_u64::<LittleEndian>()?,
        first_usable: reader.read_u64::<LittleEndian>()?,
        last_usable: reader.read_u64::<LittleEndian>()?,
        disk_guid: parse_uuid(&mut reader)?,
        part_start: reader.read_u64::<LittleEndian>()?,
        num_parts: reader.read_u32::<LittleEndian>()?,
        part_size: reader.read_u32::<LittleEndian>()?,
        crc32_parts: reader.read_u32::<LittleEndian>()?,
    };

    let mut hdr_crc = hdr;
    for crc_byte in hdr_crc.iter_mut().skip(16).take(4) {
        *crc_byte = 0;
    }
    let c = crc32::checksum_ieee(&hdr_crc);
    trace!("header CRC32: {:#x} - computed CRC32: {:#x}", h.crc32, c);
    if crc32::checksum_ieee(&hdr_crc) == h.crc32 {
        Ok(h)
    } else {
        Err(Error::new(ErrorKind::Other, "invalid CRC32 checksum"))
    }
}

pub(crate) fn find_backup_lba(f: &mut File, sector_size: disk::LogicalBlockSize) -> Result<u64> {
    trace!("querying file size to find backup header location");
    let lb_size: u64 = sector_size.into();
    let old_pos = f.seek(std::io::SeekFrom::Current(0))?;
    let len = f.seek(std::io::SeekFrom::End(0))?;
    f.seek(std::io::SeekFrom::Start(old_pos))?;
    if len <= lb_size {
        return Err(Error::new(
            ErrorKind::Other,
            "disk image too small for backup header",
        ));
    }
    let bak_offset = len.saturating_sub(lb_size);
    let bak_lba = bak_offset / lb_size;
    trace!(
        "backup header: LBA={}, bytes offset={}",
        bak_lba,
        bak_offset
    );

    Ok(bak_lba)
}

fn calculate_crc32(b: &[u8]) -> Result<u32> {
    let mut digest = crc32::Digest::new(crc32::IEEE);
    trace!("Writing buffer to digest calculator");
    digest.write(b);

    Ok(digest.sum32())
}

pub(crate) fn partentry_checksum(
    file: &mut File,
    hdr: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<u32> {
    // Seek to start of partition table.
    let start = hdr.part_start
        .checked_mul(lb_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - partition table start"))?;
    let _ = file.seek(SeekFrom::Start(start))?;

    // Read partition table.
    let pt_len = u64::from(hdr.num_parts)
        .checked_mul(hdr.part_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partition table - size"))?;
    let mut buf = vec![0; pt_len as usize];
    file.read_exact(&mut buf)?;

    // Compute CRC32 over all table bits.
    let mut digest = crc32::Digest::new(crc32::IEEE);
    digest.write(&buf);
    Ok(digest.sum32())
}

/// A helper function to create a new header and write it to disk.
/// If the uuid isn't given a random one will be generated.  Use
/// this in conjunction with Partition::write()
// TODO: Move this to Header::new() and Header::write to write it
// that will match the Partition::write() API
pub fn write_header(
    p: &Path,
    uuid: Option<uuid::Uuid>,
    sector_size: disk::LogicalBlockSize,
) -> Result<uuid::Uuid> {
    debug!("opening {} for writing", p.display());
    let mut file = OpenOptions::new().write(true).read(true).open(p)?;
    let bak = find_backup_lba(&mut file, sector_size)?;
    let guid = match uuid {
        Some(u) => u,
        None => {
            let u = uuid::Uuid::new_v4();
            debug!("Generated random uuid: {}", u);
            u
        }
    };

    let hdr = Header::compute_new(true, &[], guid, bak)?;
    debug!("new header: {:#?}", hdr);
    hdr.write_primary(&mut file, sector_size)?;

    Ok(guid)
}
<|endoftext|>


<commit_before>use std::io::{BufRead, Read, BufReader, self, Error, ErrorKind};
use std::{cmp,mem};

use super::InflateStream;

/// A DEFLATE decoder/decompressor.
///
/// This structuree implements a `BufRead` interface and takes a stream of compressed data as input,
/// provoding the decompressed data when read from.
pub struct DeflateDecoderBuf<R> {
    /// The inner reader instance
    reader: R,
    /// The raw decompressor
    decompressor: InflateStream,
    /// How many bytes of the decompressor's output buffer still need to be output.
    pending_output_bytes: usize,
    /// Total number of bytes read from the underlying reader.
    total_in: u64,
    /// Total number of bytes written in `read` calls.
    total_out: u64,
}

impl<R: BufRead> DeflateDecoderBuf<R> {
    pub fn new(reader: R) -> DeflateDecoderBuf<R> {
        DeflateDecoderBuf {
            reader: reader,
            decompressor: InflateStream::new(),
            pending_output_bytes: 0,
            total_in: 0,
            total_out: 0,
        }
    }
}

impl<R> DeflateDecoderBuf<R> {
    /// Resets the decompressor, and replaces the current inner `BufRead` instance by `r`.
    /// without doing any extra reallocations.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset(&mut self, r: R) -> R {
        self.decompressor.reset();
        mem::replace(&mut self.reader, r)
    }

    /// Resets the decoder, but continue to read from the same reader.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset_data(&mut self) {
        self.decompressor.reset()
    }

    /// Returns a reference to the underlying `BufRead` instance.
    pub fn get_ref(&self) -> &R {
        &self.reader
    }

    /// Returns a mutable reference to the underlying `BufRead` instance.
    ///
    /// Note that mutation of the reader may cause surprising results if the decoder is going to
    /// keep being used.
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.reader
    }

    /// Drops the decoder and return the inner `BufRead` instance.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn into_inner(self) -> R {
        self.reader
    }

    /// Returns the total bytes read from the underlying `BufRead` instance.
    pub fn total_in(&self) -> u64 {
        self.total_in
    }

    /// Returns the total number of bytes output from this decoder.
    pub fn total_out(&self) -> u64 {
        self.total_out
    }
}

impl<R: BufRead> Read for DeflateDecoderBuf<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut bytes_out = 0;
        // If there is still data left to ouput from the last call to `update()`, that needs to be
        // output first
        if self.pending_output_bytes != 0 {
            // Get the part of the buffer that has not been output yet.
            // The decompressor sets `pos` to 0 when it reaches the end of it's internal buffer,
            // so we have to check for that.
            let start = if self.decompressor.pos != 0 {
                self.decompressor.pos as usize - self.pending_output_bytes
            } else {
                self.decompressor.buffer.len() - self.pending_output_bytes
            };

            // Copy as much decompressed as possible to buf.
            let bytes_to_copy = cmp::min(buf.len(), self.pending_output_bytes);
            let pending_data =
                &self.decompressor.buffer[start..
                                          start + bytes_to_copy];
            buf[..bytes_to_copy].copy_from_slice(pending_data);
            bytes_out += bytes_to_copy;
            // This won't underflow since `bytes_to_copy` will be at most
            // the same value as `pending_output_bytes`.
            self.pending_output_bytes -= bytes_to_copy;
            if self.pending_output_bytes != 0 {
                self.total_out += bytes_out as u64;
                // If there is still decompressed data left that didn't
                // fit in `buf`, return what we read.
                return Ok(bytes_out);
            }
        }

        // There is space in `buf` for more data, so try to read more.
        let (input_bytes_read, remaining_bytes) = {
            self.pending_output_bytes = 0;
            let input = try!(self.reader.fill_buf());
            if input.len() == 0 {
                self.total_out += bytes_out as u64;
                //If there is nothing more to read, return.
                return Ok(bytes_out);
            }
            let (input_bytes_read, data) =
                match self.decompressor.update(&input) {
                    Ok(res) => res,
                    Err(m) => return Err(Error::new(ErrorKind::Other, m))
                };

            // Space left in `buf`
            let space_left = buf.len() - bytes_out;
            let bytes_to_copy = cmp::min(space_left, data.len());
            buf[bytes_out..bytes_out + bytes_to_copy].copy_from_slice(&data[..bytes_to_copy]);

            bytes_out += bytes_to_copy;

            // Can't underflow as bytes_to_copy is bounded by data.len().
            (input_bytes_read, data.len() - bytes_to_copy)

        };

        self.pending_output_bytes += remaining_bytes;
        self.total_in += input_bytes_read as u64;
        self.total_out += bytes_out as u64;
        self.reader.consume(input_bytes_read);

        Ok(bytes_out)
    }
}

/// A DEFLATE decoder/decompressor.
///
/// This structuree implements a `Read` interface and takes a stream of compressed data as input,
/// provoding the decompressed data when read from.
pub struct DeflateDecoder<R> {
    /// Inner DeflateDecoderBuf, with R wrapped in a `BufReader`.
    inner: DeflateDecoderBuf<BufReader<R>>
}

impl<R: Read> DeflateDecoder<R> {
    pub fn new(reader: R) -> DeflateDecoder<R> {
        DeflateDecoder {
            inner: DeflateDecoderBuf::new(BufReader::new(reader))
        }
    }

    /// Resets the decompressor, and replaces the current inner `BufRead` instance by `r`.
    /// without doing any extra reallocations.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset(&mut self, r: R) -> R {
        self.inner.reset(BufReader::new(r)).into_inner()
    }

    /// Returns a reference to the underlying reader.
    pub fn get_ref(&self) -> &R {
        self.inner.get_ref().get_ref()
    }

    /// Returns a mutable reference to the underlying reader.
    ///
    /// Note that mutation of the reader may cause surprising results if the decoder is going to
    /// keep being used.
    pub fn get_mut(&mut self) -> &mut R {
        self.inner.get_mut().get_mut()
    }

    /// Returns the total number of bytes output from this decoder.
    pub fn into_inner(self) -> R {
        self.inner.into_inner().into_inner()
    }
}

impl<R> DeflateDecoder<R> {
    /// Resets the decoder, but continue to read from the same reader.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset_data(&mut self) {
        self.inner.reset_data()
    }

    /// Returns the total bytes read from the underlying reader.
    pub fn total_in(&self) -> u64 {
        self.inner.total_in
    }

    /// Returns the total number of bytes output from this decoder.
    pub fn total_out(&self) -> u64 {
        self.inner.total_out
    }
}

impl<R: BufRead> Read for DeflateDecoder<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.inner.read(buf)
    }
}

#[cfg(test)]
mod test {
    use super::{DeflateDecoder};

    #[test]
    fn deflate_reader() {
        let encoded = vec![243, 72, 205, 201, 201, 215, 81, 40, 207, 47, 202, 73, 1, 0];
        let mut decoder = DeflateDecoder::new(encoded.as_slice());
        let mut output = Vec::new();
        decoder.read_to_end(&mut output).unwrap();
        assert_eq!(String::from_utf8(output).unwrap(), "Hello, world");
    }
}
<commit_msg>Support zlib for readers, fix 1.0 compat<commit_after>use std::io::{self, BufRead, Read, BufReader, Error, ErrorKind, Write};
use std::{cmp,mem};

use super::InflateStream;

/// Workaround for lack of copy_from_slice on pre-1.9 rust.
fn copy_from_slice(mut to: &mut [u8], from: &[u8]) {
    assert_eq!(to.len(), from.len());
    to.write_all(from).unwrap();
}

/// A DEFLATE decoder/decompressor.
///
/// This structure implements a `Read` interface and takes a stream of compressed data as input,
/// providing the decompressed data when read from.
pub struct DeflateDecoderBuf<R> {
    /// The inner reader instance
    reader: R,
    /// The raw decompressor
    decompressor: InflateStream,
    /// How many bytes of the decompressor's output buffer still need to be output.
    pending_output_bytes: usize,
    /// Total number of bytes read from the underlying reader.
    total_in: u64,
    /// Total number of bytes written in `read` calls.
    total_out: u64,
}

impl<R: BufRead> DeflateDecoderBuf<R> {
    pub fn new(reader: R) -> DeflateDecoderBuf<R> {
        DeflateDecoderBuf {
            reader: reader,
            decompressor: InflateStream::new(),
            pending_output_bytes: 0,
            total_in: 0,
            total_out: 0,
        }
    }

    pub fn from_zlib(reader: R) -> DeflateDecoderBuf<R> {
        DeflateDecoderBuf {
            reader: reader,
            decompressor: InflateStream::from_zlib(),
            pending_output_bytes: 0,
            total_in: 0,
            total_out: 0,
        }
    }
}

impl<R> DeflateDecoderBuf<R> {
    /// Resets the decompressor, and replaces the current inner `BufRead` instance by `r`.
    /// without doing any extra reallocations.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset(&mut self, r: R) -> R {
        self.decompressor.reset();
        mem::replace(&mut self.reader, r)
    }

    /// Resets the decoder, but continue to read from the same reader.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset_data(&mut self) {
        self.decompressor.reset()
    }

    /// Returns a reference to the underlying `BufRead` instance.
    pub fn get_ref(&self) -> &R {
        &self.reader
    }

    /// Returns a mutable reference to the underlying `BufRead` instance.
    ///
    /// Note that mutation of the reader may cause surprising results if the decoder is going to
    /// keep being used.
    pub fn get_mut(&mut self) -> &mut R {
        &mut self.reader
    }

    /// Drops the decoder and return the inner `BufRead` instance.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn into_inner(self) -> R {
        self.reader
    }

    /// Returns the total bytes read from the underlying `BufRead` instance.
    pub fn total_in(&self) -> u64 {
        self.total_in
    }

    /// Returns the total number of bytes output from this decoder.
    pub fn total_out(&self) -> u64 {
        self.total_out
    }
}

impl<R: BufRead> Read for DeflateDecoderBuf<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut bytes_out = 0;
        // If there is still data left to ouput from the last call to `update()`, that needs to be
        // output first
        if self.pending_output_bytes != 0 {
            // Get the part of the buffer that has not been output yet.
            // The decompressor sets `pos` to 0 when it reaches the end of it's internal buffer,
            // so we have to check for that.
            let start = if self.decompressor.pos != 0 {
                self.decompressor.pos as usize - self.pending_output_bytes
            } else {
                self.decompressor.buffer.len() - self.pending_output_bytes
            };

            // Copy as much decompressed as possible to buf.
            let bytes_to_copy = cmp::min(buf.len(), self.pending_output_bytes);
            let pending_data =
                &self.decompressor.buffer[start..
                                          start + bytes_to_copy];
            copy_from_slice(&mut buf[..bytes_to_copy],pending_data);
            bytes_out += bytes_to_copy;
            // This won't underflow since `bytes_to_copy` will be at most
            // the same value as `pending_output_bytes`.
            self.pending_output_bytes -= bytes_to_copy;
            if self.pending_output_bytes != 0 {
                self.total_out += bytes_out as u64;
                // If there is still decompressed data left that didn't
                // fit in `buf`, return what we read.
                return Ok(bytes_out);
            }
        }

        // There is space in `buf` for more data, so try to read more.
        let (input_bytes_read, remaining_bytes) = {
            self.pending_output_bytes = 0;
            let input = try!(self.reader.fill_buf());
            if input.len() == 0 {
                self.total_out += bytes_out as u64;
                //If there is nothing more to read, return.
                return Ok(bytes_out);
            }
            let (input_bytes_read, data) =
                match self.decompressor.update(&input) {
                    Ok(res) => res,
                    Err(m) => return Err(Error::new(ErrorKind::Other, m))
                };

            // Space left in `buf`
            let space_left = buf.len() - bytes_out;
            let bytes_to_copy = cmp::min(space_left, data.len());

            copy_from_slice(&mut buf[bytes_out..bytes_out + bytes_to_copy], &data[..bytes_to_copy]);

            bytes_out += bytes_to_copy;

            // Can't underflow as bytes_to_copy is bounded by data.len().
            (input_bytes_read, data.len() - bytes_to_copy)

        };

        self.pending_output_bytes += remaining_bytes;
        self.total_in += input_bytes_read as u64;
        self.total_out += bytes_out as u64;
        self.reader.consume(input_bytes_read);

        Ok(bytes_out)
    }
}

/// A DEFLATE decoder/decompressor.
///
/// This structure implements a `Read` interface and takes a stream of compressed data as input,
/// provoding the decompressed data when read from.
pub struct DeflateDecoder<R> {
    /// Inner DeflateDecoderBuf, with R wrapped in a `BufReader`.
    inner: DeflateDecoderBuf<BufReader<R>>
}

impl<R: Read> DeflateDecoder<R> {
    pub fn new(reader: R) -> DeflateDecoder<R> {
        DeflateDecoder {
            inner: DeflateDecoderBuf::new(BufReader::new(reader))
        }
    }

    pub fn from_zlib(reader: R) -> DeflateDecoder<R> {
        DeflateDecoder {
            inner: DeflateDecoderBuf::from_zlib(BufReader::new(reader))
        }
    }

    /// Resets the decompressor, and replaces the current inner `BufRead` instance by `r`.
    /// without doing any extra reallocations.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset(&mut self, r: R) -> R {
        self.inner.reset(BufReader::new(r)).into_inner()
    }

    /// Returns a reference to the underlying reader.
    pub fn get_ref(&self) -> &R {
        self.inner.get_ref().get_ref()
    }

    /// Returns a mutable reference to the underlying reader.
    ///
    /// Note that mutation of the reader may cause surprising results if the decoder is going to
    /// keep being used.
    pub fn get_mut(&mut self) -> &mut R {
        self.inner.get_mut().get_mut()
    }

    /// Returns the total number of bytes output from this decoder.
    pub fn into_inner(self) -> R {
        self.inner.into_inner().into_inner()
    }
}

impl<R> DeflateDecoder<R> {
    /// Resets the decoder, but continue to read from the same reader.
    ///
    /// Note that this function doesn't ensure that all data has been output.
    pub fn reset_data(&mut self) {
        self.inner.reset_data()
    }

    /// Returns the total bytes read from the underlying reader.
    pub fn total_in(&self) -> u64 {
        self.inner.total_in
    }

    /// Returns the total number of bytes output from this decoder.
    pub fn total_out(&self) -> u64 {
        self.inner.total_out
    }
}

impl<R: Read> Read for DeflateDecoder<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.inner.read(buf)
    }
}

#[cfg(test)]
mod test {
    use super::{DeflateDecoder};
    use std::io::Read;

    #[test]
    fn deflate_reader() {
        const TEST_STRING: &'static str = "Hello, world";
        let encoded = vec![243, 72, 205, 201, 201, 215, 81, 40, 207, 47, 202, 73, 1, 0];
        let mut decoder = DeflateDecoder::new(&encoded[..]);
        let mut output = Vec::new();
        decoder.read_to_end(&mut output).unwrap();
        assert_eq!(String::from_utf8(output).unwrap(), TEST_STRING);
        assert_eq!(decoder.total_in(), encoded.len() as u64);
        assert_eq!(decoder.total_out(), TEST_STRING.len() as u64);
    }

    #[test]
    fn zlib_reader() {
        const TEST_STRING: &'static str = "Hello, zlib!";
        let encoded = vec![120, 156, 243, 72, 205, 201, 201, 215, 81, 168, 202, 201,
                       76, 82, 4, 0, 27, 101, 4, 19];
        let mut decoder = DeflateDecoder::from_zlib(&encoded[..]);
        let mut output = Vec::new();
        decoder.read_to_end(&mut output).unwrap();
        assert_eq!(String::from_utf8(output).unwrap(), TEST_STRING);
        assert_eq!(decoder.total_in(), encoded.len() as u64);
        assert_eq!(decoder.total_out(), TEST_STRING.len() as u64);
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Terminal formatting library.
//!
//! This crate provides the `Terminal` trait, which abstracts over an [ANSI
//! Termina][ansi] to provide color printing, among other things. There are two implementations,
//! the `TerminfoTerminal`, which uses control characters from a
//! [terminfo][ti] database, and `WinConsole`, which uses the [Win32 Console
//! API][win].
//!
//! ## Example
//!
//! ```rust
//! extern crate term;
//!
//! fn main() {
//!     let mut t = term::stdout().unwrap();
//!     t.fg(term::color::GREEN).unwrap();
//!     println!("hello, ");
//!     t.fg(term::color::RED).unwrap();
//!     println!("world!");
//!     t.reset().unwrap();
//! }
//! ```
//!
//! [ansi]: https://en.wikipedia.org/wiki/ANSI_escape_code
//! [win]: http://msdn.microsoft.com/en-us/library/windows/desktop/ms682010%28v=vs.85%29.aspx
//! [ti]: https://en.wikipedia.org/wiki/Terminfo

#![crate_id = "term#0.11.0-pre"]
#![comment = "Simple ANSI color library"]
#![license = "MIT/ASL2"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://static.rust-lang.org/doc/master")]

#![feature(macro_rules, phase)]

#![deny(missing_doc)]

#[phase(syntax, link)] extern crate log;
extern crate collections;

pub use terminfo::TerminfoTerminal;
#[cfg(windows)]
pub use win::WinConsole;

use std::io::IoResult;

pub mod terminfo;

#[cfg(windows)]
mod win;

#[cfg(not(windows))]
/// Return a Terminal wrapping stdout, or None if a terminal couldn't be
/// opened.
pub fn stdout() -> Option<Box<Terminal<Box<Writer:Send>>:Send>> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stdout() as Box<Writer:Send>);
    ti.map(|t| box t as Box<Terminal<Box<Writer:Send>:Send>:Send>)
}

#[cfg(windows)]
/// Return a Terminal wrapping stdout, or None if a terminal couldn't be
/// opened.
pub fn stdout() -> Option<Box<Terminal<Box<Writer:Send>:Send>:Send>> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stdout() as Box<Writer:Send>);

    match ti {
        Some(t) => Some(box t as Box<Terminal<Box<Writer:Send>:Send>:Send>),
        None => {
            let wc: Option<WinConsole<Box<Writer:Send>>>
                = Terminal::new(box std::io::stdout() as Box<Writer:Send>);
            wc.map(|w| box w as Box<Terminal<Box<Writer:Send>:Send>:Send>)
        }
    }
}

#[cfg(not(windows))]
/// Return a Terminal wrapping stderr, or None if a terminal couldn't be
/// opened.
pub fn stderr() -> Option<Box<Terminal<Box<Writer:Send>:Send>:Send>:Send> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stderr() as Box<Writer:Send>);
    ti.map(|t| box t as Box<Terminal<Box<Writer:Send>:Send>:Send>)
}

#[cfg(windows)]
/// Return a Terminal wrapping stderr, or None if a terminal couldn't be
/// opened.
pub fn stderr() -> Option<Box<Terminal<Box<Writer:Send>:Send>:Send>> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stderr() as Box<Writer:Send>);

    match ti {
        Some(t) => Some(box t as Box<Terminal<Box<Writer:Send>:Send>:Send>),
        None => {
            let wc: Option<WinConsole<Box<Writer:Send>>>
                = Terminal::new(box std::io::stderr() as Box<Writer:Send>);
            wc.map(|w| box w as Box<Terminal<Box<Writer:Send>:Send>:Send>)
        }
    }
}

/// Terminal color definitions
pub mod color {
    /// Number for a terminal color
    pub type Color = u16;

    pub static BLACK:   Color = 0u16;
    pub static RED:     Color = 1u16;
    pub static GREEN:   Color = 2u16;
    pub static YELLOW:  Color = 3u16;
    pub static BLUE:    Color = 4u16;
    pub static MAGENTA: Color = 5u16;
    pub static CYAN:    Color = 6u16;
    pub static WHITE:   Color = 7u16;

    pub static BRIGHT_BLACK:   Color = 8u16;
    pub static BRIGHT_RED:     Color = 9u16;
    pub static BRIGHT_GREEN:   Color = 10u16;
    pub static BRIGHT_YELLOW:  Color = 11u16;
    pub static BRIGHT_BLUE:    Color = 12u16;
    pub static BRIGHT_MAGENTA: Color = 13u16;
    pub static BRIGHT_CYAN:    Color = 14u16;
    pub static BRIGHT_WHITE:   Color = 15u16;
}

/// Terminal attributes
pub mod attr {
    /// Terminal attributes for use with term.attr().
    ///
    /// Most attributes can only be turned on and must be turned off with term.reset().
    /// The ones that can be turned off explicitly take a boolean value.
    /// Color is also represented as an attribute for convenience.
    pub enum Attr {
        /// Bold (or possibly bright) mode
        Bold,
        /// Dim mode, also called faint or half-bright. Often not supported
        Dim,
        /// Italics mode. Often not supported
        Italic(bool),
        /// Underline mode
        Underline(bool),
        /// Blink mode
        Blink,
        /// Standout mode. Often implemented as Reverse, sometimes coupled with Bold
        Standout(bool),
        /// Reverse mode, inverts the foreground and background colors
        Reverse,
        /// Secure mode, also called invis mode. Hides the printed text
        Secure,
        /// Convenience attribute to set the foreground color
        ForegroundColor(super::color::Color),
        /// Convenience attribute to set the background color
        BackgroundColor(super::color::Color)
    }
}

/// A terminal with similar capabilities to an ANSI Terminal
/// (foreground/background colors etc).
pub trait Terminal<T: Writer>: Writer {
    /// Returns `None` whenever the terminal cannot be created for some
    /// reason.
    fn new(out: T) -> Option<Self>;

    /// Sets the foreground color to the given color.
    ///
    /// If the color is a bright color, but the terminal only supports 8 colors,
    /// the corresponding normal color will be used instead.
    ///
    /// Returns `Ok(true)` if the color was set, `Ok(false)` otherwise, and `Err(e)`
    /// if there was an I/O error.
    fn fg(&mut self, color: color::Color) -> IoResult<bool>;

    /// Sets the background color to the given color.
    ///
    /// If the color is a bright color, but the terminal only supports 8 colors,
    /// the corresponding normal color will be used instead.
    ///
    /// Returns `Ok(true)` if the color was set, `Ok(false)` otherwise, and `Err(e)`
    /// if there was an I/O error.
    fn bg(&mut self, color: color::Color) -> IoResult<bool>;

    /// Sets the given terminal attribute, if supported.  Returns `Ok(true)`
    /// if the attribute was supported, `Ok(false)` otherwise, and `Err(e)` if
    /// there was an I/O error.
    fn attr(&mut self, attr: attr::Attr) -> IoResult<bool>;

    /// Returns whether the given terminal attribute is supported.
    fn supports_attr(&self, attr: attr::Attr) -> bool;

    /// Resets all terminal attributes and color to the default.
    /// Returns `Ok()`.
    fn reset(&mut self) -> IoResult<()>;

    /// Returns the contained stream, destroying the `Terminal`
    fn unwrap(self) -> T;

    /// Gets an immutable reference to the stream inside
    fn get_ref<'a>(&'a self) -> &'a T;

    /// Gets a mutable reference to the stream inside
    fn get_mut<'a>(&'a mut self) -> &'a mut T;
}
<commit_msg>auto merge of #14257 : alexcrichton/rust/fix-snap, r=huonw<commit_after>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Terminal formatting library.
//!
//! This crate provides the `Terminal` trait, which abstracts over an [ANSI
//! Termina][ansi] to provide color printing, among other things. There are two implementations,
//! the `TerminfoTerminal`, which uses control characters from a
//! [terminfo][ti] database, and `WinConsole`, which uses the [Win32 Console
//! API][win].
//!
//! ## Example
//!
//! ```no_run
//! extern crate term;
//!
//! fn main() {
//!     let mut t = term::stdout().unwrap();
//!     t.fg(term::color::GREEN).unwrap();
//!     println!("hello, ");
//!     t.fg(term::color::RED).unwrap();
//!     println!("world!");
//!     t.reset().unwrap();
//! }
//! ```
//!
//! [ansi]: https://en.wikipedia.org/wiki/ANSI_escape_code
//! [win]: http://msdn.microsoft.com/en-us/library/windows/desktop/ms682010%28v=vs.85%29.aspx
//! [ti]: https://en.wikipedia.org/wiki/Terminfo

#![crate_id = "term#0.11.0-pre"]
#![comment = "Simple ANSI color library"]
#![license = "MIT/ASL2"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![doc(html_logo_url = "http://www.rust-lang.org/logos/rust-logo-128x128-blk-v2.png",
       html_favicon_url = "http://www.rust-lang.org/favicon.ico",
       html_root_url = "http://static.rust-lang.org/doc/master")]

#![feature(macro_rules, phase)]

#![deny(missing_doc)]

#[phase(syntax, link)] extern crate log;
extern crate collections;

pub use terminfo::TerminfoTerminal;
#[cfg(windows)]
pub use win::WinConsole;

use std::io::IoResult;

pub mod terminfo;

#[cfg(windows)]
mod win;

#[cfg(not(windows))]
/// Return a Terminal wrapping stdout, or None if a terminal couldn't be
/// opened.
pub fn stdout() -> Option<Box<Terminal<Box<Writer:Send>>:Send>> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stdout() as Box<Writer:Send>);
    ti.map(|t| box t as Box<Terminal<Box<Writer:Send>:Send>:Send>)
}

#[cfg(windows)]
/// Return a Terminal wrapping stdout, or None if a terminal couldn't be
/// opened.
pub fn stdout() -> Option<Box<Terminal<Box<Writer:Send>:Send>:Send>> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stdout() as Box<Writer:Send>);

    match ti {
        Some(t) => Some(box t as Box<Terminal<Box<Writer:Send>:Send>:Send>),
        None => {
            let wc: Option<WinConsole<Box<Writer:Send>>>
                = Terminal::new(box std::io::stdout() as Box<Writer:Send>);
            wc.map(|w| box w as Box<Terminal<Box<Writer:Send>:Send>:Send>)
        }
    }
}

#[cfg(not(windows))]
/// Return a Terminal wrapping stderr, or None if a terminal couldn't be
/// opened.
pub fn stderr() -> Option<Box<Terminal<Box<Writer:Send>:Send>:Send>:Send> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stderr() as Box<Writer:Send>);
    ti.map(|t| box t as Box<Terminal<Box<Writer:Send>:Send>:Send>)
}

#[cfg(windows)]
/// Return a Terminal wrapping stderr, or None if a terminal couldn't be
/// opened.
pub fn stderr() -> Option<Box<Terminal<Box<Writer:Send>:Send>:Send>> {
    let ti: Option<TerminfoTerminal<Box<Writer:Send>>>
        = Terminal::new(box std::io::stderr() as Box<Writer:Send>);

    match ti {
        Some(t) => Some(box t as Box<Terminal<Box<Writer:Send>:Send>:Send>),
        None => {
            let wc: Option<WinConsole<Box<Writer:Send>>>
                = Terminal::new(box std::io::stderr() as Box<Writer:Send>);
            wc.map(|w| box w as Box<Terminal<Box<Writer:Send>:Send>:Send>)
        }
    }
}

/// Terminal color definitions
pub mod color {
    /// Number for a terminal color
    pub type Color = u16;

    pub static BLACK:   Color = 0u16;
    pub static RED:     Color = 1u16;
    pub static GREEN:   Color = 2u16;
    pub static YELLOW:  Color = 3u16;
    pub static BLUE:    Color = 4u16;
    pub static MAGENTA: Color = 5u16;
    pub static CYAN:    Color = 6u16;
    pub static WHITE:   Color = 7u16;

    pub static BRIGHT_BLACK:   Color = 8u16;
    pub static BRIGHT_RED:     Color = 9u16;
    pub static BRIGHT_GREEN:   Color = 10u16;
    pub static BRIGHT_YELLOW:  Color = 11u16;
    pub static BRIGHT_BLUE:    Color = 12u16;
    pub static BRIGHT_MAGENTA: Color = 13u16;
    pub static BRIGHT_CYAN:    Color = 14u16;
    pub static BRIGHT_WHITE:   Color = 15u16;
}

/// Terminal attributes
pub mod attr {
    /// Terminal attributes for use with term.attr().
    ///
    /// Most attributes can only be turned on and must be turned off with term.reset().
    /// The ones that can be turned off explicitly take a boolean value.
    /// Color is also represented as an attribute for convenience.
    pub enum Attr {
        /// Bold (or possibly bright) mode
        Bold,
        /// Dim mode, also called faint or half-bright. Often not supported
        Dim,
        /// Italics mode. Often not supported
        Italic(bool),
        /// Underline mode
        Underline(bool),
        /// Blink mode
        Blink,
        /// Standout mode. Often implemented as Reverse, sometimes coupled with Bold
        Standout(bool),
        /// Reverse mode, inverts the foreground and background colors
        Reverse,
        /// Secure mode, also called invis mode. Hides the printed text
        Secure,
        /// Convenience attribute to set the foreground color
        ForegroundColor(super::color::Color),
        /// Convenience attribute to set the background color
        BackgroundColor(super::color::Color)
    }
}

/// A terminal with similar capabilities to an ANSI Terminal
/// (foreground/background colors etc).
pub trait Terminal<T: Writer>: Writer {
    /// Returns `None` whenever the terminal cannot be created for some
    /// reason.
    fn new(out: T) -> Option<Self>;

    /// Sets the foreground color to the given color.
    ///
    /// If the color is a bright color, but the terminal only supports 8 colors,
    /// the corresponding normal color will be used instead.
    ///
    /// Returns `Ok(true)` if the color was set, `Ok(false)` otherwise, and `Err(e)`
    /// if there was an I/O error.
    fn fg(&mut self, color: color::Color) -> IoResult<bool>;

    /// Sets the background color to the given color.
    ///
    /// If the color is a bright color, but the terminal only supports 8 colors,
    /// the corresponding normal color will be used instead.
    ///
    /// Returns `Ok(true)` if the color was set, `Ok(false)` otherwise, and `Err(e)`
    /// if there was an I/O error.
    fn bg(&mut self, color: color::Color) -> IoResult<bool>;

    /// Sets the given terminal attribute, if supported.  Returns `Ok(true)`
    /// if the attribute was supported, `Ok(false)` otherwise, and `Err(e)` if
    /// there was an I/O error.
    fn attr(&mut self, attr: attr::Attr) -> IoResult<bool>;

    /// Returns whether the given terminal attribute is supported.
    fn supports_attr(&self, attr: attr::Attr) -> bool;

    /// Resets all terminal attributes and color to the default.
    /// Returns `Ok()`.
    fn reset(&mut self) -> IoResult<()>;

    /// Returns the contained stream, destroying the `Terminal`
    fn unwrap(self) -> T;

    /// Gets an immutable reference to the stream inside
    fn get_ref<'a>(&'a self) -> &'a T;

    /// Gets a mutable reference to the stream inside
    fn get_mut<'a>(&'a mut self) -> &'a mut T;
}
<|endoftext|>