<commit_before>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;
use uuid::Uuid;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;

use Provider;
use Track;
use soundcloud;
use youtube;
use error::Error;

use url::percent_encoding::{percent_decode};

static YOUTUBE_EMBED:       &'static str = r"www.youtube.com/embed/([^/\?&{videoseries}].+)";
static YOUTUBE_LIST:        &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:       &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:    &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST: &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:     &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";

pub fn extract_tracks(url: &str) -> Result<Vec<Track>, Error> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut tracks  = Vec::new();
        walk(0, dom.document, &mut tracks);
        Ok(tracks)
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        Err(Error::NotFound)
    }
}

// This is not proper HTML serialization, of course.
fn walk(indent: usize, handle: Handle, tracks: &mut Vec<Track>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let ts: Vec<Track> = extract_tracks_from_tag(tag_name, attrs);
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(indent+4, child.clone(), tracks);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => match cap.at(1) {
                Some(str) => {
                    let strs: Vec<&str> = str.split('?').collect();
                    return Some(strs[0].to_string())
                },
                None => None
            },
            None => None
        },
        Err(_) => None
    }
}

pub fn extract_tracks_from_tag(tag_name: &str,
                               attrs: &Vec<Attribute>) -> Vec<Track> {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_tracks_from_url(src.to_string()),
            None => vec![]
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_tracks_from_url(href.to_string()),
            None => vec![]
        }
    } else {
        vec![]
    }
}

fn extract_tracks_from_url(url: String) -> Vec<Track> {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => {
            return match youtube::fetch_playlist(&identifier) {
                Ok(res) => res.items.iter()
                                    .map(|ref i| Track::from_yt_playlist_item(i))
                                    .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::SoundCloud,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => {
            return match soundcloud::fetch_playlist(&identifier) {
                Ok(playlist) => playlist.tracks
                                        .iter()
                                        .map(|ref t| Track::from_sc_track(t))
                                        .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => {
            print!("user {}", identifier);
            return match soundcloud::fetch_user_tracks(&identifier) {
                Ok(tracks) => tracks.iter()
                                    .map(|ref t| Track::from_sc_track(t))
                                    .collect::<Vec<_>>(),
                Err(_)     => vec![]
            }
        },
        None => ()
    }
    return vec![]
}

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

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }

        match extract_identifier(youtube_list, super::YOUTUBE_EMBED) {
            Some(_) => assert!(false),
            None    => assert!(true)
        }
    }
}
<commit_msg>Fix bug of extracting youtube identifer: ignore video id = videoseires<commit_after>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;
use uuid::Uuid;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;

use Provider;
use Track;
use soundcloud;
use youtube;
use error::Error;

use url::percent_encoding::{percent_decode};

static YOUTUBE_EMBED:       &'static str = r"www.youtube.com/embed/([a-zA-Z0-9_-].+)";
static YOUTUBE_LIST:        &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:       &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:    &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST: &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:     &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";

pub fn extract_tracks(url: &str) -> Result<Vec<Track>, Error> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut tracks  = Vec::new();
        walk(0, dom.document, &mut tracks);
        Ok(tracks)
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        Err(Error::NotFound)
    }
}

// This is not proper HTML serialization, of course.
fn walk(indent: usize, handle: Handle, tracks: &mut Vec<Track>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let ts: Vec<Track> = extract_tracks_from_tag(tag_name, attrs);
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(indent+4, child.clone(), tracks);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => match cap.at(1) {
                Some(str) => {
                    let strs: Vec<&str> = str.split('?').collect();
                    return Some(strs[0].to_string())
                },
                None => None
            },
            None => None
        },
        Err(_) => None
    }
}

pub fn extract_tracks_from_tag(tag_name: &str,
                               attrs: &Vec<Attribute>) -> Vec<Track> {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_tracks_from_url(src.to_string()),
            None => vec![]
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_tracks_from_url(href.to_string()),
            None => vec![]
        }
    } else {
        vec![]
    }
}

fn extract_tracks_from_url(url: String) -> Vec<Track> {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => {
            return match youtube::fetch_playlist(&identifier) {
                Ok(res) => res.items.iter()
                                    .map(|ref i| Track::from_yt_playlist_item(i))
                                    .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::SoundCloud,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => {
            return match soundcloud::fetch_playlist(&identifier) {
                Ok(playlist) => playlist.tracks
                                        .iter()
                                        .map(|ref t| Track::from_sc_track(t))
                                        .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => {
            print!("user {}", identifier);
            return match soundcloud::fetch_user_tracks(&identifier) {
                Ok(tracks) => tracks.iter()
                                    .map(|ref t| Track::from_sc_track(t))
                                    .collect::<Vec<_>>(),
                Err(_)     => vec![]
            }
        },
        None => ()
    }
    return vec![]
}

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

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }
    }
}
<|endoftext|>


<commit_before>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;
use uuid::Uuid;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;

use Provider;
use Track;
use soundcloud;
use youtube;

use url::percent_encoding::{percent_decode};

static YOUTUBE_EMBED:    &'static str = r"www.youtube.com/embed/([^\?&{videoseries}].+)";
static YOUTUBE_LIST:         &'static str = r"www.youtube.com/embed/videoseries\?list=([^\?]+)";
static YOUTUBE_WATCH:        &'static str = r"www.youtube.com/watch\?v=([^\?&]+)";
static SOUNDCLOUD_TRACK:     &'static str = r"api.soundcloud.com/tracks/([^\?&]+)";
static SOUNDCLOUD_PLAYLIST:  &'static str = r"api.soundcloud.com/playlists/([^\?&]+)";
static SOUNDCLOUD_USER:      &'static str = r"api.soundcloud.com/users/([^\?&]+)";

pub fn extract_tracks(url: &str) -> Option<Vec<Track>> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut tracks  = Vec::new();
        walk(0, dom.document, &mut tracks);
        return Some(tracks)
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        return None
    }
}

// This is not proper HTML serialization, of course.
fn walk(indent: usize, handle: Handle, tracks: &mut Vec<Track>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let ts: Vec<Track> = extract_tracks_from_tag(tag_name, attrs);
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(indent+4, child.clone(), tracks);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => match cap.at(1) {
                Some(str) => {
                    let strs: Vec<&str> = str.split('?').collect();
                    return Some(strs[0].to_string())
                },
                None => None
            },
            None => None
        },
        Err(_) => None
    }
}

pub fn extract_tracks_from_tag(tag_name: &str,
                               attrs: &Vec<Attribute>) -> Vec<Track> {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_tracks_from_url(src.to_string()),
            None => vec![]
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_tracks_from_url(href.to_string()),
            None => vec![]
        }
    } else {
        vec![]
    }
}

fn extract_tracks_from_url(url: String) -> Vec<Track> {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => {
            return match youtube::fetch_playlist(&identifier) {
                Ok(res) => res.items.iter()
                                    .map(|ref i| Track::from_yt_playlist_item(i))
                                    .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::SoundCloud,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => {
            return match soundcloud::fetch_playlist(&identifier) {
                Ok(playlist) => playlist.tracks
                                        .iter()
                                        .map(|ref t| Track::from_sc_track(t))
                                        .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => {
            print!("user {}", identifier);
            return match soundcloud::fetch_user_tracks(&identifier) {
                Ok(tracks) => tracks.iter()
                                    .map(|ref t| Track::from_sc_track(t))
                                    .collect::<Vec<_>>(),
                Err(_)     => vec![]
            }
        },
        None => ()
    }
    return vec![]
}

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

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }

        match extract_identifier(youtube_list, super::YOUTUBE_EMBED) {
            Some(identifier) => assert!(false),
            None             => assert!(true)
        }
    }
}
<commit_msg>Fix track extractor. Use white list regex in order to remove `/stream` from sound cloud identifier<commit_after>extern crate html5ever;
extern crate tendril;
extern crate regex;
extern crate hyper;
extern crate string_cache;
extern crate url;

use html5ever::rcdom::{Document, Doctype, Text, Comment, Element};
use html5ever::rcdom::{RcDom, Handle};
use html5ever::{parse_document, Attribute};
use tendril::stream::TendrilSink;
use std::default::Default;
use uuid::Uuid;

use regex::Regex;
use hyper::Client;
use hyper::header::Connection;
use hyper::header::ConnectionOption;

use Provider;
use Track;
use soundcloud;
use youtube;

use url::percent_encoding::{percent_decode};

static YOUTUBE_EMBED:       &'static str = r"www.youtube.com/embed/([^/\?&{videoseries}].+)";
static YOUTUBE_LIST:        &'static str = r"www.youtube.com/embed/videoseries\?list=([a-zA-Z0-9_-]+)";
static YOUTUBE_WATCH:       &'static str = r"www.youtube.com/watch\?v=([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_TRACK:    &'static str = r"api.soundcloud.com/tracks/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_PLAYLIST: &'static str = r"api.soundcloud.com/playlists/([a-zA-Z0-9_-]+)";
static SOUNDCLOUD_USER:     &'static str = r"api.soundcloud.com/users/([a-zA-Z0-9_-]+)";

pub fn extract_tracks(url: &str) -> Option<Vec<Track>> {
    let client = Client::new();
    let mut res = client.get(url)
        .header(Connection(vec![ConnectionOption::Close]))
        .send()
        .unwrap();
    if res.status.is_success() {
        let dom = parse_document(RcDom::default(), Default::default())
            .from_utf8()
            .read_from(&mut res)
            .unwrap();
        let mut tracks  = Vec::new();
        walk(0, dom.document, &mut tracks);
        return Some(tracks)
    } else {
        println!("Failed to get entry html {}: {}", res.status, url);
        return None
    }
}

// This is not proper HTML serialization, of course.
fn walk(indent: usize, handle: Handle, tracks: &mut Vec<Track>) {
    let node = handle.borrow();
    match node.node {
        Document         => (),
        Doctype(_, _, _) => (),
        Text(_)          => (),
        Comment(_)       => (),
        Element(ref name, _, ref attrs) => {
            let tag_name = name.local.as_ref();
            let ts: Vec<Track> = extract_tracks_from_tag(tag_name, attrs);
            for track in ts.iter().cloned() {
                if !(tracks).iter().any(|t| track == *t) {
                    (*tracks).push(track)
                }
            }
        }
    }
    for child in node.children.iter() {
        walk(indent+4, child.clone(), tracks);
    }
}

fn attr(attr_name: &str, attrs: &Vec<Attribute>) -> Option<String> {
    for attr in attrs.iter() {
        if attr.name.local.as_ref() == attr_name {
            return Some(attr.value.to_string())
        }
    }
    None
}

fn extract_identifier(value: &str, regex_str: &str) -> Option<String> {
    match Regex::new(regex_str) {
        Ok(re) => match re.captures(value) {
            Some(cap) => match cap.at(1) {
                Some(str) => {
                    let strs: Vec<&str> = str.split('?').collect();
                    return Some(strs[0].to_string())
                },
                None => None
            },
            None => None
        },
        Err(_) => None
    }
}

pub fn extract_tracks_from_tag(tag_name: &str,
                               attrs: &Vec<Attribute>) -> Vec<Track> {
    if tag_name == "iframe" {
        match attr("src", attrs) {
            Some(ref src) => extract_tracks_from_url(src.to_string()),
            None => vec![]
        }
    } else if tag_name == "a" || tag_name == "link" {
        match attr("href", attrs) {
            Some(ref href) => extract_tracks_from_url(href.to_string()),
            None => vec![]
        }
    } else {
        vec![]
    }
}

fn extract_tracks_from_url(url: String) -> Vec<Track> {
    let decoded = percent_decode(url.as_bytes()).decode_utf8_lossy().into_owned();
    match extract_identifier(&decoded, YOUTUBE_EMBED) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_WATCH) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::YouTube,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, YOUTUBE_LIST) {
        Some(identifier) => {
            return match youtube::fetch_playlist(&identifier) {
                Ok(res) => res.items.iter()
                                    .map(|ref i| Track::from_yt_playlist_item(i))
                                    .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_TRACK) {
        Some(identifier) => {
            return vec![Track {
                        id: Uuid::new_v4(),
                  provider: Provider::SoundCloud,
                     title: "".to_string(),
                       url: url.to_string(),
                identifier: identifier
            }]
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_PLAYLIST) {
        Some(identifier) => {
            return match soundcloud::fetch_playlist(&identifier) {
                Ok(playlist) => playlist.tracks
                                        .iter()
                                        .map(|ref t| Track::from_sc_track(t))
                                        .collect::<Vec<_>>(),
                Err(_)       => vec![]
            }
        },
        None => ()
    }
    match extract_identifier(&decoded, SOUNDCLOUD_USER) {
        Some(identifier) => {
            print!("user {}", identifier);
            return match soundcloud::fetch_user_tracks(&identifier) {
                Ok(tracks) => tracks.iter()
                                    .map(|ref t| Track::from_sc_track(t))
                                    .collect::<Vec<_>>(),
                Err(_)     => vec![]
            }
        },
        None => ()
    }
    return vec![]
}

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

    #[test]
    fn test_extract_identifier() {
        let soundcloud_src = "https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/tracks/195425494/stream&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&visual=true";
        match extract_identifier(soundcloud_src, super::SOUNDCLOUD_TRACK) {
            Some(identifier) => assert_eq!(identifier, "195425494".to_string()),
            None             => assert!(false)
        }
        let youtube_embed = "https://www.youtube.com/embed/X8tOngmlES0?rel=0";
        match extract_identifier(youtube_embed, super::YOUTUBE_EMBED) {
            Some(identifier) => assert_eq!(identifier, "X8tOngmlES0".to_string()),
            None             => assert!(false)
        }

        let youtube_watch = "https://www.youtube.com/watch?v=oDuif301F-8";
        match extract_identifier(youtube_watch, super::YOUTUBE_WATCH) {
            Some(identifier) => assert_eq!(identifier, "oDuif301F-8".to_string()),
            None             => assert!(false)
        }

        let youtube_list = "https://www.youtube.com/embed/videoseries?list=PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX";
        match extract_identifier(youtube_list, super::YOUTUBE_LIST) {
            Some(identifier) => assert_eq!(identifier, "PLy8LZ8FM-o0ViuGAF68RAaXkQ8V-3dbTX".to_string()),
            None             => assert!(false)
        }

        match extract_identifier(youtube_list, super::YOUTUBE_EMBED) {
            Some(_) => assert!(false),
            None    => assert!(true)
        }
    }
}
<|endoftext|>


<commit_before>//! Iterators that are sources (produce elements from parameters,
//! not from another iterator).

use std::fmt;
use std::mem;

/// See [`repeat_call`](../fn.repeat_call.html) for more information.
pub struct RepeatCall<F> {
    f: F,
}

impl<F> fmt::Debug for RepeatCall<F>
{
    debug_fmt_fields!(RepeatCall, );
}

/// An iterator source that produces elements indefinitely by calling
/// a given closure.
///
/// Iterator element type is the return type of the closure.
///
/// ```
/// use itertools::repeat_call;
/// use itertools::Itertools;
/// use std::collections::BinaryHeap;
///
/// let mut heap = BinaryHeap::from(vec![2, 5, 3, 7, 8]);
///
/// // extract each element in sorted order
/// for element in repeat_call(|| heap.pop()).while_some() {
///     print!("{}", element);
/// }
///
/// itertools::assert_equal(
///     repeat_call(|| 1).take(5),
///     vec![1, 1, 1, 1, 1]
/// );
/// ```
pub fn repeat_call<F>(function: F) -> RepeatCall<F> {
    RepeatCall { f: function }
}

impl Iterator for RepeatCall<F>
    where F: FnMut() -> A
{
    type Item = A;

    #[inline]
    fn next(&mut self) -> Option {
        Some((self.f)())
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::max_value(), None)
    }
}

/// Creates a new unfold source with the specified closure as the "iterator
/// function" and an initial state to eventually pass to the closure
///
/// `unfold` is a general iterator builder: it has a mutable state value,
/// and a closure with access to the state that produces the next value.
///
/// This more or less equivalent to a regular struct with an `Iterator`
/// implementation, and is useful for one-off iterators.
///
/// ```
/// // an iterator that yields sequential Fibonacci numbers,
/// // and stops at the maximum representable value.
///
/// use itertools::unfold;
///
/// let (mut x1, mut x2) = (1u32, 1u32);
/// let mut fibonacci = unfold((), move |_| {
///     // Attempt to get the next Fibonacci number
///     let next = x1.saturating_add(x2);
///
///     // Shift left: ret <- x1 <- x2 <- next
///     let ret = x1;
///     x1 = x2;
///     x2 = next;
///
///     // If addition has saturated at the maximum, we are finished
///     if ret == x1 && ret > 1 {
///         return None;
///     }
///
///     Some(ret)
/// });
///
/// itertools::assert_equal(fibonacci.by_ref().take(8),
///                         vec![1, 1, 2, 3, 5, 8, 13, 21]);
/// assert_eq!(fibonacci.last(), Some(2_971_215_073))
/// ```
pub fn unfold(initial_state: St, f: F) -> Unfold<St, F>
    where F: FnMut(&mut St) -> Option
{
    Unfold {
        f: f,
        state: initial_state,
    }
}

impl<St, F> fmt::Debug for Unfold<St, F>
    where St: fmt::Debug,
{
    debug_fmt_fields!(Unfold, state);
}

/// See [`unfold`](../fn.unfold.html) for more information.
#[derive(Clone)]
pub struct Unfold<St, F> {
    f: F,
    /// Internal state that will be passed to the closure on the next iteration
    pub state: St,
}

impl Iterator for Unfold<St, F>
    where F: FnMut(&mut St) -> Option
{
    type Item = A;

    #[inline]
    fn next(&mut self) -> Option {
        (self.f)(&mut self.state)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        // no possible known bounds at this point
        (0, None)
    }
}

/// An iterator that infinitely applies function to value and yields results.
///
/// This `struct` is created by the [`iterate()`] function. See its documentation for more.
///
/// [`iterate()`]: ../fn.iterate.html
#[derive(Clone)]
pub struct Iterate<St, F> {
    state: St,
    f: F,
}

impl<St, F> fmt::Debug for Iterate<St, F>
    where St: fmt::Debug,
{
    debug_fmt_fields!(Iterate, state);
}

impl<St, F> Iterator for Iterate<St, F>
    where F: FnMut(&St) -> St
{
    type Item = St;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let next_state = (self.f)(&self.state);
        Some(mem::replace(&mut self.state, next_state))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::max_value(), None)
    }
}

/// Creates a new iterator that infinitely applies function to value and yields results.
///
/// ```
/// use itertools::iterate;
///
/// itertools::assert_equal(iterate(1, |&i| i * 3).take(5), vec![1, 3, 9, 27, 81]);
/// ```
pub fn iterate<St, F>(initial_value: St, f: F) -> Iterate<St, F>
    where F: FnMut(&St) -> St
{
    Iterate {
        state: initial_value,
        f: f,
    }
}
<commit_msg>API: Add closure bound to repeat_call<commit_after>//! Iterators that are sources (produce elements from parameters,
//! not from another iterator).

use std::fmt;
use std::mem;

/// See [`repeat_call`](../fn.repeat_call.html) for more information.
pub struct RepeatCall<F> {
    f: F,
}

impl<F> fmt::Debug for RepeatCall<F>
{
    debug_fmt_fields!(RepeatCall, );
}

/// An iterator source that produces elements indefinitely by calling
/// a given closure.
///
/// Iterator element type is the return type of the closure.
///
/// ```
/// use itertools::repeat_call;
/// use itertools::Itertools;
/// use std::collections::BinaryHeap;
///
/// let mut heap = BinaryHeap::from(vec![2, 5, 3, 7, 8]);
///
/// // extract each element in sorted order
/// for element in repeat_call(|| heap.pop()).while_some() {
///     print!("{}", element);
/// }
///
/// itertools::assert_equal(
///     repeat_call(|| 1).take(5),
///     vec![1, 1, 1, 1, 1]
/// );
/// ```
pub fn repeat_call<F, A>(function: F) -> RepeatCall<F>
    where F: FnMut() -> A
{
    RepeatCall { f: function }
}

impl Iterator for RepeatCall<F>
    where F: FnMut() -> A
{
    type Item = A;

    #[inline]
    fn next(&mut self) -> Option {
        Some((self.f)())
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::max_value(), None)
    }
}

/// Creates a new unfold source with the specified closure as the "iterator
/// function" and an initial state to eventually pass to the closure
///
/// `unfold` is a general iterator builder: it has a mutable state value,
/// and a closure with access to the state that produces the next value.
///
/// This more or less equivalent to a regular struct with an `Iterator`
/// implementation, and is useful for one-off iterators.
///
/// ```
/// // an iterator that yields sequential Fibonacci numbers,
/// // and stops at the maximum representable value.
///
/// use itertools::unfold;
///
/// let (mut x1, mut x2) = (1u32, 1u32);
/// let mut fibonacci = unfold((), move |_| {
///     // Attempt to get the next Fibonacci number
///     let next = x1.saturating_add(x2);
///
///     // Shift left: ret <- x1 <- x2 <- next
///     let ret = x1;
///     x1 = x2;
///     x2 = next;
///
///     // If addition has saturated at the maximum, we are finished
///     if ret == x1 && ret > 1 {
///         return None;
///     }
///
///     Some(ret)
/// });
///
/// itertools::assert_equal(fibonacci.by_ref().take(8),
///                         vec![1, 1, 2, 3, 5, 8, 13, 21]);
/// assert_eq!(fibonacci.last(), Some(2_971_215_073))
/// ```
pub fn unfold(initial_state: St, f: F) -> Unfold<St, F>
    where F: FnMut(&mut St) -> Option
{
    Unfold {
        f: f,
        state: initial_state,
    }
}

impl<St, F> fmt::Debug for Unfold<St, F>
    where St: fmt::Debug,
{
    debug_fmt_fields!(Unfold, state);
}

/// See [`unfold`](../fn.unfold.html) for more information.
#[derive(Clone)]
pub struct Unfold<St, F> {
    f: F,
    /// Internal state that will be passed to the closure on the next iteration
    pub state: St,
}

impl Iterator for Unfold<St, F>
    where F: FnMut(&mut St) -> Option
{
    type Item = A;

    #[inline]
    fn next(&mut self) -> Option {
        (self.f)(&mut self.state)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        // no possible known bounds at this point
        (0, None)
    }
}

/// An iterator that infinitely applies function to value and yields results.
///
/// This `struct` is created by the [`iterate()`] function. See its documentation for more.
///
/// [`iterate()`]: ../fn.iterate.html
#[derive(Clone)]
pub struct Iterate<St, F> {
    state: St,
    f: F,
}

impl<St, F> fmt::Debug for Iterate<St, F>
    where St: fmt::Debug,
{
    debug_fmt_fields!(Iterate, state);
}

impl<St, F> Iterator for Iterate<St, F>
    where F: FnMut(&St) -> St
{
    type Item = St;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        let next_state = (self.f)(&self.state);
        Some(mem::replace(&mut self.state, next_state))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (usize::max_value(), None)
    }
}

/// Creates a new iterator that infinitely applies function to value and yields results.
///
/// ```
/// use itertools::iterate;
///
/// itertools::assert_equal(iterate(1, |&i| i * 3).take(5), vec![1, 3, 9, 27, 81]);
/// ```
pub fn iterate<St, F>(initial_value: St, f: F) -> Iterate<St, F>
    where F: FnMut(&St) -> St
{
    Iterate {
        state: initial_value,
        f: f,
    }
}
<|endoftext|>


<commit_before>use std::env;

fn main() {
    let mut newline = true;
    for arg in env::args().skip(1) {
        if arg == "-n" {
            newline = false
            continue;
        }
        print!("{} ", arg);
    }
    if newline {
        print!("\n");
    }
}
<commit_msg>Fix compile error<commit_after>use std::env;

fn main() {
    let mut newline = true;
    for arg in env::args().skip(1) {
        if arg == "-n" {
            newline = false;
            continue;
        }
        print!("{} ", arg);
    }
    if newline {
        print!("\n");
    }
}
<|endoftext|>


<commit_before>use std::io::process::ExitStatus;

use cargo::ops;
use cargo::core::MultiShell;
use cargo::util::{CliResult, CliError, CargoError};
use cargo::util::important_paths::{find_root_manifest_for_cwd};

#[deriving(Decodable)]
struct Options {
    arg_args: Vec<String>,
    flag_features: Vec<String>,
    flag_jobs: Option<uint>,
    flag_manifest_path: Option<String>,
    flag_name: Option<String>,
    flag_no_default_features: bool,
    flag_no_run: bool,
    flag_package: Option<String>,
    flag_target: Option<String>,
    flag_verbose: bool,
}

pub const USAGE: &'static str = "
Execute all unit and integration tests of a local package

Usage:
    cargo test [options] [--] [<args>...]

Options:
    -h, --help               Print this message
    --name NAME              Name of the test to run
    --no-run                 Compile, but don't run tests
    -p SPEC, --package SPEC  Package to run tests for
    -j N, --jobs N           The number of jobs to run in parallel
    --features FEATURES      Space-separated list of features to also build
    --no-default-features    Do not build the `default` feature
    --target TRIPLE          Build for the target triple
    --manifest-path PATH     Path to the manifest to build tests for
    -v, --verbose            Use verbose output

All of the trailing arguments are passed to the test binaries generated for
filtering tests and generally providing options configuring how they run.

If the --package argument is given, then SPEC is a package id specification
which indicates which package should be tested. If it is not given, then the
current package is tested. For more information on SPEC and its format, see the
`cargo help pkgid` command.
";

pub fn execute(options: Options, shell: &mut MultiShell) -> CliResult<Option<()>> {
    let root = try!(find_root_manifest_for_cwd(options.flag_manifest_path));
    shell.set_verbose(options.flag_verbose);

    let mut ops = ops::TestOptions {
        name: options.flag_name.as_ref().map(|s| s.as_slice()),
        no_run: options.flag_no_run,
        compile_opts: ops::CompileOptions {
            env: "test",
            shell: shell,
            jobs: options.flag_jobs,
            target: options.flag_target.as_ref().map(|s| s.as_slice()),
            dev_deps: true,
            features: options.flag_features.as_slice(),
            no_default_features: options.flag_no_default_features,
            spec: options.flag_package.as_ref().map(|s| s.as_slice()),
        },
    };

    let err = try!(ops::run_tests(&root, &mut ops,
                                  options.arg_args.as_slice()).map_err(|err| {
        CliError::from_boxed(err, 101)
    }));
    match err {
        None => Ok(None),
        Some(err) => {
            Err(match err.exit {
                Some(ExitStatus(i)) => CliError::new("", i as uint),
                _ => CliError::from_boxed(err.concrete().mark_human(), 101)
            })
        }
    }
}
<commit_msg>auto merge of #816 : alexcrichton/cargo/issue-812, r=brson<commit_after>use std::io::process::ExitStatus;

use cargo::ops;
use cargo::core::MultiShell;
use cargo::util::{CliResult, CliError, CargoError};
use cargo::util::important_paths::{find_root_manifest_for_cwd};

#[deriving(Decodable)]
struct Options {
    arg_args: Vec<String>,
    flag_features: Vec<String>,
    flag_jobs: Option<uint>,
    flag_manifest_path: Option<String>,
    flag_name: Option<String>,
    flag_no_default_features: bool,
    flag_no_run: bool,
    flag_package: Option<String>,
    flag_target: Option<String>,
    flag_verbose: bool,
}

pub const USAGE: &'static str = "
Execute all unit and integration tests of a local package

Usage:
    cargo test [options] [--] [<args>...]

Options:
    -h, --help               Print this message
    --name NAME              Name of the test executable to run
    --no-run                 Compile, but don't run tests
    -p SPEC, --package SPEC  Package to run tests for
    -j N, --jobs N           The number of jobs to run in parallel
    --features FEATURES      Space-separated list of features to also build
    --no-default-features    Do not build the `default` feature
    --target TRIPLE          Build for the target triple
    --manifest-path PATH     Path to the manifest to build tests for
    -v, --verbose            Use verbose output

All of the trailing arguments are passed to the test binaries generated for
filtering tests and generally providing options configuring how they run. For
example, this will run all tests with the name `foo` in their name:

    cargo test foo

If the --package argument is given, then SPEC is a package id specification
which indicates which package should be tested. If it is not given, then the
current package is tested. For more information on SPEC and its format, see the
`cargo help pkgid` command.
";

pub fn execute(options: Options, shell: &mut MultiShell) -> CliResult<Option<()>> {
    let root = try!(find_root_manifest_for_cwd(options.flag_manifest_path));
    shell.set_verbose(options.flag_verbose);

    let mut ops = ops::TestOptions {
        name: options.flag_name.as_ref().map(|s| s.as_slice()),
        no_run: options.flag_no_run,
        compile_opts: ops::CompileOptions {
            env: "test",
            shell: shell,
            jobs: options.flag_jobs,
            target: options.flag_target.as_ref().map(|s| s.as_slice()),
            dev_deps: true,
            features: options.flag_features.as_slice(),
            no_default_features: options.flag_no_default_features,
            spec: options.flag_package.as_ref().map(|s| s.as_slice()),
        },
    };

    let err = try!(ops::run_tests(&root, &mut ops,
                                  options.arg_args.as_slice()).map_err(|err| {
        CliError::from_boxed(err, 101)
    }));
    match err {
        None => Ok(None),
        Some(err) => {
            Err(match err.exit {
                Some(ExitStatus(i)) => CliError::new("", i as uint),
                _ => CliError::from_boxed(err.concrete().mark_human(), 101)
            })
        }
    }
}
<|endoftext|>


<commit_before>//! Kafka Consumer
//!
//! A simple consumer based on KafkaClient. Accepts an instance of KafkaClient, a group and a
//! topic. Partitions can be specfied using builder pattern (Assumes all partitions if not
//! specfied).
//!
//! # Example
//!
//! ```no_run
//! let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_owned()));
//! let res = client.load_metadata_all();
//! let con = kafka::consumer::Consumer::new(client, "test-group".to_owned(), "my-topic".to_owned())
//!             .partition(0);
//! for msg in con {
//!     println!("{:?}", msg);
//! }
//! ```
//!
//! Consumer auto-commits the offsets afer consuming COMMIT_INTERVAL messages (Currently set at
//! 100)
//!
//! Consumer implements Iterator.

use std::collections::HashMap;
use error::Result;
use utils::{TopicMessage, TopicPartitionOffset};
use client::KafkaClient;

const COMMIT_INTERVAL: i32 = 100; // Commit after every 100 message

#[derive(Default, Debug)]
pub struct Consumer {
    client: KafkaClient,
    group: String,
    topic: String,
    partitions: Vec<i32>,
    initialized: bool,
    messages: Vec<TopicMessage>,
    index: usize,
    offsets: HashMap<i32, i64>,
    consumed: i32
}

impl Consumer {
    /// Constructor
    ///
    /// Create a new consumer. Expects a KafkaClient, group, and topic as arguments.
    pub fn new(client: KafkaClient, group: String, topic: String) -> Consumer {
        Consumer{
            client: client,
            group: group,
            topic: topic,
            initialized: false,
            index: 0,
            ..Consumer::default()
        }
    }

    /// Set the partitions of this consumer.
    ///
    /// If this function is never called, all partitions are assumed.
    /// This function call be called multiple times to add more than 1 partitions.
    pub fn partition(mut self, partition: i32) -> Consumer {
        self.partitions.push(partition);
        self
    }

    fn commit_offsets(&mut self) -> Result<()> {
        let tpos = self.offsets.iter()
                        .map(|(p, o)| TopicPartitionOffset{
                                topic: self.topic.clone(),
                                partition: p.clone(),
                                offset: o.clone()
                            })
                        .collect();
        self.client.commit_offsets(self.group.clone(), tpos)
    }

    fn fetch_offsets(&mut self) -> Result<()> {
        let tpos = try!(self.client.fetch_group_topic_offset(self.group.clone(), self.topic.clone()));
        if self.partitions.len() == 0 {
            self.partitions = self.client.topic_partitions.get(&self.topic).unwrap_or(&vec!()).clone();
        }
        for tpo in tpos {
            if self.partitions.contains(&tpo.partition) {
                self.offsets.insert(tpo.partition, tpo.offset);
            }
        }
        Ok(())
    }

    fn make_request(&mut self) -> Result<()>{
        if ! self.initialized {
            try!(self.fetch_offsets());
        }
        let tpos = self.offsets.iter()
                        .map(|(p, o)| TopicPartitionOffset{
                                topic: self.topic.clone(),
                                partition: p.clone(),
                                offset: o.clone()
                            })
                        .collect();
        self.messages = try!(self.client.fetch_messages_multi(tpos));
        self.initialized = true;
        self.index = 0;
        Ok(())
    }
}

impl Iterator for Consumer {
    type Item = TopicMessage;

    fn next(&mut self) -> Option<TopicMessage> {
        if self.initialized {
            self.index += 1;
            self.consumed += 1;
            if self.consumed % COMMIT_INTERVAL == 0 {
                let _ = self.commit_offsets();
            }
            if self.index <= self.messages.len() {
                if self.messages[self.index-1].error.is_none() {
                    let curr = self.offsets.entry(self.messages[self.index-1].partition).or_insert(0);
                    *curr = *curr+1;
                    return Some(self.messages[self.index-1].clone());
                }
                return None;
            }
            let _ = self.commit_offsets();
            if self.messages.len() == 0 {
                return None;
            }
        }
        match self.make_request() {
            Err(_) => None,
            Ok(_) => self.next()
        }
    }
}
<commit_msg>#! Fix initializing start offsets when missing for consumer-group<commit_after>//! Kafka Consumer
//!
//! A simple consumer based on KafkaClient. Accepts an instance of KafkaClient, a group and a
//! topic. Partitions can be specified using builder pattern (Assumes all partitions if not
//! specified).
//!
//! # Example
//!
//! ```no_run
//! let mut client = kafka::client::KafkaClient::new(vec!("localhost:9092".to_owned()));
//! let res = client.load_metadata_all();
//! let con = kafka::consumer::Consumer::new(client, "test-group".to_owned(), "my-topic".to_owned())
//!             .partition(0)
//!             .fallback_offset(-2);
//! for msg in con {
//!     println!("{:?}", msg);
//! }
//! ```
//!
//! Consumer auto-commits the offsets after consuming COMMIT_INTERVAL messages (Currently set at
//! 100)
//!
//! Consumer implements Iterator.

use std::collections::HashMap;
use error::{Error, Result};
use utils::{TopicMessage, TopicPartitionOffset, TopicPartitionOffsetError};
use client::KafkaClient;

const COMMIT_INTERVAL: i32 = 100; // Commit after every 100 message

#[derive(Default, Debug)]
pub struct Consumer {
    client: KafkaClient,
    group: String,
    topic: String,
    partitions: Vec<i32>,
    initialized: bool,
    messages: Vec<TopicMessage>,
    index: usize,
    offsets: HashMap<i32, i64>,
    consumed: i32,

    fallback_offset: Option<i64>,
}

impl Consumer {

    /// Constructor
    ///
    /// Create a new consumer. Expects a KafkaClient, group, and topic as arguments.
    pub fn new(client: KafkaClient, group: String, topic: String) -> Consumer {
        Consumer{
            client: client,
            group: group,
            topic: topic,
            initialized: false,
            index: 0,
            ..Consumer::default()
        }
    }

    /// Set the partitions of this consumer.
    ///
    /// If this function is never called, all partitions are assumed.
    /// This function call be called multiple times to add more than 1 partitions.
    pub fn partition(mut self, partition: i32) -> Consumer {
        self.partitions.push(partition);
        self
    }

    /// Specify the offset to use when none was committed for the
    /// underlying group yet.
    ///
    /// Running the underlying group for the first time against a
    /// topic results in the question where to start reading from the
    /// topic?  It might already contain a lot of messages.  Common
    /// strategies are starting at the earliest available message
    /// (thereby consuming whatever is currently in the topic) or at
    /// the latest one (thereby staring to consume only newly arriving
    /// messages.)  The parameter here corresponds to `time` in
    /// `KafkaClient::fetch_offsets`.
    ///
    /// Unless this method is called and there is no offset committed
    /// for the underlying group yet, this consumer will _not_ retrieve
    /// any messages from the underlying topic.
    pub fn fallback_offset(mut self, fallback_offset_time: i64) -> Consumer {
        self.fallback_offset = Some(fallback_offset_time);
        self
    }

    fn commit_offsets(&mut self) -> Result<()> {
        let tpos = self.offsets.iter()
                        .map(|(p, o)| TopicPartitionOffset{
                                topic: self.topic.clone(),
                                partition: p.clone(),
                                offset: o.clone()
                            })
                        .collect();
        self.client.commit_offsets(self.group.clone(), tpos)
    }

    fn fetch_offsets(&mut self) -> Result<()> {
        // ~ fetch the so far commited group offsets
        let mut tpos = try!(self.client.fetch_group_topic_offset(self.group.clone(), self.topic.clone()));
        if self.partitions.len() == 0 {
            self.partitions = self.client.topic_partitions.get(&self.topic).unwrap_or(&vec!()).clone();
        }

        // ~ it might well that there were no group offsets committed
        // yet ... fallback to default offsets.
        try!(self.set_fallback_offsets(&mut tpos));

        // ~ now initialized from the fetched offsets
        for tpo in &tpos {
            if self.partitions.contains(&tpo.partition) {
                self.offsets.insert(tpo.partition, tpo.offset);
            }
        }
        Ok(())
    }

    /// Try setting the "fallback offsets" for all of `tpos` where
    /// `offset == -1`. Fails if retrieving the fallback offset is not
    /// possible for some reason for the affected elements from
    /// `tpos`.
    fn set_fallback_offsets(&mut self, tpos: &mut [TopicPartitionOffsetError]) -> Result<()> {
        // ~ it looks like kafka (0.8.2.1) is sending an error code
        // (even though it documents it won't: https://cwiki.apache.org/confluence/display/KAFKA/A+Guide+To+The+Kafka+Protocol#AGuideToTheKafkaProtocol-OffsetFetchResponse)
        // so let's check only for `offset == -1` (and don't verify `error == None`)
        let has_missing_offs = tpos.iter().filter(|tpo| tpo.offset == -1).next().is_some();
        if has_missing_offs {
            // ~ if the user specified a "fallback offset" strategy
            // will try to fetch the current offset for that strategy
            // and start consuming starting at those
            if let Some(fallback_offset) = self.fallback_offset {
                // ~ now fetch the offset according to the specified strategy
                let new_offs = {
                    let mut r = try!(self.client.fetch_topic_offset(self.topic.clone(), fallback_offset));
                    match r.remove(&self.topic) {
                        None => return Err(Error::UnknownTopicOrPartition),
                        Some(pts) => pts,
                    }
                };
                // ehm ... not really fast (O(n^2))
                for tpo in tpos.iter_mut() {
                    if tpo.offset == -1 {
                        if let Some(new_off) = new_offs.iter().find(|pt| pt.partition == tpo.partition) {
                            tpo.offset = new_off.offset;
                            tpo.error = None;
                        }
                    }
                }
            } else {
                // XXX might want to produce some log message and return a dedicated error code
                return Err(Error::Unknown);
            }
        }
        Ok(())
    }

    fn make_request(&mut self) -> Result<()> {
        if ! self.initialized {
            try!(self.fetch_offsets());
        }
        let tpos = self.offsets.iter()
                        .map(|(p, o)| TopicPartitionOffset{
                                topic: self.topic.clone(),
                                partition: p.clone(),
                                offset: o.clone()
                            })
                        .collect();
        self.messages = try!(self.client.fetch_messages_multi(tpos));
        self.initialized = true;
        self.index = 0;
        Ok(())
    }
}

impl Iterator for Consumer {
    type Item = TopicMessage;

    fn next(&mut self) -> Option<TopicMessage> {
        if self.initialized {
            self.index += 1;
            self.consumed += 1;
            if self.consumed % COMMIT_INTERVAL == 0 {
                let _ = self.commit_offsets();
            }
            if self.index <= self.messages.len() {
                if self.messages[self.index-1].error.is_none() {
                    let curr = self.offsets.entry(self.messages[self.index-1].partition).or_insert(0);
                    *curr = *curr+1;
                    return Some(self.messages[self.index-1].clone());
                }
                return None;
            }
            let _ = self.commit_offsets();
            if self.messages.len() == 0 {
                return None;
            }
        }
        match self.make_request() {
            Err(_) => None,
            Ok(_) => self.next()
        }
    }
}
<|endoftext|>


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

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

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

const BLOCK_SIZE: usize = 16384;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<commit_msg>Storing the Block data to PeerConnection.downloaded<commit_after>use std::{convert, io};
use std::io::{Read, Write};
use std::net::{Ipv4Addr, TcpStream};

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

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

const BLOCK_SIZE: usize = 16384;
type Block = Vec<u8>;

struct PeerConnection {
    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 PeerConnection {
    fn connect(peer: &Peer, metainfo: &Metainfo) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let num_pieces = metainfo.info.pieces.len();
        let mut conn = PeerConnection {
            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.handshake(&metainfo.info_hash));
        loop {
            let message = try!(conn.receive_message());
            println!("{:?}", message);
            try!(conn.process(message));
        }
        Ok(conn)
    }

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

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

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

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

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

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

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

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

#[derive(Debug)]
enum Message {
    KeepAlive,
    Choke,
    Unchoke,
    Interested,
    NotInterested,
    Have(usize),
    Bitfield(Vec<u8>),
    Request, // TODO add params
    Piece(usize, usize, 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(convert_big_endian_to_integer(body)),
            5 => Message::Bitfield(body.to_owned()),
            6 => Message::Request,
            7 => {
                let index = convert_big_endian_to_integer(&body[0..4]);
                let begin = convert_big_endian_to_integer(&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)
        }
    }
}

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

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

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

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

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


<commit_before>// Defined note types for GNU systems.

#[cfg(feature = "log")]
use log::debug;
#[cfg(feature = "alloc")]
use scroll::{IOread, IOwrite, Pread, Pwrite, SizeWith};

/// ABI information.
///
/// The descriptor consists of words:
///  * word 0: OS descriptor
///  * word 1: major version of the ABI
///  * word 2: minor version of the ABI
///  * word 3: subminor version of the ABI
pub const NT_GNU_ABI_TAG: u32 = 1;

/// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
// `NT_GNU_ABI_TAG` note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

/// Synthetic `hwcap` information.
///
/// The descriptor begins with two words:
///  * word 0: number of entries
///  * word 1: bitmask of enabled entries
///
/// Then follow variable-length entries, one byte followed by a '\0'-terminated
/// `hwcap` name string.  The byte gives the bit number to test if enabled,
/// `(1U << bit) & bitmask`.
pub const NT_GNU_HWCAP: u32 = 2;

/// Build ID bits as generated by ld --build-id.
///
/// The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

/// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

///Contains copy of prstatus struct.
pub const NT_PRSTATUS: u32 = 1;

///Contains copy of prpsinfo struct.
pub const NT_PRPSINFO: u32 = 3;

///Fields of siginfo_t.
pub const NT_SIGINFO: u32 = 0x5349_4749;

///Description of mapped files.
pub const NT_FILE: u32 = 0x4649_4c45;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_alloc! {
    use crate::error;
    use crate::container;
    use scroll::ctx;
    use alloc::vec::Vec;

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, usize), Self::Error> {
            let offset = &mut 0;
            let mut alignment = alignment;
            if alignment < 4 {
                alignment = 4;
            }
            let header: NoteHeader = {
                match alignment {
                    4|8 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz.saturating_sub(1)))?;
            if header.n_namesz > 0 {
                *offset += 1;
            }
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<commit_msg>elf: Implement Plain for Nhdr{32,64} (#317)<commit_after>// Defined note types for GNU systems.

#[cfg(feature = "log")]
use log::debug;
#[cfg(feature = "alloc")]
use scroll::{IOread, IOwrite, Pread, Pwrite, SizeWith};

/// ABI information.
///
/// The descriptor consists of words:
///  * word 0: OS descriptor
///  * word 1: major version of the ABI
///  * word 2: minor version of the ABI
///  * word 3: subminor version of the ABI
pub const NT_GNU_ABI_TAG: u32 = 1;

/// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
// `NT_GNU_ABI_TAG` note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

/// Synthetic `hwcap` information.
///
/// The descriptor begins with two words:
///  * word 0: number of entries
///  * word 1: bitmask of enabled entries
///
/// Then follow variable-length entries, one byte followed by a '\0'-terminated
/// `hwcap` name string.  The byte gives the bit number to test if enabled,
/// `(1U << bit) & bitmask`.
pub const NT_GNU_HWCAP: u32 = 2;

/// Build ID bits as generated by ld --build-id.
///
/// The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

/// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

///Contains copy of prstatus struct.
pub const NT_PRSTATUS: u32 = 1;

///Contains copy of prpsinfo struct.
pub const NT_PRPSINFO: u32 = 3;

///Fields of siginfo_t.
pub const NT_SIGINFO: u32 = 0x5349_4749;

///Description of mapped files.
pub const NT_FILE: u32 = 0x4649_4c45;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

// Declare that this is a plain type.
unsafe impl plain::Plain for Nhdr32 {}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

// Declare that this is a plain type.
unsafe impl plain::Plain for Nhdr64 {}

if_alloc! {
    use crate::error;
    use crate::container;
    use scroll::ctx;
    use alloc::vec::Vec;

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, usize), Self::Error> {
            let offset = &mut 0;
            let mut alignment = alignment;
            if alignment < 4 {
                alignment = 4;
            }
            let header: NoteHeader = {
                match alignment {
                    4|8 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz.saturating_sub(1)))?;
            if header.n_namesz > 0 {
                *offset += 1;
            }
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<|endoftext|>