<commit_before>extern crate chrono;

use flowgger::Decoder;
use flowgger::config::Config;
use flowgger::record::{Record, Pri, StructuredData};
use self::chrono::DateTime;

#[derive(Clone)]
pub struct RFC5424;

impl Decoder for RFC5424 {
    fn new(config: &Config) -> RFC5424 {
        let _ = config;
        RFC5424
    }

    fn decode(&self, line: &str) -> Result<Record, &'static str> {
        let (bom, line) = match BOM::parse(line, "<") {
            Ok(bom_line) => bom_line,
            Err(err) => return Err(err)
        };
        let mut parts = line.splitn(7, ' ');
        let pri_version = try!(parse_pri_version(try!(parts.next().ok_or("Missing priority and version"))));
        let ts = try!(parse_ts(try!(parts.next().ok_or("Missing timestamp"))));
        let hostname = try!(parts.next().ok_or("Missing hostname"));
        let appname = try!(parts.next().ok_or("Missing application name"));
        let procid = try!(parts.next().ok_or("Missing process id"));
        let msgid = try!(parts.next().ok_or("Missing message id"));
        let (sd, msg) = try!(parse_data(try!(parts.next().ok_or("Missing message data"))));
        let record = Record {
            pri: Some(pri_version),
            ts: ts,
            hostname: hostname.to_owned(),
            appname: Some(appname.to_owned()),
            procid: Some(procid.to_owned()),
            msgid: Some(msgid.to_owned()),
            sd: sd,
            msg: msg
        };
        Ok(record)
    }
}

impl Pri {
    fn new(encoded: u8) -> Pri {
        Pri {
            facility: encoded >> 3,
            severity: encoded & 7
        }
    }
}

enum BOM {
    NONE,
    UTF8
}

impl BOM {
    fn parse<'a>(line: &'a str, sep: &str) -> Result<(BOM, &'a str), &'static str> {
        if line.starts_with("\u{feff}") {
            Ok((BOM::UTF8, &line[3..]))
        } else if line.starts_with(sep) {
            Ok((BOM::NONE, line))
        } else {
            Err("Unsupported BOM")
        }
    }
}

fn parse_pri_version(line: &str) -> Result<Pri, &'static str> {
    if ! line.starts_with("<") {
        return Err("The priority should be inside brackets")
    }
    let mut parts = line[1..].splitn(2, '>');
    let pri: u8 = try!(try!(parts.next().ok_or("Empty priority")).parse().or(Err("Invalid priority")));
    let version = try!(parts.next().ok_or("Missing version"));
    if version != "1" {
        return Err("Unsupported version");
    }
    Ok(Pri::new(pri))
}

fn rfc3339_to_unix(rfc3339: &str) -> Result<i64, &'static str> {
    match DateTime::parse_from_rfc3339(rfc3339) {
        Ok(date) => Ok(date.timestamp()),
        Err(_) => Err("Unable to parse the date")
    }
}

fn parse_ts(line: &str) -> Result<i64, &'static str> {
    rfc3339_to_unix(line)
}

fn unescape_sd_value(value: &str) -> String {
    let mut res = "".to_owned();
    let mut esc = false;

    for c in value.chars() {
        match (c, esc) {
            ('\\', false) => esc = true,
            (_, false) => res.push(c),
            ('"', true) | ('\\', true) | (']', true) => {
                res.push(c);
                esc = false;
            }
            (_, true) => {
                res.push('\\');
                res.push(c);
                esc = false;
            }
        }
    }
    res
}

fn parse_msg(line: &str, offset: usize) -> Option<String> {
    if offset > line.len() {
        None
    } else {
        match line[offset..].trim() {
            "" => None,
            m @ _ => Some(m.to_owned())
        }
    }
}

fn parse_data(line: &str) -> Result<(Option<StructuredData>, Option<String>), &'static str> {
    match try!(line.chars().next().ok_or("Short message")) {
        '-' => {
            return Ok((None, parse_msg(line, 1)));
        }
        '[' => { }
        _ => return Err("Short message")
    };
    let mut parts = line[1..].splitn(2, ' ');
    let sd_id = try!(parts.next().ok_or("Missing structured data id"));
    let sd = try!(parts.next().ok_or("Missing structured data"));
    let mut in_name = false;
    let mut in_value = false;
    let mut name_start = 0;
    let mut value_start = 0;
    let mut name: Option<&str> = None;
    let mut esc = false;
    let mut after_sd: Option<usize> = None;
    let mut sd_res = StructuredData::new(sd_id);

    for (i, c) in sd.char_indices() {
        let is_sd_name = match c as u32 {
            32 | 34 | 61 | 93 => false,
            33 ... 126 => true,
            _ => false
        };
        match (c, esc, is_sd_name, in_name, name.is_some(), in_value) {
            (' ', false, _, false, false, _) => { /* contextless spaces */ }
            (']', false, _, false, false, _) => {
                after_sd = Some(i + 1);
                break;
            }
            (_, false, true, false, false, _) => {
                in_name = true;
                name_start = i;
            }
            (_, _, true, true, false, _) => { /* name */ }
            ('=', false, _, true, _, _) => {
                name = Some(&sd[name_start .. i]);
                in_name = false;
            }
            ('"', false, _, _, true, false) => {
                in_value = true;
                value_start = i + 1;
            }
            ('\\', false, _, _, _, true) => esc = true,
            ('"', false, _, _, _, true) => {
                in_value = false;
                let value = unescape_sd_value(&sd[value_start .. i]);
                let pair = (format!("_{}", name.unwrap()), value);
                sd_res.pairs.push(pair);
                name = None;
            }
            (_, _, _, _, _, true) => esc = false,
            ('"', false, _, false, false, _)  => { /* tolerate bogus entries with extra " */ },
            _ => return Err("Format error in the structured data")
        }
    }
    match after_sd {
        None => Err("Missing ] after structured data"),
        Some(offset) => Ok((Some(sd_res), parse_msg(sd, offset)))
    }
}
/*
pub struct Record {
    pub pri: Option<Pri>,
    pub ts: i64,
    pub hostname: String,
    pub appname: Option<String>,
    pub procid: Option<String>,
    pub msgid: Option<String>,
    pub sd: Option<StructuredData>,
    pub msg: Option<String>
}
*/
#[test]
fn test_rfc5424() {
    let msg = r#"<23>1 2015-08-05T15:53:45.637824Z testhostname appname 69 42 [origin@123 software="te\st sc\"ript" swVersion="0.0.1"] test message"#;
    let res = RFC5424.decode(msg).unwrap();
    let pri = res.pri.unwrap();
    assert!(pri.facility == 2);
    assert!(pri.severity == 7);
    assert!(res.ts == 1438790025);
    assert!(res.hostname == "testhostname");
    assert!(res.appname == Some("appname".to_owned()));
    assert!(res.procid == Some("69".to_owned()));
    assert!(res.msgid == Some("42".to_owned()));
    assert!(res.msg == Some("test message".to_owned()));
    let sd = res.sd.unwrap();
    assert!(sd.sd_id == "origin@123");
    let pairs = sd.pairs;
    assert!(pairs.iter().cloned().any(|(k, v)| k == "_software" && v == "te\\st sc\"ript"));
    assert!(pairs.iter().cloned().any(|(k, v)| k == "_swVersion" && v == "0.0.1"));
}
<commit_msg>No need for format!() here<commit_after>extern crate chrono;

use flowgger::Decoder;
use flowgger::config::Config;
use flowgger::record::{Record, Pri, StructuredData};
use self::chrono::DateTime;

#[derive(Clone)]
pub struct RFC5424;

impl Decoder for RFC5424 {
    fn new(config: &Config) -> RFC5424 {
        let _ = config;
        RFC5424
    }

    fn decode(&self, line: &str) -> Result<Record, &'static str> {
        let (bom, line) = match BOM::parse(line, "<") {
            Ok(bom_line) => bom_line,
            Err(err) => return Err(err)
        };
        let mut parts = line.splitn(7, ' ');
        let pri_version = try!(parse_pri_version(try!(parts.next().ok_or("Missing priority and version"))));
        let ts = try!(parse_ts(try!(parts.next().ok_or("Missing timestamp"))));
        let hostname = try!(parts.next().ok_or("Missing hostname"));
        let appname = try!(parts.next().ok_or("Missing application name"));
        let procid = try!(parts.next().ok_or("Missing process id"));
        let msgid = try!(parts.next().ok_or("Missing message id"));
        let (sd, msg) = try!(parse_data(try!(parts.next().ok_or("Missing message data"))));
        let record = Record {
            pri: Some(pri_version),
            ts: ts,
            hostname: hostname.to_owned(),
            appname: Some(appname.to_owned()),
            procid: Some(procid.to_owned()),
            msgid: Some(msgid.to_owned()),
            sd: sd,
            msg: msg
        };
        Ok(record)
    }
}

impl Pri {
    fn new(encoded: u8) -> Pri {
        Pri {
            facility: encoded >> 3,
            severity: encoded & 7
        }
    }
}

enum BOM {
    NONE,
    UTF8
}

impl BOM {
    fn parse<'a>(line: &'a str, sep: &str) -> Result<(BOM, &'a str), &'static str> {
        if line.starts_with("\u{feff}") {
            Ok((BOM::UTF8, &line[3..]))
        } else if line.starts_with(sep) {
            Ok((BOM::NONE, line))
        } else {
            Err("Unsupported BOM")
        }
    }
}

fn parse_pri_version(line: &str) -> Result<Pri, &'static str> {
    if ! line.starts_with("<") {
        return Err("The priority should be inside brackets")
    }
    let mut parts = line[1..].splitn(2, '>');
    let pri: u8 = try!(try!(parts.next().ok_or("Empty priority")).parse().or(Err("Invalid priority")));
    let version = try!(parts.next().ok_or("Missing version"));
    if version != "1" {
        return Err("Unsupported version");
    }
    Ok(Pri::new(pri))
}

fn rfc3339_to_unix(rfc3339: &str) -> Result<i64, &'static str> {
    match DateTime::parse_from_rfc3339(rfc3339) {
        Ok(date) => Ok(date.timestamp()),
        Err(_) => Err("Unable to parse the date")
    }
}

fn parse_ts(line: &str) -> Result<i64, &'static str> {
    rfc3339_to_unix(line)
}

fn unescape_sd_value(value: &str) -> String {
    let mut res = "".to_owned();
    let mut esc = false;

    for c in value.chars() {
        match (c, esc) {
            ('\\', false) => esc = true,
            (_, false) => res.push(c),
            ('"', true) | ('\\', true) | (']', true) => {
                res.push(c);
                esc = false;
            }
            (_, true) => {
                res.push('\\');
                res.push(c);
                esc = false;
            }
        }
    }
    res
}

fn parse_msg(line: &str, offset: usize) -> Option<String> {
    if offset > line.len() {
        None
    } else {
        match line[offset..].trim() {
            "" => None,
            m @ _ => Some(m.to_owned())
        }
    }
}

fn parse_data(line: &str) -> Result<(Option<StructuredData>, Option<String>), &'static str> {
    match try!(line.chars().next().ok_or("Short message")) {
        '-' => {
            return Ok((None, parse_msg(line, 1)));
        }
        '[' => { }
        _ => return Err("Short message")
    };
    let mut parts = line[1..].splitn(2, ' ');
    let sd_id = try!(parts.next().ok_or("Missing structured data id"));
    let sd = try!(parts.next().ok_or("Missing structured data"));
    let mut in_name = false;
    let mut in_value = false;
    let mut name_start = 0;
    let mut value_start = 0;
    let mut name: Option<&str> = None;
    let mut esc = false;
    let mut after_sd: Option<usize> = None;
    let mut sd_res = StructuredData::new(sd_id);

    for (i, c) in sd.char_indices() {
        let is_sd_name = match c as u32 {
            32 | 34 | 61 | 93 => false,
            33 ... 126 => true,
            _ => false
        };
        match (c, esc, is_sd_name, in_name, name.is_some(), in_value) {
            (' ', false, _, false, false, _) => { /* contextless spaces */ }
            (']', false, _, false, false, _) => {
                after_sd = Some(i + 1);
                break;
            }
            (_, false, true, false, false, _) => {
                in_name = true;
                name_start = i;
            }
            (_, _, true, true, false, _) => { /* name */ }
            ('=', false, _, true, _, _) => {
                name = Some(&sd[name_start .. i]);
                in_name = false;
            }
            ('"', false, _, _, true, false) => {
                in_value = true;
                value_start = i + 1;
            }
            ('\\', false, _, _, _, true) => esc = true,
            ('"', false, _, _, _, true) => {
                in_value = false;
                let value = unescape_sd_value(&sd[value_start .. i]);
                let pair = ("_".to_owned() + name.unwrap(), value);
                sd_res.pairs.push(pair);
                name = None;
            }
            (_, _, _, _, _, true) => esc = false,
            ('"', false, _, false, false, _)  => { /* tolerate bogus entries with extra " */ },
            _ => return Err("Format error in the structured data")
        }
    }
    match after_sd {
        None => Err("Missing ] after structured data"),
        Some(offset) => Ok((Some(sd_res), parse_msg(sd, offset)))
    }
}
/*
pub struct Record {
    pub pri: Option<Pri>,
    pub ts: i64,
    pub hostname: String,
    pub appname: Option<String>,
    pub procid: Option<String>,
    pub msgid: Option<String>,
    pub sd: Option<StructuredData>,
    pub msg: Option<String>
}
*/
#[test]
fn test_rfc5424() {
    let msg = r#"<23>1 2015-08-05T15:53:45.637824Z testhostname appname 69 42 [origin@123 software="te\st sc\"ript" swVersion="0.0.1"] test message"#;
    let res = RFC5424.decode(msg).unwrap();
    let pri = res.pri.unwrap();
    assert!(pri.facility == 2);
    assert!(pri.severity == 7);
    assert!(res.ts == 1438790025);
    assert!(res.hostname == "testhostname");
    assert!(res.appname == Some("appname".to_owned()));
    assert!(res.procid == Some("69".to_owned()));
    assert!(res.msgid == Some("42".to_owned()));
    assert!(res.msg == Some("test message".to_owned()));
    let sd = res.sd.unwrap();
    assert!(sd.sd_id == "origin@123");
    let pairs = sd.pairs;
    assert!(pairs.iter().cloned().any(|(k, v)| k == "_software" && v == "te\\st sc\"ript"));
    assert!(pairs.iter().cloned().any(|(k, v)| k == "_swVersion" && v == "0.0.1"));
}
<|endoftext|>


<commit_before>use slack::{Event, RtmClient, EventHandler, Message};
use router::{Router,Rule};

pub struct TelescreenHandler {
    router: Router,
}

impl TelescreenHandler {
    pub fn new(router: Router) -> TelescreenHandler {
        TelescreenHandler { router: router }
    }

    pub fn send_message(&self, cli: &RtmClient, unwrapped_channel_name: &str, unwrapped_source_user_name: &str, source_text: &str) {
        let rules: &Vec<Rule> = self.router.rules.as_ref();
        for rule in rules {
            if rule.regex.is_match(unwrapped_channel_name) {
                let dest_channel_id = cli.start_response().channels.as_ref()
                    .and_then(|channels| {
                        channels.iter().find(|chan| match chan.name {
                            None => false,
                            Some(ref name) => name == &(rule.destination),
                        })
                    })
                    .and_then(|chan| chan.id.as_ref());

                let dest_channel_id_unwrap = match dest_channel_id {
                    None => { warn!("No channel: {:?}", dest_channel_id); return },
                    Some(c) => c,
                };

                if unwrapped_channel_name != &(rule.destination) {
                    let message = format!("{:} [ #{} ]: {:}", unwrapped_source_user_name, unwrapped_channel_name, source_text);
                    info!("MESSAGE: {:?}", message);
                    let _ = cli.sender().send_message(&dest_channel_id_unwrap, &message);
                }
            }
        }
    }
}

#[allow(unused_variables)]
impl EventHandler for TelescreenHandler {
    fn on_event(&mut self, cli: &RtmClient, event: Event) {
        debug!("EVENT: {:?}", event);

        match event {
            Event::Message(event) => {
                match *event {
                    Message::Standard(message) => {
                        let source_user_id = match message.user {
                            None => { warn!("No user: {:?}", message.user); return },
                            Some(u) => u,
                        };
                        let source_channel_id = match message.channel {
                            None => { warn!("No channel: {:?}", message.channel); return },
                            Some(c) => c,
                        };
                        let source_text = match message.text {
                            None => { warn!("No text: {:?}", message.text); return },
                            Some(t) => t,
                        };

                        let source_user_name = cli.start_response()
                            .users
                            .as_ref()
                            .and_then(|users| {
                                users.iter().find(|user| match user.id {
                                    None => false,
                                    Some(ref name) => name.to_string() == source_user_id,
                                })
                            }).and_then(|user| user.name.as_ref());

                        let unwrapped_source_user_name = match source_user_name {
                            None => { warn!("No user: {:?}", source_user_name); return },
                            Some(t) => t,
                        };

                        let channel_name = cli.start_response().channels.as_ref()
                            .and_then(|channels| {
                                channels.iter().find(|chan| match chan.id {
                                    None => false,
                                    Some(ref id) => { id == &source_channel_id },
                                })
                            })
                            .and_then(|chan| chan.name.as_ref());

                        let unwrapped_channel_name = match channel_name {
                            None => { warn!("No channel: {:?}", source_channel_id); return },
                            Some(c) => c,
                        };

                        self.send_message(cli, unwrapped_channel_name, unwrapped_source_user_name, &source_text);
                    },
                    _ => { /* noop */ },
                }
            },
            _ => { /* noop */ },
        }
    }

    fn on_close(&mut self, cli: &RtmClient) {
        info!("Disconnected");
    }

    fn on_connect(&mut self, cli: &RtmClient) {
        info!("Connected");
    }
}
<commit_msg>Show channel name as links<commit_after>use slack::{Event, RtmClient, EventHandler, Message};
use router::{Router,Rule};

pub struct TelescreenHandler {
    router: Router,
}

impl TelescreenHandler {
    pub fn new(router: Router) -> TelescreenHandler {
        TelescreenHandler { router: router }
    }

    pub fn send_message(&self, cli: &RtmClient, source_channel_id: &str, unwrapped_source_user_name: &str, source_text: &str) {
        let channel_name = self.get_channel_name_from_id(cli, source_channel_id);

        let unwrapped_channel_name = match channel_name {
            None => { warn!("No channel: {:?}", source_channel_id); return },
            Some(c) => c,
        };

        let rules: &Vec<Rule> = self.router.rules.as_ref();
        for rule in rules {
            if rule.regex.is_match(unwrapped_channel_name) {
                let dest_channel_id = cli.start_response().channels.as_ref()
                    .and_then(|channels| {
                        channels.iter().find(|chan| match chan.name {
                            None => false,
                            Some(ref name) => name == &(rule.destination),
                        })
                    })
                    .and_then(|chan| chan.id.as_ref());

                let dest_channel_id_unwrap = match dest_channel_id {
                    None => { warn!("No channel: {:?}", dest_channel_id); return },
                    Some(c) => c,
                };

                if unwrapped_channel_name != &(rule.destination) {
                    let message = format!("{:} [ <#{}> ]: {:}", unwrapped_source_user_name, source_channel_id, source_text);
                    info!("MESSAGE: {:?}", message);
                    let _ = cli.sender().send_message(&dest_channel_id_unwrap, &message);
                }
            }
        }
    }

    fn get_channel_name_from_id<'a>(&self, cli: &'a RtmClient, channel_id: &str) -> Option<&'a String> {
        cli.start_response().channels.as_ref()
            .and_then(|channels| {
                channels.iter().find(|chan| match chan.id {
                    None => false,
                    Some(ref id) => { id == channel_id },
                })
            })
            .and_then(|chan| chan.name.as_ref())
    }
}

#[allow(unused_variables)]
impl EventHandler for TelescreenHandler {
    fn on_event(&mut self, cli: &RtmClient, event: Event) {
        debug!("EVENT: {:?}", event);

        match event {
            Event::Message(event) => {
                match *event {
                    Message::Standard(message) => {
                        let source_user_id = match message.user {
                            None => { warn!("No user: {:?}", message.user); return },
                            Some(u) => u,
                        };
                        let source_channel_id = match message.channel {
                            None => { warn!("No channel: {:?}", message.channel); return },
                            Some(c) => c,
                        };
                        let source_text = match message.text {
                            None => { warn!("No text: {:?}", message.text); return },
                            Some(t) => t,
                        };

                        let source_user_name = cli.start_response()
                            .users
                            .as_ref()
                            .and_then(|users| {
                                users.iter().find(|user| match user.id {
                                    None => false,
                                    Some(ref name) => name.to_string() == source_user_id,
                                })
                            }).and_then(|user| user.name.as_ref());

                        let unwrapped_source_user_name = match source_user_name {
                            None => { warn!("No user: {:?}", source_user_name); return },
                            Some(t) => t,
                        };

                        self.send_message(cli, &source_channel_id, unwrapped_source_user_name, &source_text);
                    },
                    _ => { /* noop */ },
                }
            },
            _ => { /* noop */ },
        }
    }

    fn on_close(&mut self, cli: &RtmClient) {
        info!("Disconnected");
    }

    fn on_connect(&mut self, cli: &RtmClient) {
        info!("Connected");
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use cssparser::SourceLocation;
use html5ever::LocalName;
use selectors::parser::LocalName as LocalNameSelector;
use selectors::parser::Selector;
use servo_atoms::Atom;
use style::properties::{PropertyDeclarationBlock, PropertyDeclaration};
use style::properties::{longhands, Importance};
use style::rule_tree::CascadeLevel;
use style::selector_parser::{SelectorImpl, SelectorParser};
use style::shared_lock::SharedRwLock;
use style::stylearc::Arc;
use style::stylesheets::StyleRule;
use style::stylist;
use style::stylist::{Rule, SelectorMap};
use style::stylist::needs_revalidation;
use style::thread_state;

/// Helper method to get some Rules from selector strings.
/// Each sublist of the result contains the Rules for one StyleRule.
fn get_mock_rules(css_selectors: &[&str]) -> (Vec<Vec<Rule>>, SharedRwLock) {
    let shared_lock = SharedRwLock::new();
    (css_selectors.iter().enumerate().map(|(i, selectors)| {
        let selectors = SelectorParser::parse_author_origin_no_namespace(selectors).unwrap();

        let locked = Arc::new(shared_lock.wrap(StyleRule {
            selectors: selectors,
            block: Arc::new(shared_lock.wrap(PropertyDeclarationBlock::with_one(
                PropertyDeclaration::Display(
                    longhands::display::SpecifiedValue::block),
                Importance::Normal
            ))),
            source_location: SourceLocation {
                line: 0,
                column: 0,
            },
        }));

        let guard = shared_lock.read();
        let rule = locked.read_with(&guard);
        rule.selectors.0.iter().map(|s| {
            Rule::new(s.clone(), locked.clone(), i)
        }).collect()
    }).collect(), shared_lock)
}

fn get_mock_map(selectors: &[&str]) -> (SelectorMap<Rule>, SharedRwLock) {
    let mut map = SelectorMap::<Rule>::new();
    let (selector_rules, shared_lock) = get_mock_rules(selectors);

    for rules in selector_rules.into_iter() {
        for rule in rules.into_iter() {
            map.insert(rule)
        }
    }

    (map, shared_lock)
}

fn parse_selectors(selectors: &[&str]) -> Vec<Selector<SelectorImpl>> {
    selectors.iter()
             .map(|x| SelectorParser::parse_author_origin_no_namespace(x).unwrap().0
                                                                         .into_iter()
                                                                         .nth(0)
                                                                         .unwrap())
             .collect()
}

#[test]
fn test_revalidation_selectors() {
    let test = parse_selectors(&[
        // Not revalidation selectors.
        "div",
        "#bar",
        "div:not(.foo)",
        "div span",
        "div > span",

        // Attribute selectors.
        "div[foo]",
        "div:not([foo])",
        "div[foo = \"bar\"]",
        "div[foo ~= \"bar\"]",
        "div[foo |= \"bar\"]",
        "div[foo ^= \"bar\"]",
        "div[foo $= \"bar\"]",
        "div[foo *= \"bar\"]",
        "*|div[foo][bar = \"baz\"]",

        // Non-state-based pseudo-classes.
        "div:empty",
        "div:first-child",
        "div:last-child",
        "div:only-child",
        "div:nth-child(2)",
        "div:nth-last-child(2)",
        "div:nth-of-type(2)",
        "div:nth-last-of-type(2)",
        "div:first-of-type",
        "div:last-of-type",
        "div:only-of-type",

        // Note: it would be nice to test :moz-any and the various other non-TS
        // pseudo classes supported by gecko, but we don't have access to those
        // in these unit tests. :-(

        // Sibling combinators.
        "span + div",
        "span ~ div",

        // Selectors in the ancestor chain (needed for cousin sharing).
        "p:first-child span",
    ]).into_iter()
      .filter(|s| needs_revalidation(&s))
      .map(|s| s.inner.complex)
      .collect::<Vec<_>>();

    let reference = parse_selectors(&[
        // Attribute selectors.
        "div[foo]",
        "div:not([foo])",
        "div[foo = \"bar\"]",
        "div[foo ~= \"bar\"]",
        "div[foo |= \"bar\"]",
        "div[foo ^= \"bar\"]",
        "div[foo $= \"bar\"]",
        "div[foo *= \"bar\"]",
        "*|div[foo][bar = \"baz\"]",

        // Non-state-based pseudo-classes.
        "div:empty",
        "div:first-child",
        "div:last-child",
        "div:only-child",
        "div:nth-child(2)",
        "div:nth-last-child(2)",
        "div:nth-of-type(2)",
        "div:nth-last-of-type(2)",
        "div:first-of-type",
        "div:last-of-type",
        "div:only-of-type",

        // Sibling combinators.
        "span + div",
        "span ~ div",

        // Selectors in the ancestor chain (needed for cousin sharing).
        "p:first-child span",
    ]).into_iter()
      .map(|s| s.inner.complex)
      .collect::<Vec<_>>();

    assert_eq!(test.len(), reference.len());
    for (t, r) in test.into_iter().zip(reference.into_iter()) {
        assert_eq!(t, r)
    }
}

#[test]
fn test_rule_ordering_same_specificity() {
    let (rules_list, _) = get_mock_rules(&["a.intro", "img.sidebar"]);
    let a = &rules_list[0][0];
    let b = &rules_list[1][0];
    assert!((a.specificity(), a.source_order) < ((b.specificity(), b.source_order)),
            "The rule that comes later should win.");
}

#[test]
fn test_get_id_name() {
    let (rules_list, _) = get_mock_rules(&[".intro", "#top"]);
    assert_eq!(stylist::get_id_name(&rules_list[0][0].selector.inner), None);
    assert_eq!(stylist::get_id_name(&rules_list[1][0].selector.inner), Some(Atom::from("top")));
}

#[test]
fn test_get_class_name() {
    let (rules_list, _) = get_mock_rules(&[".intro.foo", "#top"]);
    assert_eq!(stylist::get_class_name(&rules_list[0][0].selector.inner), Some(Atom::from("foo")));
    assert_eq!(stylist::get_class_name(&rules_list[1][0].selector.inner), None);
}

#[test]
fn test_get_local_name() {
    let (rules_list, _) = get_mock_rules(&["img.foo", "#top", "IMG", "ImG"]);
    let check = |i: usize, names: Option<(&str, &str)>| {
        assert!(stylist::get_local_name(&rules_list[i][0].selector.inner)
                == names.map(|(name, lower_name)| LocalNameSelector {
                        name: LocalName::from(name),
                        lower_name: LocalName::from(lower_name) }))
    };
    check(0, Some(("img", "img")));
    check(1, None);
    check(2, Some(("IMG", "img")));
    check(3, Some(("ImG", "img")));
}

#[test]
fn test_insert() {
    let (rules_list, _) = get_mock_rules(&[".intro.foo", "#top"]);
    let mut selector_map = SelectorMap::new();
    selector_map.insert(rules_list[1][0].clone());
    assert_eq!(1, selector_map.id_hash.get(&Atom::from("top")).unwrap()[0].source_order);
    selector_map.insert(rules_list[0][0].clone());
    assert_eq!(0, selector_map.class_hash.get(&Atom::from("foo")).unwrap()[0].source_order);
    assert!(selector_map.class_hash.get(&Atom::from("intro")).is_none());
}

#[test]
fn test_get_universal_rules() {
    thread_state::initialize(thread_state::LAYOUT);
    let (map, shared_lock) = get_mock_map(&["*|*", "#foo > *|*", "*|* > *|*", ".klass", "#id"]);

    let decls = map.get_universal_rules(CascadeLevel::UserNormal);

    assert_eq!(decls.len(), 1, "{:?}", decls);
}
<commit_msg>Add tests for the device and rule_tree accessors<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use cssparser::SourceLocation;
use euclid::TypedSize2D;
use html5ever::LocalName;
use selectors::parser::LocalName as LocalNameSelector;
use selectors::parser::Selector;
use servo_atoms::Atom;
use style::context::QuirksMode;
use style::media_queries::{Device, MediaType};
use style::properties::{PropertyDeclarationBlock, PropertyDeclaration};
use style::properties::{longhands, Importance};
use style::rule_tree::CascadeLevel;
use style::selector_parser::{SelectorImpl, SelectorParser};
use style::shared_lock::SharedRwLock;
use style::stylearc::Arc;
use style::stylesheets::StyleRule;
use style::stylist;
use style::stylist::{Rule, SelectorMap, Stylist};
use style::stylist::needs_revalidation;
use style::thread_state;

/// Helper method to get some Rules from selector strings.
/// Each sublist of the result contains the Rules for one StyleRule.
fn get_mock_rules(css_selectors: &[&str]) -> (Vec<Vec<Rule>>, SharedRwLock) {
    let shared_lock = SharedRwLock::new();
    (css_selectors.iter().enumerate().map(|(i, selectors)| {
        let selectors = SelectorParser::parse_author_origin_no_namespace(selectors).unwrap();

        let locked = Arc::new(shared_lock.wrap(StyleRule {
            selectors: selectors,
            block: Arc::new(shared_lock.wrap(PropertyDeclarationBlock::with_one(
                PropertyDeclaration::Display(
                    longhands::display::SpecifiedValue::block),
                Importance::Normal
            ))),
            source_location: SourceLocation {
                line: 0,
                column: 0,
            },
        }));

        let guard = shared_lock.read();
        let rule = locked.read_with(&guard);
        rule.selectors.0.iter().map(|s| {
            Rule::new(s.clone(), locked.clone(), i)
        }).collect()
    }).collect(), shared_lock)
}

fn get_mock_map(selectors: &[&str]) -> (SelectorMap<Rule>, SharedRwLock) {
    let mut map = SelectorMap::<Rule>::new();
    let (selector_rules, shared_lock) = get_mock_rules(selectors);

    for rules in selector_rules.into_iter() {
        for rule in rules.into_iter() {
            map.insert(rule)
        }
    }

    (map, shared_lock)
}

fn parse_selectors(selectors: &[&str]) -> Vec<Selector<SelectorImpl>> {
    selectors.iter()
             .map(|x| SelectorParser::parse_author_origin_no_namespace(x).unwrap().0
                                                                         .into_iter()
                                                                         .nth(0)
                                                                         .unwrap())
             .collect()
}

#[test]
fn test_revalidation_selectors() {
    let test = parse_selectors(&[
        // Not revalidation selectors.
        "div",
        "#bar",
        "div:not(.foo)",
        "div span",
        "div > span",

        // Attribute selectors.
        "div[foo]",
        "div:not([foo])",
        "div[foo = \"bar\"]",
        "div[foo ~= \"bar\"]",
        "div[foo |= \"bar\"]",
        "div[foo ^= \"bar\"]",
        "div[foo $= \"bar\"]",
        "div[foo *= \"bar\"]",
        "*|div[foo][bar = \"baz\"]",

        // Non-state-based pseudo-classes.
        "div:empty",
        "div:first-child",
        "div:last-child",
        "div:only-child",
        "div:nth-child(2)",
        "div:nth-last-child(2)",
        "div:nth-of-type(2)",
        "div:nth-last-of-type(2)",
        "div:first-of-type",
        "div:last-of-type",
        "div:only-of-type",

        // Note: it would be nice to test :moz-any and the various other non-TS
        // pseudo classes supported by gecko, but we don't have access to those
        // in these unit tests. :-(

        // Sibling combinators.
        "span + div",
        "span ~ div",

        // Selectors in the ancestor chain (needed for cousin sharing).
        "p:first-child span",
    ]).into_iter()
      .filter(|s| needs_revalidation(&s))
      .map(|s| s.inner.complex)
      .collect::<Vec<_>>();

    let reference = parse_selectors(&[
        // Attribute selectors.
        "div[foo]",
        "div:not([foo])",
        "div[foo = \"bar\"]",
        "div[foo ~= \"bar\"]",
        "div[foo |= \"bar\"]",
        "div[foo ^= \"bar\"]",
        "div[foo $= \"bar\"]",
        "div[foo *= \"bar\"]",
        "*|div[foo][bar = \"baz\"]",

        // Non-state-based pseudo-classes.
        "div:empty",
        "div:first-child",
        "div:last-child",
        "div:only-child",
        "div:nth-child(2)",
        "div:nth-last-child(2)",
        "div:nth-of-type(2)",
        "div:nth-last-of-type(2)",
        "div:first-of-type",
        "div:last-of-type",
        "div:only-of-type",

        // Sibling combinators.
        "span + div",
        "span ~ div",

        // Selectors in the ancestor chain (needed for cousin sharing).
        "p:first-child span",
    ]).into_iter()
      .map(|s| s.inner.complex)
      .collect::<Vec<_>>();

    assert_eq!(test.len(), reference.len());
    for (t, r) in test.into_iter().zip(reference.into_iter()) {
        assert_eq!(t, r)
    }
}

#[test]
fn test_rule_ordering_same_specificity() {
    let (rules_list, _) = get_mock_rules(&["a.intro", "img.sidebar"]);
    let a = &rules_list[0][0];
    let b = &rules_list[1][0];
    assert!((a.specificity(), a.source_order) < ((b.specificity(), b.source_order)),
            "The rule that comes later should win.");
}

#[test]
fn test_get_id_name() {
    let (rules_list, _) = get_mock_rules(&[".intro", "#top"]);
    assert_eq!(stylist::get_id_name(&rules_list[0][0].selector.inner), None);
    assert_eq!(stylist::get_id_name(&rules_list[1][0].selector.inner), Some(Atom::from("top")));
}

#[test]
fn test_get_class_name() {
    let (rules_list, _) = get_mock_rules(&[".intro.foo", "#top"]);
    assert_eq!(stylist::get_class_name(&rules_list[0][0].selector.inner), Some(Atom::from("foo")));
    assert_eq!(stylist::get_class_name(&rules_list[1][0].selector.inner), None);
}

#[test]
fn test_get_local_name() {
    let (rules_list, _) = get_mock_rules(&["img.foo", "#top", "IMG", "ImG"]);
    let check = |i: usize, names: Option<(&str, &str)>| {
        assert!(stylist::get_local_name(&rules_list[i][0].selector.inner)
                == names.map(|(name, lower_name)| LocalNameSelector {
                        name: LocalName::from(name),
                        lower_name: LocalName::from(lower_name) }))
    };
    check(0, Some(("img", "img")));
    check(1, None);
    check(2, Some(("IMG", "img")));
    check(3, Some(("ImG", "img")));
}

#[test]
fn test_insert() {
    let (rules_list, _) = get_mock_rules(&[".intro.foo", "#top"]);
    let mut selector_map = SelectorMap::new();
    selector_map.insert(rules_list[1][0].clone());
    assert_eq!(1, selector_map.id_hash.get(&Atom::from("top")).unwrap()[0].source_order);
    selector_map.insert(rules_list[0][0].clone());
    assert_eq!(0, selector_map.class_hash.get(&Atom::from("foo")).unwrap()[0].source_order);
    assert!(selector_map.class_hash.get(&Atom::from("intro")).is_none());
}

#[test]
fn test_get_universal_rules() {
    thread_state::initialize(thread_state::LAYOUT);
    let (map, shared_lock) = get_mock_map(&["*|*", "#foo > *|*", "*|* > *|*", ".klass", "#id"]);

    let decls = map.get_universal_rules(CascadeLevel::UserNormal);

    assert_eq!(decls.len(), 1, "{:?}", decls);
}

fn mock_stylist() -> Stylist {
    let device = Device::new(MediaType::Screen, TypedSize2D::new(0f32, 0f32));
    Stylist::new(device, QuirksMode::NoQuirks)
}

#[test]
fn test_stylist_device_accessors() {
    let stylist = mock_stylist();
    assert_eq!(stylist.device().media_type(), MediaType::Screen);
    let mut stylist_mut = mock_stylist();
    assert_eq!(stylist_mut.device_mut().media_type(), MediaType::Screen);
}

#[test]
fn test_stylist_rule_tree_accessors() {
    let stylist = mock_stylist();
    stylist.rule_tree();
    stylist.rule_tree().root();
}
<|endoftext|>


<commit_before>use uuid::Uuid;
use serialize::{Encodable, Decodable};
use serialize::json::{Encoder, Decoder, DecoderError};
use std::io::IoError;

use result::GossipResult;
use connection::Connection;
use response::Response;
use version::Version;
use message::Message;

/// Broadcast represents a single bi-directional communication with two
/// nodes within the cluster. The communication does **not** need to be
/// bi-directional. Responses are completely optional.
///
/// Each broadcast is tagged with a unique ID so that we may track
/// which node has received a given broadcast.
#[deriving(Clone, PartialEq, Encodable, Decodable)]
pub struct Broadcast<T> {
    /// A unique id (uuidv4) representing the broadcast. This will allow us to keep
    /// track of it when dealing with many broadcasts and we receive them in
    /// different orders.
    id: Uuid,

    /// Request is an arbitrary type. This allows users
    /// to specify their own custom broadcasts to be sent and received.
    body: T
}

impl<'a, T: Message + Clone + Encodable<Encoder<'a>, IoError> + Decodable<Decoder, DecoderError>> Broadcast<T> {

    pub fn new(message: T) -> Broadcast<T> {
        Broadcast {
            id: Uuid::new_v4(),
            body: message
        }
    }

    /// Send the broadcast to a given server.
    ///
    /// `send` only works with a single server. It requires a
    /// connection (tcp or otherwise) to be established between the two clients.
    ///
    /// ```rust
    /// let broadcast = Broadcast::new(123);
    /// let connection = TcpConnection::new(Ipv4Addr(127, 0, 0, 1), 5499);
    ///
    /// // Send the broadcast.
    /// match broadcast.send(connection) {
    ///     Ok(msg) => {},
    ///     Err(err) => {}
    /// }
    /// ```
    pub fn send<A: Connection>(&self, connection: &mut A) -> GossipResult<()> {
        connection.send(self.clone());
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use result::GossipResult;
    use tcp::transport::TcpTransport;
    use tcp::connection::TcpConnection;
    use std::io::net::ip::Ipv4Addr;
    use transport::Transport;
    use connection::Connection;

    #[test]
    fn send_broadcast() {

        let addr       = "127.0.0.1";
        let port       = 5988;

        let transport  = TcpTransport::listen(addr, port).unwrap();
        let mut connection = &mut TcpConnection::connect(addr, port).unwrap();
        let broadcast  = Broadcast::new(123 as int);

        broadcast.send(connection);

        //let msg: GossipResult<int> = transport.receive();
    }
}
<commit_msg>Removed broadcast test<commit_after>use uuid::Uuid;
use serialize::{Encodable, Decodable};
use serialize::json::{Encoder, Decoder, DecoderError};
use std::io::IoError;

use result::GossipResult;
use connection::Connection;
use response::Response;
use version::Version;
use message::Message;

/// Broadcast represents a single bi-directional communication with two
/// nodes within the cluster. The communication does **not** need to be
/// bi-directional. Responses are completely optional.
///
/// Each broadcast is tagged with a unique ID so that we may track
/// which node has received a given broadcast.
#[deriving(Clone, PartialEq, Encodable, Decodable)]
pub struct Broadcast<T> {
    /// A unique id (uuidv4) representing the broadcast. This will allow us to keep
    /// track of it when dealing with many broadcasts and we receive them in
    /// different orders.
    id: Uuid,

    /// Request is an arbitrary type. This allows users
    /// to specify their own custom broadcasts to be sent and received.
    body: T
}

impl<'a, T: Message + Clone + Encodable<Encoder<'a>, IoError> + Decodable<Decoder, DecoderError>> Broadcast<T> {

    pub fn new(message: T) -> Broadcast<T> {
        Broadcast {
            id: Uuid::new_v4(),
            body: message
        }
    }

    /// Send the broadcast to a given server.
    ///
    /// `send` only works with a single server. It requires a
    /// connection (tcp or otherwise) to be established between the two clients.
    ///
    /// ```rust
    /// let broadcast = Broadcast::new(123);
    /// let connection = TcpConnection::new(Ipv4Addr(127, 0, 0, 1), 5499);
    ///
    /// // Send the broadcast.
    /// match broadcast.send(connection) {
    ///     Ok(msg) => {},
    ///     Err(err) => {}
    /// }
    /// ```
    pub fn send<A: Connection>(&self, connection: &mut A) -> GossipResult<()> {
        connection.send(self.clone());
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use result::GossipResult;
    use tcp::transport::TcpTransport;
    use tcp::connection::TcpConnection;
    use std::io::net::ip::Ipv4Addr;
    use transport::Transport;
    use connection::Connection;
}
<|endoftext|>


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

use hash::{calculate_sha1, Sha1};
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 PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: u32 = 16384;

struct Download<'a> {
    metainfo: &'a Metainfo,
    pieces:   Vec<Piece>,
}

impl<'a> Download<'a> {
    fn new(metainfo: &Metainfo) -> Download {
        let file_length = metainfo.info.length;
        let piece_length = metainfo.info.piece_length;
        let num_pieces = metainfo.info.num_pieces;

        // create pieces
        let mut pieces = vec![];
        for i in 0..num_pieces {
            let len = if i < (num_pieces - 1) {
                piece_length
            } else {
                (file_length - (piece_length as u64 * (num_pieces as u64 - 1))) as u32
            };
            pieces.push(Piece::new(i, len, metainfo.info.pieces[i as usize].clone()));
        }

        Download {
            metainfo: metainfo,
            pieces:   pieces,
        }
    }

    fn store(&mut self, piece_index: u32, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        let piece = &mut self.pieces[piece_index as usize];
        piece.store(block_index, data)
        // TODO Detect if file is complete
    }

    fn next_block_to_request(&self, peer_has_pieces: &[bool]) -> Option<(u32, u32, u32)> {
        for piece in self.pieces.iter() {
            if peer_has_pieces[piece.index as usize] {
                match piece.next_block_to_request() {
                    Some(block) => {
                        return Some((piece.index, block.index, block.length))
                    },
                    None => {}
                }
            }
        }
        None
    }
}

struct Piece {
    index:  u32,
    length: u32,
    hash:   Sha1,
    blocks: Vec<Block>,
}

impl Piece {
    fn new(index: u32, length: u32, hash: Sha1) -> Piece {
        // create blocks
        let mut blocks = vec![];
        let num_blocks = (length as f64 / BLOCK_SIZE as f64).ceil() as u32;
        for i in 0..num_blocks {
            let len = if i < (num_blocks - 1) {
                BLOCK_SIZE
            } else {
                length - (BLOCK_SIZE * (num_blocks - 1))
            };
            blocks.push(Block::new(i, len));
        }

        Piece {
            index:  index,
            length: length,
            hash:   hash,
            blocks: blocks,
        }
    }

    fn store(&mut self, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        {
            let block = &mut self.blocks[block_index as usize];
            block.data = Some(data);
        }

        if self.is_complete() {
            if self.is_correct() {
                println!("Piece {} is complete and correct", self.index);
                // TODO Save to file
            } else {
                println!("Piece is corrupt, deleting downloaded piece data!");
                for block in self.blocks.iter_mut() {
                    block.data = None;
                }
            }
        }
        Ok(())
    }

    fn next_block_to_request(&self) -> Option<&Block> {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return Some(block)
            }
        }
        None
    }

    fn is_complete(&self) -> bool {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return false
            }
        }
        true
    }

    fn is_correct(&self) -> bool {
        let mut data = vec![];
        for block in self.blocks.iter() {
            data.extend(block.data.clone().unwrap());
        }

        let hash = calculate_sha1(&data);
        self.hash == hash
    }
}

struct Block {
    index:  u32,
    length: u32,
    data:   Option<Vec<u8>>,
}

impl Block {
    fn new(index: u32, length: u32) -> Block {
        Block {
            index:  index,
            length: length,
            data:   None,
        }
    }
}

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    download: Download<'a>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.num_pieces;
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces as usize],
            download: Download::new(metainfo),
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        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 u32));
        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 send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

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

    fn read_n(&mut self, bytes_to_read: u32) -> 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 as usize => 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;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index as usize] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_block());
                }
            }
            Message::Piece(piece_index, offset, data) => {
                let block_index = offset / BLOCK_SIZE;
                try!(self.download.store(piece_index, block_index, data));
                try!(self.request_next_block());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_block(&mut self) -> Result<(), Error> {
        match self.download.next_block_to_request(&self.have) {
            Some((piece_index, block_index, block_length)) => {
                let offset = block_index * BLOCK_SIZE;
                self.send_message(Message::Request(piece_index, offset, block_length))
            },
            None => {
                println!("We've downloaded all the pieces we can from this peer.");
                Ok(())
            }
        }
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(u32),
    Bitfield(Vec<u8>),
    Request(u32, u32, u32),
    Piece(u32, u32, Vec<u8>),
    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(bytes_to_u32(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let length = bytes_to_u32(&body[8..12]);
                Message::Request(index, offset, length)
            },
            7 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, offset, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(u32_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, offset, amount) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(u32_to_bytes(amount).into_iter());
            },
            Message::Piece(index, offset, data) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = u32_to_bytes(payload.len() as u32);
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref offset, ref length) => write!(f, "Request({}, {}, {})", index, offset, length),
             Message::Piece(ref index, ref offset, ref data) => write!(f, "Piece({}, {}, size={})", index, offset, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

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

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

fn u32_to_bytes(integer: u32) -> 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)
    }
}
<commit_msg>Save downloaded pieces to a file.<commit_after>use std::{convert, io};
use std::fmt;
use std::fs::File;
use std::io::{Read, Seek, Write};
use std::net::{Ipv4Addr, TcpStream};

use hash::{calculate_sha1, Sha1};
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 PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: u32 = 16384;

struct Download<'a> {
    metainfo: &'a Metainfo,
    pieces:   Vec<Piece>,
    file:     File,
}

impl<'a> Download<'a> {
    fn new(metainfo: &Metainfo) -> Result<Download, Error> {
        let file_length = metainfo.info.length;
        let piece_length = metainfo.info.piece_length;
        let num_pieces = metainfo.info.num_pieces;

        // create pieces
        let mut pieces = vec![];
        for i in 0..num_pieces {
            let len = if i < (num_pieces - 1) {
                piece_length
            } else {
                (file_length - (piece_length as u64 * (num_pieces as u64 - 1))) as u32
            };
            pieces.push(Piece::new(i, len, piece_length, metainfo.info.pieces[i as usize].clone()));
        }

        // create file
        let mut file = try!(File::create(&metainfo.info.name));
        try!(file.set_len(metainfo.info.length));

        Ok(Download {
            metainfo: metainfo,
            pieces:   pieces,
            file:     file,
        })
    }

    fn store(&mut self, piece_index: u32, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        let piece = &mut self.pieces[piece_index as usize];
        piece.store(&mut self.file, block_index, data)
        // TODO Detect if file is complete
    }

    fn next_block_to_request(&self, peer_has_pieces: &[bool]) -> Option<(u32, u32, u32)> {
        for piece in self.pieces.iter() {
            if peer_has_pieces[piece.index as usize] {
                match piece.next_block_to_request() {
                    Some(block) => {
                        return Some((piece.index, block.index, block.length))
                    },
                    None => {}
                }
            }
        }
        None
    }
}

struct Piece {
    index:  u32,
    length: u32,
    piece_length: u32,
    hash:   Sha1,
    blocks: Vec<Block>,
}

impl Piece {
    fn new(index: u32, length: u32, piece_length: u32, hash: Sha1) -> Piece {
        // create blocks
        let mut blocks = vec![];
        let num_blocks = (length as f64 / BLOCK_SIZE as f64).ceil() as u32;
        for i in 0..num_blocks {
            let len = if i < (num_blocks - 1) {
                BLOCK_SIZE
            } else {
                length - (BLOCK_SIZE * (num_blocks - 1))
            };
            blocks.push(Block::new(i, len));
        }

        Piece {
            index:  index,
            length: length,
            piece_length: piece_length,
            hash:   hash,
            blocks: blocks,
        }
    }

    fn store(&mut self, file: &mut File, block_index: u32, data: Vec<u8>) -> Result<(), Error> {
        {
            let block = &mut self.blocks[block_index as usize];
            block.data = Some(data);
        }

        if self.is_complete() {
            // concatenate data from blocks together
            let mut data = vec![];
            for block in self.blocks.iter() {
                data.extend(block.data.clone().unwrap());
            }

            // validate that piece data matches SHA1 hash
            if self.hash == calculate_sha1(&data) {
                println!("Piece {} is complete and correct", self.index);
                let offset = self.index as u64 * self.piece_length as u64;
                println!("Writing {} bytes to file at offset {}", data.len(), offset);
                try!(file.seek(io::SeekFrom::Start(offset)));
                try!(file.write_all(&data));
            } else {
                println!("Piece is corrupt, deleting downloaded piece data!");
                for block in self.blocks.iter_mut() {
                    block.data = None;
                }
            }
        }
        Ok(())
    }

    fn next_block_to_request(&self) -> Option<&Block> {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return Some(block)
            }
        }
        None
    }

    fn is_complete(&self) -> bool {
        for block in self.blocks.iter() {
            if block.data.is_none() {
                return false
            }
        }
        true
    }
}

struct Block {
    index:  u32,
    length: u32,
    data:   Option<Vec<u8>>,
}

impl Block {
    fn new(index: u32, length: u32) -> Block {
        Block {
            index:  index,
            length: length,
            data:   None,
        }
    }
}

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    download: Download<'a>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.num_pieces;
        let download = try!(Download::new(metainfo));
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces as usize],
            download: download,
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        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 u32));
        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 send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

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

    fn read_n(&mut self, bytes_to_read: u32) -> 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 as usize => 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;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index as usize] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_block());
                }
            }
            Message::Piece(piece_index, offset, data) => {
                let block_index = offset / BLOCK_SIZE;
                try!(self.download.store(piece_index, block_index, data));
                try!(self.request_next_block());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_block(&mut self) -> Result<(), Error> {
        match self.download.next_block_to_request(&self.have) {
            Some((piece_index, block_index, block_length)) => {
                let offset = block_index * BLOCK_SIZE;
                self.send_message(Message::Request(piece_index, offset, block_length))
            },
            None => {
                println!("We've downloaded all the pieces we can from this peer.");
                Ok(())
            }
        }
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(u32),
    Bitfield(Vec<u8>),
    Request(u32, u32, u32),
    Piece(u32, u32, Vec<u8>),
    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(bytes_to_u32(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let length = bytes_to_u32(&body[8..12]);
                Message::Request(index, offset, length)
            },
            7 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, offset, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(u32_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, offset, amount) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(u32_to_bytes(amount).into_iter());
            },
            Message::Piece(index, offset, data) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = u32_to_bytes(payload.len() as u32);
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref offset, ref length) => write!(f, "Request({}, {}, {})", index, offset, length),
             Message::Piece(ref index, ref offset, ref data) => write!(f, "Piece({}, {}, size={})", index, offset, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

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

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

fn u32_to_bytes(integer: u32) -> 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|>


<commit_before>use std::{convert, io};
use std::fmt;
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 PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: u32 = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        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 u32));
        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 send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

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

    fn read_n(&mut self, bytes_to_read: u32) -> 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 as usize => 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;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index as usize] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, begin, data) => {
                self.downloaded[index as usize] = Some(data);
                try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<u32> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i as u32)
            }
        }
        println!("Done downloading file!");
        None
    }

    fn send_request(&mut self, piece: u32) -> Result<(), Error> {
        let num_pieces = self.downloaded.len() as u32;
        let request_size = if piece == num_pieces - 1 {
            (self.metainfo.info.length - (self.metainfo.info.piece_length as u64 * (num_pieces as u64 - 1))) as u32
        } else {
            BLOCK_SIZE
        };
        self.send_message(Message::Request(piece, 0, request_size))
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(u32),
    Bitfield(Vec<u8>),
    Request(u32, u32, u32),
    Piece(u32, u32, Block),
    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(bytes_to_u32(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_u32(&body[0..4]);
                let begin = bytes_to_u32(&body[4..8]);
                let offset = bytes_to_u32(&body[8..12]);
                Message::Request(index, begin, offset)
            },
            7 => {
                let index = bytes_to_u32(&body[0..4]);
                let begin = bytes_to_u32(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, begin, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(u32_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, begin, amount) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(begin).into_iter());
                payload.extend(u32_to_bytes(amount).into_iter());
            },
            Message::Piece(index, begin, data) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(begin).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = u32_to_bytes(payload.len() as u32);
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref begin, ref offset) => write!(f, "Request({}, {}, {})", index, begin, offset),
             Message::Piece(ref index, ref begin, ref data) => write!(f, "Piece({}, {}, size={})", index, begin, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

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

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

fn u32_to_bytes(integer: u32) -> 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)
    }
}
<commit_msg>Renamed offset to length and begin to offset in requests and pieces.<commit_after>use std::{convert, io};
use std::fmt;
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 PROTOCOL: &'static str = "BitTorrent protocol";
const BLOCK_SIZE: u32 = 16384;
type Block = Vec<u8>;

struct PeerConnection<'a> {
    metainfo: &'a Metainfo,
    stream: TcpStream,
    have: Vec<bool>,
    downloaded: Vec<Option<Block>>,
    am_i_choked: bool,
    am_i_interested: bool,
    are_they_choked: bool,
    are_they_interested: bool,
}

impl<'a> PeerConnection<'a> {
    fn connect(peer: &Peer, metainfo: &'a Metainfo) -> Result<PeerConnection<'a>, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            metainfo: metainfo,
            stream: stream,
            have: vec![false; num_pieces],
            downloaded: vec![None; num_pieces],
            am_i_choked: true,
            am_i_interested: false,
            are_they_choked: true,
            are_they_interested: false,
        };
        try!(conn.run());
        Ok(conn)
    }

    fn run(&mut self) -> Result<(), Error> {
        try!(self.send_handshake());
        try!(self.receive_handshake());
        loop {
            let message = try!(self.receive_message());
            println!("Recieved: {:?}", message);
            try!(self.process(message));
        }
        Ok(())
    }

    fn send_handshake(&mut self) -> Result<(), Error> {
        let mut message = vec![];
        message.push(PROTOCOL.len() as u8);
        message.extend(PROTOCOL.bytes());
        message.extend(vec![0; 8].into_iter());
        message.extend(self.metainfo.info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".bytes());
        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 u32));
        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 send_message(&mut self, message: Message) -> Result<(), Error> {
        println!("Sending: {:?}", message);
        try!(self.stream.write_all(&message.serialize()));
        Ok(())
    }

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

    fn read_n(&mut self, bytes_to_read: u32) -> 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 as usize => 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;
                };
                try!(self.send_interested());
            },
            Message::Have(have_index) => {
                self.have[have_index as usize] = true;
                try!(self.send_interested());
            },
            Message::Unchoke => {
                if self.am_i_choked {
                    self.am_i_choked = false;
                    try!(self.request_next_piece());
                }
            }
            Message::Piece(index, offset, data) => {
                self.downloaded[index as usize] = Some(data);
                try!(self.request_next_piece());
            }
            _ => panic!("Need to process message: {:?}", message)
        };
        Ok(())
    }

    fn send_interested(&mut self) -> Result<(), Error> {
        if self.am_i_interested == false {
            self.am_i_interested = true;
            try!(self.send_message(Message::Interested));
        }
        Ok(())
    }

    fn request_next_piece(&mut self) -> Result<(), Error> {
        match self.next_piece_to_request() {
            Some(piece_index) => self.send_request(piece_index),
            None => Ok(())
        }
    }

    fn next_piece_to_request(&self) -> Option<u32> {
        for i in 0..self.downloaded.len() {
            if self.downloaded[i].is_none() {
                return Some(i as u32)
            }
        }
        println!("Done downloading file!");
        None
    }

    fn send_request(&mut self, piece: u32) -> Result<(), Error> {
        let num_pieces = self.downloaded.len() as u32;
        let request_size = if piece == num_pieces - 1 {
            (self.metainfo.info.length - (self.metainfo.info.piece_length as u64 * (num_pieces as u64 - 1))) as u32
        } else {
            BLOCK_SIZE
        };
        self.send_message(Message::Request(piece, 0, request_size))
    }
}

enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(u32),
    Bitfield(Vec<u8>),
    Request(u32, u32, u32),
    Piece(u32, u32, Block),
    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(bytes_to_u32(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let length = bytes_to_u32(&body[8..12]);
                Message::Request(index, offset, length)
            },
            7 => {
                let index = bytes_to_u32(&body[0..4]);
                let offset = bytes_to_u32(&body[4..8]);
                let data = body[8..].to_owned();
                Message::Piece(index, offset, data)
            },
            8 => Message::Cancel,
            9 => Message::Port,
            _ => panic!("Bad message id: {}", id)
        }
    }

    fn serialize(self) -> Vec<u8> {
        let mut payload = vec![];
        match self {
            Message::KeepAlive => {},
            Message::Choke => payload.push(0),
            Message::Unchoke => payload.push(1),
            Message::Interested => payload.push(2),
            Message::NotInterested => payload.push(3),
            Message::Have(index) => {
                payload.push(4);
                payload.extend(u32_to_bytes(index).into_iter());
            },
            Message::Bitfield(bytes) => {
                payload.push(5);
                payload.extend(bytes);
            },
            Message::Request(index, offset, amount) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(u32_to_bytes(amount).into_iter());
            },
            Message::Piece(index, offset, data) => {
                payload.push(6);
                payload.extend(u32_to_bytes(index).into_iter());
                payload.extend(u32_to_bytes(offset).into_iter());
                payload.extend(data);
            },
            Message::Cancel => payload.push(8),
            Message::Port => payload.push(9),
        };

        // prepend size
        let mut size = u32_to_bytes(payload.len() as u32);
        size.extend(payload);
        size
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
             Message::KeepAlive => write!(f, "KeepAlive"),
             Message::Choke => write!(f, "Choke"),
             Message::Unchoke => write!(f, "Unchoke"),
             Message::Interested => write!(f, "Interested"),
             Message::NotInterested => write!(f, "NotInterested"),
             Message::Have(ref index) => write!(f, "Have({})", index),
             Message::Bitfield(ref bytes) => write!(f, "Bitfield({:?})", bytes),
             Message::Request(ref index, ref offset, ref length) => write!(f, "Request({}, {}, {})", index, offset, length),
             Message::Piece(ref index, ref offset, ref data) => write!(f, "Piece({}, {}, size={})", index, offset, data.len()),
             Message::Cancel => write!(f, "Cancel"),
             Message::Port => write!(f, "Port"),
        }
    }
}

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

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

fn u32_to_bytes(integer: u32) -> 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|>