<commit_before>//! The Divide & Conquer recursive strategy.
//!
//! From [Wikipedia][1]:
//! > A divide and conquer algorithm works by recursively breaking down a
//! > problem into two or more sub-problems of the same (or related) type
//! > (**divide**) until these become simple enough to be solved directly
//! > (**conquer**). The solution to the sub-problems are then combined
//! > to give a solution to the original problem.
//!
//! To use this strategy you need to implement the [`DacProblem`][2] trait and
//! everything else is handled for you.
//!
//! [1]: https://en.wikipedia.org/wiki/Divide_and_conquer_algorithms
//! [2]: trait.DacProblem.html

use std::collections::HashMap;
use std::hash::Hash;
use std::marker::PhantomData;

/// Divide & Conquer Problem.
///
/// You have to implement this trait in order to use the algorithm. Simply fill
/// all the methods and you're all set. We have two generic types:
///
/// - `E`: A partial solution, wich we work with.
/// - `S`: The final solution, the output of the algorithm.
///
/// Sometimes `E` and `S` are the same type because we don't need conversion
/// between the partial and the final solution (e.g. when calculating a
/// factorial), but if we are working with a list or an array, `E` could simply
/// be a pointer to an element and `S` the actual element.
///
/// Once you have defined the problem, use [`DacAlgorithm`][1] or
/// [`DacMemAlgorithm`][2] to solve it.
///
/// [1]: struct.DacAlgorithm.html
/// [2]: struct.DacMemAlgorithm.html
pub trait DacProblem<S, E> {
    /// The size of your problem. Usually referred to as `n`.
    ///
    /// It should decrease as you split the problem, otherwise the algorithm
    /// will never stop, and we don't want that.
    fn size(&self) -> usize;

    /// The moment when the problem is simple enough to be solved directly.
    ///
    /// Here `DacProblem::size` is the smallest.
    fn is_base_case(&self) -> bool;

    /// The base case's solution.
    fn base_case_solution(&self) -> E;

    /// The number of subproblems the problem has.
    fn subproblem_count(&self) -> usize;

    /// Return the subproblem corresponding to the given `i`. (i.e. the first
    /// subproblem, the second...).
    ///
    /// `i` is in the range [0, `DacProblem::subproblem_count`).
    fn get_subproblem(&self, i: usize) -> Self;

    /// How to combine the subproblems' solutions into the current problem's
    /// solution.
    fn combine(&self, solutions: Vec<E>) -> E;

    /// Transforms the partial solution `E`, if it is possible, into a final
    /// solution `S`.
    fn get_solution(&self, partial_solution: &E) -> Option<S>;
}

/// The problem solver.
///
/// Solves a divide & conquer problem without memory. Useful if we don't need
/// to store partial solutions.
pub struct DacAlgorithm<S, E, P: DacProblem<S, E>> {
    phan: PhantomData<S>,
    partial_solution: E,
    problem: P,
}

impl<S, E, P: DacProblem<S, E>> DacAlgorithm<S, E, P> {
    /// Solve the `problem` problem.
    pub fn new(problem: P) -> Self {
        let e = Self::solve(&problem);

        DacAlgorithm {
            phan: PhantomData,
            partial_solution: e,
            problem: problem,
        }
    }

    fn solve(problem: &P) -> E {
        if problem.is_base_case() {
            problem.base_case_solution()

        } else {
            let solutions = (0..problem.subproblem_count())
                                .map(|i| problem.get_subproblem(i))
                                .map(|p| Self::solve(&p))
                                .collect::<Vec<E>>();

            problem.combine(solutions)
        }
    }

    /// Get the final solution.
    pub fn get_solution(&self) -> Option<S> {
        self.problem.get_solution(&self.partial_solution)
    }
}

/// The problem solver, with memory.
///
/// # Explanation
///
/// Solves a divide & conquer problem with memory: it stores the partial
/// solutions of all the subproblems, in case they are needed by another
/// subproblem. An usual example is the [Fibonacci sequence][1]:
///
///     fn fib(n: u64) -> u64 {
///         match n {
///             0 => 0,
///             1 => 1,
///             n => fib(n - 1) + fib(n - 2),
///         }
///     }
///
/// As you can see, `fib(5)` needs `fib(4)` and `fib(3)`. But then `fib(4)` also
/// needs `fib(3)`. Memory here is useful so we don't recalculate values.
///
/// # Usage
///
/// Using this is the same as using [`DacAlgorithm`][2] but you problem has to
/// implement some traits:
///
/// - `DacProblem` has to implement the following traits:
///     - `Eq` (and therefore `PartialEq`)
///     - `Hash`
///     - `Clone`
/// - `E` has to be `Clone` too.
///
/// Which is not complicated to do:
///
///     #[derive(Eq, PartialEq, Hash, Clone)]
///     struct Fibonacci(u64);
///
/// [1]: https://en.wikipedia.org/wiki/Fibonacci_number
/// [2]: struct.DacAlgorithm.html#usage
pub struct DacMemAlgorithm<P, S, E>
    where P: DacProblem<S, E> + Eq + Hash + Clone,
          E: Clone
{
    phan: PhantomData<S>,
    solutions: HashMap<P, E>,
    problem: P,
}

impl<P, S, E> DacMemAlgorithm<P, S, E>
    where P: DacProblem<S, E> + Eq + Hash + Clone,
          E: Clone
{
    pub fn new(problem: P) -> Self {
        let mut map = HashMap::new();
        Self::solve(&problem, &mut map);

        DacMemAlgorithm {
            phan: PhantomData,
            solutions: map,
            problem: problem,
        }
    }

    fn solve(problem: &P, mut solutions: &mut HashMap<P, E>) -> E {
        if problem.is_base_case() {
            problem.base_case_solution()

        } else if solutions.contains_key(&problem) {
            solutions.get(&problem).unwrap().clone()

        } else {
            let solution = {
                let solutions = (0..problem.subproblem_count())
                                    .map(|i| problem.get_subproblem(i))
                                    .map(|p| Self::solve(&p, &mut solutions))
                                    .collect::<Vec<E>>();

                problem.combine(solutions)
            };

            solutions.insert(problem.clone(), solution);
            solutions.get(&problem).unwrap().clone()
        }
    }

    pub fn get_solution(&self) -> Option<S> {
        self.solutions
            .get(&self.problem)
            .and_then(|e| self.problem.get_solution(e))
    }
}
<commit_msg>Clear DacAlgorithm generics type signature<commit_after>//! The Divide & Conquer recursive strategy.
//!
//! From [Wikipedia][1]:
//! > A divide and conquer algorithm works by recursively breaking down a
//! > problem into two or more sub-problems of the same (or related) type
//! > (**divide**) until these become simple enough to be solved directly
//! > (**conquer**). The solution to the sub-problems are then combined
//! > to give a solution to the original problem.
//!
//! To use this strategy you need to implement the [`DacProblem`][2] trait and
//! everything else is handled for you.
//!
//! [1]: https://en.wikipedia.org/wiki/Divide_and_conquer_algorithms
//! [2]: trait.DacProblem.html

use std::collections::HashMap;
use std::hash::Hash;
use std::marker::PhantomData;

/// Divide & Conquer Problem.
///
/// You have to implement this trait in order to use the algorithm. Simply fill
/// all the methods and you're all set. We have two generic types:
///
/// - `E`: A partial solution, wich we work with.
/// - `S`: The final solution, the output of the algorithm.
///
/// Sometimes `E` and `S` are the same type because we don't need conversion
/// between the partial and the final solution (e.g. when calculating a
/// factorial), but if we are working with a list or an array, `E` could simply
/// be a pointer to an element and `S` the actual element.
///
/// Once you have defined the problem, use [`DacAlgorithm`][1] or
/// [`DacMemAlgorithm`][2] to solve it.
///
/// [1]: struct.DacAlgorithm.html
/// [2]: struct.DacMemAlgorithm.html
pub trait DacProblem<S, E> {
    /// The size of your problem. Usually referred to as `n`.
    ///
    /// It should decrease as you split the problem, otherwise the algorithm
    /// will never stop, and we don't want that.
    fn size(&self) -> usize;

    /// The moment when the problem is simple enough to be solved directly.
    ///
    /// Here `DacProblem::size` is the smallest.
    fn is_base_case(&self) -> bool;

    /// The base case's solution.
    fn base_case_solution(&self) -> E;

    /// The number of subproblems the problem has.
    fn subproblem_count(&self) -> usize;

    /// Return the subproblem corresponding to the given `i`. (i.e. the first
    /// subproblem, the second...).
    ///
    /// `i` is in the range [0, `DacProblem::subproblem_count`).
    fn get_subproblem(&self, i: usize) -> Self;

    /// How to combine the subproblems' solutions into the current problem's
    /// solution.
    fn combine(&self, solutions: Vec<E>) -> E;

    /// Transforms the partial solution `E`, if it is possible, into a final
    /// solution `S`.
    fn get_solution(&self, partial_solution: &E) -> Option<S>;
}

/// The problem solver.
///
/// Solves a divide & conquer problem without memory. Useful if we don't need
/// to store partial solutions.
pub struct DacAlgorithm<P, S, E>
    where P: DacProblem<S, E>
{
    phan: PhantomData<S>,
    partial_solution: E,
    problem: P,
}

impl<P, S, E> DacAlgorithm<P, S, E>
    where P: DacProblem<S, E>
{
    /// Solve the `problem` problem.
    pub fn new(problem: P) -> Self {
        let e = Self::solve(&problem);

        DacAlgorithm {
            phan: PhantomData,
            partial_solution: e,
            problem: problem,
        }
    }

    fn solve(problem: &P) -> E {
        if problem.is_base_case() {
            problem.base_case_solution()

        } else {
            let solutions = (0..problem.subproblem_count())
                                .map(|i| problem.get_subproblem(i))
                                .map(|p| Self::solve(&p))
                                .collect::<Vec<E>>();

            problem.combine(solutions)
        }
    }

    /// Get the final solution.
    pub fn get_solution(&self) -> Option<S> {
        self.problem.get_solution(&self.partial_solution)
    }
}

/// The problem solver, with memory.
///
/// # Explanation
///
/// Solves a divide & conquer problem with memory: it stores the partial
/// solutions of all the subproblems, in case they are needed by another
/// subproblem. An usual example is the [Fibonacci sequence][1]:
///
///     fn fib(n: u64) -> u64 {
///         match n {
///             0 => 0,
///             1 => 1,
///             n => fib(n - 1) + fib(n - 2),
///         }
///     }
///
/// As you can see, `fib(5)` needs `fib(4)` and `fib(3)`. But then `fib(4)` also
/// needs `fib(3)`. Memory here is useful so we don't recalculate values.
///
/// # Usage
///
/// Using this is the same as using [`DacAlgorithm`][2] but you problem has to
/// implement some traits:
///
/// - `DacProblem` has to implement the following traits:
///     - `Eq` (and therefore `PartialEq`)
///     - `Hash`
///     - `Clone`
/// - `E` has to be `Clone` too.
///
/// Which is not complicated to do:
///
///     #[derive(Eq, PartialEq, Hash, Clone)]
///     struct Fibonacci(u64);
///
/// [1]: https://en.wikipedia.org/wiki/Fibonacci_number
/// [2]: struct.DacAlgorithm.html#usage
pub struct DacMemAlgorithm<P, S, E>
    where P: DacProblem<S, E> + Eq + Hash + Clone,
          E: Clone
{
    phan: PhantomData<S>,
    solutions: HashMap<P, E>,
    problem: P,
}

impl<P, S, E> DacMemAlgorithm<P, S, E>
    where P: DacProblem<S, E> + Eq + Hash + Clone,
          E: Clone
{
    pub fn new(problem: P) -> Self {
        let mut map = HashMap::new();
        Self::solve(&problem, &mut map);

        DacMemAlgorithm {
            phan: PhantomData,
            solutions: map,
            problem: problem,
        }
    }

    fn solve(problem: &P, mut solutions: &mut HashMap<P, E>) -> E {
        if problem.is_base_case() {
            problem.base_case_solution()

        } else if solutions.contains_key(&problem) {
            solutions.get(&problem).unwrap().clone()

        } else {
            let solution = {
                let solutions = (0..problem.subproblem_count())
                                    .map(|i| problem.get_subproblem(i))
                                    .map(|p| Self::solve(&p, &mut solutions))
                                    .collect::<Vec<E>>();

                problem.combine(solutions)
            };

            solutions.insert(problem.clone(), solution);
            solutions.get(&problem).unwrap().clone()
        }
    }

    pub fn get_solution(&self) -> Option<S> {
        self.solutions
            .get(&self.problem)
            .and_then(|e| self.problem.get_solution(e))
    }
}
<|endoftext|>


<commit_before>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::demux::{Demuxer, Event};
use av_format::demuxer::context::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational32;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement, simple_block};
use nom::{self, Err, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    //self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn open(&mut self) {}

    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| d as u64);
                if let Some(ref t) = self.tracks {
                    info.streams = t.tracks.iter().map(|tr| track_to_stream(tr)).collect();
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
             Err(Err::Incomplete(_)) => Err(ErrorKind::MoreDataNeeded.into()),
            e => {
                println!("error reading headers: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }

    fn read_packet(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            println!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets();
                            self.queue.extend(packets.drain(..));
                            /*
                        for block in c.simple_block.iter() {
                          println!("got simple block of size {}", block.len());
                          if let IResult::Done(_,o) = simple_block(block) {
                            println!("parsing simple block: {:?}", o);
                          } else {
                            println!("error parsing simple block");
                          }
                          self.blockstream.extend(*block);
                        }

                        for block_group in c.block_group.iter() {
                          println!("got block group of size {}", block_group.block.len());
                          self.blockstream.extend(block_group.block);
                        }
                        */
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        el => {
                            println!("got element: {:#?}", el);
                        }
                    }

                    Ok((seek, Event::MoreDataNeeded))
                }
                 Err(Err::Incomplete(_)) => Ok((SeekFrom::Current(0), Event::MoreDataNeeded)),
                e => {
                    println!("parsing issue: {:?}", e);
                    Err(ErrorKind::InvalidData.into())
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9" => Some(CodecID::VP9),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

pub fn track_to_stream(t: &TrackEntry) -> Stream {
    Stream {
        id: t.track_uid as usize,
        index: t.track_number as usize,
        start: None,
        duration: t.default_duration,
        timebase: Rational32::from_integer(1),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                //println!("parsing simple block: {:?}", block);
                let packet = Packet {
                    data: i.into(),
                    pts: Some(block.timecode as i64),
                    dts: None,
                    pos: None,
                    stream_index: block.track_number as isize,
                    is_key: block.keyframe,
                    is_corrupted: false,
                };

                v.push(Event::NewPacket(packet));
            } else {
                println!("error parsing simple block");
            }
        }

        v
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::Offset;
    use av_format::demuxer::context::DemuxerContext;

    const webm: &'static [u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        println!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                println!("offset: {} bytes", webm.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
        panic!();
    }

    #[test]
    fn context() {
        let mut context = DemuxerContext::new(Box::new(MkvDemuxer::new()),
                                              Box::new(Cursor::new(webm)));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT streams: {:?}", context.info.streams);
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        panic!();
    }
}
<commit_msg>Update to the new Packet and Stream<commit_after>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_data::timeinfo::TimeInfo;
use av_format::stream::*;
use av_format::buffer::Buffered;
use av_format::demuxer::demux::{Demuxer, Event};
use av_format::demuxer::context::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational64;

use ebml::{ebml_header, EBMLHeader};
use elements::{segment, segment_element, Cluster, SeekHead, Info, Tracks, TrackEntry,
               SegmentElement, simple_block};
use nom::{self, Err, IResult, Offset};

#[derive(Debug, Clone, PartialEq)]
pub struct MkvDemuxer {
    pub header: Option<EBMLHeader>,
    pub seek_head: Option<SeekHead>,
    pub info: Option<Info>,
    pub tracks: Option<Tracks>,
    pub queue: VecDeque<Event>,
    pub blockstream: Vec<u8>,
}

impl MkvDemuxer {
    pub fn new() -> MkvDemuxer {
        MkvDemuxer {
            header: None,
            seek_head: None,
            info: None,
            tracks: None,
            queue: VecDeque::new(),
            blockstream: Vec::new(),
        }
    }

    pub fn parse_until_tracks<'a>(&mut self, original_input: &'a [u8]) -> IResult<&'a [u8], ()> {
        let (i1, header) = try_parse!(original_input, ebml_header);

        self.header = Some(header);

        let (mut input, _) = try_parse!(i1, segment);

        loop {
            if self.seek_head.is_some() && self.info.is_some() && self.tracks.is_some() {
                return Ok((input, ()));
            }

            println!("offset: {}", original_input.offset(input));

            let (i3, element) = try_parse!(input, segment_element);

            match element {
                SegmentElement::SeekHead(s) => {
                    println!("got seek head: {:#?}", s);
                    if self.seek_head.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.seek_head = Some(s);
                    }
                }
                SegmentElement::Info(i) => {
                    println!("got info: {:#?}", i);
                    if self.info.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.info = Some(i);
                    }
                }
                SegmentElement::Tracks(t) => {
                    println!("got tracks: {:#?}", t);
                    if self.tracks.is_some() {
                        return Err(Err::Error(error_position!(nom::ErrorKind::Custom(1), input)));
                    } else {
                        self.tracks = Some(t);
                    }
                }
                SegmentElement::Cluster(c) => {
                    println!("got a cluster: {:#?}", c);
                    //self.clusters.push(c);
                }
                el => {
                    println!("got element: {:#?}", el);
                }
            }

            input = i3;
        }
    }
}

impl Demuxer for MkvDemuxer {
    fn open(&mut self) {}

    fn read_headers(&mut self, buf: &Box<Buffered>, info: &mut GlobalInfo) -> Result<SeekFrom> {
        match self.parse_until_tracks(buf.data()) {
            Ok((i, _)) => {
                info.duration = self.info.as_ref().and_then(|info| info.duration).map(|d| d as u64);
                if let Some(ref t) = self.tracks {
                    info.streams = t.tracks.iter().map(|tr| track_to_stream(tr)).collect();
                }
                Ok(SeekFrom::Current(buf.data().offset(i) as i64))
            }
             Err(Err::Incomplete(_)) => Err(ErrorKind::MoreDataNeeded.into()),
            e => {
                println!("error reading headers: {:?}", e);
                Err(ErrorKind::InvalidData.into())
            }
        }
    }

    fn read_packet(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            println!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets();
                            self.queue.extend(packets.drain(..));
                            /*
                        for block in c.simple_block.iter() {
                          println!("got simple block of size {}", block.len());
                          if let IResult::Done(_,o) = simple_block(block) {
                            println!("parsing simple block: {:?}", o);
                          } else {
                            println!("error parsing simple block");
                          }
                          self.blockstream.extend(*block);
                        }

                        for block_group in c.block_group.iter() {
                          println!("got block group of size {}", block_group.block.len());
                          self.blockstream.extend(block_group.block);
                        }
                        */
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        el => {
                            println!("got element: {:#?}", el);
                        }
                    }

                    Ok((seek, Event::MoreDataNeeded))
                }
                 Err(Err::Incomplete(_)) => Ok((SeekFrom::Current(0), Event::MoreDataNeeded)),
                e => {
                    println!("parsing issue: {:?}", e);
                    Err(ErrorKind::InvalidData.into())
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<CodecID> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS" => Some(CodecID::Opus),
        "V_VP9" => Some(CodecID::VP9),
        _ => None,
    }
}

fn track_entry_video_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref video) = t.video {
        let v = VideoInfo {
            width: video.pixel_width as usize,
            height: video.pixel_height as usize,
            // TODO parse Colour and/or CodecPrivate to extract the format
            format: None,
        };
        Some(MediaKind::Video(v))
    } else {
        None
    }
}

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let a = AudioInfo {
            samples: 0,
            rate: rate as usize,
            map: None,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

fn track_entry_media_kind(t: &TrackEntry) -> Option<MediaKind> {
    // TODO: Use an enum for the track type
    match t.track_type {
        0x1 => track_entry_video_kind(t),
        0x2 => track_entry_audio_kind(t),
        _ => None,
    }
}

pub fn track_to_stream(t: &TrackEntry) -> Stream {
    Stream {
        id: t.track_uid as usize,
        index: t.track_number as usize,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::from_integer(1),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
    }
}

impl<'a> Cluster<'a> {
    pub fn generate_packets(&self) -> Vec<Event> {
        let mut v = Vec::new();

        for block_data in self.simple_block.iter() {
            if let Ok((i, block)) = simple_block(block_data) {
                //println!("parsing simple block: {:?}", block);
                let packet = Packet {
                    data: i.into(),
                    t: TimeInfo {
                        pts: Some(block.timecode as i64),
                        dts: None,
                        duration: None,
                        timebase: None,
                    },
                    pos: None,
                    stream_index: block.track_number as isize,
                    is_key: block.keyframe,
                    is_corrupted: false,
                };

                v.push(Event::NewPacket(packet));
            } else {
                println!("error parsing simple block");
            }
        }

        v
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::Offset;
    use av_format::demuxer::context::DemuxerContext;

    const webm: &'static [u8] = include_bytes!("../assets/bbb-vp9-opus.webm");

    #[test]
    fn parse_headers() {
        let mut demuxer = MkvDemuxer::new();

        let res = demuxer.parse_until_tracks(webm);
        println!("got parsing res: {:?}", res);
        match res {
            Ok((i, _)) => {
                println!("offset: {} bytes", webm.offset(i));
            }
            e => {
                println!("could not parse: {:?}", e);
            }
        }

        println!("demuxer: {:#?}", demuxer);
        panic!();
    }

    #[test]
    fn context() {
        let mut context = DemuxerContext::new(Box::new(MkvDemuxer::new()),
                                              Box::new(Cursor::new(webm)));
        println!("DEMUXER CONTEXT read headers: {:?}", context.read_headers());
        println!("DEMUXER CONTEXT streams: {:?}", context.info.streams);
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        println!("DEMUXER CONTEXT event: {:?}", context.read_packet());
        panic!();
    }
}
<|endoftext|>


<commit_before>//! Provides the parser for WSplit splits files.

use crate::{settings::Image, RealTime, Run, Segment, TimeSpan};
use snafu::ResultExt;
use std::io::{self, BufRead};
use std::num::{ParseFloatError, ParseIntError};
use std::result::Result as StdResult;

/// The Error type for splits files that couldn't be parsed by the WSplit
/// Parser.
#[derive(Debug, snafu::Snafu)]
pub enum Error {
    /// Expected the name of the segment, but didn't find it.
    ExpectedSegmentName,
    /// Expected the old time, but didn't find it.
    ExpectedOldTime,
    /// Expected the split time, but didn't find it.
    ExpectedPbTime,
    /// Expected the best segment time, but didn't find it.
    ExpectedBestTime,
    /// Failed to parse the amount of attempts.
    Attempt {
        /// The underlying error.
        source: ParseIntError,
    },
    /// Failed to parse the time the timer starts at.
    Offset {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to parse the personal best split time of a segment.
    PbTime {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to parse the best segment time of a segment.
    BestSegment {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to parse the time the "Old Run" of a segment.
    OldTime {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to read the next line.
    Line {
        /// The underlying error.
        source: io::Error,
    },
}

/// The Result type for the WSplit Parser.
pub type Result<T> = StdResult<T, Error>;

/// Attempts to parse a WSplit splits file. In addition to the source to parse,
/// you need to specify if additional files for the icons should be loaded from
/// the file system. If you are using livesplit-core in a server-like
/// environment, set this to `false`. Only client-side applications should set
/// this to `true`.
pub fn parse<R: BufRead>(source: R, load_icons: bool) -> Result<Run> {
    let mut run = Run::new();
    let mut icon_buf = Vec::new();
    let mut icons_list = Vec::new();
    let mut old_run_exists = false;
    let mut goal = None;

    for line in source.lines() {
        let line = line.context(Line)?;
        if !line.is_empty() {
            if line.starts_with("Title=") {
                run.set_category_name(&line["Title=".len()..]);
            } else if line.starts_with("Attempts=") {
                run.set_attempt_count(line["Attempts=".len()..].parse().context(Attempt)?);
            } else if line.starts_with("Offset=") {
                let offset = &line["Offset=".len()..];
                if !offset.is_empty() {
                    run.set_offset(TimeSpan::from_milliseconds(
                        -offset.parse::<f64>().context(Offset)?,
                    ));
                }
            } else if line.starts_with("Size=") {
                // Ignore
            } else if line.starts_with("Icons=") {
                if load_icons {
                    let icons = &line["Icons=".len()..];
                    icons_list.clear();
                    for path in icons.split(',') {
                        if path.len() >= 2 {
                            let path = &path[1..path.len() - 1];
                            if let Ok(image) = Image::from_file(path, &mut icon_buf) {
                                icons_list.push(image);
                                continue;
                            }
                        }
                        icons_list.push(Image::default());
                    }
                }
            } else if line.starts_with("Goal=") {
                goal = Some(line["Goal=".len()..].to_owned());
            } else {
                // must be a split Kappa
                let mut split_info = line.split(',');

                let segment_name = split_info.next().ok_or(Error::ExpectedSegmentName)?;
                let old_time = split_info.next().ok_or(Error::ExpectedOldTime)?;
                let pb_time = split_info.next().ok_or(Error::ExpectedPbTime)?;
                let best_time = split_info.next().ok_or(Error::ExpectedBestTime)?;

                let pb_time = TimeSpan::from_seconds(pb_time.parse().context(PbTime)?);
                let best_time = TimeSpan::from_seconds(best_time.parse().context(BestSegment)?);
                let old_time = TimeSpan::from_seconds(old_time.parse().context(OldTime)?);

                let mut segment = Segment::new(segment_name);

                if pb_time != TimeSpan::zero() {
                    segment.set_personal_best_split_time(RealTime(Some(pb_time)).into());
                }
                if best_time != TimeSpan::zero() {
                    segment.set_best_segment_time(RealTime(Some(best_time)).into());
                }
                if old_time != TimeSpan::zero() {
                    *segment.comparison_mut("Old Run") = RealTime(Some(old_time)).into();
                    old_run_exists = true;
                }

                run.push_segment(segment);
            }
        }
    }

    if let Some(goal) = goal {
        if run.category_name().is_empty() {
            run.set_category_name(goal);
        } else {
            let category_name = format!("{} (Goal: {})", run.category_name(), goal);
            run.set_category_name(category_name);
        }
    }

    if old_run_exists {
        run.add_custom_comparison("Old Run")
            .expect("WSplit: Old Run");
    }

    for (icon, segment) in icons_list.into_iter().zip(run.segments_mut().iter_mut()) {
        segment.set_icon(icon);
    }

    Ok(run)
}

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

    #[test]
    fn goal_parsing() {
        const RUN: &[u8] = br#"Title=WarioWare, Inc
Attempts=1
Offset=0
Size=374,61
Goal=sub 2h
Introduction,0,85.48,85.48
Jimmy,0,219.68,134.2
"#;

        let run = parse(RUN, false).unwrap();
        assert_eq!(run.category_name(), "WarioWare, Inc (Goal: sub 2h)");
        assert_eq!(run.len(), 2);
    }
}
<commit_msg>Don't include WSplit Goal if it's empty<commit_after>//! Provides the parser for WSplit splits files.

use crate::{settings::Image, RealTime, Run, Segment, TimeSpan};
use snafu::ResultExt;
use std::io::{self, BufRead};
use std::num::{ParseFloatError, ParseIntError};
use std::result::Result as StdResult;

/// The Error type for splits files that couldn't be parsed by the WSplit
/// Parser.
#[derive(Debug, snafu::Snafu)]
pub enum Error {
    /// Expected the name of the segment, but didn't find it.
    ExpectedSegmentName,
    /// Expected the old time, but didn't find it.
    ExpectedOldTime,
    /// Expected the split time, but didn't find it.
    ExpectedPbTime,
    /// Expected the best segment time, but didn't find it.
    ExpectedBestTime,
    /// Failed to parse the amount of attempts.
    Attempt {
        /// The underlying error.
        source: ParseIntError,
    },
    /// Failed to parse the time the timer starts at.
    Offset {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to parse the personal best split time of a segment.
    PbTime {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to parse the best segment time of a segment.
    BestSegment {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to parse the time the "Old Run" of a segment.
    OldTime {
        /// The underlying error.
        source: ParseFloatError,
    },
    /// Failed to read the next line.
    Line {
        /// The underlying error.
        source: io::Error,
    },
}

/// The Result type for the WSplit Parser.
pub type Result<T> = StdResult<T, Error>;

/// Attempts to parse a WSplit splits file. In addition to the source to parse,
/// you need to specify if additional files for the icons should be loaded from
/// the file system. If you are using livesplit-core in a server-like
/// environment, set this to `false`. Only client-side applications should set
/// this to `true`.
pub fn parse<R: BufRead>(source: R, load_icons: bool) -> Result<Run> {
    let mut run = Run::new();
    let mut icon_buf = Vec::new();
    let mut icons_list = Vec::new();
    let mut old_run_exists = false;
    let mut goal = None;

    for line in source.lines() {
        let line = line.context(Line)?;
        if !line.is_empty() {
            if line.starts_with("Title=") {
                run.set_category_name(&line["Title=".len()..]);
            } else if line.starts_with("Attempts=") {
                run.set_attempt_count(line["Attempts=".len()..].parse().context(Attempt)?);
            } else if line.starts_with("Offset=") {
                let offset = &line["Offset=".len()..];
                if !offset.is_empty() {
                    run.set_offset(TimeSpan::from_milliseconds(
                        -offset.parse::<f64>().context(Offset)?,
                    ));
                }
            } else if line.starts_with("Size=") {
                // Ignore
            } else if line.starts_with("Icons=") {
                if load_icons {
                    let icons = &line["Icons=".len()..];
                    icons_list.clear();
                    for path in icons.split(',') {
                        if path.len() >= 2 {
                            let path = &path[1..path.len() - 1];
                            if let Ok(image) = Image::from_file(path, &mut icon_buf) {
                                icons_list.push(image);
                                continue;
                            }
                        }
                        icons_list.push(Image::default());
                    }
                }
            } else if line.starts_with("Goal=") {
                let goal_value = line["Goal=".len()..].trim();
                if !goal_value.is_empty() {
                    goal = Some(goal_value.to_owned());
                }
            } else {
                // must be a split Kappa
                let mut split_info = line.split(',');

                let segment_name = split_info.next().ok_or(Error::ExpectedSegmentName)?;
                let old_time = split_info.next().ok_or(Error::ExpectedOldTime)?;
                let pb_time = split_info.next().ok_or(Error::ExpectedPbTime)?;
                let best_time = split_info.next().ok_or(Error::ExpectedBestTime)?;

                let pb_time = TimeSpan::from_seconds(pb_time.parse().context(PbTime)?);
                let best_time = TimeSpan::from_seconds(best_time.parse().context(BestSegment)?);
                let old_time = TimeSpan::from_seconds(old_time.parse().context(OldTime)?);

                let mut segment = Segment::new(segment_name);

                if pb_time != TimeSpan::zero() {
                    segment.set_personal_best_split_time(RealTime(Some(pb_time)).into());
                }
                if best_time != TimeSpan::zero() {
                    segment.set_best_segment_time(RealTime(Some(best_time)).into());
                }
                if old_time != TimeSpan::zero() {
                    *segment.comparison_mut("Old Run") = RealTime(Some(old_time)).into();
                    old_run_exists = true;
                }

                run.push_segment(segment);
            }
        }
    }

    if let Some(goal) = goal {
        if run.category_name().is_empty() {
            run.set_category_name(goal);
        } else {
            let category_name = format!("{} (Goal: {})", run.category_name(), goal);
            run.set_category_name(category_name);
        }
    }

    if old_run_exists {
        run.add_custom_comparison("Old Run")
            .expect("WSplit: Old Run");
    }

    for (icon, segment) in icons_list.into_iter().zip(run.segments_mut().iter_mut()) {
        segment.set_icon(icon);
    }

    Ok(run)
}

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

    #[test]
    fn goal_parsing() {
        const RUN: &[u8] = br#"Title=WarioWare, Inc
Attempts=1
Offset=0
Size=374,61
Goal=sub 2h
Introduction,0,85.48,85.48
Jimmy,0,219.68,134.2
"#;

        let run = parse(RUN, false).unwrap();
        assert_eq!(run.category_name(), "WarioWare, Inc (Goal: sub 2h)");
        assert_eq!(run.len(), 2);
    }

    #[test]
    fn skip_goal_if_its_empty() {
        const RUN: &[u8] = br#"Title=WarioWare, Inc
Attempts=1
Offset=0
Size=374,61
Goal=
Introduction,0,85.48,85.48
Jimmy,0,219.68,134.2
"#;

        let run = parse(RUN, false).unwrap();
        assert_eq!(run.category_name(), "WarioWare, Inc");
        assert_eq!(run.len(), 2);
    }
}
<|endoftext|>


<commit_before>// The MIT License (MIT)

// Copyright (c) 2015 Y. T. Chung <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

//! Processing unit of a thread

use std::cell::UnsafeCell;
use std::io;
#[cfg(target_os = "linux")]
use std::os::unix::io::AsRawFd;
#[cfg(target_os = "linux")]
use std::convert::From;
// use std::sync::Arc;
use std::thread;
use std::mem;
use std::sync::mpsc::{self, Sender, Receiver};
use std::boxed::FnBox;

use mio::{EventLoop, Evented, Handler, Token, EventSet, PollOpt};
use mio::util::Slab;
#[cfg(target_os = "linux")]
use mio::Io;

// use mio::util::BoundedQueue;

use deque::{BufferPool, Stolen, Worker, Stealer};

use rand;

use scheduler::Scheduler;
use coroutine::{self, Coroutine, State, Handle, SendableCoroutinePtr};
use options::Options;

const MAX_TOKEN_NUM: usize = 102400;

thread_local!(static PROCESSOR: UnsafeCell<Processor> = UnsafeCell::new(Processor::new()));

/// Processing unit of a thread
pub struct Processor {
    event_loop: EventLoop<Processor>,

    main_coro: Handle,
    cur_running: Option<*mut Coroutine>,
    last_result: Option<coroutine::Result<State>>,

    queue_worker: Worker<SendableCoroutinePtr>,
    queue_stealer: Stealer<SendableCoroutinePtr>,
    neighbor_stealers: Vec<Stealer<SendableCoroutinePtr>>,

    chan_sender: Sender<ProcMessage>,
    chan_receiver: Receiver<ProcMessage>,

    is_scheduling: bool,
    has_ready_tasks: bool,

    #[cfg(any(target_os = "linux",
              target_os = "android"))]
    io_slabs: Slab<(*mut Coroutine, Io)>,
    #[cfg(any(target_os = "macos",
              target_os = "freebsd",
              target_os = "dragonfly",
              target_os = "ios",
              target_os = "bitrig",
              target_os = "openbsd"))]
    io_slabs: Slab<*mut Coroutine>,

    timer_slabs: Slab<*mut Coroutine>,
}

impl Processor {
    fn new() -> Processor {
        Processor::new_with_neighbors(Vec::new())
    }

    fn new_with_neighbors(neigh: Vec<Stealer<SendableCoroutinePtr>>) -> Processor {
        let main_coro = unsafe {
            Coroutine::empty()
        };

        let (worker, stealer) = BufferPool::new().deque();
        let (tx, rx) = mpsc::channel();

        Processor {
            event_loop: EventLoop::new().expect("Unable to create the eventloop"),

            main_coro: main_coro,
            cur_running: None,
            last_result: None,

            queue_worker: worker,
            queue_stealer: stealer,
            neighbor_stealers: neigh,

            chan_sender: tx,
            chan_receiver: rx,

            is_scheduling: false,
            has_ready_tasks: false,

            io_slabs: Slab::new(MAX_TOKEN_NUM),

            timer_slabs: Slab::new(MAX_TOKEN_NUM),
        }
    }

    #[inline]
    // Get the current running coroutine
    pub unsafe fn running(&mut self) -> Option<*mut Coroutine> {
        self.cur_running
    }

    #[inline]
    pub fn set_neighbors(&mut self, neigh: Vec<Stealer<SendableCoroutinePtr>>) {
        self.neighbor_stealers = neigh;
    }

    /// Get the thread local processor
    #[inline]
    pub fn current() -> &'static mut Processor {
        PROCESSOR.with(|p| unsafe { &mut *p.get() })
    }

    #[inline]
    pub fn stealer(&self) -> Stealer<SendableCoroutinePtr> {
        self.queue_stealer.clone()
    }

    #[inline]
    pub fn handle(&self) -> Sender<ProcMessage> {
        self.chan_sender.clone()
    }

    #[inline]
    // Call by scheduler
    pub unsafe fn ready(&mut self, coro_ptr: *mut Coroutine) {
        self.has_ready_tasks = true;
        self.queue_worker.push(SendableCoroutinePtr(coro_ptr));
    }

    #[inline]
    pub fn spawn_opts<F>(&mut self, f: Box<F>, opts: Options)
        where F: FnBox() + 'static
    {
        let coro = Coroutine::spawn_opts(f, opts);

        unsafe {
            self.ready(mem::transmute(coro));
        }
        self.sched();
    }

    #[inline]
    fn set_last_result(&mut self, r: coroutine::Result<State>) {
        self.last_result = Some(r);
    }

    #[inline]
    unsafe fn run_with_all_local_tasks(&mut self, coro_ptr: *mut Coroutine) {
        let mut hdl = coro_ptr;
        loop {
            let is_suspended = match self.resume(hdl) {
                Ok(State::Suspended) => {
                    true
                },
                Ok(State::Finished) | Err(..) => {
                    Scheduler::finished(hdl);
                    false
                },
                Ok(..) => {
                    false
                }
            };

            // Try to fetch one task from the local queue
            let next_hdl = self.queue_worker.pop();
            if is_suspended {
                // If the current task has to be suspended, then
                // push it back to the local work queue
                self.ready(hdl);
            }

            match next_hdl {
                Some(h) => hdl = h.0,
                None => break
            }
        }
    }

    /// Run the processor
    pub fn schedule(&mut self) -> io::Result<()> {
        self.is_scheduling = true;

        'outerloop:
        loop {
            // 1. Run all tasks in local queue
            if let Some(hdl) = self.queue_worker.pop() {
                unsafe {
                    self.run_with_all_local_tasks(hdl.0);
                }
            } else {
                self.has_ready_tasks = false;
            }

            // 2. Get work from timer
            //    In favor of mio's timer implementation

            // 3. Well, check if there are any works could be waked up
            if self.io_slabs.count() != 0 || self.timer_slabs.count() != 0 {
                // Make the borrow checker happy.
                let proc_ptr: *mut Processor = self;
                if let Err(err) = self.event_loop.run_once(unsafe { &mut *proc_ptr }) {
                    self.is_scheduling = false;
                    error!("EventLoop failed with {:?}", err);
                    return Err(err);
                }

                if self.has_ready_tasks {
                    continue 'outerloop;
                }
            } else if Scheduler::instance().work_count() == 0 {
                break;
            } else {
                // We don't have active tasks in the local queue
                // And we don't have any activated tasks from event loop
            }

            // 4. Randomly steal from neighbors
            let rand_idx = rand::random::<usize>();
            let total_stealers = self.neighbor_stealers.len();
            for idx in (0..self.neighbor_stealers.len()).map(|x| (x + rand_idx) % total_stealers) {
                if let Stolen::Data(hdl) = self.neighbor_stealers[idx].steal() {
                    unsafe {
                        self.run_with_all_local_tasks(hdl.0);
                    }
                    continue 'outerloop;
                }
            }

            // 5. Check the mailbox
            while let Ok(msg) = self.chan_receiver.try_recv() {
                match msg {
                    ProcMessage::NewNeighbor(nei) => self.neighbor_stealers.push(nei),
                }
            }

            thread::sleep_ms(100);
        }

        self.is_scheduling = false;

        Ok(())
    }

    #[inline]
    pub fn resume(&mut self, coro_ptr: *mut Coroutine) -> coroutine::Result<State> {
        self.cur_running = Some(coro_ptr);
        unsafe {
            self.main_coro.yield_to(&mut *coro_ptr);
        }

        match self.last_result.take() {
            None => Ok(State::Suspended),
            Some(r) => r,
        }
    }

    /// Block the current coroutine until the specific time
    #[inline]
    pub fn sleep_ms(&mut self, delay: u64) {
        if let Some(coro_ptr) = unsafe { self.running() } {
            // Set into the timer
            loop {
                match self.timer_slabs.insert(coro_ptr) {
                    Ok(token) => {
                        self.event_loop.timeout_ms(token, delay).unwrap();
                        break;
                    },
                    Err(..) => {}
                }
            }

            // Just block it, the handle has already been sent to
            // the timer.
            self.block();
        } else {
            warn!("Called `wait_until` without running coroutine");
        }
    }

    /// Suspended the current running coroutine, equivalent to `Scheduler::sched`
    #[inline]
    pub fn sched(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ptr) => unsafe {
                self.set_last_result(Ok(State::Suspended));
                (&mut *coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    /// Block the current running coroutine, equivalent to `Scheduler::block`
    #[inline]
    pub fn block(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ptr) => unsafe {
                self.set_last_result(Ok(State::Blocked));
                (&mut *coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    /// Yield the current running coroutine with specified result
    #[inline]
    pub fn yield_with(&mut self, r: coroutine::Result<State>) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ptr) => unsafe {
                self.set_last_result(r);
                (&mut *coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
impl Processor {
    /// Register and wait I/O
    pub fn wait_event<E: Evented + AsRawFd>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.io_slabs.insert((unsafe { Processor::current().running().unwrap() },
                                               From::from(fd.as_raw_fd()))).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        self.block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

#[cfg(any(target_os = "linux",
          target_os = "android"))]
impl Handler for Processor {
    type Timeout = Token;
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.io_slabs.remove(token) {
            Some((hdl, fd)) => {
                // Linux EPoll needs to explicit EPOLL_CTL_DEL the fd
                event_loop.deregister(&fd).unwrap();
                mem::forget(fd);
                unsafe {
                    self.ready(hdl);
                }
            },
            None => {
                warn!("No coroutine is waiting on readable {:?}", token);
            }
        }
    }

    fn timeout(&mut self, _: &mut EventLoop<Self>, token: Token) {
        match self.timer_slabs.remove(token) {
            Some(coro_ptr) => unsafe {
                self.ready(coro_ptr);
            },
            None => {
                warn!("Timer token {:?} was awaited without waiting coroutines", token);
            }
        }
    }
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
impl Processor {
    /// Register and wait I/O
    pub fn wait_event<E: Evented>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.io_slabs.insert(unsafe { Processor::current().running().unwrap() }).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        self.block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

#[cfg(any(target_os = "macos",
          target_os = "freebsd",
          target_os = "dragonfly",
          target_os = "ios",
          target_os = "bitrig",
          target_os = "openbsd"))]
impl Handler for Processor {
    type Timeout = Token;
    type Message = ();

    fn ready(&mut self, _: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.io_slabs.remove(token) {
            Some(hdl) => {
                unsafe {
                    self.ready(hdl);
                }
            },
            None => {
                warn!("No coroutine is waiting on readable {:?}", token);
            }
        }
    }

    fn timeout(&mut self, _: &mut EventLoop<Self>, token: Token) {
        match self.timer_slabs.remove(token) {
            Some(coro_ptr) => unsafe {
                self.ready(coro_ptr);
            },
            None => {
                warn!("Timer token {:?} was awaited without waiting coroutines", token);
            }
        }
    }
}

pub enum ProcMessage {
    NewNeighbor(Stealer<SendableCoroutinePtr>),
}
<commit_msg>combine different impl together<commit_after>// The MIT License (MIT)

// Copyright (c) 2015 Y. T. Chung <zonyitoo@gmail.com>

// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

//! Processing unit of a thread

use std::cell::UnsafeCell;
use std::io;
#[cfg(target_os = "linux")]
use std::os::unix::io::AsRawFd;
#[cfg(target_os = "linux")]
use std::convert::From;
// use std::sync::Arc;
use std::thread;
use std::mem;
use std::sync::mpsc::{self, Sender, Receiver};
use std::boxed::FnBox;

use mio::{EventLoop, Evented, Handler, Token, EventSet, PollOpt};
use mio::util::Slab;
#[cfg(target_os = "linux")]
use mio::Io;

// use mio::util::BoundedQueue;

use deque::{BufferPool, Stolen, Worker, Stealer};

use rand;

use scheduler::Scheduler;
use coroutine::{self, Coroutine, State, Handle, SendableCoroutinePtr};
use options::Options;

const MAX_TOKEN_NUM: usize = 102400;

thread_local!(static PROCESSOR: UnsafeCell<Processor> = UnsafeCell::new(Processor::new()));

/// Processing unit of a thread
pub struct Processor {
    event_loop: EventLoop<Processor>,

    main_coro: Handle,
    cur_running: Option<*mut Coroutine>,
    last_result: Option<coroutine::Result<State>>,

    queue_worker: Worker<SendableCoroutinePtr>,
    queue_stealer: Stealer<SendableCoroutinePtr>,
    neighbor_stealers: Vec<Stealer<SendableCoroutinePtr>>,

    chan_sender: Sender<ProcMessage>,
    chan_receiver: Receiver<ProcMessage>,

    is_scheduling: bool,
    has_ready_tasks: bool,

    #[cfg(any(target_os = "linux",
              target_os = "android"))]
    io_slabs: Slab<(*mut Coroutine, Io)>,
    #[cfg(any(target_os = "macos",
              target_os = "freebsd",
              target_os = "dragonfly",
              target_os = "ios",
              target_os = "bitrig",
              target_os = "openbsd"))]
    io_slabs: Slab<*mut Coroutine>,

    timer_slabs: Slab<*mut Coroutine>,
}

impl Processor {
    fn new() -> Processor {
        Processor::new_with_neighbors(Vec::new())
    }

    fn new_with_neighbors(neigh: Vec<Stealer<SendableCoroutinePtr>>) -> Processor {
        let main_coro = unsafe {
            Coroutine::empty()
        };

        let (worker, stealer) = BufferPool::new().deque();
        let (tx, rx) = mpsc::channel();

        Processor {
            event_loop: EventLoop::new().expect("Unable to create the eventloop"),

            main_coro: main_coro,
            cur_running: None,
            last_result: None,

            queue_worker: worker,
            queue_stealer: stealer,
            neighbor_stealers: neigh,

            chan_sender: tx,
            chan_receiver: rx,

            is_scheduling: false,
            has_ready_tasks: false,

            io_slabs: Slab::new(MAX_TOKEN_NUM),

            timer_slabs: Slab::new(MAX_TOKEN_NUM),
        }
    }

    #[inline]
    // Get the current running coroutine
    pub unsafe fn running(&mut self) -> Option<*mut Coroutine> {
        self.cur_running
    }

    #[inline]
    pub fn set_neighbors(&mut self, neigh: Vec<Stealer<SendableCoroutinePtr>>) {
        self.neighbor_stealers = neigh;
    }

    /// Get the thread local processor
    #[inline]
    pub fn current() -> &'static mut Processor {
        PROCESSOR.with(|p| unsafe { &mut *p.get() })
    }

    #[inline]
    pub fn stealer(&self) -> Stealer<SendableCoroutinePtr> {
        self.queue_stealer.clone()
    }

    #[inline]
    pub fn handle(&self) -> Sender<ProcMessage> {
        self.chan_sender.clone()
    }

    #[inline]
    // Call by scheduler
    pub unsafe fn ready(&mut self, coro_ptr: *mut Coroutine) {
        self.has_ready_tasks = true;
        self.queue_worker.push(SendableCoroutinePtr(coro_ptr));
    }

    #[inline]
    pub fn spawn_opts<F>(&mut self, f: Box<F>, opts: Options)
        where F: FnBox() + 'static
    {
        let coro = Coroutine::spawn_opts(f, opts);

        unsafe {
            self.ready(mem::transmute(coro));
        }
        self.sched();
    }

    #[inline]
    fn set_last_result(&mut self, r: coroutine::Result<State>) {
        self.last_result = Some(r);
    }

    #[inline]
    unsafe fn run_with_all_local_tasks(&mut self, coro_ptr: *mut Coroutine) {
        let mut hdl = coro_ptr;
        loop {
            let is_suspended = match self.resume(hdl) {
                Ok(State::Suspended) => {
                    true
                },
                Ok(State::Finished) | Err(..) => {
                    Scheduler::finished(hdl);
                    false
                },
                Ok(..) => {
                    false
                }
            };

            // Try to fetch one task from the local queue
            let next_hdl = self.queue_worker.pop();
            if is_suspended {
                // If the current task has to be suspended, then
                // push it back to the local work queue
                self.ready(hdl);
            }

            match next_hdl {
                Some(h) => hdl = h.0,
                None => break
            }
        }
    }

    /// Run the processor
    pub fn schedule(&mut self) -> io::Result<()> {
        self.is_scheduling = true;

        'outerloop:
        loop {
            // 1. Run all tasks in local queue
            if let Some(hdl) = self.queue_worker.pop() {
                unsafe {
                    self.run_with_all_local_tasks(hdl.0);
                }
            } else {
                self.has_ready_tasks = false;
            }

            // 2. Get work from timer
            //    In favor of mio's timer implementation

            // 3. Well, check if there are any works could be waked up
            if self.io_slabs.count() != 0 || self.timer_slabs.count() != 0 {
                // Make the borrow checker happy.
                let proc_ptr: *mut Processor = self;
                if let Err(err) = self.event_loop.run_once(unsafe { &mut *proc_ptr }) {
                    self.is_scheduling = false;
                    error!("EventLoop failed with {:?}", err);
                    return Err(err);
                }

                if self.has_ready_tasks {
                    continue 'outerloop;
                }
            } else if Scheduler::instance().work_count() == 0 {
                break;
            } else {
                // We don't have active tasks in the local queue
                // And we don't have any activated tasks from event loop
            }

            // 4. Randomly steal from neighbors
            let rand_idx = rand::random::<usize>();
            let total_stealers = self.neighbor_stealers.len();
            for idx in (0..self.neighbor_stealers.len()).map(|x| (x + rand_idx) % total_stealers) {
                if let Stolen::Data(hdl) = self.neighbor_stealers[idx].steal() {
                    unsafe {
                        self.run_with_all_local_tasks(hdl.0);
                    }
                    continue 'outerloop;
                }
            }

            // 5. Check the mailbox
            while let Ok(msg) = self.chan_receiver.try_recv() {
                match msg {
                    ProcMessage::NewNeighbor(nei) => self.neighbor_stealers.push(nei),
                }
            }

            thread::sleep_ms(100);
        }

        self.is_scheduling = false;

        Ok(())
    }

    #[inline]
    pub fn resume(&mut self, coro_ptr: *mut Coroutine) -> coroutine::Result<State> {
        self.cur_running = Some(coro_ptr);
        unsafe {
            self.main_coro.yield_to(&mut *coro_ptr);
        }

        match self.last_result.take() {
            None => Ok(State::Suspended),
            Some(r) => r,
        }
    }

    /// Block the current coroutine until the specific time
    #[inline]
    pub fn sleep_ms(&mut self, delay: u64) {
        if let Some(coro_ptr) = unsafe { self.running() } {
            // Set into the timer
            loop {
                match self.timer_slabs.insert(coro_ptr) {
                    Ok(token) => {
                        self.event_loop.timeout_ms(token, delay).unwrap();
                        break;
                    },
                    Err(..) => {}
                }
            }

            // Just block it, the handle has already been sent to
            // the timer.
            self.block();
        } else {
            warn!("Called `wait_until` without running coroutine");
        }
    }

    /// Suspended the current running coroutine, equivalent to `Scheduler::sched`
    #[inline]
    pub fn sched(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ptr) => unsafe {
                self.set_last_result(Ok(State::Suspended));
                (&mut *coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    /// Block the current running coroutine, equivalent to `Scheduler::block`
    #[inline]
    pub fn block(&mut self) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ptr) => unsafe {
                self.set_last_result(Ok(State::Blocked));
                (&mut *coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }

    /// Yield the current running coroutine with specified result
    #[inline]
    pub fn yield_with(&mut self, r: coroutine::Result<State>) {
        match self.cur_running.take() {
            None => {},
            Some(coro_ptr) => unsafe {
                self.set_last_result(r);
                (&mut *coro_ptr).yield_to(&*self.main_coro)
            }
        }
    }
}

impl Processor {
    /// Register and wait I/O
    #[cfg(any(target_os = "linux",
              target_os = "android"))]
    pub fn wait_event<E: Evented + AsRawFd>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.io_slabs.insert((unsafe { Processor::current().running().unwrap() },
                                               From::from(fd.as_raw_fd()))).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        self.block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }

    #[cfg(any(target_os = "macos",
              target_os = "freebsd",
              target_os = "dragonfly",
              target_os = "ios",
              target_os = "bitrig",
              target_os = "openbsd"))]
    pub fn wait_event<E: Evented>(&mut self, fd: &E, interest: EventSet) -> io::Result<()> {
        let token = self.io_slabs.insert(unsafe { Processor::current().running().unwrap() }).unwrap();
        try!(self.event_loop.register_opt(fd, token, interest,
                                          PollOpt::edge()|PollOpt::oneshot()));

        debug!("wait_event: Blocked current Coroutine ...; token={:?}", token);
        self.block();
        debug!("wait_event: Waked up; token={:?}", token);

        Ok(())
    }
}

impl Handler for Processor {
    type Timeout = Token;
    type Message = ();

    #[cfg(any(target_os = "linux",
              target_os = "android"))]
    fn ready(&mut self, event_loop: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.io_slabs.remove(token) {
            Some((hdl, fd)) => {
                // Linux EPoll needs to explicit EPOLL_CTL_DEL the fd
                event_loop.deregister(&fd).unwrap();
                mem::forget(fd);
                unsafe {
                    self.ready(hdl);
                }
            },
            None => {
                warn!("No coroutine is waiting on token {:?}", token);
            }
        }
    }

    #[cfg(any(target_os = "macos",
              target_os = "freebsd",
              target_os = "dragonfly",
              target_os = "ios",
              target_os = "bitrig",
              target_os = "openbsd"))]
    fn ready(&mut self, _: &mut EventLoop<Self>, token: Token, events: EventSet) {
        debug!("Got {:?} for {:?}", events, token);

        match self.io_slabs.remove(token) {
            Some(hdl) => {
                unsafe {
                    self.ready(hdl);
                }
            },
            None => {
                warn!("No coroutine is waiting on token {:?}", token);
            }
          }
    }

    fn timeout(&mut self, _: &mut EventLoop<Self>, token: Token) {
        match self.timer_slabs.remove(token) {
            Some(coro_ptr) => unsafe {
                self.ready(coro_ptr);
            },
            None => {
                warn!("Timer token {:?} was awaited without waiting coroutines", token);
            }
        }
    }
}

pub enum ProcMessage {
    NewNeighbor(Stealer<SendableCoroutinePtr>),
}
<|endoftext|>


<commit_before>//! Fairing implementation
use std::str::FromStr;

use rocket::{self, Request, Outcome};
use rocket::http::{self, Status, Header};

use {Cors, Error, validate, preflight_response, actual_request_response, origin, request_headers};

/// An injected header to quickly give the result of CORS
static CORS_HEADER: &str = "ROCKET-CORS";
enum InjectedHeader {
    Success,
    Failure,
}

impl InjectedHeader {
    fn to_str(&self) -> &'static str {
        match *self {
            InjectedHeader::Success => "Success",
            InjectedHeader::Failure => "Failure",
        }
    }
}

impl FromStr for InjectedHeader {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self, Error> {
        match s {
            "Success" => Ok(InjectedHeader::Success),
            "Failure" => Ok(InjectedHeader::Failure),
            other => {
                error_!(
                    "Unknown injected header encountered: {}\nThis is probably a bug.",
                    other
                );
                Err(Error::UnknownInjectedHeader)
            }
        }
    }
}

/// Route for Fairing error handling
pub(crate) fn fairing_error_route<'r>(
    request: &'r Request,
    _: rocket::Data,
) -> rocket::handler::Outcome<'r> {
    let status = request.get_param::<u16>(0).unwrap_or_else(|e| {
        error_!("Fairing Error Handling Route error: {:?}", e);
        500
    });
    let status = Status::from_code(status).unwrap_or_else(|| Status::InternalServerError);
    Outcome::Failure(status)
}

/// Create a new `Route` for Fairing handling
fn fairing_route() -> rocket::Route {
    rocket::Route::new(http::Method::Get, "/<status>", fairing_error_route)
}

/// Modifies a `Request` to route to Fairing error handler
fn route_to_fairing_error_handler(options: &Cors, status: u16, request: &mut Request) {
    request.set_method(http::Method::Get);
    request.set_uri(format!("{}/{}", options.fairing_route_base, status));
}

/// Inject a header into the Request with result
fn inject_request_header(header: &InjectedHeader, request: &mut Request) {
    request.replace_header(Header::new(CORS_HEADER, header.to_str()));
}

fn on_response_wrapper(
    options: &Cors,
    request: &Request,
    response: &mut rocket::Response,
) -> Result<(), Error> {
    let origin = match origin(request)? {
        None => {
            // Not a CORS request
            return Ok(());
        }
        Some(origin) => origin,
    };

    // Get validation result from injected header
    let injected_header = request.headers().get_one(CORS_HEADER).ok_or_else(|| {
        Error::MissingInjectedHeader
    })?;
    let result = InjectedHeader::from_str(injected_header)?;

    if let InjectedHeader::Failure = result {
        // Nothing else for us to do
        return Ok(());
    }

    let cors_response = if request.method() == http::Method::Options {
        let headers = request_headers(request)?;
        preflight_response(options, &origin, headers.as_ref())
    } else {
        actual_request_response(options, &origin)
    };

    cors_response.merge(response);

    // If this was an OPTIONS request and no route can be found, we should turn this
    // into a HTTP 204 with no content body.
    // This allows the user to not have to specify an OPTIONS route for everything.
    //
    // TODO: Is there anyway we can make this smarter? Only modify status codes for
    // requests where an actual route exist?
    if request.method() == http::Method::Options && request.method() == http::Method::Options &&
        request.route().is_none()
    {
        info_!(
            "CORS Fairing: Turned missing route {} into an OPTIONS pre-flight request",
            request
        );
        response.set_status(Status::NoContent);
        let _ = response.take_body();
    }
    Ok(())
}

impl rocket::fairing::Fairing for Cors {
    fn info(&self) -> rocket::fairing::Info {
        rocket::fairing::Info {
            name: "CORS",
            kind: rocket::fairing::Kind::Attach | rocket::fairing::Kind::Request |
                rocket::fairing::Kind::Response,
        }
    }

    fn on_attach(&self, rocket: rocket::Rocket) -> Result<rocket::Rocket, rocket::Rocket> {
        match self.validate() {
            Ok(()) => {
                Ok(rocket.mount(&self.fairing_route_base, vec![fairing_route()]))
            }
            Err(e) => {
                error_!("Error attaching CORS fairing: {}", e);
                Err(rocket)
            }
        }
    }

    fn on_request(&self, request: &mut Request, _: &rocket::Data) {
        let injected_header = match validate(self, request) {
            Ok(_) => InjectedHeader::Success,
            Err(err) => {
                error_!("CORS Error: {}", err);
                let status = err.status();
                route_to_fairing_error_handler(self, status.code, request);
                InjectedHeader::Failure
            }
        };

        inject_request_header(&injected_header, request);
    }

    fn on_response(&self, request: &Request, response: &mut rocket::Response) {
        if let Err(err) = on_response_wrapper(self, request, response) {
            error_!("Fairings on_response error: {}\nMost likely a bug", err);
            response.set_status(Status::InternalServerError);
            let _ = response.take_body();
        }
    }
}

#[cfg(test)]
mod tests {
    use rocket::Rocket;
    use rocket::http::{Method, Status};
    use rocket::local::Client;

    use {Cors, AllOrSome, AllowedOrigins, AllowedHeaders};

    const CORS_ROOT: &'static str = "/my_cors";

    fn make_cors_options() -> Cors {
        let (allowed_origins, failed_origins) = AllowedOrigins::some(&["https://www.acme.com"]);
        assert!(failed_origins.is_empty());

        Cors {
            allowed_origins: allowed_origins,
            allowed_methods: vec![Method::Get].into_iter().map(From::from).collect(),
            allowed_headers: AllowedHeaders::some(&["Authorization", "Accept"]),
            allow_credentials: true,
            fairing_route_base: CORS_ROOT.to_string(),

            ..Default::default()
        }
    }

    fn rocket(fairing: Cors) -> Rocket {
        Rocket::ignite().attach(fairing)
    }

    #[test]
    fn fairing_error_route_returns_passed_in_status() {
        let client = Client::new(rocket(make_cors_options())).expect("to not fail");
        let request = client.get(format!("{}/403", CORS_ROOT));
        let response = request.dispatch();
        assert_eq!(Status::Forbidden, response.status());
    }

    #[test]
    fn fairing_error_route_returns_500_for_unknown_status() {
        let client = Client::new(rocket(make_cors_options())).expect("to not fail");
        let request = client.get(format!("{}/999", CORS_ROOT));
        let response = request.dispatch();
        assert_eq!(Status::InternalServerError, response.status());
    }

    #[test]
    fn error_route_is_mounted_on_attach() {
        let rocket = rocket(make_cors_options());

        let expected_uri = format!("{}/<status>", CORS_ROOT);
        let error_route = rocket.routes().find(|r| {
            r.method == Method::Get && r.uri.as_str() == expected_uri
        });
        assert!(error_route.is_some());
    }

    #[test]
    #[should_panic(expected = "launch fairing failure")]
    fn options_are_validated_on_attach() {
        let mut options = make_cors_options();
        options.allowed_origins = AllOrSome::All;
        options.send_wildcard = true;

        let _ = rocket(options).launch();
    }

    // Rest of the things can only be tested in integration tests
}
<commit_msg>update fairing.rs removed duplicate condition (#27)<commit_after>//! Fairing implementation
use std::str::FromStr;

use rocket::{self, Request, Outcome};
use rocket::http::{self, Status, Header};

use {Cors, Error, validate, preflight_response, actual_request_response, origin, request_headers};

/// An injected header to quickly give the result of CORS
static CORS_HEADER: &str = "ROCKET-CORS";
enum InjectedHeader {
    Success,
    Failure,
}

impl InjectedHeader {
    fn to_str(&self) -> &'static str {
        match *self {
            InjectedHeader::Success => "Success",
            InjectedHeader::Failure => "Failure",
        }
    }
}

impl FromStr for InjectedHeader {
    type Err = Error;
    fn from_str(s: &str) -> Result<Self, Error> {
        match s {
            "Success" => Ok(InjectedHeader::Success),
            "Failure" => Ok(InjectedHeader::Failure),
            other => {
                error_!(
                    "Unknown injected header encountered: {}\nThis is probably a bug.",
                    other
                );
                Err(Error::UnknownInjectedHeader)
            }
        }
    }
}

/// Route for Fairing error handling
pub(crate) fn fairing_error_route<'r>(
    request: &'r Request,
    _: rocket::Data,
) -> rocket::handler::Outcome<'r> {
    let status = request.get_param::<u16>(0).unwrap_or_else(|e| {
        error_!("Fairing Error Handling Route error: {:?}", e);
        500
    });
    let status = Status::from_code(status).unwrap_or_else(|| Status::InternalServerError);
    Outcome::Failure(status)
}

/// Create a new `Route` for Fairing handling
fn fairing_route() -> rocket::Route {
    rocket::Route::new(http::Method::Get, "/<status>", fairing_error_route)
}

/// Modifies a `Request` to route to Fairing error handler
fn route_to_fairing_error_handler(options: &Cors, status: u16, request: &mut Request) {
    request.set_method(http::Method::Get);
    request.set_uri(format!("{}/{}", options.fairing_route_base, status));
}

/// Inject a header into the Request with result
fn inject_request_header(header: &InjectedHeader, request: &mut Request) {
    request.replace_header(Header::new(CORS_HEADER, header.to_str()));
}

fn on_response_wrapper(
    options: &Cors,
    request: &Request,
    response: &mut rocket::Response,
) -> Result<(), Error> {
    let origin = match origin(request)? {
        None => {
            // Not a CORS request
            return Ok(());
        }
        Some(origin) => origin,
    };

    // Get validation result from injected header
    let injected_header = request.headers().get_one(CORS_HEADER).ok_or_else(|| {
        Error::MissingInjectedHeader
    })?;
    let result = InjectedHeader::from_str(injected_header)?;

    if let InjectedHeader::Failure = result {
        // Nothing else for us to do
        return Ok(());
    }

    let cors_response = if request.method() == http::Method::Options {
        let headers = request_headers(request)?;
        preflight_response(options, &origin, headers.as_ref())
    } else {
        actual_request_response(options, &origin)
    };

    cors_response.merge(response);

    // If this was an OPTIONS request and no route can be found, we should turn this
    // into a HTTP 204 with no content body.
    // This allows the user to not have to specify an OPTIONS route for everything.
    //
    // TODO: Is there anyway we can make this smarter? Only modify status codes for
    // requests where an actual route exist?
    if request.method() == http::Method::Options &&
        request.route().is_none()
    {
        info_!(
            "CORS Fairing: Turned missing route {} into an OPTIONS pre-flight request",
            request
        );
        response.set_status(Status::NoContent);
        let _ = response.take_body();
    }
    Ok(())
}

impl rocket::fairing::Fairing for Cors {
    fn info(&self) -> rocket::fairing::Info {
        rocket::fairing::Info {
            name: "CORS",
            kind: rocket::fairing::Kind::Attach | rocket::fairing::Kind::Request |
                rocket::fairing::Kind::Response,
        }
    }

    fn on_attach(&self, rocket: rocket::Rocket) -> Result<rocket::Rocket, rocket::Rocket> {
        match self.validate() {
            Ok(()) => {
                Ok(rocket.mount(&self.fairing_route_base, vec![fairing_route()]))
            }
            Err(e) => {
                error_!("Error attaching CORS fairing: {}", e);
                Err(rocket)
            }
        }
    }

    fn on_request(&self, request: &mut Request, _: &rocket::Data) {
        let injected_header = match validate(self, request) {
            Ok(_) => InjectedHeader::Success,
            Err(err) => {
                error_!("CORS Error: {}", err);
                let status = err.status();
                route_to_fairing_error_handler(self, status.code, request);
                InjectedHeader::Failure
            }
        };

        inject_request_header(&injected_header, request);
    }

    fn on_response(&self, request: &Request, response: &mut rocket::Response) {
        if let Err(err) = on_response_wrapper(self, request, response) {
            error_!("Fairings on_response error: {}\nMost likely a bug", err);
            response.set_status(Status::InternalServerError);
            let _ = response.take_body();
        }
    }
}

#[cfg(test)]
mod tests {
    use rocket::Rocket;
    use rocket::http::{Method, Status};
    use rocket::local::Client;

    use {Cors, AllOrSome, AllowedOrigins, AllowedHeaders};

    const CORS_ROOT: &'static str = "/my_cors";

    fn make_cors_options() -> Cors {
        let (allowed_origins, failed_origins) = AllowedOrigins::some(&["https://www.acme.com"]);
        assert!(failed_origins.is_empty());

        Cors {
            allowed_origins: allowed_origins,
            allowed_methods: vec![Method::Get].into_iter().map(From::from).collect(),
            allowed_headers: AllowedHeaders::some(&["Authorization", "Accept"]),
            allow_credentials: true,
            fairing_route_base: CORS_ROOT.to_string(),

            ..Default::default()
        }
    }

    fn rocket(fairing: Cors) -> Rocket {
        Rocket::ignite().attach(fairing)
    }

    #[test]
    fn fairing_error_route_returns_passed_in_status() {
        let client = Client::new(rocket(make_cors_options())).expect("to not fail");
        let request = client.get(format!("{}/403", CORS_ROOT));
        let response = request.dispatch();
        assert_eq!(Status::Forbidden, response.status());
    }

    #[test]
    fn fairing_error_route_returns_500_for_unknown_status() {
        let client = Client::new(rocket(make_cors_options())).expect("to not fail");
        let request = client.get(format!("{}/999", CORS_ROOT));
        let response = request.dispatch();
        assert_eq!(Status::InternalServerError, response.status());
    }

    #[test]
    fn error_route_is_mounted_on_attach() {
        let rocket = rocket(make_cors_options());

        let expected_uri = format!("{}/<status>", CORS_ROOT);
        let error_route = rocket.routes().find(|r| {
            r.method == Method::Get && r.uri.as_str() == expected_uri
        });
        assert!(error_route.is_some());
    }

    #[test]
    #[should_panic(expected = "launch fairing failure")]
    fn options_are_validated_on_attach() {
        let mut options = make_cors_options();
        options.allowed_origins = AllOrSome::All;
        options.send_wildcard = true;

        let _ = rocket(options).launch();
    }

    // Rest of the things can only be tested in integration tests
}
<|endoftext|>


<commit_before>extern crate irc;
use self::irc::client::prelude::*;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use chat::ChatMessage;
use chat::ReplyMessage;
use std::path::Path;
use std::thread::sleep;
use std::time::Duration;

extern crate itertools;
use self::itertools::Itertools;

const IRC_RECONNECT_WAIT: u64 = 1;

pub fn connect<P>(sender: Sender<ReplyMessage>, path: P)
    where P: AsRef<Path>
{
    let config = Config::load(path)
        .unwrap_or_else(|e| panic!("IRC Fatal: Couldn't load config file: {}", e));
    loop {
        let server = IrcServer::from_config(config.clone())
            .unwrap_or_else(|e| panic!("IRC Fatal: Couldn't make server: {}", e));
        server.identify().unwrap_or_else(|e| panic!("IRC Fatal: Failed to identify server: {}", e));
        for message in server.iter() {
            let message = message.unwrap_or_else(|e| panic!("IRC Fatal: Failed to get message: {}", e));
            match message.command {
                Command::PRIVMSG(target, msg) => {
                    let name = message.prefix.unwrap_or_else(|| panic!("IRC Fatal: Unable to get msg name"));
                    let nick = server.config().nickname().to_string();
                    let chat_message = ChatMessage {
                        source: target.clone(),
                        author: name,
                        message: msg.split(' ').filter(|e| e == &nick).join(" "),
                    };
                    if msg.contains(&nick) {
                        let (reply_sender, reply_reciever) = channel();
                        sender.send(ReplyMessage(chat_message, Some(reply_sender)))
                            .unwrap_or_else(|e| panic!("IRC Fatal: Message sender closed: {}", e));
                        let reply_message = reply_reciever.recv()
                            .unwrap_or_else(|e| panic!("IRC Fatal: Reply receiver failed: {}", e));
                        if let Some(reply_string) = reply_message {
                            if let Some(fc) = reply_string.chars().next() {
                                if fc != '.' && fc != '/' {
                                    server.send_privmsg(target.as_str(), reply_string.as_str())
                                        .unwrap_or_else(|e| panic!("IRC Fatal: Failed to send message: {}", e));
                                }
                            }
                        }
                    } else {
                        sender.send(ReplyMessage(chat_message, None))
                            .unwrap_or_else(|e| panic!("IRC Fatal: Message sender closed: {}", e));
                    }
                },
                _ => {},
            }
        }
        // Connection ended, so we need to wait an amount of time before trying again
        sleep(Duration::from_secs(IRC_RECONNECT_WAIT));
    }
}
<commit_msg>umm, remove everything BUT its name...oops!<commit_after>extern crate irc;
use self::irc::client::prelude::*;
use std::sync::mpsc::channel;
use std::sync::mpsc::Sender;
use chat::ChatMessage;
use chat::ReplyMessage;
use std::path::Path;
use std::thread::sleep;
use std::time::Duration;

extern crate itertools;
use self::itertools::Itertools;

const IRC_RECONNECT_WAIT: u64 = 1;

pub fn connect<P>(sender: Sender<ReplyMessage>, path: P)
    where P: AsRef<Path>
{
    let config = Config::load(path)
        .unwrap_or_else(|e| panic!("IRC Fatal: Couldn't load config file: {}", e));
    loop {
        let server = IrcServer::from_config(config.clone())
            .unwrap_or_else(|e| panic!("IRC Fatal: Couldn't make server: {}", e));
        server.identify().unwrap_or_else(|e| panic!("IRC Fatal: Failed to identify server: {}", e));
        for message in server.iter() {
            let message = message.unwrap_or_else(|e| panic!("IRC Fatal: Failed to get message: {}", e));
            match message.command {
                Command::PRIVMSG(target, msg) => {
                    let name = message.prefix.unwrap_or_else(|| panic!("IRC Fatal: Unable to get msg name"));
                    let nick = server.config().nickname().to_string();
                    let chat_message = ChatMessage {
                        source: target.clone(),
                        author: name,
                        message: msg.split(' ').filter(|e| e != &nick).join(" "),
                    };
                    if msg.contains(&nick) {
                        let (reply_sender, reply_reciever) = channel();
                        sender.send(ReplyMessage(chat_message, Some(reply_sender)))
                            .unwrap_or_else(|e| panic!("IRC Fatal: Message sender closed: {}", e));
                        let reply_message = reply_reciever.recv()
                            .unwrap_or_else(|e| panic!("IRC Fatal: Reply receiver failed: {}", e));
                        if let Some(reply_string) = reply_message {
                            if let Some(fc) = reply_string.chars().next() {
                                if fc != '.' && fc != '/' {
                                    server.send_privmsg(target.as_str(), reply_string.as_str())
                                        .unwrap_or_else(|e| panic!("IRC Fatal: Failed to send message: {}", e));
                                }
                            }
                        }
                    } else {
                        sender.send(ReplyMessage(chat_message, None))
                            .unwrap_or_else(|e| panic!("IRC Fatal: Message sender closed: {}", e));
                    }
                },
                _ => {},
            }
        }
        // Connection ended, so we need to wait an amount of time before trying again
        sleep(Duration::from_secs(IRC_RECONNECT_WAIT));
    }
}
<|endoftext|>