<commit_before>extern crate pgn_reader;
extern crate memmap;
extern crate madvise;
extern crate btoi;
extern crate rand;
extern crate reqwest;

use std::env;
use std::mem;
use std::str;
use std::cmp::min;
use std::fs::File;
use std::io::Read;
use std::time::Duration;

use memmap::Mmap;
use madvise::{AccessPattern, AdviseMemory};
use pgn_reader::{Visitor, Skip, Reader, San};
use btoi::ParseIntegerError;
use rand::{thread_rng, Rng};
use rand::distributions::OpenClosed01;

const BATCH_SIZE: usize = 50;

const MAX_PLIES: usize = 50;

#[derive(Debug)]
enum TimeControl {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

#[derive(Debug)]
struct TimeControlError;

impl From<ParseIntegerError> for TimeControlError {
    fn from(_: ParseIntegerError) -> TimeControlError {
        TimeControlError { }
    }
}

impl TimeControl {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> TimeControl {
        let total = seconds + 40 * increment;

        if total < 30 {
            TimeControl::UltraBullet
        } else if total < 180 {
            TimeControl::Bullet
        } else if total < 480 {
            TimeControl::Blitz
        } else if total < 1500 {
            TimeControl::Rapid
        } else if total < 21_600 {
            TimeControl::Classical
        } else {
            TimeControl::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<TimeControl, TimeControlError> {
        if bytes == b"-" {
            return Ok(TimeControl::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(TimeControlError)?)?;
        let increment = btoi::btou(parts.next().ok_or(TimeControlError)?)?;
        Ok(TimeControl::from_seconds_and_increment(seconds, increment))
    }
}

struct Indexer<'pgn> {
    client: reqwest::blocking::Client,

    filename: String,
    date: &'pgn [u8],

    white_elo: i16,
    black_elo: i16,
    time_control: TimeControl,
    skip: bool,

    current_game: Vec<u8>,
    plies: usize,
    standard: bool,

    batch: Vec<u8>,
    batch_size: usize,
}

impl<'pgn> Indexer<'pgn> {
    fn new(filename: &str) -> Indexer {
        Indexer {
            client: reqwest::blocking::Client::builder().timeout(Duration::from_secs(60)).build().expect("client"),

            filename: filename.into(),
            date: b"0000.00.00",

            white_elo: 0,
            black_elo: 0,
            time_control: TimeControl::Correspondence,
            skip: true,

            current_game: Vec::new(),
            plies: 0,
            standard: true,

            batch: Vec::new(),
            batch_size: 0,
        }
    }

    fn send(&mut self) {
        if self.batch_size > 0 {
            self.batch_size = 0;

            let mut res = self.client
                .put("http://127.0.0.1:9000/import/lichess")
                .body(mem::replace(&mut self.batch, Vec::new()))
                .send().expect("send batch");

            let mut answer = String::new();
            res.read_to_string(&mut answer).expect("decode response");
            println!("{} d: {} t: {}", self.filename, str::from_utf8(self.date).expect("date is valid utf8"), answer);
            assert!(res.status().is_success());
        }
    }
}

impl<'pgn> Visitor<'pgn> for Indexer<'pgn> {
    type Result = ();

    fn begin_game(&mut self) {
        self.current_game.clear();
        self.plies = 0;
    }

    fn begin_headers(&mut self) {
        self.white_elo = 0;
        self.black_elo = 0;
        self.time_control = TimeControl::Correspondence;
        self.standard = true;
    }

    fn header(&mut self, key: &'pgn [u8], value: &'pgn [u8]) {
        if key == b"WhiteElo" {
            self.white_elo = if value == b"?" { 0 } else { btoi::btoi(value).expect("WhiteElo") };
        } else if key == b"BlackElo" {
            self.black_elo = if value == b"?" { 0 } else { btoi::btoi(value).expect("BlackElo") };
        } else if key == b"TimeControl" {
            self.time_control = TimeControl::from_bytes(value).expect("TimeControl");
        } else if key == b"Variant" {
            self.standard = value == b"Standard";
            if self.standard {
                return; // we add this unconditionally later
            }
        } else if key == b"Date" || key == b"UTCDate" {
            self.date = value;
        }

        let (key, value) = if key == b"Site" {
            (&b"LichessID"[..], value.rsplitn(2, |ch| *ch == b'/').next().expect("Site"))
        } else {
            (key, value)
        };

        self.current_game.push(b'[');
        self.current_game.extend(key);
        self.current_game.extend(b" \"");
        self.current_game.extend(value);
        self.current_game.extend(b"\"]\n");
    }

    fn end_headers(&mut self) -> Skip {
        let rating = (self.white_elo + self.black_elo) / 2;

        let probability = if self.standard {
            self.current_game.extend(b"[Variant \"Standard\"]\n");

            match self.time_control {
                TimeControl::Correspondence | TimeControl::Classical => 1.0,
                TimeControl::Rapid if rating >= 2000 => 1.0,
                TimeControl::Rapid if rating >= 1800 => 2.0 / 5.0,
                TimeControl::Rapid => 1.0 / 8.0,
                TimeControl::Blitz if rating >= 2000 => 1.0,
                TimeControl::Blitz if rating >= 1800 => 1.0 / 4.0,
                TimeControl::Blitz => 1.0 / 15.0,
                TimeControl::Bullet if rating >= 2300 => 1.0,
                TimeControl::Bullet if rating >= 2200 => 4.0 / 5.0,
                TimeControl::Bullet if rating >= 2000 => 1.0 / 4.0,
                TimeControl::Bullet if rating >= 1800 => 1.0 / 7.0,
                _ => 1.0 / 20.0,
            }
        } else {
            if rating >= 1600 { 1.0 } else { 0.5 } // variant games
        };

        self.current_game.push(b'\n');

        let rnd = thread_rng().sample(OpenClosed01);
        let accept = min(self.white_elo, self.black_elo) >= 1500 && probability >= rnd;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: San) {
        if self.plies < MAX_PLIES {
            if self.plies > 0 {
                self.current_game.push(b' ');
            }

            self.current_game.extend(san.to_string().as_bytes());
            self.plies += 1;
        }
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self, _game: &'pgn [u8]) {
        if !self.skip && self.plies > 8 {
            if self.batch_size >= BATCH_SIZE {
                self.send();
            }

            if self.batch_size > 0 {
                self.batch.extend(b"\n\n\n");
            }

            self.batch.extend(&self.current_game);
            self.batch_size += 1;
        }
    }
}

fn main() {
    for arg in env::args().skip(1) {
        eprintln!("% indexing lichess games from {} ...", arg);
        let file = File::open(&arg).expect("fopen");
        if file.metadata().expect("meta").len() != 0 {
            let pgn = unsafe { Mmap::map(&file).expect("mmap") };
            pgn.advise_memory_access(AccessPattern::Sequential).expect("madvise");

            let mut indexer = Indexer::new(&arg);
            Reader::new(&mut indexer, &pgn[..]).read_all();
            indexer.send(); // send last
        }
    }
}
<commit_msg>specify body charset also in index-lichess<commit_after>extern crate pgn_reader;
extern crate memmap;
extern crate madvise;
extern crate btoi;
extern crate rand;
extern crate reqwest;

use std::env;
use std::mem;
use std::str;
use std::cmp::min;
use std::fs::File;
use std::io::Read;
use std::time::Duration;

use memmap::Mmap;
use madvise::{AccessPattern, AdviseMemory};
use pgn_reader::{Visitor, Skip, Reader, San};
use btoi::ParseIntegerError;
use rand::{thread_rng, Rng};
use rand::distributions::OpenClosed01;

const BATCH_SIZE: usize = 50;

const MAX_PLIES: usize = 50;

#[derive(Debug)]
enum TimeControl {
    UltraBullet,
    Bullet,
    Blitz,
    Rapid,
    Classical,
    Correspondence,
}

#[derive(Debug)]
struct TimeControlError;

impl From<ParseIntegerError> for TimeControlError {
    fn from(_: ParseIntegerError) -> TimeControlError {
        TimeControlError { }
    }
}

impl TimeControl {
    fn from_seconds_and_increment(seconds: u64, increment: u64) -> TimeControl {
        let total = seconds + 40 * increment;

        if total < 30 {
            TimeControl::UltraBullet
        } else if total < 180 {
            TimeControl::Bullet
        } else if total < 480 {
            TimeControl::Blitz
        } else if total < 1500 {
            TimeControl::Rapid
        } else if total < 21_600 {
            TimeControl::Classical
        } else {
            TimeControl::Correspondence
        }
    }

    fn from_bytes(bytes: &[u8]) -> Result<TimeControl, TimeControlError> {
        if bytes == b"-" {
            return Ok(TimeControl::Correspondence);
        }

        let mut parts = bytes.splitn(2, |ch| *ch == b'+');
        let seconds = btoi::btou(parts.next().ok_or(TimeControlError)?)?;
        let increment = btoi::btou(parts.next().ok_or(TimeControlError)?)?;
        Ok(TimeControl::from_seconds_and_increment(seconds, increment))
    }
}

struct Indexer<'pgn> {
    client: reqwest::blocking::Client,

    filename: String,
    date: &'pgn [u8],

    white_elo: i16,
    black_elo: i16,
    time_control: TimeControl,
    skip: bool,

    current_game: Vec<u8>,
    plies: usize,
    standard: bool,

    batch: Vec<u8>,
    batch_size: usize,
}

impl<'pgn> Indexer<'pgn> {
    fn new(filename: &str) -> Indexer {
        Indexer {
            client: reqwest::blocking::Client::builder().timeout(Duration::from_secs(60)).build().expect("client"),

            filename: filename.into(),
            date: b"0000.00.00",

            white_elo: 0,
            black_elo: 0,
            time_control: TimeControl::Correspondence,
            skip: true,

            current_game: Vec::new(),
            plies: 0,
            standard: true,

            batch: Vec::new(),
            batch_size: 0,
        }
    }

    fn send(&mut self) {
        if self.batch_size > 0 {
            self.batch_size = 0;

            let mut res = self.client
                .put("http://127.0.0.1:9000/import/lichess")
                .header("Content-Type", "application/vnd.chess-pgn;charset=utf-8")
                .body(mem::replace(&mut self.batch, Vec::new()))
                .send().expect("send batch");

            let mut answer = String::new();
            res.read_to_string(&mut answer).expect("decode response");
            println!("{} d: {} t: {}", self.filename, str::from_utf8(self.date).expect("date is valid utf8"), answer);
            assert!(res.status().is_success());
        }
    }
}

impl<'pgn> Visitor<'pgn> for Indexer<'pgn> {
    type Result = ();

    fn begin_game(&mut self) {
        self.current_game.clear();
        self.plies = 0;
    }

    fn begin_headers(&mut self) {
        self.white_elo = 0;
        self.black_elo = 0;
        self.time_control = TimeControl::Correspondence;
        self.standard = true;
    }

    fn header(&mut self, key: &'pgn [u8], value: &'pgn [u8]) {
        if key == b"WhiteElo" {
            self.white_elo = if value == b"?" { 0 } else { btoi::btoi(value).expect("WhiteElo") };
        } else if key == b"BlackElo" {
            self.black_elo = if value == b"?" { 0 } else { btoi::btoi(value).expect("BlackElo") };
        } else if key == b"TimeControl" {
            self.time_control = TimeControl::from_bytes(value).expect("TimeControl");
        } else if key == b"Variant" {
            self.standard = value == b"Standard";
            if self.standard {
                return; // we add this unconditionally later
            }
        } else if key == b"Date" || key == b"UTCDate" {
            self.date = value;
        }

        let (key, value) = if key == b"Site" {
            (&b"LichessID"[..], value.rsplitn(2, |ch| *ch == b'/').next().expect("Site"))
        } else {
            (key, value)
        };

        self.current_game.push(b'[');
        self.current_game.extend(key);
        self.current_game.extend(b" \"");
        self.current_game.extend(value);
        self.current_game.extend(b"\"]\n");
    }

    fn end_headers(&mut self) -> Skip {
        let rating = (self.white_elo + self.black_elo) / 2;

        let probability = if self.standard {
            self.current_game.extend(b"[Variant \"Standard\"]\n");

            match self.time_control {
                TimeControl::Correspondence | TimeControl::Classical => 1.0,
                TimeControl::Rapid if rating >= 2000 => 1.0,
                TimeControl::Rapid if rating >= 1800 => 2.0 / 5.0,
                TimeControl::Rapid => 1.0 / 8.0,
                TimeControl::Blitz if rating >= 2000 => 1.0,
                TimeControl::Blitz if rating >= 1800 => 1.0 / 4.0,
                TimeControl::Blitz => 1.0 / 15.0,
                TimeControl::Bullet if rating >= 2300 => 1.0,
                TimeControl::Bullet if rating >= 2200 => 4.0 / 5.0,
                TimeControl::Bullet if rating >= 2000 => 1.0 / 4.0,
                TimeControl::Bullet if rating >= 1800 => 1.0 / 7.0,
                _ => 1.0 / 20.0,
            }
        } else {
            if rating >= 1600 { 1.0 } else { 0.5 } // variant games
        };

        self.current_game.push(b'\n');

        let rnd = thread_rng().sample(OpenClosed01);
        let accept = min(self.white_elo, self.black_elo) >= 1500 && probability >= rnd;

        self.skip = !accept;
        Skip(self.skip)
    }

    fn san(&mut self, san: San) {
        if self.plies < MAX_PLIES {
            if self.plies > 0 {
                self.current_game.push(b' ');
            }

            self.current_game.extend(san.to_string().as_bytes());
            self.plies += 1;
        }
    }

    fn begin_variation(&mut self) -> Skip {
        Skip(true) // stay in the mainline
    }

    fn end_game(&mut self, _game: &'pgn [u8]) {
        if !self.skip && self.plies > 8 {
            if self.batch_size >= BATCH_SIZE {
                self.send();
            }

            if self.batch_size > 0 {
                self.batch.extend(b"\n\n\n");
            }

            self.batch.extend(&self.current_game);
            self.batch_size += 1;
        }
    }
}

fn main() {
    for arg in env::args().skip(1) {
        eprintln!("% indexing lichess games from {} ...", arg);
        let file = File::open(&arg).expect("fopen");
        if file.metadata().expect("meta").len() != 0 {
            let pgn = unsafe { Mmap::map(&file).expect("mmap") };
            pgn.advise_memory_access(AccessPattern::Sequential).expect("madvise");

            let mut indexer = Indexer::new(&arg);
            Reader::new(&mut indexer, &pgn[..]).read_all();
            indexer.send(); // send last
        }
    }
}
<|endoftext|>


<commit_before>extern crate textwrap;

use textwrap::wrap;

fn main() {
    let example = "
Memory safety without garbage collection.
Concurrency without data races.
Zero-cost abstractions.
";
    let mut prev_lines = vec![];
    let mut wrapper = Wrapper::new(0);
    for width in 15..60 {
        wrapper.width = width;
        let lines = wrapper.wrap(example);
        if lines != prev_lines {
            let title = format!(" Width: {} ", width);
            println!(".{:-^1$}.", title, width + 2);
            for line in &lines {
                println!("| {:1$} |", line, width);
            }
            prev_lines = lines;
        }
    }
}
<commit_msg>layout: use US English language corpus when wrapping<commit_after>extern crate hyphenation;
extern crate textwrap;

use hyphenation::Language;
use textwrap::Wrapper;

fn main() {
    let corpus = hyphenation::load(Language::English_US).unwrap();
    let example = "
Memory safety without garbage collection.
Concurrency without data races.
Zero-cost abstractions.
";
    let mut prev_lines = vec![];
    let mut wrapper = Wrapper::new(0);
    wrapper.corpus = Some(&corpus);
    for width in 15..60 {
        wrapper.width = width;
        let lines = wrapper.wrap(example);
        if lines != prev_lines {
            let title = format!(" Width: {} ", width);
            println!(".{:-^1$}.", title, width + 2);
            for line in &lines {
                println!("| {:1$} |", line, width);
            }
            prev_lines = lines;
        }
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]

extern crate github;

use github::Client;
use github::activity::events::*;

fn main() {
    let client = &Client::new("glendc");
    match list_events(client) {
        Ok((events, resp)) => {
            println!("listing public events succesfull, we have {} requsts remaining of {}. Limit resets @ {}...",
                resp.rate.remaining, resp.rate.limit, resp.rate.reset);
            for event in events {
                println!("event #{} at {} by {}...",
                    event.id, event.created_at, event.actor.login);
            }
        }
        Err(e) => {
            println!("list_events => {}", e);
        }
    }
}
<commit_msg>expanded public eexample<commit_after>#![deny(warnings)]

extern crate github;

use github::Client;
use github::error::*;
use github::activity::events::*;

fn main() {
    let client = &Client::new("glendc");
    match list_events(client) {
        Ok((events, resp)) => {
            println!("listing public events succesfull, we have {} requsts remaining of {}. Limit resets @ {}...",
                resp.rate.remaining, resp.rate.limit, resp.rate.reset);
            for event in events {
                println!("event #{} at {} by {}...",
                    event.id, event.created_at, event.actor.login);
            }
        }
        Err(err) => {
            println!("list_events => {}", err);
            if let ClientError::Http(http_error) = err {
                for error in http_error.errors {
                    println!("    {}", error);;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>///! Definition and description of the yet output json format to webui
///
use super::{
    super::super::ctrl::{ForwardNetMessage, NetMessages, Status},
    CollectionPathAlive, InternalUiMsg,
};
use serde_json;
use std::fmt;

/// some internal shall fail because they are not supposed to be sent out,
/// so it results in a correct MyJson or the attribute that shall
/// not be send out but was tried to do ...
pub fn convert_internal_message(internal_msg: &InternalUiMsg) -> Result<WSJsonOut, String> {
    match internal_msg {
        InternalUiMsg::Update(ref forward_net_message) => match forward_net_message.net {
            NetMessages::Debug => Err("No debug messages".to_string()),
            NetMessages::ShowNewHost => Ok(WSJsonOut::update(ViewData::host(
                forward_net_message.cnt.clone(),
            ))),
            NetMessages::ShowStats { show: stats } => Ok(WSJsonOut::nothing()),
        },
        InternalUiMsg::StartAnimate(paths_alive, status) => match paths_alive {
            CollectionPathAlive::BusyPath(nr) => Ok(WSJsonOut::searching(AnimateData::cnt(
                RefreshData::path { nr: *nr },
                match status {
                    Status::ON => true,
                    Status::OFF => false,
                },
            ))),
            CollectionPathAlive::HostSearch => Ok(WSJsonOut::searching(AnimateData::cnt(
                RefreshData::net,
                match status {
                    Status::ON => true,
                    Status::OFF => false,
                },
            ))),
        },
        InternalUiMsg::StepAndAnimate(paths_alive) => match paths_alive {
            CollectionPathAlive::BusyPath(nr) => {
                Ok(WSJsonOut::refresh(RefreshData::path { nr: *nr }))
            }
            CollectionPathAlive::HostSearch => Ok(WSJsonOut::refresh(RefreshData::net)),
        },
        InternalUiMsg::Terminate => Err("terminate is not really of interest, is it?".to_string()),
    }
}

pub fn convert_external_message(input_data: &str) -> Result<WSJsonIn, String> {
    serde_json::from_str(input_data).map_err(|_| match serde_json::from_str(input_data) {
        Ok::<serde_json::Value, serde_json::error::Error>(good_json) => format!(
            "error: this is a json but not as expected: {}",
            good_json.to_string()
        ),
        Err(_bad_json) => format!("error: this is not even a json: {}", input_data),
    })
}

pub fn generate_init_data(paths: &Vec<String>) -> WSJsonOut {
    WSJsonOut::init(InitData {
        paths: paths
            .iter()
            .enumerate()
            .map(|(e, b)| PathData {
                nr: e.clone(),
                name: b.clone(),
            })
            .collect(),
    })
}

//////////////////////////////////////////////////////////////////////
// help on that here https://serde.rs/enum-representations.html

// todo: recreate InternalUiMessages as well to use more structs rather than enums, just like here

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "type", content = "cnt")]
pub enum RefreshData {
    path { nr: usize },
    net,
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(untagged)]
pub enum AnimateData {
    cnt(RefreshData, bool),
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct PathData {
    nr: usize,
    name: String,
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct InitData {
    paths: Vec<PathData>,
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "view", content = "cnt")]
pub enum ViewData {
    host(String),
}

// This is the critical part here, "event" and "data" to work with
// "ws_events_dispatcher.js" !!!!!!!!
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "event", content = "data")]
#[allow(non_camel_case_types)]
pub enum WSJsonOut {
    refresh(RefreshData),
    searching(AnimateData),
    init(InitData),
    update(ViewData),
    nothing(),
}

impl fmt::Display for WSJsonOut {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", serde_json::to_string(&self).unwrap())
    }
}

// This is the critical part here, "event" and "data" to work with
// "ws_events_dispatcher.js" !!!!!!!!
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "event", content = "data")]
#[allow(non_camel_case_types)]
pub enum WSJsonIn {
    start,
}
<commit_msg>[clean-up] removed unnecessary import<commit_after>///! Definition and description of the yet output json format to webui
///
use super::{
    super::super::ctrl::{NetMessages, Status},
    CollectionPathAlive, InternalUiMsg,
};
use serde_json;
use std::fmt;

/// some internal shall fail because they are not supposed to be sent out,
/// so it results in a correct MyJson or the attribute that shall
/// not be send out but was tried to do ...
pub fn convert_internal_message(internal_msg: &InternalUiMsg) -> Result<WSJsonOut, String> {
    match internal_msg {
        InternalUiMsg::Update(ref forward_net_message) => match forward_net_message.net {
            NetMessages::Debug => Err("No debug messages".to_string()),
            NetMessages::ShowNewHost => Ok(WSJsonOut::update(ViewData::host(
                forward_net_message.cnt.clone(),
            ))),
            NetMessages::ShowStats { show: stats } => Ok(WSJsonOut::nothing()),
        },
        InternalUiMsg::StartAnimate(paths_alive, status) => match paths_alive {
            CollectionPathAlive::BusyPath(nr) => Ok(WSJsonOut::searching(AnimateData::cnt(
                RefreshData::path { nr: *nr },
                match status {
                    Status::ON => true,
                    Status::OFF => false,
                },
            ))),
            CollectionPathAlive::HostSearch => Ok(WSJsonOut::searching(AnimateData::cnt(
                RefreshData::net,
                match status {
                    Status::ON => true,
                    Status::OFF => false,
                },
            ))),
        },
        InternalUiMsg::StepAndAnimate(paths_alive) => match paths_alive {
            CollectionPathAlive::BusyPath(nr) => {
                Ok(WSJsonOut::refresh(RefreshData::path { nr: *nr }))
            }
            CollectionPathAlive::HostSearch => Ok(WSJsonOut::refresh(RefreshData::net)),
        },
        InternalUiMsg::Terminate => Err("terminate is not really of interest, is it?".to_string()),
    }
}

pub fn convert_external_message(input_data: &str) -> Result<WSJsonIn, String> {
    serde_json::from_str(input_data).map_err(|_| match serde_json::from_str(input_data) {
        Ok::<serde_json::Value, serde_json::error::Error>(good_json) => format!(
            "error: this is a json but not as expected: {}",
            good_json.to_string()
        ),
        Err(_bad_json) => format!("error: this is not even a json: {}", input_data),
    })
}

pub fn generate_init_data(paths: &Vec<String>) -> WSJsonOut {
    WSJsonOut::init(InitData {
        paths: paths
            .iter()
            .enumerate()
            .map(|(e, b)| PathData {
                nr: e.clone(),
                name: b.clone(),
            })
            .collect(),
    })
}

//////////////////////////////////////////////////////////////////////
// help on that here https://serde.rs/enum-representations.html

// todo: recreate InternalUiMessages as well to use more structs rather than enums, just like here

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "type", content = "cnt")]
pub enum RefreshData {
    path { nr: usize },
    net,
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(untagged)]
pub enum AnimateData {
    cnt(RefreshData, bool),
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct PathData {
    nr: usize,
    name: String,
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct InitData {
    paths: Vec<PathData>,
}

#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "view", content = "cnt")]
pub enum ViewData {
    host(String),
}

// This is the critical part here, "event" and "data" to work with
// "ws_events_dispatcher.js" !!!!!!!!
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "event", content = "data")]
#[allow(non_camel_case_types)]
pub enum WSJsonOut {
    refresh(RefreshData),
    searching(AnimateData),
    init(InitData),
    update(ViewData),
    nothing(),
}

impl fmt::Display for WSJsonOut {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", serde_json::to_string(&self).unwrap())
    }
}

// This is the critical part here, "event" and "data" to work with
// "ws_events_dispatcher.js" !!!!!!!!
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(tag = "event", content = "data")]
#[allow(non_camel_case_types)]
pub enum WSJsonIn {
    start,
}
<|endoftext|>


<commit_before>use futures::channel::mpsc;
use futures::executor::block_on;
use futures::future::join3;
use futures::sink::SinkExt;
use futures::stream::{self, StreamExt};

#[test]
fn it_works() {
    let (tx1, rx1) = mpsc::channel(1);
    let (tx2, rx2) = mpsc::channel(2);
    let tx = tx1.fanout(tx2).sink_map_err(|_| ());

    let src = stream::iter((0..10).map(Ok));
    let fwd = src.forward(tx);

    let collect_fut1 = rx1.collect::<Vec<_>>();
    let collect_fut2 = rx2.collect::<Vec<_>>();
    let (_, vec1, vec2) = block_on(join3(fwd, collect_fut1, collect_fut2));

    let expected = (0..10).collect::<Vec<_>>();

    assert_eq!(vec1, expected);
    assert_eq!(vec2, expected);
}
<commit_msg>futures: tests/sink_fanout.rs: Don't break w/o features<commit_after>#[cfg(all(feature = "alloc", feature="std", feature="executor"))] // channel::mpsc, executor::
#[test]
fn it_works() {
    use futures::channel::mpsc;
    use futures::executor::block_on;
    use futures::future::join3;
    use futures::sink::SinkExt;
    use futures::stream::{self, StreamExt};

    let (tx1, rx1) = mpsc::channel(1);
    let (tx2, rx2) = mpsc::channel(2);
    let tx = tx1.fanout(tx2).sink_map_err(|_| ());

    let src = stream::iter((0..10).map(Ok));
    let fwd = src.forward(tx);

    let collect_fut1 = rx1.collect::<Vec<_>>();
    let collect_fut2 = rx2.collect::<Vec<_>>();
    let (_, vec1, vec2) = block_on(join3(fwd, collect_fut1, collect_fut2));

    let expected = (0..10).collect::<Vec<_>>();

    assert_eq!(vec1, expected);
    assert_eq!(vec2, expected);
}
<|endoftext|>


<commit_before>use ws;
use time;
use std::fmt;
use uuid::Uuid;
use std::sync::{Arc,Mutex};
use std::collections::{HashMap,HashSet};

#[derive(PartialEq, Eq, Hash, Debug, Clone)]
pub enum Item {
    Potatoes,
    Berries,
    TreeSap,
}

#[derive(Debug, Clone)]
pub struct PlayerState {
    inventory: HashSet<Item>,
}

#[derive(Debug, Clone)]
pub enum Player {
    AnonymousPlayer {
        state : PlayerState,
    },
    RegisteredPlayer {
        id    : i32,
        name  : String,
        state : PlayerState,
    }
}

mod AnonymousPlayer {
    use super::*;
    pub fn new () -> Player {
        let mut new_player_inventory = HashSet::new();
        new_player_inventory.insert(Item::Potatoes);

        Player::AnonymousPlayer {
            state: PlayerState {
                inventory: new_player_inventory
            }
        }
    }
}

mod RegisteredPlayer {
    use super::*;
    pub fn new (id: i32, name: String) -> Player {
        let mut new_player_inventory = HashSet::new();
        new_player_inventory.insert(Item::Potatoes);

        Player::RegisteredPlayer {
            id: id,
            name: name,
            state: PlayerState {
                inventory: new_player_inventory
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct Connection {
    uuid   : Uuid,
    player : Player,
}

impl Connection {
    fn new(u: Uuid, p: Player) -> Connection {
        Connection {
            uuid  : u,
            player: p
        }
    }
}

impl fmt::Display for Connection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "conn{{{}}}", self.uuid)
    }
}

fn parse_cookies (req: &ws::Request) -> HashMap<String, String> {
    req.header("cookie")
        .and_then(|cookies_bytes| String::from_utf8(cookies_bytes.to_vec()).ok())
        .unwrap_or(String::from(""))
        .rsplit(";")
        .filter_map(|cookie_string| {
            let mut cookie_pair = cookie_string.split("=");
            match (cookie_pair.next(), cookie_pair.next()) {
                (Some(name), Some(value)) => {
                    Some((String::from(name.trim()), String::from(value.trim())))
                },
                _ => None
            }
        })
        .collect()
}

fn put_cookie (name: String, value: String, resp: &mut ws::Response) {
    let headers = resp.headers_mut();
    let cookie_bytes = format!("{}={}", name, value).as_bytes().to_vec();
    headers.push((String::from("Set-Cookie"), cookie_bytes));
}

type ConnectionMap = Arc<Mutex<HashMap<Uuid, Connection>>>;

struct WSServer {
    out: ws::Sender,
    connections: ConnectionMap,
}

fn parse_message (msg: ws::Message) -> Option<(String, String)> {
    msg.into_text()
        .ok()
        .and_then(|message_blob: String| {
            let mut message_parts = message_blob.split(":");
            match (message_parts.next(), message_parts.next()) {
                (Some(ticket), Some(msg_type)) => Some((ticket.to_owned(), msg_type.to_owned())),
                _ => None
            }
        })
}

impl ws::Handler for WSServer {
    fn on_request(&mut self, req: &ws::Request) -> ws::Result<ws::Response> {
        let mut resp = ws::Response::from_request(req).unwrap();

        let mut conn_map = self.connections.lock().unwrap();

        println!("ws:req[{}]", time::precise_time_ns());

        let cookies : HashMap<String, String> = parse_cookies(req);

        let mut cookie_existed = true;
        let ticket = cookies.get("bzwf_anon_wstx")
            .and_then(|uuid_string| Uuid::parse_str(uuid_string.as_str()).ok())
            .unwrap_or_else(|| {
                println!("ws:req[{}]: no bzwf_anon_wstx cookie found", time::precise_time_ns());
                cookie_existed = false;
                Uuid::new_v4()
            });

        let users_count = conn_map.len();

        if conn_map.contains_key(&ticket) {
            println!("ws:req[{}]: reconnect: uuid {}", time::precise_time_ns(), ticket);
            println!("{} connected users", users_count);
        } else {
            println!("ws:req[{}]: new connection with uuid {}", time::precise_time_ns(), ticket);

            let cookie_op_str = if cookie_existed {
                "replacing persistence cookie"
            } else {
                "creating persistence cookie"
            };

            println!("ws:req[{}]: {}: bzwf_anon_wstx={}", time::precise_time_ns(), cookie_op_str, ticket);
            println!("{} connected users", users_count + 1);
            put_cookie(String::from("bzwf_anon_wstx"), ticket.to_string(), &mut resp);
            conn_map.insert(ticket, Connection::new(ticket, AnonymousPlayer::new()));
        }

        Ok(resp)
    }

    fn on_message(&mut self, msg: ws::Message) -> ws::Result<()> {
        println!("ws:rcv[{}]: user with uuid {} sent ws msg {}",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]",
                 msg);
        match parse_message(msg) {
            // TODO
            _ => self.out.send("I hear you loud and clear").unwrap()
        };
        Ok(())
    }

    fn on_open(&mut self, _: ws::Handshake) -> ws::Result<()> {
        // TODO
        println!("ws:opn[{}]: user with uuid {} opened ws cxn",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]");
        Ok(())
    }

    fn on_close(&mut self, code: ws::CloseCode, reason: &str) {
        println!("ws:cls[{}]: user with uuid {} closed ws cxn\n\tCode [{:?}] reason: {}",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]",
                 code,
                 reason);
        // TODO
    }

    fn on_error(&mut self, err: ws::Error) {
        println!("ws:err[{}]: user with uuid {} got error {}",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]",
                 err);
        // TODO
    }
}

pub fn server (domain: String, port: i32) -> () {
    let connections : ConnectionMap = Arc::new(Mutex::new(HashMap::new()));
    let addr = format!("{}:{}", domain, port);
    println!("WebSockets server listening on {}", addr);
    ws::listen(addr, |out| WSServer {
        out: out,
        connections: connections.clone(),
    }).unwrap();
}
<commit_msg>WIP far more advanced message parsing<commit_after>use ws;
use time;
use std::fmt;
use uuid::Uuid;
use std::sync::{Arc,Mutex};
use std::collections::{HashMap,HashSet};

#[derive(PartialEq, Eq, Hash, Debug, Clone)]
pub enum Item {
    Potatoes,
    Berries,
    TreeSap,
}

trait ParseFrom<T, Out = Self> {
    fn parse (from: T) -> Result<Out, String>;
}

impl ParseFrom<i32> for Item {
    fn parse (i: i32) -> Result<Self, String> {
        match i {
            0 => Ok(Item::Potatoes),
            1 => Ok(Item::Berries),
            2 => Ok(Item::TreeSap),
            _ => Err(format!("Item.parse<i32> failure: unrecognized Item {}", i)),
        }
    }
}

#[derive(Debug)]
pub enum Action {
    addItemToInventory(Item),
    default,
}

#[derive(Debug, Clone)]
pub struct PlayerState {
    inventory: HashSet<Item>,
}

#[derive(Debug, Clone)]
pub enum Player {
    AnonymousPlayer {
        state : PlayerState,
    },
    RegisteredPlayer {
        id    : i32,
        name  : String,
        state : PlayerState,
    }
}

mod AnonymousPlayer {
    use super::*;
    pub fn new () -> Player {
        let mut new_player_inventory = HashSet::new();
        new_player_inventory.insert(Item::Potatoes);

        Player::AnonymousPlayer {
            state: PlayerState {
                inventory: new_player_inventory
            }
        }
    }
}

mod RegisteredPlayer {
    use super::*;
    pub fn new (id: i32, name: String) -> Player {
        let mut new_player_inventory = HashSet::new();
        new_player_inventory.insert(Item::Potatoes);

        Player::RegisteredPlayer {
            id: id,
            name: name,
            state: PlayerState {
                inventory: new_player_inventory
            }
        }
    }
}

#[derive(Debug, Clone)]
pub struct Connection {
    uuid   : Uuid,
    player : Player,
}

impl Connection {
    fn new(u: Uuid, p: Player) -> Connection {
        Connection {
            uuid  : u,
            player: p
        }
    }
}

impl fmt::Display for Connection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "conn{{{}}}", self.uuid)
    }
}

fn parse_cookies (req: &ws::Request) -> HashMap<String, String> {
    req.header("cookie")
        .and_then(|cookies_bytes| String::from_utf8(cookies_bytes.to_vec()).ok())
        .unwrap_or(String::from(""))
        .rsplit(";")
        .filter_map(|cookie_string| {
            let mut cookie_pair = cookie_string.split("=");
            match (cookie_pair.next(), cookie_pair.next()) {
                (Some(name), Some(value)) => {
                    Some((String::from(name.trim()), String::from(value.trim())))
                },
                _ => None
            }
        })
        .collect()
}

fn put_cookie (name: String, value: String, resp: &mut ws::Response) {
    let headers = resp.headers_mut();
    let cookie_bytes = format!("{}={}", name, value).as_bytes().to_vec();
    headers.push((String::from("Set-Cookie"), cookie_bytes));
}

fn authorize_ticket_dumb (ticket: Uuid, msg_type: Vec<String>, conn_map: ConnectionMap) -> Result<Vec<String>, String> {
   if conn_map.lock().unwrap().contains_key(&ticket) {
        Ok(msg_type)
   } else {
        Err(String::from("authorize_ticket_dumb authorization failed"))
   }
}

type ConnectionMap = Arc<Mutex<HashMap<Uuid, Connection>>>;

struct WSServer {
    out: ws::Sender,
    connections: ConnectionMap,
}

fn parse_message (msg: ws::Message) -> Result<(Uuid, Vec<String>), String> {
    fn parse_error_msg (reason: &str) -> String {
        format!("parse_message failure: {}", reason)
    }

    msg.into_text()
        .or(Err(parse_error_msg("cannot get text from Message")))
        .and_then(|message_blob: String| {
            let mut message_parts = message_blob.split(":");

            let msg_ticket   = message_parts.next();
            let msg_contents = message_parts.map(|p| p.to_string()).collect();

            msg_ticket
                .ok_or(parse_error_msg("no ticket"))
                .and_then(|uuid_string| {
                    Uuid::parse_str(uuid_string)
                        .or(Err(parse_error_msg("ticket is invalid uuidv4")))
                })
                .map(|uuid| (uuid, msg_contents))
        })
}

impl ws::Handler for WSServer {
    fn on_request(&mut self, req: &ws::Request) -> ws::Result<ws::Response> {
        let mut resp = ws::Response::from_request(req).unwrap();

        let mut conn_map = self.connections.lock().unwrap();

        println!("ws:req[{}]", time::precise_time_ns());

        let cookies : HashMap<String, String> = parse_cookies(req);

        let mut cookie_existed = true;
        let ticket = cookies.get("bzwf_anon_wstx")
            .and_then(|uuid_string| Uuid::parse_str(uuid_string.as_str()).ok())
            .unwrap_or_else(|| {
                println!("ws:req[{}]: no bzwf_anon_wstx cookie found", time::precise_time_ns());
                cookie_existed = false;
                Uuid::new_v4()
            });

        let users_count = conn_map.len();

        if conn_map.contains_key(&ticket) {
            println!("ws:req[{}]: reconnect: uuid {}", time::precise_time_ns(), ticket);
            println!("{} connected users", users_count);
        } else {
            println!("ws:req[{}]: new connection with uuid {}", time::precise_time_ns(), ticket);

            let cookie_op_str = if cookie_existed {
                "replacing persistence cookie"
            } else {
                "creating persistence cookie"
            };

            println!("ws:req[{}]: {}: bzwf_anon_wstx={}", time::precise_time_ns(), cookie_op_str, ticket);
            println!("{} connected users", users_count + 1);
            put_cookie(String::from("bzwf_anon_wstx"), ticket.to_string(), &mut resp);
            conn_map.insert(ticket, Connection::new(ticket, AnonymousPlayer::new()));
        }

        Ok(resp)
    }

    fn on_message(&mut self, msg: ws::Message) -> ws::Result<()> {
        println!("ws:rcv[{}]: user with uuid {} sent ws msg {}",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]",
                 msg);

        parse_message(msg)
            .and_then(|(msg_ticket, msg_contents)| {
                println!("authorizing ticket {}", msg_ticket);
                authorize_ticket_dumb(msg_ticket, msg_contents, self.connections.clone())
            })
            .and_then(|msg_contents| {
                // FIXME possible out of bounds error
                let ref msg_type  = msg_contents[0];
                let ref msg_param = msg_contents[1];

                println!("trying to read message contents of {}({})", msg_type, msg_param);

                match (msg_type.as_str(), msg_param) {
                    ("addItemToInventory", item) => {
                        if let Ok(item_num) = item.parse::<i32>() {
                            Item::parse(item_num)
                                .map(|item| Action::addItemToInventory(item))
                        } else {
                            Err(String::from("parse_message failure: addItemToInventory: invalid item code"))
                        }
                    },
                    _ => Err(String::from("parse_message failure: unrecognized message type"))
                }
            })
            .map(|action| {
                self.out.send(format!("gotcha, you want to {:?}", action))
            })
            .unwrap_or_else(|err| {
                println!("{}", err);
                self.out.send("got your message, but not sure what it meant")
            });

        Ok(())
    }

    fn on_open(&mut self, _: ws::Handshake) -> ws::Result<()> {
        // TODO
        println!("ws:opn[{}]: user with uuid {} opened ws cxn",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]");
        Ok(())
    }

    fn on_close(&mut self, code: ws::CloseCode, reason: &str) {
        println!("ws:cls[{}]: user with uuid {} closed ws cxn\n\tCode [{:?}] reason: {}",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]",
                 code,
                 reason);
        // TODO
    }

    fn on_error(&mut self, err: ws::Error) {
        println!("ws:err[{}]: user with uuid {} got error {}",
                 time::precise_time_ns(),
                 "[[[ we don't have uuids in ws yet ]]]",
                 err);
        // TODO
    }
}

pub fn server (domain: String, port: i32) -> () {
    let connections : ConnectionMap = Arc::new(Mutex::new(HashMap::new()));
    let addr = format!("{}:{}", domain, port);
    println!("WebSockets server listening on {}", addr);
    ws::listen(addr, |out| WSServer {
        out: out,
        connections: connections.clone(),
    }).unwrap();
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#[macro_use]
extern crate log;
extern crate docopt;
extern crate quiche;
extern crate env_logger;

use std::fs;
use std::net;
use std::path;
use std::collections::hash_map;
use std::collections::HashMap;

use docopt::Docopt;

use quiche::packet;
use quiche::rand;

const LOCAL_CONN_ID_LEN: usize = 16;

const TRANSPORT_PARAMS: quiche::TransportParams = quiche::TransportParams {
    idle_timeout: 30,
    initial_max_data: 10_000_000,
    initial_max_bidi_streams: 100,
    initial_max_uni_streams: 100,
    max_packet_size: 1500,
    ack_delay_exponent: 3,
    disable_migration: true,
    max_ack_delay: 25,
    initial_max_stream_data_bidi_local: 1_000_000,
    initial_max_stream_data_bidi_remote: 1_000_000,
    initial_max_stream_data_uni: 1_000_000,
    stateless_reset_token_present: true,
    stateless_reset_token: [0xba; 16],
};

const USAGE: &str = "Usage: server [options]

Options:
  -h --help         Show this screen.
  --listen <addr>   Listen on the given IP:port [default: 127.0.0.1:4433]
  --cert <file>     TLS certificate path [default: examples/cert.crt]
  --key <file>      TLS certificate key path [default: examples/cert.key]
  --root <dir>      Root directory [default: examples/root/]
  --name <str>      Name of the server [default: quic.tech]
";

fn main() {
    let mut buf = [0; TRANSPORT_PARAMS.max_packet_size as usize];
    let mut out = [0; TRANSPORT_PARAMS.max_packet_size as usize];

    env_logger::init();

    let args = Docopt::new(USAGE)
                      .and_then(|dopt| dopt.parse())
                      .unwrap_or_else(|e| e.exit());

    let socket = net::UdpSocket::bind(args.get_str("--listen")).unwrap();

    let mut connections: HashMap<net::SocketAddr, Box<quiche::Conn>> = HashMap::new();

    loop {
        // Garbage collect closed connections.
        connections.retain(|_, ref mut c| {
            debug!("Collecting garbage");

            if c.is_closed() {
                debug!("{} connection collected", c.trace_id());
            }

            !c.is_closed()
        });

        let (len, src) = socket.recv_from(&mut buf).unwrap();
        debug!("Got {} bytes from {}", len, src);

        let buf = &mut buf[..len];

        let hdr = if packet::has_long_header(buf[0]) {
            packet::Header::decode_long(buf)
        } else {
            packet::Header::decode_short(buf, LOCAL_CONN_ID_LEN)
        };

        let hdr = match hdr {
            Ok(v) => v,

            Err(e) => {
                error!("Parsing packet header failed: {:?}", e);
                continue
            }
        };

        if hdr.ty == packet::Type::VersionNegotiation {
            error!("Version negotiation invalid on the server");
            continue;
        }

        let conn = match connections.entry(src) {
            hash_map::Entry::Vacant(v) =>{
                if hdr.version != quiche::VERSION_DRAFT15 {
                    warn!("Doing version negotiation");

                    let len = packet::negotiate_version(&hdr, &mut out).unwrap();
                    let out = &out[..len];

                    socket.send_to(out, &src).unwrap();
                    continue;
                }

                if hdr.ty != packet::Type::Initial {
                    error!("Packet is not Initial");
                    continue;
                }

                let mut scid: [u8; LOCAL_CONN_ID_LEN] = [0; LOCAL_CONN_ID_LEN];
                rand::rand_bytes(&mut scid[..]);

                let config = quiche::Config {
                    version: quiche::VERSION_DRAFT15,

                    local_conn_id: &scid,

                    local_transport_params: &TRANSPORT_PARAMS,

                    tls_server_name: args.get_str("--name"),
                    tls_certificate: args.get_str("--cert"),
                    tls_certificate_key: args.get_str("--key"),
                };

                debug!("New connection: dcid={} scid={} lcid={}",
                       hex_dump(&hdr.dcid),
                       hex_dump(&hdr.scid),
                       hex_dump(&scid));

                let conn = quiche::Conn::new(config, true).unwrap();

                v.insert(conn)
            },

            hash_map::Entry::Occupied(v) => v.into_mut(),
        };

        let mut left = len;

        // Process potentially coalesced packets.
        while left > 0 {
            let read = match conn.recv(&mut buf[len - left..len]) {
                Ok(v)  => v,
                Err(e) => panic!("{} recv failed: {:?}", conn.trace_id(), e),
            };

            left -= read;

            debug!("{} read {} bytes", conn.trace_id(), read);
        }

        let streams: Vec<u64> = conn.stream_iter().collect();
        for s in streams {
            info!("{} stream {} is readable", conn.trace_id(), s);
            handle_stream(conn, s, &args);
        }

        let mut written = 0;
        while written < buf.len() * 10 {
            let write = match conn.send(&mut out) {
                Ok(v) => v,

                Err(quiche::Error::NothingToDo) => {
                    debug!("{} done writing", conn.trace_id());
                    break;
                },

                Err(e) => panic!("{} socket send failed: {:?}",
                                 conn.trace_id(), e),
            };

            // TODO: coalesce packets.
            socket.send_to(&out[..write], &src).unwrap();
            written += write;

            debug!("{} written {} bytes", conn.trace_id(), write);
        }
    }
}

fn handle_stream(conn: &mut quiche::Conn, stream: u64, args: &docopt::ArgvMap) {
    let stream_data = match conn.stream_recv(stream) {
        Ok(v) => v,
        Err(e) => panic!("{} stream recv failed {:?}",
                         conn.trace_id(), e),
    };

    if &stream_data[..4] == b"GET " {
        let uri = &stream_data[4..stream_data.len()];
        let uri = String::from_utf8(uri.to_vec()).unwrap();
        let uri = String::from(uri.lines().next().unwrap());
        let uri = path::Path::new(&uri);
        let mut path = path::PathBuf::from(args.get_str("--root"));

        for c in uri.components() {
            if let path::Component::Normal(v) = c {
                path.push(v)
            }
        }

        info!("{} got GET request for {:?} on stream {}",
              conn.trace_id(), path, stream);

        let data = fs::read(path.as_path())
                      .unwrap_or_else(|_| Vec::from(String::from("Not Found!")));

        info!("{} sending response of size {} on stream {}",
              conn.trace_id(), data.len(), stream);

        if let Err(e) = conn.stream_send(stream, &data, true) {
            panic!("{} stream send failed {:?}", conn.trace_id(), e);
        }
    }
}

fn hex_dump(buf: &[u8]) -> String {
    let vec: Vec<String> = buf.iter()
                              .map(|b| format!("{:02x}", b))
                              .collect();

    vec.join("")
}
<commit_msg>server: remove static flight size limit<commit_after>// Copyright (c) 2018, Alessandro Ghedini
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#[macro_use]
extern crate log;
extern crate docopt;
extern crate quiche;
extern crate env_logger;

use std::fs;
use std::net;
use std::path;
use std::collections::hash_map;
use std::collections::HashMap;

use docopt::Docopt;

use quiche::packet;
use quiche::rand;

const LOCAL_CONN_ID_LEN: usize = 16;

const TRANSPORT_PARAMS: quiche::TransportParams = quiche::TransportParams {
    idle_timeout: 30,
    initial_max_data: 10_000_000,
    initial_max_bidi_streams: 100,
    initial_max_uni_streams: 100,
    max_packet_size: 1500,
    ack_delay_exponent: 3,
    disable_migration: true,
    max_ack_delay: 25,
    initial_max_stream_data_bidi_local: 1_000_000,
    initial_max_stream_data_bidi_remote: 1_000_000,
    initial_max_stream_data_uni: 1_000_000,
    stateless_reset_token_present: true,
    stateless_reset_token: [0xba; 16],
};

const USAGE: &str = "Usage: server [options]

Options:
  -h --help         Show this screen.
  --listen <addr>   Listen on the given IP:port [default: 127.0.0.1:4433]
  --cert <file>     TLS certificate path [default: examples/cert.crt]
  --key <file>      TLS certificate key path [default: examples/cert.key]
  --root <dir>      Root directory [default: examples/root/]
  --name <str>      Name of the server [default: quic.tech]
";

fn main() {
    let mut buf = [0; TRANSPORT_PARAMS.max_packet_size as usize];
    let mut out = [0; TRANSPORT_PARAMS.max_packet_size as usize];

    env_logger::init();

    let args = Docopt::new(USAGE)
                      .and_then(|dopt| dopt.parse())
                      .unwrap_or_else(|e| e.exit());

    let socket = net::UdpSocket::bind(args.get_str("--listen")).unwrap();

    let mut connections: HashMap<net::SocketAddr, Box<quiche::Conn>> = HashMap::new();

    loop {
        // Garbage collect closed connections.
        connections.retain(|_, ref mut c| {
            debug!("Collecting garbage");

            if c.is_closed() {
                debug!("{} connection collected", c.trace_id());
            }

            !c.is_closed()
        });

        let (len, src) = socket.recv_from(&mut buf).unwrap();
        debug!("Got {} bytes from {}", len, src);

        let buf = &mut buf[..len];

        let hdr = if packet::has_long_header(buf[0]) {
            packet::Header::decode_long(buf)
        } else {
            packet::Header::decode_short(buf, LOCAL_CONN_ID_LEN)
        };

        let hdr = match hdr {
            Ok(v) => v,

            Err(e) => {
                error!("Parsing packet header failed: {:?}", e);
                continue
            }
        };

        if hdr.ty == packet::Type::VersionNegotiation {
            error!("Version negotiation invalid on the server");
            continue;
        }

        let conn = match connections.entry(src) {
            hash_map::Entry::Vacant(v) =>{
                if hdr.version != quiche::VERSION_DRAFT15 {
                    warn!("Doing version negotiation");

                    let len = packet::negotiate_version(&hdr, &mut out).unwrap();
                    let out = &out[..len];

                    socket.send_to(out, &src).unwrap();
                    continue;
                }

                if hdr.ty != packet::Type::Initial {
                    error!("Packet is not Initial");
                    continue;
                }

                let mut scid: [u8; LOCAL_CONN_ID_LEN] = [0; LOCAL_CONN_ID_LEN];
                rand::rand_bytes(&mut scid[..]);

                let config = quiche::Config {
                    version: quiche::VERSION_DRAFT15,

                    local_conn_id: &scid,

                    local_transport_params: &TRANSPORT_PARAMS,

                    tls_server_name: args.get_str("--name"),
                    tls_certificate: args.get_str("--cert"),
                    tls_certificate_key: args.get_str("--key"),
                };

                debug!("New connection: dcid={} scid={} lcid={}",
                       hex_dump(&hdr.dcid),
                       hex_dump(&hdr.scid),
                       hex_dump(&scid));

                let conn = quiche::Conn::new(config, true).unwrap();

                v.insert(conn)
            },

            hash_map::Entry::Occupied(v) => v.into_mut(),
        };

        let mut left = len;

        // Process potentially coalesced packets.
        while left > 0 {
            let read = match conn.recv(&mut buf[len - left..len]) {
                Ok(v)  => v,
                Err(e) => panic!("{} recv failed: {:?}", conn.trace_id(), e),
            };

            left -= read;

            debug!("{} read {} bytes", conn.trace_id(), read);
        }

        let streams: Vec<u64> = conn.stream_iter().collect();
        for s in streams {
            info!("{} stream {} is readable", conn.trace_id(), s);
            handle_stream(conn, s, &args);
        }

        loop {
            let write = match conn.send(&mut out) {
                Ok(v) => v,

                Err(quiche::Error::NothingToDo) => {
                    debug!("{} done writing", conn.trace_id());
                    break;
                },

                Err(e) => panic!("{} socket send failed: {:?}",
                                 conn.trace_id(), e),
            };

            // TODO: coalesce packets.
            socket.send_to(&out[..write], &src).unwrap();

            debug!("{} written {} bytes", conn.trace_id(), write);
        }
    }
}

fn handle_stream(conn: &mut quiche::Conn, stream: u64, args: &docopt::ArgvMap) {
    let stream_data = match conn.stream_recv(stream) {
        Ok(v) => v,
        Err(e) => panic!("{} stream recv failed {:?}",
                         conn.trace_id(), e),
    };

    if &stream_data[..4] == b"GET " {
        let uri = &stream_data[4..stream_data.len()];
        let uri = String::from_utf8(uri.to_vec()).unwrap();
        let uri = String::from(uri.lines().next().unwrap());
        let uri = path::Path::new(&uri);
        let mut path = path::PathBuf::from(args.get_str("--root"));

        for c in uri.components() {
            if let path::Component::Normal(v) = c {
                path.push(v)
            }
        }

        info!("{} got GET request for {:?} on stream {}",
              conn.trace_id(), path, stream);

        let data = fs::read(path.as_path())
                      .unwrap_or_else(|_| Vec::from(String::from("Not Found!")));

        info!("{} sending response of size {} on stream {}",
              conn.trace_id(), data.len(), stream);

        if let Err(e) = conn.stream_send(stream, &data, true) {
            panic!("{} stream send failed {:?}", conn.trace_id(), e);
        }
    }
}

fn hex_dump(buf: &[u8]) -> String {
    let vec: Vec<String> = buf.iter()
                              .map(|b| format!("{:02x}", b))
                              .collect();

    vec.join("")
}
<|endoftext|>


<commit_before>extern crate rust_rocket;

use rust_rocket::{Rocket, Event};

fn main() {
    let mut rocket = Rocket::new().unwrap();
    rocket.get_track("test");
    rocket.get_track("test2");
    rocket.get_track("a:test2");

    let mut current_row = 0;
    loop {
        if let Some(event) = rocket.poll_events() {
            match event {
                Event::SetRow(row) => {
                    println!("SetRow (row: {:?})", row);
                    current_row = row;
                },
                Event::Pause(_) => {
                    {
                        let track1 = rocket.get_track("test");
                        println!("Pause (value: {:?}) (row: {:?})", track1.get_value(current_row as f32), current_row);
                    }
                },
                _ => (),
            }
            println!("{:?}", event);
        }
        std::thread::sleep_ms(1);
    }
}
<commit_msg>Make example work again<commit_after>extern crate rust_rocket;

use rust_rocket::{Rocket, Event};

fn main() {
    let mut rocket = Rocket::new().unwrap();
    rocket.get_track_mut("test");
    rocket.get_track_mut("test2");
    rocket.get_track_mut("a:test2");

    let mut current_row = 0;
    loop {
        if let Some(event) = rocket.poll_events() {
            match event {
                Event::SetRow(row) => {
                    println!("SetRow (row: {:?})", row);
                    current_row = row;
                },
                Event::Pause(_) => {
                    {
                        let track1 = rocket.get_track("test");
                        println!("Pause (value: {:?}) (row: {:?})", track1.get_value(current_row as f32), current_row);
                    }
                },
                _ => (),
            }
            println!("{:?}", event);
        }
        std::thread::sleep_ms(1);
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Jonathan Eyolfson

#![feature(asm)]

extern crate wayland;

extern crate libc;

use std::ptr;

static SYSCALL_MEMFD_CREATE: u32 = 319;
static MFD_CLOEXEC: u32 = 1;
static MFD_ALLOW_SEALING: u32 = 2;

// const WIDTH: i32 = 300;
// const HEIGHT: i32 = 200;
// const STRIDE: i32 = WIDTH * 4;

struct ShmBuffer {
    fd: i32,
    ptr: *mut libc::c_void,
    width: i32,
    height: i32,
    capacity: uint,
}

impl ShmBuffer {
    pub fn new() -> ShmBuffer {
        unsafe {
            let mut fd: i32;
            let name = b"rust-wayland-shm\x00";
            asm!("syscall"
                : "={rax}"(fd)
                 : "{rax}"(SYSCALL_MEMFD_CREATE),
                   "{rdi}"(name.as_ptr()),
                   "{rsi}"(MFD_CLOEXEC | MFD_ALLOW_SEALING)
                : "rcx", "r11", "memory"
                : "volatile"
            );
            assert!(fd >= 0);
            let width = 300i32;
            let height = 200i32;
            let capacity = width as uint
                * height as uint
                * std::mem::size_of::<u32>();
            assert!(libc::ftruncate(fd, capacity as i64) != -1);
            let ptr = libc::mmap(ptr::null_mut(), capacity as u64,
                                 libc::PROT_WRITE | libc::PROT_READ, 1, fd, 0);
            assert!(ptr != libc::MAP_FAILED);
            for i in range(0, width * height) {
                let p: *mut u32 = (ptr as *mut u32).offset(i as int);
                let x = i % width;
                let y = i / width;
                match x {
                    0...5 => std::ptr::write(&mut *p, 0x7FFF0000),
                    294...299 => std::ptr::write(&mut *p, 0x7F0000FF),
                    _ => match y {
                        0...5 => std::ptr::write(&mut *p, 0x7F00FF00),
                        194...199 => std::ptr::write(&mut *p, 0x7FFF00FF),
                        _ =>  std::ptr::write(&mut *p, 0x7F000000),
                    }
                }
            }
            ShmBuffer {
                fd: fd,
                ptr: ptr,
                width: width,
                height: height,
                capacity: capacity,
            }
        }
    }
    pub fn resize(&mut self, width: i32, height: i32) {
        if (width * height) > (self.width * self.height) {
            println!("TODO: actual");
        }
        self.width = width;
        self.height = height;
    }
    pub fn fd(&self) -> i32 {
        self.fd
    }
    pub fn capacity(&self) -> uint { self.capacity }
    pub fn width(&self) -> i32 { self.width }
    pub fn height(&self) -> i32 { self.height }
    pub fn stride(&self) -> i32 {
        self.width * std::mem::size_of::<u32>() as i32
    }
}

impl Drop for ShmBuffer {
    fn drop(&mut self) {
        unsafe {
            libc::munmap(self.ptr, self.capacity as u64);
            libc::close(self.fd);
        }
    }
}

fn main() {
    let mut display = wayland::Display::connect_to_env_or_default();
    let mut registry = wayland::Registry::new(&mut display);
    let mut shm_buffer = ShmBuffer::new();
    shm_buffer.resize(300, 200);
    let mut pool = registry.shm().create_pool(shm_buffer.fd(),
                                              shm_buffer.capacity() as i32);
    let mut surface = registry.compositor().create_surface();
    let mut buffer = pool.create_buffer(
        0, shm_buffer.width(), shm_buffer.height(), shm_buffer.stride(),
        wayland::raw::WL_SHM_FORMAT_ARGB8888
    );
    let mut shell_surface = registry.shell().get_shell_surface(&mut surface);
    shell_surface.set_toplevel();
    surface.attach(&mut buffer, 0, 0);
    surface.commit();
    loop {
        display.dispatch();
    }
}
<commit_msg>Refactor memfd_create<commit_after>// Copyright 2014 Jonathan Eyolfson

#![feature(asm)]

extern crate wayland;

extern crate libc;

use std::ptr;

static SYSCALL_MEMFD_CREATE: u32 = 319;

static MFD_CLOEXEC: u32 = 1;
static MFD_ALLOW_SEALING: u32 = 2;

unsafe fn memfd_create(name: *const u8, flags: u32) -> i32 {
    let mut fd: i32;
    asm!("syscall"
        : "={rax}"(fd)
        : "{rax}"(SYSCALL_MEMFD_CREATE), "{rdi}"(name), "{rsi}"(flags)
        : "rcx", "r11", "memory"
        : "volatile"
    );
    fd
}

struct ShmBuffer {
    fd: i32,
    ptr: *mut libc::c_void,
    width: i32,
    height: i32,
    capacity: uint,
}

impl ShmBuffer {
    pub fn new() -> ShmBuffer {
        unsafe {
            let name = b"rust-wayland-shm\x00";
            let fd = memfd_create(name.as_ptr(), MFD_CLOEXEC | MFD_ALLOW_SEALING);
            assert!(fd >= 0);
            let width = 300i32;
            let height = 200i32;
            let capacity = width as uint
                * height as uint
                * std::mem::size_of::<u32>();
            assert!(libc::ftruncate(fd, capacity as i64) != -1);
            let ptr = libc::mmap(ptr::null_mut(), capacity as u64,
                                 libc::PROT_WRITE | libc::PROT_READ, 1, fd, 0);
            assert!(ptr != libc::MAP_FAILED);
            for i in range(0, width * height) {
                let p: *mut u32 = (ptr as *mut u32).offset(i as int);
                let x = i % width;
                let y = i / width;
                match x {
                    0...5 => std::ptr::write(&mut *p, 0x7FFF0000),
                    294...299 => std::ptr::write(&mut *p, 0x7F0000FF),
                    _ => match y {
                        0...5 => std::ptr::write(&mut *p, 0x7F00FF00),
                        194...199 => std::ptr::write(&mut *p, 0x7FFF00FF),
                        _ =>  std::ptr::write(&mut *p, 0x7F000000),
                    }
                }
            }
            ShmBuffer {
                fd: fd,
                ptr: ptr,
                width: width,
                height: height,
                capacity: capacity,
            }
        }
    }
    pub fn resize(&mut self, width: i32, height: i32) {
        if (width * height) > (self.width * self.height) {
            println!("TODO: actual");
        }
        self.width = width;
        self.height = height;
    }
    pub fn fd(&self) -> i32 {
        self.fd
    }
    pub fn capacity(&self) -> uint { self.capacity }
    pub fn width(&self) -> i32 { self.width }
    pub fn height(&self) -> i32 { self.height }
    pub fn stride(&self) -> i32 {
        self.width * std::mem::size_of::<u32>() as i32
    }
}

impl Drop for ShmBuffer {
    fn drop(&mut self) {
        unsafe {
            libc::munmap(self.ptr, self.capacity as u64);
            libc::close(self.fd);
        }
    }
}

fn main() {
    let mut display = wayland::Display::connect_to_env_or_default();
    let mut registry = wayland::Registry::new(&mut display);
    // Create the shell surface
    let mut surface = registry.compositor().create_surface();
    let mut shell_surface = registry.shell().get_shell_surface(&mut surface);
    shell_surface.set_toplevel();
    // Create the buffer
    let mut shm_buffer = ShmBuffer::new();
    shm_buffer.resize(300, 200);
    let mut pool = registry.shm().create_pool(shm_buffer.fd(),
                                              shm_buffer.capacity() as i32);
    let mut buffer = pool.create_buffer(
        0, shm_buffer.width(), shm_buffer.height(), shm_buffer.stride(),
        wayland::raw::WL_SHM_FORMAT_ARGB8888
    );
    surface.attach(&mut buffer, 0, 0);
    surface.commit();
    loop {
        display.dispatch();
    }
}
<|endoftext|>


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

extern crate egg_mode;

mod common;

use common::tokio_core::reactor;

fn main() {
    let mut core = reactor::Core::new().unwrap();

    let config = common::Config::load(&mut core);
    let handle = core.handle();
    let tweet_id = 766678057788829697;

    println!("");
    println!("Load up an individual tweet:");
    let status = core.run(egg_mode::tweet::show(tweet_id, &config.token, &handle)).unwrap();
    common::print_tweet(&status);

    println!("");
    println!("Loading retweets of an individual tweet:");
    for rt in &core.run(egg_mode::tweet::retweets_of(tweet_id, 5, &config.token, &handle)).unwrap() {
        if let Some(ref user) = rt.user {
            println!("{} (@{})", user.name, user.screen_name);
        }
    }

    println!("");
    println!("Loading the user's home timeline:");
    let mut home = egg_mode::tweet::home_timeline(&config.token, &handle).with_page_size(5);
    for status in &core.run(home.start()).unwrap() {
        common::print_tweet(&status);
        println!("");
    }

    println!("");
    println!("Loading the user's mentions timeline:");
    let mut home = egg_mode::tweet::mentions_timeline(&config.token, &handle).with_page_size(5);
    for status in &core.run(home.start()).unwrap() {
        common::print_tweet(&status);
        println!("");
    }

    println!("");
    println!("Loading the user's timeline:");
    let mut home = egg_mode::tweet::user_timeline(config.user_id, true, true,
                                                  &config.token, &handle).with_page_size(5);
    for status in &core.run(home.start()).unwrap() {
        common::print_tweet(&status);
        println!("");
    }
}
<commit_msg>fix examples/tweets.rs<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

extern crate egg_mode;

mod common;

use common::tokio_core::reactor;

fn main() {
    let mut core = reactor::Core::new().unwrap();

    let config = common::Config::load(&mut core);
    let handle = core.handle();
    let tweet_id = 766678057788829697;

    println!("");
    println!("Load up an individual tweet:");
    let status = core.run(egg_mode::tweet::show(tweet_id, &config.token, &handle)).unwrap();
    common::print_tweet(&status);

    println!("");
    println!("Loading retweets of an individual tweet:");
    for rt in &core.run(egg_mode::tweet::retweets_of(tweet_id, 5, &config.token, &handle)).unwrap() {
        if let Some(ref user) = rt.user {
            println!("{} (@{})", user.name, user.screen_name);
        }
    }

    println!("");
    println!("Loading the user's home timeline:");
    let home = egg_mode::tweet::home_timeline(&config.token, &handle).with_page_size(5);
    let (_home, feed) = core.run(home.start()).unwrap();
    for status in feed {
        common::print_tweet(&status);
        println!("");
    }

    println!("");
    println!("Loading the user's mentions timeline:");
    let mentions = egg_mode::tweet::mentions_timeline(&config.token, &handle).with_page_size(5);
    let (_mentions, feed) = core.run(mentions.start()).unwrap();
    for status in feed {
        common::print_tweet(&status);
        println!("");
    }

    println!("");
    println!("Loading the user's timeline:");
    let user = egg_mode::tweet::user_timeline(config.user_id, true, true,
                                              &config.token, &handle).with_page_size(5);
    let (_user, feed) = core.run(user.start()).unwrap();
    for status in feed {
        common::print_tweet(&status);
        println!("");
    }
}
<|endoftext|>


<commit_before>use num::{FromPrimitive, ToPrimitive};
use std::marker::PhantomData;

use byteorder::{ByteOrder, LittleEndian};
use log::{debug};
use goblin::{Object};
use failure::{Error, Fail};

use super::super::super::arch::Arch;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

#[derive(Debug, Fail)]
pub enum PtrAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct PtrAnalyzer<A: Arch> {
    // This Analyzer must have a type parameter for it
    //  to implement Analyzer<A>.
    // however, it doesn't actually use this type itself.
    // so, we use a phantom data marker which has zero type,
    //  to ensure there is not an unused type parameter,
    //  which is a compile error.
    _phantom: PhantomData<A>,
}

impl<A: Arch> PtrAnalyzer<A> {
    pub fn new() -> PtrAnalyzer<A> {
        PtrAnalyzer {
            _phantom: PhantomData {},
        }
    }
}

struct Section<A: Arch> {
    start: A::RVA,
    end: A::RVA,
}

fn is_in_insn<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    let start: usize = rva.to_usize().unwrap();
    // TODO: remove harded max insn length
    // TODO: underflow
    let end: usize = rva.to_usize().unwrap() - 0x10;

    for i in (start..end).rev() {
        let i = A::RVA::from_usize(i).unwrap();
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + A::RVA::from_u8(len).unwrap() > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    IMAGE_REL_BASED_ABSOLUTE,
    IMAGE_REL_BASED_HIGH,
    IMAGE_REL_BASED_LOW,
    IMAGE_REL_BASED_HIGHLOW,
    IMAGE_REL_BASED_HIGHADJ,

    // IMAGE_REL_BASED_MIPS_JMPADDR,
    // IMAGE_REL_BASED_ARM_MOV32,
    // IMAGE_REL_BASED_RISCV_HIGH20,
    IMAGE_REL_ARCH1,

    IMAGE_REL_RESERVED,

    // IMAGE_REL_BASED_THUMB_MOV32,
    // IMAGE_REL_BASED_RISCV_LOW12I,
    IMAGE_REL_ARCH2,

    IMAGE_REL_BASED_RISCV_LOW12S,
    IMAGE_REL_BASED_MIPS_JMPADDR16,
    IMAGE_REL_BASED_DIR64,
}

pub struct Reloc<A: Arch> {
    pub typ: RelocationType,
    pub offset: A::RVA,
}

fn parse_reloc<A: Arch>(base: A::RVA, entry: u16) -> Result<Reloc<A>, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::IMAGE_REL_BASED_ABSOLUTE,
        1 => RelocationType::IMAGE_REL_BASED_HIGH,
        2 => RelocationType::IMAGE_REL_BASED_LOW,
        3 => RelocationType::IMAGE_REL_BASED_HIGHLOW,
        4 => RelocationType::IMAGE_REL_BASED_HIGHADJ,
        5 => RelocationType::IMAGE_REL_ARCH1,
        6 => RelocationType::IMAGE_REL_RESERVED,
        7 => RelocationType::IMAGE_REL_ARCH2,
        8 => RelocationType::IMAGE_REL_BASED_RISCV_LOW12S,
        9 => RelocationType::IMAGE_REL_BASED_MIPS_JMPADDR16,
        10 => RelocationType::IMAGE_REL_BASED_DIR64,
        _ => return Err(PtrAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + A::RVA::from_u16(reloc_offset).unwrap(),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::ptrs;
///
/// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = ptrs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   0x76008);
/// assert_eq!(relocs[277].offset, 0xA8000);
/// ```
pub fn get_relocs<A: Arch + 'static>(ws: &Workspace<A>) -> Result<Vec<Reloc<A>>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = A::RVA::from_u32(reloc_directory.virtual_address).unwrap();
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (A::RVA::from_u32(*page_rva).unwrap(), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc::<A>(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc::<A>(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

impl<A: Arch + 'static> Analyzer<A> for PtrAnalyzer<A> {
    fn get_name(&self) -> String {
        "pointer analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::PtrAnalyzer;
    ///
    /// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = PtrAnalyzer::<Arch64>::new();
    /// anal.analyze(&mut ws).unwrap();
    /// let meta = ws.get_meta(0xC7F0).unwrap();
    /// assert!(meta.is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace<A>)-> Result<(), Error> {
        let text_section = match ws.module.sections.iter()
            // currently limited to the text section.
            // TODO: accept any writable section.
            .filter(|&sec| sec.name == ".text")
            .next() {
                None => return Ok(()),
                Some(s) => s,
            };

        let text_bounds = Section::<A> {
            start: text_section.addr,
            end: text_section.addr + A::RVA::from_usize(text_section.buf.len()).unwrap(),
        };

        // scan the relocations
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.
        let o: Vec<A::RVA> = get_relocs(ws)?
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va))
            .filter(|&rva| text_bounds.start <= rva)
            .filter(|&rva| rva < text_bounds.end)
            .filter(|&rva| !is_in_insn(ws, rva))
            .filter(|&rva| !is_ptr(ws, rva))
            // TODO: maybe ensure that the insn decodes.
            .collect();

        o.iter().for_each(|&rva| {
            debug!("found ptr from .text section to .text section at {:#x}", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<commit_msg>analysis/pe/ptrs: linting<commit_after>use num::{FromPrimitive, ToPrimitive};
use std::marker::PhantomData;

use byteorder::{ByteOrder, LittleEndian};
use log::{debug};
use goblin::{Object};
use failure::{Error, Fail};

use super::super::super::arch::Arch;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

#[derive(Debug, Fail)]
pub enum PtrAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct PtrAnalyzer<A: Arch> {
    // This Analyzer must have a type parameter for it
    //  to implement Analyzer<A>.
    // however, it doesn't actually use this type itself.
    // so, we use a phantom data marker which has zero type,
    //  to ensure there is not an unused type parameter,
    //  which is a compile error.
    _phantom: PhantomData<A>,
}

impl<A: Arch> PtrAnalyzer<A> {
    pub fn new() -> PtrAnalyzer<A> {
        PtrAnalyzer {
            _phantom: PhantomData {},
        }
    }
}

struct Section<A: Arch> {
    start: A::RVA,
    end: A::RVA,
}

fn is_in_insn<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    let start: usize = rva.to_usize().unwrap();
    // TODO: remove harded max insn length
    // TODO: underflow
    let end: usize = rva.to_usize().unwrap() - 0x10;

    for i in (start..end).rev() {
        let i = A::RVA::from_usize(i).unwrap();
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + A::RVA::from_u8(len).unwrap() > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr<A: Arch + 'static>(ws: &Workspace<A>, rva: A::RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc<A: Arch> {
    pub typ: RelocationType,
    pub offset: A::RVA,
}

fn parse_reloc<A: Arch>(base: A::RVA, entry: u16) -> Result<Reloc<A>, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(PtrAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + A::RVA::from_u16(reloc_offset).unwrap(),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::ptrs;
///
/// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = ptrs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   0x76008);
/// assert_eq!(relocs[277].offset, 0xA8000);
/// ```
pub fn get_relocs<A: Arch + 'static>(ws: &Workspace<A>) -> Result<Vec<Reloc<A>>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = A::RVA::from_u32(reloc_directory.virtual_address).unwrap();
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (A::RVA::from_u32(*page_rva).unwrap(), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc::<A>(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc::<A>(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

impl<A: Arch + 'static> Analyzer<A> for PtrAnalyzer<A> {
    fn get_name(&self) -> String {
        "pointer analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::PtrAnalyzer;
    ///
    /// let mut ws = Workspace::<Arch64>::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = PtrAnalyzer::<Arch64>::new();
    /// anal.analyze(&mut ws).unwrap();
    /// let meta = ws.get_meta(0xC7F0).unwrap();
    /// assert!(meta.is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace<A>)-> Result<(), Error> {
        let text_section = match ws.module.sections.iter()
            // currently limited to the text section.
            // TODO: accept any writable section.
            .filter(|&sec| sec.name == ".text")
            .next() {
                None => return Ok(()),
                Some(s) => s,
            };

        let text_bounds = Section::<A> {
            start: text_section.addr,
            end: text_section.addr + A::RVA::from_usize(text_section.buf.len()).unwrap(),
        };

        // scan the relocations
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.
        let o: Vec<A::RVA> = get_relocs(ws)?
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va))
            .filter(|&rva| text_bounds.start <= rva)
            .filter(|&rva| rva < text_bounds.end)
            .filter(|&rva| !is_in_insn(ws, rva))
            .filter(|&rva| !is_ptr(ws, rva))
            // TODO: maybe ensure that the insn decodes.
            .collect();

        o.iter().for_each(|&rva| {
            debug!("found ptr from .text section to .text section at {:#x}", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<|endoftext|>