<commit_before>use std::io;

use args::AnyArg;
use args::settings::ArgSettings;
use term;

const TAB: &'static str = "    ";

pub struct HelpWriter<'a, A> where A: 'a {
    a: &'a A,
    l: usize,
    nlh: bool,
    pub skip_pv: bool,
    term_w: Option<usize>,
}

impl<'a, 'n, 'e, A> HelpWriter<'a, A> where A: AnyArg<'n, 'e>  {
    pub fn new(a: &'a A, l: usize, nlh: bool) -> Self {
        HelpWriter {
            a: a,
            l: l,
            nlh: nlh,
            skip_pv: false,
            term_w: term::dimensions().map(|(w, _)| w),
        }
    }
    pub fn write_to<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
        debugln!("fn=write_to;");
        try!(self.short(w));
        try!(self.long(w));
        try!(self.val(w));
        try!(self.help(w));
        write!(w, "\n")
    }

    fn short<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=short;");
        try!(write!(w, "{}", TAB));
        if let Some(s) = self.a.short() {
            write!(w, "-{}", s)
        } else if self.a.has_switch() {
            write!(w, "{}", TAB)
        } else {
            Ok(())
        }
    }

    fn long<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=long;");
        if !self.a.has_switch() {
            return Ok(());
        }
        if self.a.takes_value() {
            if let Some(l) = self.a.long() {
                try!(write!(w, "{}--{}", if self.a.short().is_some() { ", " } else { "" }, l));
            }
            try!(write!(w, " "));
        } else {
            if let Some(l) = self.a.long() {
                try!(write!(w, "{}--{}", if self.a.short().is_some() { ", " } else { "" }, l));
                if !self.nlh || !self.a.is_set(ArgSettings::NextLineHelp) {
                    write_spaces!((self.l + 4) - (l.len() + 2), w);
                }
            } else {
                if !self.nlh || !self.a.is_set(ArgSettings::NextLineHelp) {
                    // 6 is tab (4) + -- (2)
                    write_spaces!((self.l + 6), w);
                }
            }
        }
        Ok(())
    }

    fn val<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=val;");
        if !self.a.takes_value() {
            return Ok(());
        }
        if let Some(ref vec) = self.a.val_names() {
            let mut it = vec.iter().peekable();
            while let Some((_, val)) = it.next() {
                try!(write!(w, "<{}>", val));
                if it.peek().is_some() { try!(write!(w, " ")); }
            }
            let num = vec.len();
            if self.a.is_set(ArgSettings::Multiple) && num == 1 {
                try!(write!(w, "..."));
            }
        } else if let Some(num) = self.a.num_vals() {
            let mut it = (0..num).peekable();
            while let Some(_) = it.next() {
                try!(write!(w, "<{}>", self.a.name()));
                if it.peek().is_some() { try!(write!(w, " ")); }
            }
        } else {
            try!(write!(w, "<{}>{}", self.a.name(), if self.a.is_set(ArgSettings::Multiple) { "..." } else { "" }));
        }
        if self.a.has_switch() {
            if !(self.nlh || self.a.is_set(ArgSettings::NextLineHelp)) {
                let self_len = self.a.to_string().len();
                // subtract ourself
                let mut spcs = self.l - self_len;
                // Since we're writing spaces from the tab point we first need to know if we
                // had a long and short, or just short
                if self.a.long().is_some() {
                    // Only account 4 after the val
                    spcs += 4;
                } else {
                    // Only account for ', --' + 4 after the val
                    spcs += 8;
                }
                write_spaces!(spcs, w);
            }
        } else {
            if !(self.nlh || self.a.is_set(ArgSettings::NextLineHelp)) {
                write_spaces!(self.l + 4 - (self.a.to_string().len()), w);
            }
        }
        Ok(())
    }

    fn help<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=help;");
        let spec_vals = self.spec_vals();
        let mut help = String::new();
        let h = self.a.help().unwrap_or("");
        let spcs = if self.nlh || self.a.is_set(ArgSettings::NextLineHelp) {
            8 // "tab" + "tab"
        } else {
            self.l + 12
        };
        // determine if our help fits or needs to wrap
        let too_long = self.term_w.is_some() && (spcs + h.len() + spec_vals.len() >= self.term_w.unwrap_or(0));

        // Is help on next line, if so newline + 2x tab
        if self.nlh || self.a.is_set(ArgSettings::NextLineHelp) {
            try!(write!(w, "\n{}{}", TAB, TAB));
        }

        debug!("Too long...");
        if too_long {
            sdebugln!("Yes");
            if let Some(width) = self.term_w {
                help.push_str(h);
                help.push_str(&*spec_vals);
                debugln!("term width: {}", width);
                debugln!("help: {}", help);
                debugln!("help len: {}", help.len());
                // Determine how many newlines we need to insert
                let avail_chars = width - spcs;
                debugln!("Usable space: {}", avail_chars);
                let mut indices = vec![];
                let mut idx = 0;
                loop {
                    idx += avail_chars - 1;
                    if idx >= help.len() { break; }
                    // 'a' arbitrary non space char
                    if help.chars().nth(idx).unwrap_or('a') != ' ' {
                        idx = find_idx_of_space(&*help, idx);
                    }
                    debugln!("Adding idx: {}", idx);
                    debugln!("At {}: {:?}", idx, help.chars().nth(idx));
                    indices.push(idx);
                    if &help[idx..].len() <= &avail_chars {
                        break;
                    }
                }
                for (i, idx) in indices.iter().enumerate() {
                    debugln!("iter;i={},idx={}", i, idx);
                    let j = idx+(2*i);
                    debugln!("removing: {}", j);
                    debugln!("at {}: {:?}", j, help.chars().nth(j));
                    help.remove(j);
                    help.insert(j, '{');
                    help.insert(j + 1 , 'n');
                    help.insert(j + 2, '}');
                }
            }
        } else { sdebugln!("No"); }
        let help = if !help.is_empty() {
            &*help
        } else if !spec_vals.is_empty() {
            help.push_str(h);
            help.push_str(&*spec_vals);
            &*help
        } else {
            h
        };
        if help.contains("{n}") {
            if let Some(part) = help.split("{n}").next() {
                try!(write!(w, "{}", part));
            }
            for part in help.split("{n}").skip(1) {
                try!(write!(w, "\n"));
                if self.nlh || self.a.is_set(ArgSettings::NextLineHelp) {
                    try!(write!(w, "{}{}", TAB, TAB));
                } else {
                    if self.a.has_switch() {
                        write_spaces!(self.l + 12, w);
                    } else {
                        write_spaces!(self.l + 8, w);
                    }
                }
                try!(write!(w, "{}", part));
            }
        } else {
            try!(write!(w, "{}", help));
        }
        Ok(())
    }

    fn spec_vals(&self) -> String {
        debugln!("fn=spec_vals;");
        if let Some(ref pv) = self.a.default_val() {
            debugln!("Writing defaults");
            return format!(" [default: {}] {}", pv,
                if !self.skip_pv {
                    if let Some(ref pv) = self.a.possible_vals() {
                        format!(" [values: {}]", pv.join(", "))
                    } else { "".into() }
                } else { "".into() }
            );
        } else if !self.skip_pv {
            debugln!("Writing values");
            if let Some(ref pv) = self.a.possible_vals() {
                debugln!("Possible vals...{:?}", pv);
                return format!(" [values: {}]", pv.join(", "));
            }
        }
        String::new()
    }
}

fn find_idx_of_space(full: &str, start: usize) -> usize {
    debugln!("fn=find_idx_of_space;");
    let haystack = &full[..start];
    debugln!("haystack: {}", haystack);
    for (i, c) in haystack.chars().rev().enumerate() {
        debugln!("iter;c={},i={}", c, i);
        if c == ' ' {
            debugln!("Found space returning start-i...{}", start - (i+1));
            return start - (i+1);
        }
    }
    0
}
<commit_msg>fix(Help Message): fixes a bug where small terminal sizes causing a loop<commit_after>use std::io;

use args::AnyArg;
use args::settings::ArgSettings;
use term;

const TAB: &'static str = "    ";

pub struct HelpWriter<'a, A> where A: 'a {
    a: &'a A,
    l: usize,
    nlh: bool,
    pub skip_pv: bool,
    term_w: Option<usize>,
}

impl<'a, 'n, 'e, A> HelpWriter<'a, A> where A: AnyArg<'n, 'e>  {
    pub fn new(a: &'a A, l: usize, nlh: bool) -> Self {
        HelpWriter {
            a: a,
            l: l,
            nlh: nlh,
            skip_pv: false,
            term_w: term::dimensions().map(|(w, _)| w),
        }
    }
    pub fn write_to<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
        debugln!("fn=write_to;");
        try!(self.short(w));
        try!(self.long(w));
        try!(self.val(w));
        try!(self.help(w));
        write!(w, "\n")
    }

    fn short<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=short;");
        try!(write!(w, "{}", TAB));
        if let Some(s) = self.a.short() {
            write!(w, "-{}", s)
        } else if self.a.has_switch() {
            write!(w, "{}", TAB)
        } else {
            Ok(())
        }
    }

    fn long<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=long;");
        if !self.a.has_switch() {
            return Ok(());
        }
        if self.a.takes_value() {
            if let Some(l) = self.a.long() {
                try!(write!(w, "{}--{}", if self.a.short().is_some() { ", " } else { "" }, l));
            }
            try!(write!(w, " "));
        } else {
            if let Some(l) = self.a.long() {
                try!(write!(w, "{}--{}", if self.a.short().is_some() { ", " } else { "" }, l));
                if !self.nlh || !self.a.is_set(ArgSettings::NextLineHelp) {
                    write_spaces!((self.l + 4) - (l.len() + 2), w);
                }
            } else {
                if !self.nlh || !self.a.is_set(ArgSettings::NextLineHelp) {
                    // 6 is tab (4) + -- (2)
                    write_spaces!((self.l + 6), w);
                }
            }
        }
        Ok(())
    }

    fn val<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=val;");
        if !self.a.takes_value() {
            return Ok(());
        }
        if let Some(ref vec) = self.a.val_names() {
            let mut it = vec.iter().peekable();
            while let Some((_, val)) = it.next() {
                try!(write!(w, "<{}>", val));
                if it.peek().is_some() { try!(write!(w, " ")); }
            }
            let num = vec.len();
            if self.a.is_set(ArgSettings::Multiple) && num == 1 {
                try!(write!(w, "..."));
            }
        } else if let Some(num) = self.a.num_vals() {
            let mut it = (0..num).peekable();
            while let Some(_) = it.next() {
                try!(write!(w, "<{}>", self.a.name()));
                if it.peek().is_some() { try!(write!(w, " ")); }
            }
        } else {
            try!(write!(w, "<{}>{}", self.a.name(), if self.a.is_set(ArgSettings::Multiple) { "..." } else { "" }));
        }
        if self.a.has_switch() {
            if !(self.nlh || self.a.is_set(ArgSettings::NextLineHelp)) {
                let self_len = self.a.to_string().len();
                // subtract ourself
                let mut spcs = self.l - self_len;
                // Since we're writing spaces from the tab point we first need to know if we
                // had a long and short, or just short
                if self.a.long().is_some() {
                    // Only account 4 after the val
                    spcs += 4;
                } else {
                    // Only account for ', --' + 4 after the val
                    spcs += 8;
                }
                write_spaces!(spcs, w);
            }
        } else {
            if !(self.nlh || self.a.is_set(ArgSettings::NextLineHelp)) {
                write_spaces!(self.l + 4 - (self.a.to_string().len()), w);
            }
        }
        Ok(())
    }

    fn help<W>(&self, w: &mut W) -> io::Result<()>
        where W: io::Write
    {
        debugln!("fn=help;");
        let spec_vals = self.spec_vals();
        let mut help = String::new();
        let h = self.a.help().unwrap_or("");
        let spcs = if self.nlh || self.a.is_set(ArgSettings::NextLineHelp) {
            8 // "tab" + "tab"
        } else {
            self.l + 12
        };
        // determine if our help fits or needs to wrap
        let width = self.term_w.unwrap_or(0);
        debugln!("Term width...{}", width);
        let too_long = self.term_w.is_some() && (spcs + h.len() + spec_vals.len() >= width);
        debugln!("Too long...{:?}", too_long);

        // Is help on next line, if so newline + 2x tab
        if self.nlh || self.a.is_set(ArgSettings::NextLineHelp) {
            try!(write!(w, "\n{}{}", TAB, TAB));
        }

        debug!("Too long...");
        if too_long {
            sdebugln!("Yes");
            help.push_str(h);
            help.push_str(&*spec_vals);
            debugln!("help: {}", help);
            debugln!("help len: {}", help.len());
            // Determine how many newlines we need to insert
            let avail_chars = width - spcs;
            debugln!("Usable space: {}", avail_chars);
            let longest_w = {
                let mut lw = 0;
                for l in help.split(' ').map(|s| s.len()) {
                    if l > lw {
                        lw = l;
                    }
                }
                lw
            };
            debugln!("Longest word...{}", longest_w);
            debug!("Enough space...");
            if longest_w < avail_chars {
                sdebugln!("Yes");
                let mut indices = vec![];
                let mut idx = 0;
                loop {
                    idx += avail_chars - 1;
                    if idx >= help.len() { break; }
                    // 'a' arbitrary non space char
                    if help.chars().nth(idx).unwrap_or('a') != ' ' {
                        idx = find_idx_of_space(&*help, idx);
                    }
                    debugln!("Adding idx: {}", idx);
                    debugln!("At {}: {:?}", idx, help.chars().nth(idx));
                    indices.push(idx);
                    if &help[idx..].len() <= &avail_chars {
                        break;
                    }
                }
                for (i, idx) in indices.iter().enumerate() {
                    debugln!("iter;i={},idx={}", i, idx);
                    let j = idx+(2*i);
                    debugln!("removing: {}", j);
                    debugln!("at {}: {:?}", j, help.chars().nth(j));
                    help.remove(j);
                    help.insert(j, '{');
                    help.insert(j + 1 , 'n');
                    help.insert(j + 2, '}');
                }
            } else {
                sdebugln!("No");
            }
        } else { sdebugln!("No"); }
        let help = if !help.is_empty() {
            &*help
        } else if !spec_vals.is_empty() {
            help.push_str(h);
            help.push_str(&*spec_vals);
            &*help
        } else {
            h
        };
        if help.contains("{n}") {
            if let Some(part) = help.split("{n}").next() {
                try!(write!(w, "{}", part));
            }
            for part in help.split("{n}").skip(1) {
                try!(write!(w, "\n"));
                if self.nlh || self.a.is_set(ArgSettings::NextLineHelp) {
                    try!(write!(w, "{}{}", TAB, TAB));
                } else {
                    if self.a.has_switch() {
                        write_spaces!(self.l + 12, w);
                    } else {
                        write_spaces!(self.l + 8, w);
                    }
                }
                try!(write!(w, "{}", part));
            }
        } else {
            try!(write!(w, "{}", help));
        }
        Ok(())
    }

    fn spec_vals(&self) -> String {
        debugln!("fn=spec_vals;");
        if let Some(ref pv) = self.a.default_val() {
            debugln!("Writing defaults");
            return format!(" [default: {}] {}", pv,
                if !self.skip_pv {
                    if let Some(ref pv) = self.a.possible_vals() {
                        format!(" [values: {}]", pv.join(", "))
                    } else { "".into() }
                } else { "".into() }
            );
        } else if !self.skip_pv {
            debugln!("Writing values");
            if let Some(ref pv) = self.a.possible_vals() {
                debugln!("Possible vals...{:?}", pv);
                return format!(" [values: {}]", pv.join(", "));
            }
        }
        String::new()
    }
}

fn find_idx_of_space(full: &str, start: usize) -> usize {
    debugln!("fn=find_idx_of_space;");
    let haystack = &full[..start];
    debugln!("haystack: {}", haystack);
    for (i, c) in haystack.chars().rev().enumerate() {
        debugln!("iter;c={},i={}", c, i);
        if c == ' ' {
            debugln!("Found space returning start-i...{}", start - (i+1));
            return start - (i+1);
        }
    }
    0
}
<|endoftext|>


<commit_before>use crate::descriptor::MethodDescriptorProto;
use crate::descriptor::ServiceDescriptorProto;
use crate::reflect::FileDescriptor;
use crate::reflect::MessageDescriptor;

/// Dynamic representation of message type.
///
/// Used for reflection.
#[derive(Clone, Eq, PartialEq)]
pub struct ServiceDescriptor {
    file_descriptor: FileDescriptor,
    index: usize,
}

impl ServiceDescriptor {
    pub(crate) fn new(file_descriptor: FileDescriptor, index: usize) -> ServiceDescriptor {
        ServiceDescriptor {
            file_descriptor,
            index,
        }
    }

    /// Proto snippet describing this service.
    pub fn get_proto(&self) -> &ServiceDescriptorProto {
        &self.file_descriptor.proto().service[self.index]
    }

    /// Method descriptors of this service.
    pub fn methods(&self) -> Vec<MethodDescriptor> {
        let value_len = self.get_proto().method.len();
        (0..value_len)
            .map(move |index| MethodDescriptor {
                service_descriptor: self.clone(),
                index,
            })
            .collect()
    }
}

/// Service method descriptor.
pub struct MethodDescriptor {
    service_descriptor: ServiceDescriptor,
    index: usize,
}

impl MethodDescriptor {
    /// Proto snippet describing this method.
    pub fn get_proto(&self) -> &MethodDescriptorProto {
        &self.service_descriptor.get_proto().method[self.index]
    }

    /// Method input type.
    pub fn input_type(&self) -> MessageDescriptor {
        // TODO
        unimplemented!()
    }

    /// Method output type.
    pub fn output_type(&self) -> MessageDescriptor {
        // TODO
        unimplemented!()
    }
}
<commit_msg>Some docs<commit_after>use crate::descriptor::MethodDescriptorProto;
use crate::descriptor::ServiceDescriptorProto;
use crate::reflect::FileDescriptor;
use crate::reflect::MessageDescriptor;

/// Dynamic representation of service type.
///
/// Rust-protobuf does not support services (it is not an RPC library),
/// but it support querying service description. Which might be useful
/// for example to generate source files for the services.
/// or to perform invocations dynamically.
#[derive(Clone, Eq, PartialEq)]
pub struct ServiceDescriptor {
    file_descriptor: FileDescriptor,
    index: usize,
}

impl ServiceDescriptor {
    pub(crate) fn new(file_descriptor: FileDescriptor, index: usize) -> ServiceDescriptor {
        ServiceDescriptor {
            file_descriptor,
            index,
        }
    }

    /// Proto snippet describing this service.
    pub fn get_proto(&self) -> &ServiceDescriptorProto {
        &self.file_descriptor.proto().service[self.index]
    }

    /// Method descriptors of this service.
    pub fn methods(&self) -> Vec<MethodDescriptor> {
        let value_len = self.get_proto().method.len();
        (0..value_len)
            .map(move |index| MethodDescriptor {
                service_descriptor: self.clone(),
                index,
            })
            .collect()
    }
}

/// Service method descriptor.
pub struct MethodDescriptor {
    service_descriptor: ServiceDescriptor,
    index: usize,
}

impl MethodDescriptor {
    /// Proto snippet describing this method.
    pub fn get_proto(&self) -> &MethodDescriptorProto {
        &self.service_descriptor.get_proto().method[self.index]
    }

    /// Method input type.
    pub fn input_type(&self) -> MessageDescriptor {
        // TODO
        unimplemented!()
    }

    /// Method output type.
    pub fn output_type(&self) -> MessageDescriptor {
        // TODO
        unimplemented!()
    }
}
<|endoftext|>


<commit_before>use crate::comment::CommentStyle;
use std::fmt::{self, Display};

/// Formats a string as a doc comment using the given [`CommentStyle`].
#[derive(new)]
pub(super) struct DocCommentFormatter<'a> {
    literal: &'a str,
    style: CommentStyle<'a>,
}

impl Display for DocCommentFormatter<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let opener = self.style.opener().trim_end();
        let mut lines = self.literal.lines().peekable();
        while let Some(line) = lines.next() {
            let is_last_line = lines.peek().is_none();
            if is_last_line {
                write!(formatter, "{}{}", opener, line)?;
            } else {
                writeln!(formatter, "{}{}", opener, line)?;
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn literal_controls_leading_spaces() {
        test_doc_comment_is_formatted_correctly(
            "    Lorem ipsum",
            "///    Lorem ipsum",
            CommentStyle::TripleSlash,
        );
    }

    #[test]
    fn single_line_doc_comment_is_formatted_correctly() {
        test_doc_comment_is_formatted_correctly(
            "Lorem ipsum",
            "///Lorem ipsum",
            CommentStyle::TripleSlash,
        );
    }

    #[test]
    fn multi_line_doc_comment_is_formatted_correctly() {
        test_doc_comment_is_formatted_correctly(
            "Lorem ipsum\nDolor sit amet",
            "///Lorem ipsum\n///Dolor sit amet",
            CommentStyle::TripleSlash,
        );
    }

    #[test]
    fn whitespace_within_lines_is_preserved() {
        test_doc_comment_is_formatted_correctly(
            " Lorem ipsum \n Dolor sit amet ",
            "/// Lorem ipsum \n/// Dolor sit amet ",
            CommentStyle::TripleSlash,
        );
    }

    fn test_doc_comment_is_formatted_correctly(
        literal: &str,
        expected_comment: &str,
        style: CommentStyle<'_>,
    ) {
        assert_eq!(
            expected_comment,
            format!("{}", DocCommentFormatter::new(&literal, style))
        );
    }
}
<commit_msg>Insert an empty line when normalizing `#[doc = ""]`<commit_after>use crate::comment::CommentStyle;
use std::fmt::{self, Display};

/// Formats a string as a doc comment using the given [`CommentStyle`].
#[derive(new)]
pub(super) struct DocCommentFormatter<'a> {
    literal: &'a str,
    style: CommentStyle<'a>,
}

impl Display for DocCommentFormatter<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        let opener = self.style.opener().trim_end();
        let mut lines = self.literal.lines().peekable();

        // Handle `#[doc = ""]`.
        if lines.peek().is_none() {
            return write!(formatter, "{}", opener);
        }

        while let Some(line) = lines.next() {
            let is_last_line = lines.peek().is_none();
            if is_last_line {
                write!(formatter, "{}{}", opener, line)?;
            } else {
                writeln!(formatter, "{}{}", opener, line)?;
            }
        }
        Ok(())
    }
}

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

    #[test]
    fn literal_controls_leading_spaces() {
        test_doc_comment_is_formatted_correctly(
            "    Lorem ipsum",
            "///    Lorem ipsum",
            CommentStyle::TripleSlash,
        );
    }

    #[test]
    fn single_line_doc_comment_is_formatted_correctly() {
        test_doc_comment_is_formatted_correctly(
            "Lorem ipsum",
            "///Lorem ipsum",
            CommentStyle::TripleSlash,
        );
    }

    #[test]
    fn multi_line_doc_comment_is_formatted_correctly() {
        test_doc_comment_is_formatted_correctly(
            "Lorem ipsum\nDolor sit amet",
            "///Lorem ipsum\n///Dolor sit amet",
            CommentStyle::TripleSlash,
        );
    }

    #[test]
    fn whitespace_within_lines_is_preserved() {
        test_doc_comment_is_formatted_correctly(
            " Lorem ipsum \n Dolor sit amet ",
            "/// Lorem ipsum \n/// Dolor sit amet ",
            CommentStyle::TripleSlash,
        );
    }

    fn test_doc_comment_is_formatted_correctly(
        literal: &str,
        expected_comment: &str,
        style: CommentStyle<'_>,
    ) {
        assert_eq!(
            expected_comment,
            format!("{}", DocCommentFormatter::new(&literal, style))
        );
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::ffi::{OsStr, OsString};
use std::fs::{self, File, OpenOptions};
use std::io;
use std::io::prelude::*;
use std::path::{Component, Path, PathBuf};
use std::iter;

use filetime::FileTime;

use util::{internal, CargoResult};
use util::errors::{CargoError, CargoResultExt, Internal};

pub fn join_paths<T: AsRef<OsStr>>(paths: &[T], env: &str) -> CargoResult<OsString> {
    let err = match env::join_paths(paths.iter()) {
        Ok(paths) => return Ok(paths),
        Err(e) => e,
    };
    let paths = paths.iter().map(Path::new).collect::<Vec<_>>();
    let err = CargoError::from(err);
    let explain = Internal::new(format_err!("failed to join path array: {:?}", paths));
    let err = CargoError::from(err.context(explain));
    let more_explain = format!(
        "failed to join search paths together\n\
         Does ${} have an unterminated quote character?",
        env
    );
    Err(err.context(more_explain).into())
}

pub fn dylib_path_envvar() -> &'static str {
    if cfg!(windows) {
        "PATH"
    } else if cfg!(target_os = "macos") {
        "DYLD_LIBRARY_PATH"
    } else {
        "LD_LIBRARY_PATH"
    }
}

pub fn dylib_path() -> Vec<PathBuf> {
    match env::var_os(dylib_path_envvar()) {
        Some(var) => env::split_paths(&var).collect(),
        None => Vec::new(),
    }
}

pub fn normalize_path(path: &Path) -> PathBuf {
    let mut components = path.components().peekable();
    let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().cloned() {
        components.next();
        PathBuf::from(c.as_os_str())
    } else {
        PathBuf::new()
    };

    for component in components {
        match component {
            Component::Prefix(..) => unreachable!(),
            Component::RootDir => {
                ret.push(component.as_os_str());
            }
            Component::CurDir => {}
            Component::ParentDir => {
                ret.pop();
            }
            Component::Normal(c) => {
                ret.push(c);
            }
        }
    }
    ret
}

pub fn without_prefix<'a>(long_path: &'a Path, prefix: &'a Path) -> Option<&'a Path> {
    let mut a = long_path.components();
    let mut b = prefix.components();
    loop {
        match b.next() {
            Some(y) => match a.next() {
                Some(x) if x == y => continue,
                _ => return None,
            },
            None => return Some(a.as_path()),
        }
    }
}

pub fn resolve_executable(exec: &Path) -> CargoResult<PathBuf> {
    if exec.components().count() == 1 {
        let paths = env::var_os("PATH").ok_or(format_err!("no PATH"))?;
        let candidates = env::split_paths(&paths).flat_map(|path| {
            let candidate = PathBuf::from(path).join(&exec);
            let with_exe = if env::consts::EXE_EXTENSION == "" {
                None
            } else {
                Some(candidate.with_extension(env::consts::EXE_EXTENSION))
            };
            iter::once(candidate).chain(with_exe)
        });
        for candidate in candidates {
            if candidate.is_file() {
                // PATH may have a component like "." in it, so we still need to
                // canonicalize.
                return Ok(candidate.canonicalize()?);
            }
        }

        bail!("no executable for `{}` found in PATH", exec.display())
    } else {
        Ok(exec.canonicalize()?)
    }
}

pub fn read(path: &Path) -> CargoResult<String> {
    match String::from_utf8(read_bytes(path)?) {
        Ok(s) => Ok(s),
        Err(_) => bail!("path at `{}` was not valid utf-8", path.display()),
    }
}

pub fn read_bytes(path: &Path) -> CargoResult<Vec<u8>> {
    let res = (|| -> CargoResult<_> {
        let mut ret = Vec::new();
        let mut f = File::open(path)?;
        if let Ok(m) = f.metadata() {
            ret.reserve(m.len() as usize + 1);
        }
        f.read_to_end(&mut ret)?;
        Ok(ret)
    })()
        .chain_err(|| format!("failed to read `{}`", path.display()))?;
    Ok(res)
}

pub fn write(path: &Path, contents: &[u8]) -> CargoResult<()> {
    (|| -> CargoResult<()> {
        let mut f = File::create(path)?;
        f.write_all(contents)?;
        Ok(())
    })()
        .chain_err(|| format!("failed to write `{}`", path.display()))?;
    Ok(())
}

pub fn append(path: &Path, contents: &[u8]) -> CargoResult<()> {
    (|| -> CargoResult<()> {
        let mut f = OpenOptions::new()
            .write(true)
            .append(true)
            .create(true)
            .open(path)?;

        f.write_all(contents)?;
        Ok(())
    })()
        .chain_err(|| internal(format!("failed to write `{}`", path.display())))?;
    Ok(())
}

pub fn mtime(path: &Path) -> CargoResult<FileTime> {
    let meta =
        fs::metadata(path).chain_err(|| internal(format!("failed to stat `{}`", path.display())))?;
    Ok(FileTime::from_last_modification_time(&meta))
}

#[cfg(unix)]
pub fn path2bytes(path: &Path) -> CargoResult<&[u8]> {
    use std::os::unix::prelude::*;
    Ok(path.as_os_str().as_bytes())
}
#[cfg(windows)]
pub fn path2bytes(path: &Path) -> CargoResult<&[u8]> {
    match path.as_os_str().to_str() {
        Some(s) => Ok(s.as_bytes()),
        None => Err(format_err!("invalid non-unicode path: {}", path.display())),
    }
}

#[cfg(unix)]
pub fn bytes2path(bytes: &[u8]) -> CargoResult<PathBuf> {
    use std::os::unix::prelude::*;
    use std::ffi::OsStr;
    Ok(PathBuf::from(OsStr::from_bytes(bytes)))
}
#[cfg(windows)]
pub fn bytes2path(bytes: &[u8]) -> CargoResult<PathBuf> {
    use std::str;
    match str::from_utf8(bytes) {
        Ok(s) => Ok(PathBuf::from(s)),
        Err(..) => Err(format_err!("invalid non-unicode path")),
    }
}

pub fn ancestors(path: &Path) -> PathAncestors {
    PathAncestors::new(path)
}

pub struct PathAncestors<'a> {
    current: Option<&'a Path>,
    stop_at: Option<PathBuf>,
}

impl<'a> PathAncestors<'a> {
    fn new(path: &Path) -> PathAncestors {
        PathAncestors {
            current: Some(path),
            //HACK: avoid reading `~/.cargo/config` when testing Cargo itself.
            stop_at: env::var("__CARGO_TEST_ROOT").ok().map(PathBuf::from),
        }
    }
}

impl<'a> Iterator for PathAncestors<'a> {
    type Item = &'a Path;

    fn next(&mut self) -> Option<&'a Path> {
        if let Some(path) = self.current {
            self.current = path.parent();

            if let Some(ref stop_at) = self.stop_at {
                if path == stop_at {
                    self.current = None;
                }
            }

            Some(path)
        } else {
            None
        }
    }
}

pub fn remove_dir_all<P: AsRef<Path>>(p: P) -> CargoResult<()> {
    _remove_dir_all(p.as_ref())
}

fn _remove_dir_all(p: &Path) -> CargoResult<()> {
    if p.symlink_metadata()?.file_type().is_symlink() {
        return remove_file(p);
    }
    let entries = p.read_dir()
        .chain_err(|| format!("failed to read directory `{}`", p.display()))?;
    for entry in entries {
        let entry = entry?;
        let path = entry.path();
        if entry.file_type()?.is_dir() {
            remove_dir_all(&path)?;
        } else {
            remove_file(&path)?;
        }
    }
    remove_dir(&p)
}

pub fn remove_dir<P: AsRef<Path>>(p: P) -> CargoResult<()> {
    _remove_dir(p.as_ref())
}

fn _remove_dir(p: &Path) -> CargoResult<()> {
    fs::remove_dir(p).chain_err(|| format!("failed to remove directory `{}`", p.display()))?;
    Ok(())
}

pub fn remove_file<P: AsRef<Path>>(p: P) -> CargoResult<()> {
    _remove_file(p.as_ref())
}

fn _remove_file(p: &Path) -> CargoResult<()> {
    let mut err = match fs::remove_file(p) {
        Ok(()) => return Ok(()),
        Err(e) => e,
    };

    if err.kind() == io::ErrorKind::PermissionDenied && set_not_readonly(p).unwrap_or(false) {
        match fs::remove_file(p) {
            Ok(()) => return Ok(()),
            Err(e) => err = e,
        }
    }

    Err(err).chain_err(|| format!("failed to remove file `{}`", p.display()))?;
    Ok(())
}

fn set_not_readonly(p: &Path) -> io::Result<bool> {
    let mut perms = p.metadata()?.permissions();
    if !perms.readonly() {
        return Ok(false);
    }
    perms.set_readonly(false);
    fs::set_permissions(p, perms)?;
    Ok(true)
}
<commit_msg>Remove `internal` from error handling<commit_after>use std::env;
use std::ffi::{OsStr, OsString};
use std::fs::{self, File, OpenOptions};
use std::io;
use std::io::prelude::*;
use std::path::{Component, Path, PathBuf};
use std::iter;

use filetime::FileTime;

use util::errors::{CargoError, CargoResult, CargoResultExt, Internal};

pub fn join_paths<T: AsRef<OsStr>>(paths: &[T], env: &str) -> CargoResult<OsString> {
    let err = match env::join_paths(paths.iter()) {
        Ok(paths) => return Ok(paths),
        Err(e) => e,
    };
    let paths = paths.iter().map(Path::new).collect::<Vec<_>>();
    let err = CargoError::from(err);
    let explain = Internal::new(format_err!("failed to join path array: {:?}", paths));
    let err = CargoError::from(err.context(explain));
    let more_explain = format!(
        "failed to join search paths together\n\
         Does ${} have an unterminated quote character?",
        env
    );
    Err(err.context(more_explain).into())
}

pub fn dylib_path_envvar() -> &'static str {
    if cfg!(windows) {
        "PATH"
    } else if cfg!(target_os = "macos") {
        "DYLD_LIBRARY_PATH"
    } else {
        "LD_LIBRARY_PATH"
    }
}

pub fn dylib_path() -> Vec<PathBuf> {
    match env::var_os(dylib_path_envvar()) {
        Some(var) => env::split_paths(&var).collect(),
        None => Vec::new(),
    }
}

pub fn normalize_path(path: &Path) -> PathBuf {
    let mut components = path.components().peekable();
    let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().cloned() {
        components.next();
        PathBuf::from(c.as_os_str())
    } else {
        PathBuf::new()
    };

    for component in components {
        match component {
            Component::Prefix(..) => unreachable!(),
            Component::RootDir => {
                ret.push(component.as_os_str());
            }
            Component::CurDir => {}
            Component::ParentDir => {
                ret.pop();
            }
            Component::Normal(c) => {
                ret.push(c);
            }
        }
    }
    ret
}

pub fn without_prefix<'a>(long_path: &'a Path, prefix: &'a Path) -> Option<&'a Path> {
    let mut a = long_path.components();
    let mut b = prefix.components();
    loop {
        match b.next() {
            Some(y) => match a.next() {
                Some(x) if x == y => continue,
                _ => return None,
            },
            None => return Some(a.as_path()),
        }
    }
}

pub fn resolve_executable(exec: &Path) -> CargoResult<PathBuf> {
    if exec.components().count() == 1 {
        let paths = env::var_os("PATH").ok_or(format_err!("no PATH"))?;
        let candidates = env::split_paths(&paths).flat_map(|path| {
            let candidate = PathBuf::from(path).join(&exec);
            let with_exe = if env::consts::EXE_EXTENSION == "" {
                None
            } else {
                Some(candidate.with_extension(env::consts::EXE_EXTENSION))
            };
            iter::once(candidate).chain(with_exe)
        });
        for candidate in candidates {
            if candidate.is_file() {
                // PATH may have a component like "." in it, so we still need to
                // canonicalize.
                return Ok(candidate.canonicalize()?);
            }
        }

        bail!("no executable for `{}` found in PATH", exec.display())
    } else {
        Ok(exec.canonicalize()?)
    }
}

pub fn read(path: &Path) -> CargoResult<String> {
    match String::from_utf8(read_bytes(path)?) {
        Ok(s) => Ok(s),
        Err(_) => bail!("path at `{}` was not valid utf-8", path.display()),
    }
}

pub fn read_bytes(path: &Path) -> CargoResult<Vec<u8>> {
    let res = (|| -> CargoResult<_> {
        let mut ret = Vec::new();
        let mut f = File::open(path)?;
        if let Ok(m) = f.metadata() {
            ret.reserve(m.len() as usize + 1);
        }
        f.read_to_end(&mut ret)?;
        Ok(ret)
    })()
        .chain_err(|| format!("failed to read `{}`", path.display()))?;
    Ok(res)
}

pub fn write(path: &Path, contents: &[u8]) -> CargoResult<()> {
    (|| -> CargoResult<()> {
        let mut f = File::create(path)?;
        f.write_all(contents)?;
        Ok(())
    })()
        .chain_err(|| format!("failed to write `{}`", path.display()))?;
    Ok(())
}

pub fn append(path: &Path, contents: &[u8]) -> CargoResult<()> {
    (|| -> CargoResult<()> {
        let mut f = OpenOptions::new()
            .write(true)
            .append(true)
            .create(true)
            .open(path)?;

        f.write_all(contents)?;
        Ok(())
    })()
        .chain_err(|| format!("failed to write `{}`", path.display()))?;
    Ok(())
}

pub fn mtime(path: &Path) -> CargoResult<FileTime> {
    let meta =
        fs::metadata(path).chain_err(|| format!("failed to stat `{}`", path.display()))?;
    Ok(FileTime::from_last_modification_time(&meta))
}

#[cfg(unix)]
pub fn path2bytes(path: &Path) -> CargoResult<&[u8]> {
    use std::os::unix::prelude::*;
    Ok(path.as_os_str().as_bytes())
}
#[cfg(windows)]
pub fn path2bytes(path: &Path) -> CargoResult<&[u8]> {
    match path.as_os_str().to_str() {
        Some(s) => Ok(s.as_bytes()),
        None => Err(format_err!("invalid non-unicode path: {}", path.display())),
    }
}

#[cfg(unix)]
pub fn bytes2path(bytes: &[u8]) -> CargoResult<PathBuf> {
    use std::os::unix::prelude::*;
    use std::ffi::OsStr;
    Ok(PathBuf::from(OsStr::from_bytes(bytes)))
}
#[cfg(windows)]
pub fn bytes2path(bytes: &[u8]) -> CargoResult<PathBuf> {
    use std::str;
    match str::from_utf8(bytes) {
        Ok(s) => Ok(PathBuf::from(s)),
        Err(..) => Err(format_err!("invalid non-unicode path")),
    }
}

pub fn ancestors(path: &Path) -> PathAncestors {
    PathAncestors::new(path)
}

pub struct PathAncestors<'a> {
    current: Option<&'a Path>,
    stop_at: Option<PathBuf>,
}

impl<'a> PathAncestors<'a> {
    fn new(path: &Path) -> PathAncestors {
        PathAncestors {
            current: Some(path),
            //HACK: avoid reading `~/.cargo/config` when testing Cargo itself.
            stop_at: env::var("__CARGO_TEST_ROOT").ok().map(PathBuf::from),
        }
    }
}

impl<'a> Iterator for PathAncestors<'a> {
    type Item = &'a Path;

    fn next(&mut self) -> Option<&'a Path> {
        if let Some(path) = self.current {
            self.current = path.parent();

            if let Some(ref stop_at) = self.stop_at {
                if path == stop_at {
                    self.current = None;
                }
            }

            Some(path)
        } else {
            None
        }
    }
}

pub fn remove_dir_all<P: AsRef<Path>>(p: P) -> CargoResult<()> {
    _remove_dir_all(p.as_ref())
}

fn _remove_dir_all(p: &Path) -> CargoResult<()> {
    if p.symlink_metadata()?.file_type().is_symlink() {
        return remove_file(p);
    }
    let entries = p.read_dir()
        .chain_err(|| format!("failed to read directory `{}`", p.display()))?;
    for entry in entries {
        let entry = entry?;
        let path = entry.path();
        if entry.file_type()?.is_dir() {
            remove_dir_all(&path)?;
        } else {
            remove_file(&path)?;
        }
    }
    remove_dir(&p)
}

pub fn remove_dir<P: AsRef<Path>>(p: P) -> CargoResult<()> {
    _remove_dir(p.as_ref())
}

fn _remove_dir(p: &Path) -> CargoResult<()> {
    fs::remove_dir(p).chain_err(|| format!("failed to remove directory `{}`", p.display()))?;
    Ok(())
}

pub fn remove_file<P: AsRef<Path>>(p: P) -> CargoResult<()> {
    _remove_file(p.as_ref())
}

fn _remove_file(p: &Path) -> CargoResult<()> {
    let mut err = match fs::remove_file(p) {
        Ok(()) => return Ok(()),
        Err(e) => e,
    };

    if err.kind() == io::ErrorKind::PermissionDenied && set_not_readonly(p).unwrap_or(false) {
        match fs::remove_file(p) {
            Ok(()) => return Ok(()),
            Err(e) => err = e,
        }
    }

    Err(err).chain_err(|| format!("failed to remove file `{}`", p.display()))?;
    Ok(())
}

fn set_not_readonly(p: &Path) -> io::Result<bool> {
    let mut perms = p.metadata()?.permissions();
    if !perms.readonly() {
        return Ok(false);
    }
    perms.set_readonly(false);
    fs::set_permissions(p, perms)?;
    Ok(true)
}
<|endoftext|>


<commit_before>//! A flash simulator
//!
//! This module is capable of simulating the type of NOR flash commonly used in microcontrollers.
//! These generally can be written as individual bytes, but must be erased in larger units.

#[macro_use] extern crate log;
#[macro_use] extern crate error_chain;
extern crate rand;
mod pdump;

use rand::distributions::{IndependentSample, Range};
use std::fs::File;
use std::io::Write;
use std::iter::Enumerate;
use std::path::Path;
use std::slice;
use pdump::HexDump;

error_chain! {
    errors {
        OutOfBounds(t: String) {
            description("Offset is out of bounds")
            display("Offset out of bounds: {}", t)
        }
        Write(t: String) {
            description("Invalid write")
            display("Invalid write: {}", t)
        }
        SimulatedFail(t: String) {
            description("Write failed by chance")
            display("Failed write: {}", t)
        }
    }
}

pub trait Flash {
    fn erase(&mut self, offset: usize, len: usize) -> Result<()>;
    fn write(&mut self, offset: usize, payload: &[u8]) -> Result<()>;
    fn read(&self, offset: usize, data: &mut [u8]) -> Result<()>;

    fn add_bad_region(&mut self, offset: usize, len: usize, rate: f32) -> Result<()>;
    fn reset_bad_regions(&mut self);

    fn set_verify_writes(&mut self, enable: bool);

    fn sector_iter(&self) -> SectorIter;
    fn device_size(&self) -> usize;
}

fn ebounds<T: AsRef<str>>(message: T) -> ErrorKind {
    ErrorKind::OutOfBounds(message.as_ref().to_owned())
}

#[allow(dead_code)]
fn ewrite<T: AsRef<str>>(message: T) -> ErrorKind {
    ErrorKind::Write(message.as_ref().to_owned())
}

#[allow(dead_code)]
fn esimulatedwrite<T: AsRef<str>>(message: T) -> ErrorKind {
    ErrorKind::SimulatedFail(message.as_ref().to_owned())
}

/// An emulated flash device.  It is represented as a block of bytes, and a list of the sector
/// mapings.
#[derive(Clone)]
pub struct SimFlash {
    data: Vec<u8>,
    write_safe: Vec<bool>,
    sectors: Vec<usize>,
    bad_region: Vec<(usize, usize, f32)>,
    // Alignment required for writes.
    align: usize,
    verify_writes: bool,
}

impl SimFlash {
    /// Given a sector size map, construct a flash device for that.
    pub fn new(sectors: Vec<usize>, align: usize) -> SimFlash {
        // Verify that the alignment is a positive power of two.
        assert!(align > 0);
        assert!(align & (align - 1) == 0);

        let total = sectors.iter().sum();
        SimFlash {
            data: vec![0xffu8; total],
            write_safe: vec![true; total],
            sectors: sectors,
            bad_region: Vec::new(),
            align: align,
            verify_writes: true,
        }
    }

    #[allow(dead_code)]
    pub fn dump(&self) {
        self.data.dump();
    }

    /// Dump this image to the given file.
    #[allow(dead_code)]
    pub fn write_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let mut fd = File::create(path).chain_err(|| "Unable to write image file")?;
        fd.write_all(&self.data).chain_err(|| "Unable to write to image file")?;
        Ok(())
    }

    // Scan the sector map, and return the base and offset within a sector for this given byte.
    // Returns None if the value is outside of the device.
    fn get_sector(&self, offset: usize) -> Option<(usize, usize)> {
        let mut offset = offset;
        for (sector, &size) in self.sectors.iter().enumerate() {
            if offset < size {
                return Some((sector, offset));
            }
            offset -= size;
        }
        return None;
    }

}

impl Flash for SimFlash {
    /// The flash drivers tend to erase beyond the bounds of the given range.  Instead, we'll be
    /// strict, and make sure that the passed arguments are exactly at a sector boundary, otherwise
    /// return an error.
    fn erase(&mut self, offset: usize, len: usize) -> Result<()> {
        let (_start, slen) = self.get_sector(offset).ok_or_else(|| ebounds("start"))?;
        let (end, elen) = self.get_sector(offset + len - 1).ok_or_else(|| ebounds("end"))?;

        if slen != 0 {
            bail!(ebounds("offset not at start of sector"));
        }
        if elen != self.sectors[end] - 1 {
            bail!(ebounds("end not at start of sector"));
        }

        for x in &mut self.data[offset .. offset + len] {
            *x = 0xff;
        }

        for x in &mut self.write_safe[offset .. offset + len] {
            *x = true;
        }

        Ok(())
    }

    /// We restrict to only allowing writes of values that are:
    ///
    /// 1. being written to for the first time
    /// 2. being written to after being erased
    ///
    /// This emulates a flash device which starts out erased, with the
    /// added restriction that repeated writes to the same location
    /// are disallowed, even if they would be safe to do.
    fn write(&mut self, offset: usize, payload: &[u8]) -> Result<()> {
        for &(off, len, rate) in &self.bad_region {
            if offset >= off && (offset + payload.len()) <= (off + len) {
                let mut rng = rand::thread_rng();
                let between = Range::new(0., 1.);
                if between.ind_sample(&mut rng) < rate {
                    bail!(esimulatedwrite(
                        format!("Ignoring write to {:#x}-{:#x}", off, off + len)));
                }
            }
        }

        if offset + payload.len() > self.data.len() {
            panic!("Write outside of device");
        }

        // Verify the alignment (which must be a power of two).
        if offset & (self.align - 1) != 0 {
            panic!("Misaligned write address");
        }

        if payload.len() & (self.align - 1) != 0 {
            panic!("Write length not multiple of alignment");
        }

        for (i, x) in &mut self.write_safe[offset .. offset + payload.len()].iter_mut().enumerate() {
            if self.verify_writes && !(*x) {
                panic!("Write to unerased location at 0x{:x}", offset + i);
            }
            *x = false;
        }

        let mut sub = &mut self.data[offset .. offset + payload.len()];
        sub.copy_from_slice(payload);
        Ok(())
    }

    /// Read is simple.
    fn read(&self, offset: usize, data: &mut [u8]) -> Result<()> {
        if offset + data.len() > self.data.len() {
            bail!(ebounds("Read outside of device"));
        }

        let sub = &self.data[offset .. offset + data.len()];
        data.copy_from_slice(sub);
        Ok(())
    }

    /// Adds a new flash bad region. Writes to this area fail with a chance
    /// given by `rate`.
    fn add_bad_region(&mut self, offset: usize, len: usize, rate: f32) -> Result<()> {
        if rate < 0.0 || rate > 1.0 {
            bail!(ebounds("Invalid rate"));
        }

        info!("Adding new bad region {:#x}-{:#x}", offset, offset + len);
        self.bad_region.push((offset, len, rate));

        Ok(())
    }

    fn reset_bad_regions(&mut self) {
        self.bad_region.clear();
    }

    fn set_verify_writes(&mut self, enable: bool) {
        self.verify_writes = enable;
    }

    /// An iterator over each sector in the device.
    fn sector_iter(&self) -> SectorIter {
        SectorIter {
            iter: self.sectors.iter().enumerate(),
            base: 0,
        }
    }

    fn device_size(&self) -> usize {
        self.data.len()
    }
}

/// It is possible to iterate over the sectors in the device, each element returning this.
#[derive(Debug, Clone)]
pub struct Sector {
    /// Which sector is this, starting from 0.
    pub num: usize,
    /// The offset, in bytes, of the start of this sector.
    pub base: usize,
    /// The length, in bytes, of this sector.
    pub size: usize,
}

pub struct SectorIter<'a> {
    iter: Enumerate<slice::Iter<'a, usize>>,
    base: usize,
}

impl<'a> Iterator for SectorIter<'a> {
    type Item = Sector;

    fn next(&mut self) -> Option<Sector> {
        match self.iter.next() {
            None => None,
            Some((num, &size)) => {
                let base = self.base;
                self.base += size;
                Some(Sector {
                    num: num,
                    base: base,
                    size: size,
                })
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::{Flash, SimFlash, Error, ErrorKind, Result, Sector};

    #[test]
    fn test_flash() {
        // NXP-style, uniform sectors.
        let mut f1 = SimFlash::new(vec![4096usize; 256], 1);
        test_device(&mut f1);

        // STM style, non-uniform sectors
        let mut f2 = SimFlash::new(vec![16 * 1024, 16 * 1024, 16 * 1024, 64 * 1024,
                                128 * 1024, 128 * 1024, 128 * 1024], 1);
        test_device(&mut f2);
    }

    fn test_device(flash: &mut Flash) {
        let sectors: Vec<Sector> = flash.sector_iter().collect();

        flash.erase(0, sectors[0].size).unwrap();
        let flash_size = flash.device_size();
        flash.erase(0, flash_size).unwrap();
        assert!(flash.erase(0, sectors[0].size - 1).is_bounds());

        // Verify that write and erase do something.
        flash.write(0, &[0]).unwrap();
        let mut buf = [0; 4];
        flash.read(0, &mut buf).unwrap();
        assert_eq!(buf, [0, 0xff, 0xff, 0xff]);

        flash.erase(0, sectors[0].size).unwrap();
        flash.read(0, &mut buf).unwrap();
        assert_eq!(buf, [0xff; 4]);

        // Program the first and last byte of each sector, verify that has been done, and then
        // erase to verify the erase boundaries.
        for sector in &sectors {
            let byte = [(sector.num & 127) as u8];
            flash.write(sector.base, &byte).unwrap();
            flash.write(sector.base + sector.size - 1, &byte).unwrap();
        }

        // Verify the above
        let mut buf = Vec::new();
        for sector in &sectors {
            let byte = (sector.num & 127) as u8;
            buf.resize(sector.size, 0);
            flash.read(sector.base, &mut buf).unwrap();
            assert_eq!(buf.first(), Some(&byte));
            assert_eq!(buf.last(), Some(&byte));
            assert!(buf[1..buf.len()-1].iter().all(|&x| x == 0xff));
        }
    }

    // Helper checks for the result type.
    trait EChecker {
        fn is_bounds(&self) -> bool;
    }

    impl<T> EChecker for Result<T> {

        fn is_bounds(&self) -> bool {
            match *self {
                Err(Error(ErrorKind::OutOfBounds(_), _)) => true,
                _ => false,
            }
        }
    }
}
<commit_msg>sim: Remove extraneous `mut`<commit_after>//! A flash simulator
//!
//! This module is capable of simulating the type of NOR flash commonly used in microcontrollers.
//! These generally can be written as individual bytes, but must be erased in larger units.

#[macro_use] extern crate log;
#[macro_use] extern crate error_chain;
extern crate rand;
mod pdump;

use rand::distributions::{IndependentSample, Range};
use std::fs::File;
use std::io::Write;
use std::iter::Enumerate;
use std::path::Path;
use std::slice;
use pdump::HexDump;

error_chain! {
    errors {
        OutOfBounds(t: String) {
            description("Offset is out of bounds")
            display("Offset out of bounds: {}", t)
        }
        Write(t: String) {
            description("Invalid write")
            display("Invalid write: {}", t)
        }
        SimulatedFail(t: String) {
            description("Write failed by chance")
            display("Failed write: {}", t)
        }
    }
}

pub trait Flash {
    fn erase(&mut self, offset: usize, len: usize) -> Result<()>;
    fn write(&mut self, offset: usize, payload: &[u8]) -> Result<()>;
    fn read(&self, offset: usize, data: &mut [u8]) -> Result<()>;

    fn add_bad_region(&mut self, offset: usize, len: usize, rate: f32) -> Result<()>;
    fn reset_bad_regions(&mut self);

    fn set_verify_writes(&mut self, enable: bool);

    fn sector_iter(&self) -> SectorIter;
    fn device_size(&self) -> usize;
}

fn ebounds<T: AsRef<str>>(message: T) -> ErrorKind {
    ErrorKind::OutOfBounds(message.as_ref().to_owned())
}

#[allow(dead_code)]
fn ewrite<T: AsRef<str>>(message: T) -> ErrorKind {
    ErrorKind::Write(message.as_ref().to_owned())
}

#[allow(dead_code)]
fn esimulatedwrite<T: AsRef<str>>(message: T) -> ErrorKind {
    ErrorKind::SimulatedFail(message.as_ref().to_owned())
}

/// An emulated flash device.  It is represented as a block of bytes, and a list of the sector
/// mapings.
#[derive(Clone)]
pub struct SimFlash {
    data: Vec<u8>,
    write_safe: Vec<bool>,
    sectors: Vec<usize>,
    bad_region: Vec<(usize, usize, f32)>,
    // Alignment required for writes.
    align: usize,
    verify_writes: bool,
}

impl SimFlash {
    /// Given a sector size map, construct a flash device for that.
    pub fn new(sectors: Vec<usize>, align: usize) -> SimFlash {
        // Verify that the alignment is a positive power of two.
        assert!(align > 0);
        assert!(align & (align - 1) == 0);

        let total = sectors.iter().sum();
        SimFlash {
            data: vec![0xffu8; total],
            write_safe: vec![true; total],
            sectors: sectors,
            bad_region: Vec::new(),
            align: align,
            verify_writes: true,
        }
    }

    #[allow(dead_code)]
    pub fn dump(&self) {
        self.data.dump();
    }

    /// Dump this image to the given file.
    #[allow(dead_code)]
    pub fn write_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let mut fd = File::create(path).chain_err(|| "Unable to write image file")?;
        fd.write_all(&self.data).chain_err(|| "Unable to write to image file")?;
        Ok(())
    }

    // Scan the sector map, and return the base and offset within a sector for this given byte.
    // Returns None if the value is outside of the device.
    fn get_sector(&self, offset: usize) -> Option<(usize, usize)> {
        let mut offset = offset;
        for (sector, &size) in self.sectors.iter().enumerate() {
            if offset < size {
                return Some((sector, offset));
            }
            offset -= size;
        }
        return None;
    }

}

impl Flash for SimFlash {
    /// The flash drivers tend to erase beyond the bounds of the given range.  Instead, we'll be
    /// strict, and make sure that the passed arguments are exactly at a sector boundary, otherwise
    /// return an error.
    fn erase(&mut self, offset: usize, len: usize) -> Result<()> {
        let (_start, slen) = self.get_sector(offset).ok_or_else(|| ebounds("start"))?;
        let (end, elen) = self.get_sector(offset + len - 1).ok_or_else(|| ebounds("end"))?;

        if slen != 0 {
            bail!(ebounds("offset not at start of sector"));
        }
        if elen != self.sectors[end] - 1 {
            bail!(ebounds("end not at start of sector"));
        }

        for x in &mut self.data[offset .. offset + len] {
            *x = 0xff;
        }

        for x in &mut self.write_safe[offset .. offset + len] {
            *x = true;
        }

        Ok(())
    }

    /// We restrict to only allowing writes of values that are:
    ///
    /// 1. being written to for the first time
    /// 2. being written to after being erased
    ///
    /// This emulates a flash device which starts out erased, with the
    /// added restriction that repeated writes to the same location
    /// are disallowed, even if they would be safe to do.
    fn write(&mut self, offset: usize, payload: &[u8]) -> Result<()> {
        for &(off, len, rate) in &self.bad_region {
            if offset >= off && (offset + payload.len()) <= (off + len) {
                let mut rng = rand::thread_rng();
                let between = Range::new(0., 1.);
                if between.ind_sample(&mut rng) < rate {
                    bail!(esimulatedwrite(
                        format!("Ignoring write to {:#x}-{:#x}", off, off + len)));
                }
            }
        }

        if offset + payload.len() > self.data.len() {
            panic!("Write outside of device");
        }

        // Verify the alignment (which must be a power of two).
        if offset & (self.align - 1) != 0 {
            panic!("Misaligned write address");
        }

        if payload.len() & (self.align - 1) != 0 {
            panic!("Write length not multiple of alignment");
        }

        for (i, x) in &mut self.write_safe[offset .. offset + payload.len()].iter_mut().enumerate() {
            if self.verify_writes && !(*x) {
                panic!("Write to unerased location at 0x{:x}", offset + i);
            }
            *x = false;
        }

        let sub = &mut self.data[offset .. offset + payload.len()];
        sub.copy_from_slice(payload);
        Ok(())
    }

    /// Read is simple.
    fn read(&self, offset: usize, data: &mut [u8]) -> Result<()> {
        if offset + data.len() > self.data.len() {
            bail!(ebounds("Read outside of device"));
        }

        let sub = &self.data[offset .. offset + data.len()];
        data.copy_from_slice(sub);
        Ok(())
    }

    /// Adds a new flash bad region. Writes to this area fail with a chance
    /// given by `rate`.
    fn add_bad_region(&mut self, offset: usize, len: usize, rate: f32) -> Result<()> {
        if rate < 0.0 || rate > 1.0 {
            bail!(ebounds("Invalid rate"));
        }

        info!("Adding new bad region {:#x}-{:#x}", offset, offset + len);
        self.bad_region.push((offset, len, rate));

        Ok(())
    }

    fn reset_bad_regions(&mut self) {
        self.bad_region.clear();
    }

    fn set_verify_writes(&mut self, enable: bool) {
        self.verify_writes = enable;
    }

    /// An iterator over each sector in the device.
    fn sector_iter(&self) -> SectorIter {
        SectorIter {
            iter: self.sectors.iter().enumerate(),
            base: 0,
        }
    }

    fn device_size(&self) -> usize {
        self.data.len()
    }
}

/// It is possible to iterate over the sectors in the device, each element returning this.
#[derive(Debug, Clone)]
pub struct Sector {
    /// Which sector is this, starting from 0.
    pub num: usize,
    /// The offset, in bytes, of the start of this sector.
    pub base: usize,
    /// The length, in bytes, of this sector.
    pub size: usize,
}

pub struct SectorIter<'a> {
    iter: Enumerate<slice::Iter<'a, usize>>,
    base: usize,
}

impl<'a> Iterator for SectorIter<'a> {
    type Item = Sector;

    fn next(&mut self) -> Option<Sector> {
        match self.iter.next() {
            None => None,
            Some((num, &size)) => {
                let base = self.base;
                self.base += size;
                Some(Sector {
                    num: num,
                    base: base,
                    size: size,
                })
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::{Flash, SimFlash, Error, ErrorKind, Result, Sector};

    #[test]
    fn test_flash() {
        // NXP-style, uniform sectors.
        let mut f1 = SimFlash::new(vec![4096usize; 256], 1);
        test_device(&mut f1);

        // STM style, non-uniform sectors
        let mut f2 = SimFlash::new(vec![16 * 1024, 16 * 1024, 16 * 1024, 64 * 1024,
                                128 * 1024, 128 * 1024, 128 * 1024], 1);
        test_device(&mut f2);
    }

    fn test_device(flash: &mut Flash) {
        let sectors: Vec<Sector> = flash.sector_iter().collect();

        flash.erase(0, sectors[0].size).unwrap();
        let flash_size = flash.device_size();
        flash.erase(0, flash_size).unwrap();
        assert!(flash.erase(0, sectors[0].size - 1).is_bounds());

        // Verify that write and erase do something.
        flash.write(0, &[0]).unwrap();
        let mut buf = [0; 4];
        flash.read(0, &mut buf).unwrap();
        assert_eq!(buf, [0, 0xff, 0xff, 0xff]);

        flash.erase(0, sectors[0].size).unwrap();
        flash.read(0, &mut buf).unwrap();
        assert_eq!(buf, [0xff; 4]);

        // Program the first and last byte of each sector, verify that has been done, and then
        // erase to verify the erase boundaries.
        for sector in &sectors {
            let byte = [(sector.num & 127) as u8];
            flash.write(sector.base, &byte).unwrap();
            flash.write(sector.base + sector.size - 1, &byte).unwrap();
        }

        // Verify the above
        let mut buf = Vec::new();
        for sector in &sectors {
            let byte = (sector.num & 127) as u8;
            buf.resize(sector.size, 0);
            flash.read(sector.base, &mut buf).unwrap();
            assert_eq!(buf.first(), Some(&byte));
            assert_eq!(buf.last(), Some(&byte));
            assert!(buf[1..buf.len()-1].iter().all(|&x| x == 0xff));
        }
    }

    // Helper checks for the result type.
    trait EChecker {
        fn is_bounds(&self) -> bool;
    }

    impl<T> EChecker for Result<T> {

        fn is_bounds(&self) -> bool {
            match *self {
                Err(Error(ErrorKind::OutOfBounds(_), _)) => true,
                _ => false,
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::{convert, io};
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

struct PeerConnection {
    stream: TcpStream,
    have: Vec<bool>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl PeerConnection {
    fn connect(peer: &Peer, metainfo: &Metainfo) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let mut conn = PeerConnection {
            stream: stream,
            have:   vec![false; metainfo.info.pieces.len()],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                self.send_interested();
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                self.send_interested();
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                }
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }
    fn send_interested(&mut self) -> Result<(), Error>{
        if self.am_i_interested == false {
            let mut message = vec![];
            message.push(0);
            message.push(0);
            message.push(0);
            message.push(1);
            message.push(2);
            try!(self.stream.write_all(&message));
            self.am_i_interested = true;
        }
        Ok(())
    }
}

#[derive(Debug)]
enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece,   // TODO add params
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => Message::Piece,
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * 16777216 +
    bytes[1] as usize * 65536 +
    bytes[2] as usize * 256 +
    bytes[3] as usize
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Sending a Request for the first piece<commit_after>use std::{convert, io};
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

use metainfo::Metainfo;
use tracker_response::Peer;

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

const BLOCK_SIZE: usize = 16384;

struct PeerConnection {
    stream: TcpStream,
    have: Vec<bool>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl PeerConnection {
    fn connect(peer: &Peer, metainfo: &Metainfo) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let mut conn = PeerConnection {
            stream: stream,
            have:   vec![false; metainfo.info.pieces.len()],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn receive_message(&mut self) -> Result<Message, Error> {
        let message_size = convert_big_endian_to_integer(&try!(self.read_n(4)));
        if message_size > 0 {
            let message = try!(self.read_n(message_size));
            let message_id = &message[0];
            let message_body = &message[1..];
            Ok(Message::new(message_id, message_body))
        } else {
            Ok(Message::KeepAlive)
        }
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }

    fn process(&mut self, message: Message) -> Result<(), Error>{
        match message {
            Message::KeepAlive => {},
            Message::Bitfield(bytes) => {
                for have_index in 0..self.have.len() {
                    let bytes_index = have_index / 8;
                    let index_into_byte = have_index % 8;
                    let byte = bytes[bytes_index];
                    let value = (byte & (1 << (7 - index_into_byte))) != 0;
                    self.have[have_index] = value;
                };
                self.send_interested();
            },
            Message::Have(have_index) => {
                self.have[have_index] = true;
                self.send_interested();
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    let piece_index = self.next_piece_to_request();
                    self.send_request(piece_index);
                }
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }
    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            let mut message = vec![];
            message.push(0);
            message.push(0);
            message.push(0);
            message.push(1);
            message.push(2);
            try!(self.stream.write_all(&message));
            self.am_i_interested = true;
        }
        Ok(())
    }

    fn next_piece_to_request(&self) -> usize {
        0
    }

    fn send_request(&mut self, piece: usize) -> Result<(), Error> {
        let mut message = vec![];
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(13);
        message.push(6);
        message.extend(convert_usize_to_bytes(piece).iter().cloned());
        message.extend(convert_usize_to_bytes(0).iter().cloned());
        message.extend(convert_usize_to_bytes(BLOCK_SIZE).iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }
}

#[derive(Debug)]
enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece,   // TODO add params
    Cancel,  // TODO add params
    Port,    // TODO add params
}

impl Message {
    fn new(id: &u8, body: &[u8]) -> Message {
        match *id {
            0 => Message::Choke,
            1 => Message::Unchoke,
            2 => Message::Interested,
            3 => Message::NotInterested,
            4 => Message::Have(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => Message::Piece,
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }
}

const BYTE_0: usize = 256 * 256 * 256;
const BYTE_1: usize = 256 * 256;
const BYTE_2: usize = 256;
const BYTE_3: usize = 1;

fn convert_big_endian_to_integer(bytes: &[u8]) -> usize {
    bytes[0] as usize * BYTE_0 +
    bytes[1] as usize * BYTE_1 +
    bytes[2] as usize * BYTE_2 +
    bytes[3] as usize * BYTE_3
}

fn convert_usize_to_bytes(integer: usize) -> Vec<u8> {
    let mut rest = integer;
    let first = rest / BYTE_0;
    rest -= first * BYTE_0;
    let second = rest / BYTE_1;
    rest -= second * BYTE_1;
    let third = rest / BYTE_2;
    rest -= third * BYTE_2;
    let fourth = rest;
    vec![first as u8, second as u8, third as u8, fourth as u8]
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>