<commit_before>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! # Tagged stream
//!
//! A tagged stream recursively applies tags to regions of an input stream. A stream consisting of a single tag forms a tree,
//! so this stream type is useful for representing the parsed represention of an input.
//!
//! The simplest form of tagged stream is one that's just a set of untagged symbols. This can be created using the `TaggedStream::from_reader`
//! method.
//!
//! ```
//! # use concordance::*;
//! #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
//! enum SomeTag {
//!     Identifier
//! };
//!
//! let simple_tagged = TaggedStream::from_reader(&mut "Hello, world".read_symbols());
//! let with_tag      = simple_tagged.with_tags(vec![(0..5, SomeTag::Identifier)].iter().cloned());
//! ```
//!

// TODO: a possibly better way to design this would be to make it so that there's an underlying stream we read from when generating tag
// symbols so that we don't need to store the contents of the stream in memory

use std::slice::Iter;
use std::ops::Index;
use std::ops::Range;

use super::countable::*;
use super::symbol_reader::*;
use super::tokenizer::*;
use super::symbol_range_dfa::*;

///
/// Represents a symbol in a tagged stream.
///
/// A symbol can either be a sequence of untagged base symbols, or it can be a tag (which itself can contain any number of tagged symbol)
///
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum TagSymbol<Base: Clone+Ord, Tag: Clone+Ord> {
    /// An untagged element in this stream
    Untagged(Base),

    /// A tagged region of the stream
    Tagged(Tag, TaggedStream<Base, Tag>)
}

use TagSymbol::*;

///
/// Represents a stream of tagged symbols
///
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct TaggedStream<Base: Clone+Ord, Tag: Clone+Ord> {
    /// The data in this stream
    data: Vec<TagSymbol<Base, Tag>>
}

impl<Base: Ord+Clone, Tag: Ord+Clone+'static> TaggedStream<Base, Tag> {
    ///
    /// Creates a basic tagged stream from a source of the base symbol
    ///
    pub fn from_reader(reader: &mut SymbolReader<Base>) -> TaggedStream<Base, Tag> {
        // Read from the reader into an array
        let mut symbols = vec![];

        while let Some(next_symbol) = reader.next_symbol() {
            symbols.push(Untagged(next_symbol));
        }

        // Generate a simple tagged stream from the result
        TaggedStream { data: symbols }
    }

    ///
    /// The number of symbols in this stream
    ///
    pub fn len(&self) -> usize {
        self.data.len()
    }

    ///
    /// Replaces a range in this stream with a tag
    ///
    pub fn tag(&mut self, tag: Tag, range: Range<usize>) {
        // Create a tag to replace the range
        let replaced_symbols    = self.data[range.clone()].to_vec();
        let tag_symbol          = Tagged(tag, TaggedStream { data: replaced_symbols });

        // Draining seems to be for reading a range but does double duty for deleting a range?
        // I don't think rust has a way to replace a range in a vector, or at least not one that's easy to find in the docs.
        self.data.drain(range.clone());

        // Draining then inserting is inefficient compared to flat out replacing items :-/ This may be possible but the vec docs aren't very easy to read
        self.data.insert(range.start, tag_symbol);
    }

    ///
    /// Creates a tag symbol by tagging a particular range within this stream
    ///
    pub fn tag_range(&self, range: Range<usize>, tag: Tag) -> TagSymbol<Base, Tag> {
        let tag_data = self.data[range].to_vec();

        Tagged(tag, TaggedStream { data: tag_data })
    }

    ///
    /// Applies tags to the specified ranges in this stream
    ///
    /// The tag ranges must be in ascending order and must not overlap
    ///
    pub fn with_tags<I>(&self, tags: I) -> TaggedStream<Base, Tag>
        where I : Iterator<Item=(Range<usize>, Tag)> {
        // The data that will make up the new stream
        let mut new_stream = vec![];

        // The last processed range (we haven't processed any range yet, so this starts at 0..0)
        let mut last_range = 0..0;

        // Build up the result from the tag list and the last range
        for tag in tags {
            // Fetch the information for this tag and check ordering
            let (next_range, next_tag)  = tag;
            let end                     = last_range.end;

            if end > next_range.start {
                panic!("Tags for with_tags must be sorted in order");
            }

            // Sometimes need to leave some data from this stream untagged
            if end < next_range.start {
                new_stream.extend(self.data[end..next_range.start].iter().cloned());
            }

            // Push the tag
            new_stream.push(self.tag_range(next_range.clone(), next_tag));

            // Update state
            last_range = next_range;
        }

        // Append anything left in the stream
        if last_range.end < self.data.len() {
            new_stream.extend(self.data[last_range.end..self.data.len()].iter().cloned());
        }

        // Final result
        TaggedStream { data: new_stream }
    }

    ///
    /// Runs the current values of this tagged stream through a tokenizer and tags anything it matches
    ///
    /// This takes a mapping function to describe how symbols in this stream map to symbols in the DFA. Note that every symbol
    /// (tagged or untagged) must be mapped to a DFA symbol, so if only tagged or untagged symbols are being used it's necessary
    /// to decide how the other symbols are mapped (eg, to an unused symbol)
    ///
    pub fn tokenize<DfaSymbol: Ord+Countable+Clone, MapFn>(&self, token_matcher: &SymbolRangeDfa<DfaSymbol, Tag>, map_symbol: MapFn) -> TaggedStream<Base, Tag>
        where MapFn: FnMut(TagSymbol<Base, Tag>) -> DfaSymbol {
        // Generate a symbol reader with the mapping function
        let reader = self.data.read_symbols().map_symbols(map_symbol);

        // Tokenize it
        let tokenizer = Tokenizer::new_prepared(reader, token_matcher);

        // Use the tokenizer to generate the result
        self.with_tags(tokenizer)
    }
}

impl<Base: Clone+Ord, Tag: Clone+Ord> Index<usize> for TaggedStream<Base, Tag> {
    type Output = TagSymbol<Base, Tag>;

    fn index(&self, index: usize) -> &TagSymbol<Base, Tag> {
        &self.data[index]
    }
}

impl<'a, Base: Clone+Ord, Tag: Clone+Ord> SymbolSource<'a, TagSymbol<Base, Tag>> for &'a TaggedStream<Base, Tag> {
    type SymbolReader = Iter<'a, TagSymbol<Base, Tag>>;

    /// Returns a new object that can read the symbols from this one
    fn read_symbols(self) -> Self::SymbolReader {
        self.data.read_symbols()
    }
}

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

    #[test]
    fn can_tag_range() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello,
            World
        }

        let mut tagged: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());

        tagged.tag(Tags::Hello, 0..5);
        tagged.tag(Tags::World, 1..6);

        assert!(tagged.len() == 2);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('o'));
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[1] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('W'));
            assert!(stream[1] == Untagged('o'));
            assert!(stream[2] == Untagged('r'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('d'));
        } else {
            assert!(false);
        }
    }

    #[test]
    fn can_tag_everything_with_tags() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello,
            World
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(0..5, Tags::Hello), (5..10, Tags::World)].iter().cloned());

        assert!(tagged.len() == 2);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('o'));
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[1] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('W'));
            assert!(stream[1] == Untagged('o'));
            assert!(stream[2] == Untagged('r'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('d'));
        } else {
            assert!(false);
        }
    }

    #[test]
    fn with_tags_preserves_middle() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello,
            World
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(0..4, Tags::Hello), (6..10, Tags::World)].iter().cloned());

        assert!(tagged.len() == 4);
        assert!(tagged[1] == Untagged('o'));
        assert!(tagged[2] == Untagged('W'));

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 4);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[3] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 4);
            assert!(stream[0] == Untagged('o'));
            assert!(stream[1] == Untagged('r'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('d'));
        } else {
            assert!(false);
        }
    }

    #[test]
    fn with_tags_preserves_end() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(0..5, Tags::Hello)].iter().cloned());

        assert!(tagged.len() == 6);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('o'));
        } else {
            assert!(false);
        }

        assert!(tagged[1] == Untagged('W'));
        assert!(tagged[2] == Untagged('o'));
        assert!(tagged[3] == Untagged('r'));
        assert!(tagged[4] == Untagged('l'));
        assert!(tagged[5] == Untagged('d'));
    }

    #[test]
    fn with_tags_preserves_start() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            World
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(5..10, Tags::World)].iter().cloned());

        assert!(tagged.len() == 6);

        if let Tagged(ref tag, ref stream) = tagged[5] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('W'));
            assert!(stream[1] == Untagged('o'));
            assert!(stream[2] == Untagged('r'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('d'));
        } else {
            assert!(false);
        }

        assert!(tagged[0] == Untagged('H'));
        assert!(tagged[1] == Untagged('e'));
        assert!(tagged[2] == Untagged('l'));
        assert!(tagged[3] == Untagged('l'));
        assert!(tagged[4] == Untagged('o'));
    }
}
<commit_msg>Test the tokenizer<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

//!
//! # Tagged stream
//!
//! A tagged stream recursively applies tags to regions of an input stream. A stream consisting of a single tag forms a tree,
//! so this stream type is useful for representing the parsed represention of an input.
//!
//! The simplest form of tagged stream is one that's just a set of untagged symbols. This can be created using the `TaggedStream::from_reader`
//! method.
//!
//! ```
//! # use concordance::*;
//! #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
//! enum SomeTag {
//!     Identifier
//! };
//!
//! let simple_tagged = TaggedStream::from_reader(&mut "Hello, world".read_symbols());
//! let with_tag      = simple_tagged.with_tags(vec![(0..5, SomeTag::Identifier)].iter().cloned());
//! ```
//!

// TODO: a possibly better way to design this would be to make it so that there's an underlying stream we read from when generating tag
// symbols so that we don't need to store the contents of the stream in memory

use std::slice::Iter;
use std::ops::Index;
use std::ops::Range;

use super::countable::*;
use super::symbol_reader::*;
use super::tokenizer::*;
use super::symbol_range_dfa::*;

///
/// Represents a symbol in a tagged stream.
///
/// A symbol can either be a sequence of untagged base symbols, or it can be a tag (which itself can contain any number of tagged symbol)
///
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum TagSymbol<Base: Clone+Ord, Tag: Clone+Ord> {
    /// An untagged element in this stream
    Untagged(Base),

    /// A tagged region of the stream
    Tagged(Tag, TaggedStream<Base, Tag>)
}

use TagSymbol::*;

///
/// Represents a stream of tagged symbols
///
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct TaggedStream<Base: Clone+Ord, Tag: Clone+Ord> {
    /// The data in this stream
    data: Vec<TagSymbol<Base, Tag>>
}

impl<Base: Ord+Clone, Tag: Ord+Clone+'static> TaggedStream<Base, Tag> {
    ///
    /// Creates a basic tagged stream from a source of the base symbol
    ///
    pub fn from_reader(reader: &mut SymbolReader<Base>) -> TaggedStream<Base, Tag> {
        // Read from the reader into an array
        let mut symbols = vec![];

        while let Some(next_symbol) = reader.next_symbol() {
            symbols.push(Untagged(next_symbol));
        }

        // Generate a simple tagged stream from the result
        TaggedStream { data: symbols }
    }

    ///
    /// The number of symbols in this stream
    ///
    pub fn len(&self) -> usize {
        self.data.len()
    }

    ///
    /// Replaces a range in this stream with a tag
    ///
    pub fn tag(&mut self, tag: Tag, range: Range<usize>) {
        // Create a tag to replace the range
        let replaced_symbols    = self.data[range.clone()].to_vec();
        let tag_symbol          = Tagged(tag, TaggedStream { data: replaced_symbols });

        // Draining seems to be for reading a range but does double duty for deleting a range?
        // I don't think rust has a way to replace a range in a vector, or at least not one that's easy to find in the docs.
        self.data.drain(range.clone());

        // Draining then inserting is inefficient compared to flat out replacing items :-/ This may be possible but the vec docs aren't very easy to read
        self.data.insert(range.start, tag_symbol);
    }

    ///
    /// Creates a tag symbol by tagging a particular range within this stream
    ///
    pub fn tag_range(&self, range: Range<usize>, tag: Tag) -> TagSymbol<Base, Tag> {
        let tag_data = self.data[range].to_vec();

        Tagged(tag, TaggedStream { data: tag_data })
    }

    ///
    /// Applies tags to the specified ranges in this stream
    ///
    /// The tag ranges must be in ascending order and must not overlap
    ///
    pub fn with_tags<I>(&self, tags: I) -> TaggedStream<Base, Tag>
        where I : Iterator<Item=(Range<usize>, Tag)> {
        // The data that will make up the new stream
        let mut new_stream = vec![];

        // The last processed range (we haven't processed any range yet, so this starts at 0..0)
        let mut last_range = 0..0;

        // Build up the result from the tag list and the last range
        for tag in tags {
            // Fetch the information for this tag and check ordering
            let (next_range, next_tag)  = tag;
            let end                     = last_range.end;

            if end > next_range.start {
                panic!("Tags for with_tags must be sorted in order");
            }

            // Sometimes need to leave some data from this stream untagged
            if end < next_range.start {
                new_stream.extend(self.data[end..next_range.start].iter().cloned());
            }

            // Push the tag
            new_stream.push(self.tag_range(next_range.clone(), next_tag));

            // Update state
            last_range = next_range;
        }

        // Append anything left in the stream
        if last_range.end < self.data.len() {
            new_stream.extend(self.data[last_range.end..self.data.len()].iter().cloned());
        }

        // Final result
        TaggedStream { data: new_stream }
    }

    ///
    /// Runs the current values of this tagged stream through a tokenizer and tags anything it matches
    ///
    /// This takes a mapping function to describe how symbols in this stream map to symbols in the DFA. Note that every symbol
    /// (tagged or untagged) must be mapped to a DFA symbol, so if only tagged or untagged symbols are being used it's necessary
    /// to decide how the other symbols are mapped (eg, to an unused symbol)
    ///
    pub fn tokenize<DfaSymbol: Ord+Countable+Clone, MapFn>(&self, token_matcher: &SymbolRangeDfa<DfaSymbol, Tag>, map_symbol: MapFn) -> TaggedStream<Base, Tag>
        where MapFn: FnMut(TagSymbol<Base, Tag>) -> DfaSymbol {
        // Generate a symbol reader with the mapping function
        let reader = self.data.read_symbols().map_symbols(map_symbol);

        // Tokenize it
        let tokenizer = Tokenizer::new_prepared(reader, token_matcher);

        // Use the tokenizer to generate the result
        self.with_tags(tokenizer)
    }
}

impl<Base: Clone+Ord, Tag: Clone+Ord> Index<usize> for TaggedStream<Base, Tag> {
    type Output = TagSymbol<Base, Tag>;

    fn index(&self, index: usize) -> &TagSymbol<Base, Tag> {
        &self.data[index]
    }
}

impl<'a, Base: Clone+Ord, Tag: Clone+Ord> SymbolSource<'a, TagSymbol<Base, Tag>> for &'a TaggedStream<Base, Tag> {
    type SymbolReader = Iter<'a, TagSymbol<Base, Tag>>;

    /// Returns a new object that can read the symbols from this one
    fn read_symbols(self) -> Self::SymbolReader {
        self.data.read_symbols()
    }
}

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

    #[test]
    fn can_tag_range() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello,
            World
        }

        let mut tagged: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());

        tagged.tag(Tags::Hello, 0..5);
        tagged.tag(Tags::World, 1..6);

        assert!(tagged.len() == 2);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('o'));
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[1] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('W'));
            assert!(stream[1] == Untagged('o'));
            assert!(stream[2] == Untagged('r'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('d'));
        } else {
            assert!(false);
        }
    }

    #[test]
    fn can_tag_everything_with_tags() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello,
            World
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(0..5, Tags::Hello), (5..10, Tags::World)].iter().cloned());

        assert!(tagged.len() == 2);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('o'));
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[1] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('W'));
            assert!(stream[1] == Untagged('o'));
            assert!(stream[2] == Untagged('r'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('d'));
        } else {
            assert!(false);
        }
    }

    #[test]
    fn with_tags_preserves_middle() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello,
            World
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(0..4, Tags::Hello), (6..10, Tags::World)].iter().cloned());

        assert!(tagged.len() == 4);
        assert!(tagged[1] == Untagged('o'));
        assert!(tagged[2] == Untagged('W'));

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 4);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[3] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 4);
            assert!(stream[0] == Untagged('o'));
            assert!(stream[1] == Untagged('r'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('d'));
        } else {
            assert!(false);
        }
    }

    #[test]
    fn with_tags_preserves_end() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            Hello
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(0..5, Tags::Hello)].iter().cloned());

        assert!(tagged.len() == 6);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == Tags::Hello);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('H'));
            assert!(stream[1] == Untagged('e'));
            assert!(stream[2] == Untagged('l'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('o'));
        } else {
            assert!(false);
        }

        assert!(tagged[1] == Untagged('W'));
        assert!(tagged[2] == Untagged('o'));
        assert!(tagged[3] == Untagged('r'));
        assert!(tagged[4] == Untagged('l'));
        assert!(tagged[5] == Untagged('d'));
    }

    #[test]
    fn with_tags_preserves_start() {
        #[derive(Clone, PartialEq, Eq, Copy, PartialOrd, Ord)]
        enum Tags {
            World
        }

        let original: TaggedStream<char, Tags> = TaggedStream::from_reader(&mut "HelloWorld".read_symbols());
        let tagged = original.with_tags(vec![(5..10, Tags::World)].iter().cloned());

        assert!(tagged.len() == 6);

        if let Tagged(ref tag, ref stream) = tagged[5] {
            assert!(*tag == Tags::World);
            assert!(stream.len() == 5);
            assert!(stream[0] == Untagged('W'));
            assert!(stream[1] == Untagged('o'));
            assert!(stream[2] == Untagged('r'));
            assert!(stream[3] == Untagged('l'));
            assert!(stream[4] == Untagged('d'));
        } else {
            assert!(false);
        }

        assert!(tagged[0] == Untagged('H'));
        assert!(tagged[1] == Untagged('e'));
        assert!(tagged[2] == Untagged('l'));
        assert!(tagged[3] == Untagged('l'));
        assert!(tagged[4] == Untagged('o'));
    }

    #[test]
    fn can_tag_with_tokenizer() {
        #[derive(Ord, PartialOrd, Eq, PartialEq, Clone)]
        enum TestToken {
            Number,
            Whitespace
        }

        let mut token_matcher = TokenMatcher::new();
        token_matcher.add_pattern(MatchRange('0', '9').repeat_forever(0), TestToken::Number);
        token_matcher.add_pattern(exactly(" ").repeat_forever(0), TestToken::Whitespace);

        let dfa = token_matcher.prepare_to_match();

        let original = TaggedStream::from_reader(&mut "12 345  56".read_symbols());
        let tagged = original.tokenize(&dfa, | symbol: TagSymbol<char, TestToken> | {
            match symbol {
                Tagged(_, _) => ' ',
                Untagged(c)  => c
            }
        });

        assert!(tagged.len() == 5);

        if let Tagged(ref tag, ref stream) = tagged[0] {
            assert!(*tag == TestToken::Number);
            assert!(stream.len() == 2);
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[1] {
            assert!(*tag == TestToken::Whitespace);
            assert!(stream.len() == 1);
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[2] {
            assert!(*tag == TestToken::Number);
            assert!(stream.len() == 3);
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[3] {
            assert!(*tag == TestToken::Whitespace);
            assert!(stream.len() == 2);
        } else {
            assert!(false);
        }

        if let Tagged(ref tag, ref stream) = tagged[4] {
            assert!(*tag == TestToken::Number);
            assert!(stream.len() == 2);
        } else {
            assert!(false);
        }
    }
}
<|endoftext|>


<commit_before>
extern crate rand;
use std::net::{Ipv4Addr, SocketAddrV4, UdpSocket};
use file_stream_iter::FileStream;
use std::io::Error;
use tftp_specific::PacketType;
use std::str::FromStr;
use std::net::AddrParseError;
use std::error;
use std::str;

pub struct EndPoint{
    local_connection : SocketAddrV4,
    remote_connection : SocketAddrV4
}

impl EndPoint{

    pub fn new(remote : &str) -> Result<EndPoint, AddrParseError>{

        let ipv4 = try!(get_ip_from(remote));
        Ok(EndPoint{
            local_connection : SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 6900),
            remote_connection : SocketAddrV4::new(ipv4, 9990)
        })
    }

    fn start_listen(&mut self)->Result<(), Error>{
        // Bind the socket
        let socket = try!(UdpSocket::bind(self.local_connection));

        // Start servicing requests from a client
        // Read from the socket
        let mut buf = [0u8; 516]; //UDP packet is 516 bytes
        let (amt, src) = try!(socket.recv_from(&mut buf));
        match buf[1]{
            1 => { //RRQ
                //Get filename to be sent from the packet
                let mut  i = 2;
                while buf[i] != '\0' as u8{
                    i += 1;
                }
                let filename = str::from_utf8(&buf[2..i]);
                // Send data via the socket
                let mut fs = FileStream::new(String::from(filename.unwrap())).unwrap();
                let (buf, num_bytes_read) = fs.next().unwrap();
                let mut v = Vec::with_capacity(num_bytes_read + 4);
                v.push(0);
                v.push(3);
                v.push(0xFF);
                let mut cnt = 0xFF;
                v.push(cnt);
                v.extend_from_slice(&buf[0..num_bytes_read]);
                try!(socket.send_to(v.as_slice(), self.remote_connection));

            },
            _ => {panic!("unrecognized packet type")}
        }

        Ok(())

    }

    pub fn recv(&mut self) -> Result<(), Error> {
        // Define the local connection information
        //let ip = Ipv4Addr::new(127, 0, 0, 1);
        //TODO we want to use port 69 here
        //let connection = SocketAddrV4::new(ip, 6900);

        // Bind the socket
        let socket = try!(UdpSocket::bind(self.local_connection));

        // Read from the socket
        let mut buf = [0u8; 516]; //UDP packet is 516 bytes
        let (amt, src) = try!(socket.recv_from(&mut buf));
        match buf[1]{
            3 =>{
                let block_num : u16 = 0u16 | (buf[2] as u16) << 8 |  buf[3] as u16;
                let block_size = amt - 4;
                println!("{:?}", block_size);
                if block_size < 512{
                    println!("recvr recvd: {:?}", &buf[0 .. amt]);
                    println!("Last block of the file received");
                }
                //send ACK

                let low = block_num & 0x00FF;
                let high = (block_num & 0xFF00) >> 8;

                socket.send_to(&[0,PacketType::ACK as u8, high as u8, low as u8], src);
            },
            _ => {}
        }

        Ok(())
    }

    pub fn send(&mut self) -> Result<(), Error> {
        // Define the local connection (to send the data from)
        //let ip = Ipv4Addr::new(127, 0, 0, 1);
        //TODO generate port # using rand
        //let connection = SocketAddrV4::new(ip, 9992);

        // Bind the socket
        let socket = try!(UdpSocket::bind(self.local_connection));

        // Define the remote connection (to send the data to)
        //let connection2 = SocketAddrV4::new(ip, 6900);

        // Send data via the socket
        let mut fs = FileStream::new("/home/abhi/code/rust/fei/target/debug/foo.txt".to_string()).unwrap();
        let (buf, num_bytes_read) = fs.next().unwrap();
        let mut v = Vec::with_capacity(num_bytes_read + 4);
        v.push(0);
        v.push(3);
        v.push(0xFF);
        let mut cnt = 0xFF;
        v.push(cnt);
        v.extend_from_slice(&buf[0..num_bytes_read]);
        try!(socket.send_to(v.as_slice(), self.remote_connection));

        let mut buf  = [0; 10];
        socket.recv_from(&mut buf);

        println!("sender recvd {:?}", buf);
        Ok(())
    }

    fn send_file(&mut self, filename : &str) {

    }

    fn create_rrq_wrq_packet(&self, p_type: PacketType, file_name: &str, mode : &'static str)->Vec<u8>{
        let mut v = Vec::with_capacity(2+file_name.len()+1+mode.len()+1);
        v.push(0);
        v.push(p_type as u8);
        v.extend(file_name.as_bytes());
        v.push('\0' as u8);
        v.push(0);//zero terminator
        v.extend(mode.as_bytes());
        v.push(0);//zero terminator
        v
    }

    fn get(&mut self, files : &[&str], mode : &'static str) -> Result<(), Error>{
        let socket = try!(UdpSocket::bind(self.local_connection));
        let mut buf  = [0; 516];
        for file in files{
            //start with a RRQ
            let packet = self.create_rrq_wrq_packet(PacketType::RRQ, file, mode);
            try!(socket.send_to(packet.as_slice(), self.remote_connection));

            loop{
                //get the first block of the requested file
                let (amt, src) = try!(socket.recv_from(&mut buf));
                match buf[1]{
                    3 =>{
                        let block_num : u16 = 0u16 | (buf[2] as u16) << 8 |  buf[3] as u16;
                        let block_size = amt - 4;
                        println!("{:?}", block_size);
                        if block_size < 512{
                            println!("recvr recvd: {:?}", &buf[0 .. amt]);
                            println!("Last block of the file received");
                        }
                        //send ACK

                        let low = block_num & 0x00FF;
                        let high = (block_num & 0xFF00) >> 8;

                        socket.send_to(&[0,PacketType::ACK as u8, high as u8, low as u8], src);
                    },
                    _ => {}
                }
                //Flush buf for next read
                //TODO fix this:
                for i in 0..516{
                    buf[0] = 0;
                }
            }

        }

        Ok(())
    }
}

fn get_ip_from(ip : &str) -> Result<Ipv4Addr, AddrParseError>{
    Ipv4Addr::from_str(ip)
}

<commit_msg>Fix block-ids packing and add ack servicing after sending data to client<commit_after>
extern crate rand;
//extern crate time;

//use time::PreciseTime;
use std::net::{Ipv4Addr, SocketAddrV4, UdpSocket};
use file_stream_iter::FileStream;
use std::io::Error;
use tftp_specific::PacketType;
use std::str::FromStr;
use std::net::AddrParseError;
use std::error;
use std::str;

pub struct EndPoint{
    local_connection : SocketAddrV4,
    remote_connection : SocketAddrV4
}

impl EndPoint{

    pub fn new(remote : &str) -> Result<EndPoint, AddrParseError>{

        let ipv4 = try!(get_ip_from(remote));
        Ok(EndPoint{
            local_connection : SocketAddrV4::new(Ipv4Addr::new(127,0,0,1), 6900),
            remote_connection : SocketAddrV4::new(ipv4, 9990)
        })
    }

    fn start_listen(&mut self)->Result<(), Error>{
        // Bind the socket
        let socket = try!(UdpSocket::bind(self.local_connection));

        // Start servicing requests from a client
        // Read from the socket
        let mut buf = [0u8; 516]; //UDP packet is 516 bytes
        let (amt, src) = try!(socket.recv_from(&mut buf));
        let mut last_blk_id = 0u16;
        match buf[1]{
            1 => { //RRQ
                //Get filename to be sent from the packet
                let mut  i = 2;
                while buf[i] != '\0' as u8{
                    i += 1;
                }
                let filename = str::from_utf8(&buf[2..i]);
                // Send data via the socket
                let mut fs = FileStream::new(String::from(filename.unwrap())).unwrap();
                let (mut buf, num_bytes_read) = fs.next().unwrap();
                let mut v = Vec::with_capacity(num_bytes_read + 4);
                v.push(0);
                v.push(PacketType::DATA as u8);
                v.push((last_blk_id >> 8) as u8);
                v.push((last_blk_id & 0x00FF) as u8);
                v.extend_from_slice(&buf[0..num_bytes_read]);
                try!(socket.send_to(v.as_slice(), self.remote_connection));
                //TODO start timer here
                loop{
                    //Get ACK
                    let (amt, src) = try!(socket.recv_from(&mut buf));
                    if buf[1] == 4{
                        last_blk_id += 1;
                        let (buf, num_bytes_read) = fs.next().unwrap();
                        let mut v = Vec::with_capacity(num_bytes_read + 4);
                        v.push(0);
                        v.push(PacketType::DATA as u8);
                        v.push((last_blk_id >> 8) as u8);
                        v.push((last_blk_id & 0x00FF) as u8);
                        v.extend_from_slice(&buf[0..num_bytes_read]);
                        try!(socket.send_to(v.as_slice(), self.remote_connection));
                    }
                }
            },
            4 => {

            },
            _ => {panic!("unrecognized packet type")}
        }

        Ok(())

    }

    pub fn recv(&mut self) -> Result<(), Error> {
        // Define the local connection information
        //let ip = Ipv4Addr::new(127, 0, 0, 1);
        //TODO we want to use port 69 here
        //let connection = SocketAddrV4::new(ip, 6900);

        // Bind the socket
        let socket = try!(UdpSocket::bind(self.local_connection));

        // Read from the socket
        let mut buf = [0u8; 516]; //UDP packet is 516 bytes
        let (amt, src) = try!(socket.recv_from(&mut buf));
        match buf[1]{
            3 =>{
                let block_num : u16 = 0u16 | (buf[2] as u16) << 8 |  buf[3] as u16;
                let block_size = amt - 4;
                println!("{:?}", block_size);
                if block_size < 512{
                    println!("recvr recvd: {:?}", &buf[0 .. amt]);
                    println!("Last block of the file received");
                }
                //send ACK

                let low = block_num & 0x00FF;
                let high = (block_num & 0xFF00) >> 8;

                socket.send_to(&[0,PacketType::ACK as u8, high as u8, low as u8], src);
            },
            _ => {}
        }

        Ok(())
    }

    pub fn send(&mut self) -> Result<(), Error> {
        // Define the local connection (to send the data from)
        //let ip = Ipv4Addr::new(127, 0, 0, 1);
        //TODO generate port # using rand
        //let connection = SocketAddrV4::new(ip, 9992);

        // Bind the socket
        let socket = try!(UdpSocket::bind(self.local_connection));

        // Define the remote connection (to send the data to)
        //let connection2 = SocketAddrV4::new(ip, 6900);

        // Send data via the socket
        let mut fs = FileStream::new("/home/abhi/code/rust/fei/target/debug/foo.txt".to_string()).unwrap();
        let (buf, num_bytes_read) = fs.next().unwrap();
        let mut v = Vec::with_capacity(num_bytes_read + 4);
        v.push(0);
        v.push(3);
        v.push(0xFF);
        let mut cnt = 0xFF;
        v.push(cnt);
        v.extend_from_slice(&buf[0..num_bytes_read]);
        try!(socket.send_to(v.as_slice(), self.remote_connection));

        let mut buf  = [0; 10];
        socket.recv_from(&mut buf);

        println!("sender recvd {:?}", buf);
        Ok(())
    }

    fn send_file(&mut self, filename : &str) {

    }

    fn create_rrq_wrq_packet(&self, p_type: PacketType, file_name: &str, mode : &'static str)->Vec<u8>{
        let mut v = Vec::with_capacity(2+file_name.len()+1+mode.len()+1);
        v.push(0);
        v.push(p_type as u8);
        v.extend(file_name.as_bytes());
        v.push('\0' as u8);
        v.push(0);//zero terminator
        v.extend(mode.as_bytes());
        v.push(0);//zero terminator
        v
    }

    fn get(&mut self, files : &[&str], mode : &'static str) -> Result<(), Error>{
        let socket = try!(UdpSocket::bind(self.local_connection));
        let mut buf  = [0; 516];
        for file in files{
            //start with a RRQ
            let packet = self.create_rrq_wrq_packet(PacketType::RRQ, file, mode);
            try!(socket.send_to(packet.as_slice(), self.remote_connection));

            loop{
                //get the first block of the requested file
                let (amt, src) = try!(socket.recv_from(&mut buf));
                match buf[1]{
                    3 =>{
                        let block_num : u16 = 0u16 | (buf[2] as u16) << 8 |  buf[3] as u16;
                        let block_size = amt - 4;
                        println!("{:?}", block_size);

                        //send ACK
                        let low = block_num & 0x00FF;
                        let high = (block_num & 0xFF00) >> 8;

                        socket.send_to(&[0,PacketType::ACK as u8, high as u8, low as u8], src);

                        if block_size < 512{
                            println!("recvr recvd: {:?}", &buf[0 .. amt]);
                            println!("Last block of the file received");
                            break;
                        }
                    },
                    _ => {}
                }
                //Flush buf for next read
                //TODO fix this:
                for i in 0..516{
                    buf[0] = 0;
                }
            }

        }

        Ok(())
    }
}

fn get_ip_from(ip : &str) -> Result<Ipv4Addr, AddrParseError>{
    Ipv4Addr::from_str(ip)
}

<|endoftext|>


<commit_before>use nom::{Err, IResult, Needed};

use std::io;
use std::io::Read;
use std::ptr;
use std::cmp;

use super::StreamProducer;

#[derive(Debug)]
pub enum ErrorKind {
  IO(io::Error),
  Incomplete(usize),
  Consumed(usize),
  Continue,
  EndOfInput,
  Unknown,
}

/// Structure that hold a slice of bytes.
pub struct ByteStream<'a> {
  offset: usize,
  bytes: &'a [u8],
}

impl<'a> ByteStream<'a> {
  /// Construct a `ByteStream` based on the passed in byte slice.
  pub fn new(bytes: &'a [u8]) -> Self {
    ByteStream {
      offset: 0,
      bytes: bytes,
    }
  }

  /// Return the number of bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.bytes.len() - self.offset
  }

  /// Return true if the stream contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }
}

impl<'a> StreamProducer for ByteStream<'a> {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T> {
    if self.is_empty() {
      return Err(ErrorKind::EndOfInput);
    }

    match f(&self.bytes[self.offset..]) {
      IResult::Done(i, o)    => {
        self.offset += self.len() - i.len();

        Ok(o)
      }
      IResult::Incomplete(n) => {
        let mut needed = self.len();

        if let Needed::Size(size) = n {
          needed = size;
        }

        Err(ErrorKind::Incomplete(needed))
      }
      IResult::Error(e)      => {
        if let Err::Position(_, i) = e {
          let consumed = self.len() - i.len();

          self.offset += consumed;

          Err(ErrorKind::Consumed(consumed))
        } else {
          Err(ErrorKind::Unknown)
        }
      }
    }
  }
}

// Growable buffer of bytes.
//
// Mainly used to the `ReadStream` structure but can be used seperately for
// manually filling with some `Read` source.
pub struct Buffer {
  data: Vec<u8>,
  filled: usize,
  offset: usize,
}

impl Buffer {
  // Default constructor for `Buffer`
  pub fn new() -> Self {
    Self::with_capacity(1024)
  }

  // Explicitly set the buffer capacity.
  pub fn with_capacity(capacity: usize) -> Self {
    let mut buffer = Vec::with_capacity(capacity);

    unsafe {
      buffer.set_len(capacity);
    }

    Buffer {
      data: buffer,
      filled: 0,
      offset: 0,
    }
  }

  // Return the number of read bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.filled - self.offset
  }

  // Return true if buffer contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }

  // The set length of the unlining buffer.
  #[inline]
  pub fn capacity(&self) -> usize {
    self.data.len()
  }

  // Return a reference to the slice of unread bytes.
  pub fn as_slice(&self) -> &[u8] {
    &self.data[self.offset..self.filled]
  }

  // Fill the buffer with bytes from a `Read` source.
  pub fn fill<R: Read>(&mut self, reader: &mut R) -> io::Result<usize> {
    reader.read(&mut self.data[self.filled..]).map(|consumed| {
      self.filled += consumed;

      consumed
    })
  }

  // Resize the current buffer
  //
  // This will only allocate data when the size requests is larger than the
  // current capacity of the buffer, otherwise it moves the currently filled
  // data to the beginning of the buffer.
  pub fn resize(&mut self, size: usize) {
    if size > self.data.capacity() {
      self.data.reserve(size);

      let capacity = self.data.capacity();

      unsafe {
        self.data.set_len(capacity);
      }
    }

    if self.data.len() - self.filled < size  {
      let length  = self.filled - self.offset;
      let mut_ptr = self.data.as_mut_ptr();

      unsafe {
        let offset_ptr  = self.data.as_ptr().offset(self.offset as isize);

        ptr::copy(offset_ptr, mut_ptr, length);
      }

      self.filled -= self.offset;
      self.offset  = 0;
    }
  }

  // Move the offset by the amount of consumed bytes.
  pub fn consume(&mut self, consumed: usize) {
    self.offset += consumed;
  }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParserState {
  Incomplete,
  EndOfInput,
}

fn fill<R: Read>(buffer: &mut Buffer, reader: &mut R, needed: usize)
                 -> io::Result<usize> {
  let mut read = 0;

  if buffer.len() < needed {
    buffer.resize(needed);

    while buffer.len() < needed {
      let size_read = try!(buffer.fill(reader));

      if size_read > 0 {
        read += size_read;
      } else {
        break;
      }
    }
  }

  Ok(read)
}

// Structure that hold a reader for a source of bytes.
pub struct ReadStream<R: Read> {
  reader: R,
  buffer: Buffer,
  needed: usize,
  state: ParserState,
}

impl<R> ReadStream<R> where R: Read {
  // Constructor for `ReadStream` based on a `Read` source.
  pub fn new(reader: R) -> Self {
    ReadStream {
      reader: reader,
      buffer: Buffer::new(),
      needed: 0,
      state: ParserState::Incomplete,
    }
  }

  // Fill the stream with bytes from a `Read` source.
  fn fill(&mut self) -> io::Result<usize> {
    let needed = cmp::max(1, self.needed);

    fill(&mut self.buffer, &mut self.reader, needed).map(|consumed| {
      if self.buffer.len() < needed {
        self.state = ParserState::EndOfInput;
      }

      consumed
    })
  }
}

fn from_iresult<T>(buffer: &Buffer, result: IResult<&[u8], T>)
                   -> Result<(usize, T), ErrorKind> {
  match result {
    IResult::Done(i, o)    => Ok((buffer.len() - i.len(), o)),
    IResult::Incomplete(n) => {
      let mut needed = buffer.capacity();

      if let Needed::Size(size) = n {
        needed = size;
      }

      Err(ErrorKind::Incomplete(needed))
    }
    IResult::Error(e)      => {
      if let Err::Position(_, i) = e {
        Err(ErrorKind::Consumed(buffer.len() - i.len()))
      } else {
        Err(ErrorKind::Unknown)
      }
    }
  }
}

impl<R> StreamProducer for ReadStream<R> where R: Read {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T> {
    if self.state != ParserState::EndOfInput {
      try!(self.fill().map_err(ErrorKind::IO));
    }

    let mut buffer = &mut self.buffer;

    if buffer.is_empty() {
      self.state = ParserState::EndOfInput;

      return Err(ErrorKind::EndOfInput);
    }

    let result = {
      let iresult = f(buffer.as_slice());

      from_iresult(&buffer, iresult)
    };

    match result {
      Ok((consumed, o)) => {
        buffer.consume(consumed);

        Ok(o)
      }
      Err(kind)         => {
        match kind {
          ErrorKind::Incomplete(needed) => {
            self.needed = needed;

            Err(ErrorKind::Continue)
          }
          ErrorKind::Consumed(consumed) => {
            buffer.consume(consumed);

            Err(ErrorKind::Continue)
          }
          _                             => Err(kind),
        }
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use utility::StreamProducer;
  use nom::be_u32;

  #[test]
  fn test_buffer() {
    let mut buffer = Buffer::new();
    let bytes      = b"Hello World";
    let mut reader = &bytes[..];

    assert!(buffer.is_empty());
    assert_eq!(buffer.capacity(), 1024);

    let bytes_read = buffer.fill(&mut reader).unwrap_or(0);
    let bytes_len  = bytes.len();

    assert_eq!(bytes_read, bytes_len);
    assert_eq!(buffer.len(), bytes_len);
    assert_eq!(buffer.as_slice(), bytes);

    buffer.resize(512);
    assert_eq!(buffer.capacity(), 1024);
  }

  #[test]
  fn test_byte_stream() {
    let bytes      = b"Hello World";
    let mut stream = ByteStream::new(bytes);

    assert_eq!(stream.len(), bytes.len());

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444);
    assert_eq!(stream.len(), 7);
  }

  #[test]
  fn test_read_stream() {
    let bytes      = b"Hello World";
    let mut stream = ReadStream::new(&bytes[..]);

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444)
  }
}
<commit_msg>Change `ReadStream` comments to documentation<commit_after>use nom::{Err, IResult, Needed};

use std::io;
use std::io::Read;
use std::ptr;
use std::cmp;

use super::StreamProducer;

#[derive(Debug)]
pub enum ErrorKind {
  IO(io::Error),
  Incomplete(usize),
  Consumed(usize),
  Continue,
  EndOfInput,
  Unknown,
}

/// Structure that hold a slice of bytes.
pub struct ByteStream<'a> {
  offset: usize,
  bytes: &'a [u8],
}

impl<'a> ByteStream<'a> {
  /// Construct a `ByteStream` based on the passed in byte slice.
  pub fn new(bytes: &'a [u8]) -> Self {
    ByteStream {
      offset: 0,
      bytes: bytes,
    }
  }

  /// Return the number of bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.bytes.len() - self.offset
  }

  /// Return true if the stream contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }
}

impl<'a> StreamProducer for ByteStream<'a> {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T> {
    if self.is_empty() {
      return Err(ErrorKind::EndOfInput);
    }

    match f(&self.bytes[self.offset..]) {
      IResult::Done(i, o)    => {
        self.offset += self.len() - i.len();

        Ok(o)
      }
      IResult::Incomplete(n) => {
        let mut needed = self.len();

        if let Needed::Size(size) = n {
          needed = size;
        }

        Err(ErrorKind::Incomplete(needed))
      }
      IResult::Error(e)      => {
        if let Err::Position(_, i) = e {
          let consumed = self.len() - i.len();

          self.offset += consumed;

          Err(ErrorKind::Consumed(consumed))
        } else {
          Err(ErrorKind::Unknown)
        }
      }
    }
  }
}

// Growable buffer of bytes.
//
// Mainly used to the `ReadStream` structure but can be used seperately for
// manually filling with some `Read` source.
pub struct Buffer {
  data: Vec<u8>,
  filled: usize,
  offset: usize,
}

impl Buffer {
  // Default constructor for `Buffer`
  pub fn new() -> Self {
    Self::with_capacity(1024)
  }

  // Explicitly set the buffer capacity.
  pub fn with_capacity(capacity: usize) -> Self {
    let mut buffer = Vec::with_capacity(capacity);

    unsafe {
      buffer.set_len(capacity);
    }

    Buffer {
      data: buffer,
      filled: 0,
      offset: 0,
    }
  }

  // Return the number of read bytes that haven't been consumed yet.
  #[inline]
  pub fn len(&self) -> usize {
    self.filled - self.offset
  }

  // Return true if buffer contains no more bytes.
  #[inline]
  pub fn is_empty(&self) -> bool {
    self.len() == 0
  }

  // The set length of the unlining buffer.
  #[inline]
  pub fn capacity(&self) -> usize {
    self.data.len()
  }

  // Return a reference to the slice of unread bytes.
  pub fn as_slice(&self) -> &[u8] {
    &self.data[self.offset..self.filled]
  }

  // Fill the buffer with bytes from a `Read` source.
  pub fn fill<R: Read>(&mut self, reader: &mut R) -> io::Result<usize> {
    reader.read(&mut self.data[self.filled..]).map(|consumed| {
      self.filled += consumed;

      consumed
    })
  }

  // Resize the current buffer
  //
  // This will only allocate data when the size requests is larger than the
  // current capacity of the buffer, otherwise it moves the currently filled
  // data to the beginning of the buffer.
  pub fn resize(&mut self, size: usize) {
    if size > self.data.capacity() {
      self.data.reserve(size);

      let capacity = self.data.capacity();

      unsafe {
        self.data.set_len(capacity);
      }
    }

    if self.data.len() - self.filled < size  {
      let length  = self.filled - self.offset;
      let mut_ptr = self.data.as_mut_ptr();

      unsafe {
        let offset_ptr  = self.data.as_ptr().offset(self.offset as isize);

        ptr::copy(offset_ptr, mut_ptr, length);
      }

      self.filled -= self.offset;
      self.offset  = 0;
    }
  }

  // Move the offset by the amount of consumed bytes.
  pub fn consume(&mut self, consumed: usize) {
    self.offset += consumed;
  }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ParserState {
  Incomplete,
  EndOfInput,
}

fn fill<R: Read>(buffer: &mut Buffer, reader: &mut R, needed: usize)
                 -> io::Result<usize> {
  let mut read = 0;

  if buffer.len() < needed {
    buffer.resize(needed);

    while buffer.len() < needed {
      let size_read = try!(buffer.fill(reader));

      if size_read > 0 {
        read += size_read;
      } else {
        break;
      }
    }
  }

  Ok(read)
}

/// Structure that hold a reader for a source of bytes.
pub struct ReadStream<R: Read> {
  reader: R,
  buffer: Buffer,
  needed: usize,
  state: ParserState,
}

impl<R> ReadStream<R> where R: Read {
  /// Constructor for `ReadStream` based on a `Read` source.
  pub fn new(reader: R) -> Self {
    ReadStream {
      reader: reader,
      buffer: Buffer::new(),
      needed: 0,
      state: ParserState::Incomplete,
    }
  }

  // Fill the stream with bytes from a `Read` source.
  fn fill(&mut self) -> io::Result<usize> {
    let needed = cmp::max(1, self.needed);

    fill(&mut self.buffer, &mut self.reader, needed).map(|consumed| {
      if self.buffer.len() < needed {
        self.state = ParserState::EndOfInput;
      }

      consumed
    })
  }
}

fn from_iresult<T>(buffer: &Buffer, result: IResult<&[u8], T>)
                   -> Result<(usize, T), ErrorKind> {
  match result {
    IResult::Done(i, o)    => Ok((buffer.len() - i.len(), o)),
    IResult::Incomplete(n) => {
      let mut needed = buffer.capacity();

      if let Needed::Size(size) = n {
        needed = size;
      }

      Err(ErrorKind::Incomplete(needed))
    }
    IResult::Error(e)      => {
      if let Err::Position(_, i) = e {
        Err(ErrorKind::Consumed(buffer.len() - i.len()))
      } else {
        Err(ErrorKind::Unknown)
      }
    }
  }
}

impl<R> StreamProducer for ReadStream<R> where R: Read {
  fn parse<F, T>(&mut self, f: F) -> Result<T, ErrorKind>
   where F: FnOnce(&[u8]) -> IResult<&[u8], T> {
    if self.state != ParserState::EndOfInput {
      try!(self.fill().map_err(ErrorKind::IO));
    }

    let mut buffer = &mut self.buffer;

    if buffer.is_empty() {
      self.state = ParserState::EndOfInput;

      return Err(ErrorKind::EndOfInput);
    }

    let result = {
      let iresult = f(buffer.as_slice());

      from_iresult(&buffer, iresult)
    };

    match result {
      Ok((consumed, o)) => {
        buffer.consume(consumed);

        Ok(o)
      }
      Err(kind)         => {
        match kind {
          ErrorKind::Incomplete(needed) => {
            self.needed = needed;

            Err(ErrorKind::Continue)
          }
          ErrorKind::Consumed(consumed) => {
            buffer.consume(consumed);

            Err(ErrorKind::Continue)
          }
          _                             => Err(kind),
        }
      }
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use utility::StreamProducer;
  use nom::be_u32;

  #[test]
  fn test_buffer() {
    let mut buffer = Buffer::new();
    let bytes      = b"Hello World";
    let mut reader = &bytes[..];

    assert!(buffer.is_empty());
    assert_eq!(buffer.capacity(), 1024);

    let bytes_read = buffer.fill(&mut reader).unwrap_or(0);
    let bytes_len  = bytes.len();

    assert_eq!(bytes_read, bytes_len);
    assert_eq!(buffer.len(), bytes_len);
    assert_eq!(buffer.as_slice(), bytes);

    buffer.resize(512);
    assert_eq!(buffer.capacity(), 1024);
  }

  #[test]
  fn test_byte_stream() {
    let bytes      = b"Hello World";
    let mut stream = ByteStream::new(bytes);

    assert_eq!(stream.len(), bytes.len());

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444);
    assert_eq!(stream.len(), 7);
  }

  #[test]
  fn test_read_stream() {
    let bytes      = b"Hello World";
    let mut stream = ReadStream::new(&bytes[..]);

    let result = stream.parse(be_u32).unwrap_or(0);

    assert_eq!(result, 1214606444)
  }
}
<|endoftext|>


<commit_before>//! A reverse proxy application.
//!
//! This module provides a reverse proxy implementation for the Open
//! Matchmaking project.
//!
//! The purposes of the application  are handling client connections, applying
//! a middleware with checks for a token when it was specified and communicating
//! with a message broker for getting responses from microservices in the certain
//! format.
//!

use std::collections::{HashMap};
use std::error::{Error};
use std::net::{SocketAddr};
use std::sync::{Arc, Mutex, RwLock};

use cli::{CliOptions};
use futures::sync::{mpsc};
use futures::{Future, Sink};
use futures::stream::{Stream};
use tokio::net::{TcpListener};
use tokio::reactor::{Handle};
use tokio::runtime::{run};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_current_thread::{spawn};
use tokio_tungstenite::{accept_async};
use tungstenite::protocol::{Message};

use auth::middleware::{Middleware, EmptyMiddleware};
use auth::token::middleware::{JwtTokenMiddleware};
use engine::{Engine};
use error::{PathfinderError};
use rabbitmq::{RabbitMQClient};

/// A reverse proxy application.
pub struct Proxy<'a, T: 'a> {
    engine: Arc<RwLock<Box<Engine<'a, T>>>>,
    rabbitmq_client: Arc<RwLock<Box<RabbitMQClient<'a, T>>>>,
    connections: Arc<Mutex<HashMap<SocketAddr, mpsc::UnboundedSender<Message>>>>,
    auth_middleware: Arc<RwLock<Box<Middleware>>>,
}

impl<'a, T: AsyncRead + AsyncWrite + Send + Sync + 'static> Proxy<'a, T> {
    /// Returns a new instance of a reverse proxy application.
    pub fn new(cli: &CliOptions, engine: Box<Engine<'a, T>>) -> Proxy<'a, T> {
        let auth_middleware: Box<Middleware> = match cli.validate {
            true => Box::new(JwtTokenMiddleware::new(cli)),
            _ => Box::new(EmptyMiddleware::new(cli))
        };
        let rabbitmq_client = RabbitMQClient::new(cli);

        Proxy {
            engine: Arc::new(RwLock::new(engine)),
            rabbitmq_client: Arc::new(Box::new(rabbitmq_client)),
            connections: Arc::new(Mutex::new(HashMap::new())),
            auth_middleware: Arc::new(RwLock::new(auth_middleware)),
        }
    }

    fn init(&mut self) -> impl Future<Item=(), Error=PathfinderError> {
        self.rabbitmq_client.clone().write().unwrap().init();
        Ok(())
    }

    /// Run the server on the specified address and port.
    pub fn run(&self, address: SocketAddr) {
        let handle = Handle::default();
        let listener = TcpListener::bind(&address).unwrap();
        println!("Listening on: {}", address);

        let server = listener.incoming().for_each(move |stream| {
            let addr = stream.peer_addr().expect("Connected stream should have a peer address.");

            let engine_local = self.engine.clone();
            let connections_local = self.connections.clone();
            let auth_middleware_local = self.auth_middleware.clone();
            let handle_local = handle.clone();

            accept_async(stream)
                // Process the messages
                .and_then(move |ws_stream| {
                    // Create a channel for the stream, which other sockets will use to
                    // send us messages. It could be used for broadcasting your data to
                    // another users in the future.
                    let (tx, rx) = mpsc::unbounded();
                    connections_local.lock().unwrap().insert(addr, tx);

                    // Split the WebSocket stream so that it will be possible to work
                    // with the reading and writing halves separately.
                    let (sink, stream) = ws_stream.split();

                    // Read and process each message
                    let handle_inner = handle_local.clone();
                    let connections_inner = connections_local.clone();
                    let ws_reader = stream.for_each(move |message: Message| {

                        // Get references to required components
                        let addr_nested = addr.clone();
                        let engine_nested = engine_local.clone();
                        let connections_nested = connections_inner.clone();
                        let transmitter_nested = &connections_nested.clone().lock().unwrap()[&addr_nested];
                        let transmitter_nested2 = transmitter_nested.clone();

                        // 1. Deserialize message into JSON
                        let json_message = match engine_local.read().unwrap().deserialize_message(&message) {
                            Ok(json_message) => json_message,
                            Err(err) => {
                                let formatted_error = format!("{}", err);
                                let error_message = engine_nested.read().unwrap().wrap_an_error(formatted_error.as_str());
                                transmitter_nested.unbounded_send(error_message).unwrap();
                                return Ok(())
                            }
                        };

                        // 2. Apply a middleware to each incoming message
                        let auth_future = auth_middleware_local.read().unwrap()
                            .process_request(json_message.clone());

                        // 3. Put request into a queue in RabbitMQ and receive the response
                        let rabbitmq_future = engine_local.read().unwrap().handle(
                            json_message.clone(), transmitter_nested.clone(), &handle_inner
                        );

                        let processing_request_future = auth_future
                            .and_then(move |_| rabbitmq_future)
                            .map_err(move |err| {
                                let formatted_error = format!("{}", err);
                                let error_message = engine_nested.read().unwrap().wrap_an_error(formatted_error.as_str());
                                transmitter_nested2.unbounded_send(error_message).unwrap();
                                ()
                            });

                        spawn(processing_request_future);
                        Ok(())
                    });

                    // Write back prepared responses
                    let ws_writer = rx.fold(sink, |mut sink, msg| {
                        sink.start_send(msg).unwrap();
                        Ok(sink)
                    });

                    // Wait for either half to be done to tear down the other
                    let connection = ws_reader.map(|_| ()).map_err(|_| ())
                        .select(ws_writer.map(|_| ()).map_err(|_| ()));

                    // Close the connection after using
                    spawn(connection.then(move |_| {
                        connections_inner.lock().unwrap().remove(&addr);
                        debug!("Connection {} closed.", addr);
                        Ok(())
                    }));

                    Ok(())
                })
                // An error occurred during the WebSocket handshake
                .or_else(|err| {
                    debug!("{}", err.description());
                    Ok(())
                })
        });

        // Run the server
        run(self.init.and_then(|_| server.map_err(|_e| ())));
    }
}
<commit_msg>Fixed init future for proxy<commit_after>//! A reverse proxy application.
//!
//! This module provides a reverse proxy implementation for the Open
//! Matchmaking project.
//!
//! The purposes of the application  are handling client connections, applying
//! a middleware with checks for a token when it was specified and communicating
//! with a message broker for getting responses from microservices in the certain
//! format.
//!

use std::collections::{HashMap};
use std::error::{Error};
use std::net::{SocketAddr};
use std::sync::{Arc, Mutex, RwLock};

use cli::{CliOptions};
use futures::sync::{mpsc};
use futures::{Future, Sink};
use futures::future::{self, FutureResult};
use futures::stream::{Stream};
use tokio::net::{TcpListener};
use tokio::reactor::{Handle};
use tokio::runtime::{run};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_current_thread::{spawn};
use tokio_tungstenite::{accept_async};
use tungstenite::protocol::{Message};

use auth::middleware::{Middleware, EmptyMiddleware};
use auth::token::middleware::{JwtTokenMiddleware};
use engine::{Engine};
use error::{PathfinderError};
use rabbitmq::{RabbitMQClient};

/// A reverse proxy application.
pub struct Proxy<'a, T: 'a> {
    engine: Arc<RwLock<Box<Engine<'a, T>>>>,
    rabbitmq_client: Arc<RwLock<Box<RabbitMQClient<'a, T>>>>,
    connections: Arc<Mutex<HashMap<SocketAddr, mpsc::UnboundedSender<Message>>>>,
    auth_middleware: Arc<RwLock<Box<Middleware>>>,
}

impl<'a, T: AsyncRead + AsyncWrite + Send + Sync + 'static> Proxy<'a, T> {
    /// Returns a new instance of a reverse proxy application.
    pub fn new(cli: &CliOptions, engine: Box<Engine<'a, T>>) -> Proxy<'a, T> {
        let auth_middleware: Box<Middleware> = match cli.validate {
            true => Box::new(JwtTokenMiddleware::new(cli)),
            _ => Box::new(EmptyMiddleware::new(cli))
        };
        let rabbitmq_client = RabbitMQClient::new(cli);

        Proxy {
            engine: Arc::new(RwLock::new(engine)),
            rabbitmq_client: Arc::new(Box::new(rabbitmq_client)),
            connections: Arc::new(Mutex::new(HashMap::new())),
            auth_middleware: Arc::new(RwLock::new(auth_middleware)),
        }
    }

    /// Run the server on the specified address and port.
    pub fn run(&self, address: SocketAddr) {
        let handle = Handle::default();
        let listener = TcpListener::bind(&address).unwrap();
        println!("Listening on: {}", address);

        let init = future::lazy(|_| -> FutureResult<(), PathfinderError> {
            self.rabbitmq_client.clone().write().unwrap().init();
            Ok(())
        });

        let server = listener.incoming().for_each(move |stream| {
            let addr = stream.peer_addr().expect("Connected stream should have a peer address.");

            let engine_local = self.engine.clone();
            let connections_local = self.connections.clone();
            let auth_middleware_local = self.auth_middleware.clone();
            let handle_local = handle.clone();

            accept_async(stream)
                // Process the messages
                .and_then(move |ws_stream| {
                    // Create a channel for the stream, which other sockets will use to
                    // send us messages. It could be used for broadcasting your data to
                    // another users in the future.
                    let (tx, rx) = mpsc::unbounded();
                    connections_local.lock().unwrap().insert(addr, tx);

                    // Split the WebSocket stream so that it will be possible to work
                    // with the reading and writing halves separately.
                    let (sink, stream) = ws_stream.split();

                    // Read and process each message
                    let handle_inner = handle_local.clone();
                    let connections_inner = connections_local.clone();
                    let ws_reader = stream.for_each(move |message: Message| {

                        // Get references to required components
                        let addr_nested = addr.clone();
                        let engine_nested = engine_local.clone();
                        let connections_nested = connections_inner.clone();
                        let transmitter_nested = &connections_nested.clone().lock().unwrap()[&addr_nested];
                        let transmitter_nested2 = transmitter_nested.clone();

                        // 1. Deserialize message into JSON
                        let json_message = match engine_local.read().unwrap().deserialize_message(&message) {
                            Ok(json_message) => json_message,
                            Err(err) => {
                                let formatted_error = format!("{}", err);
                                let error_message = engine_nested.read().unwrap().wrap_an_error(formatted_error.as_str());
                                transmitter_nested.unbounded_send(error_message).unwrap();
                                return Ok(())
                            }
                        };

                        // 2. Apply a middleware to each incoming message
                        let auth_future = auth_middleware_local.read().unwrap()
                            .process_request(json_message.clone());

                        // 3. Put request into a queue in RabbitMQ and receive the response
                        let rabbitmq_future = engine_local.read().unwrap().handle(
                            json_message.clone(), transmitter_nested.clone(), &handle_inner
                        );

                        let processing_request_future = auth_future
                            .and_then(move |_| rabbitmq_future)
                            .map_err(move |err| {
                                let formatted_error = format!("{}", err);
                                let error_message = engine_nested.read().unwrap().wrap_an_error(formatted_error.as_str());
                                transmitter_nested2.unbounded_send(error_message).unwrap();
                                ()
                            });

                        spawn(processing_request_future);
                        Ok(())
                    });

                    // Write back prepared responses
                    let ws_writer = rx.fold(sink, |mut sink, msg| {
                        sink.start_send(msg).unwrap();
                        Ok(sink)
                    });

                    // Wait for either half to be done to tear down the other
                    let connection = ws_reader.map(|_| ()).map_err(|_| ())
                        .select(ws_writer.map(|_| ()).map_err(|_| ()));

                    // Close the connection after using
                    spawn(connection.then(move |_| {
                        connections_inner.lock().unwrap().remove(&addr);
                        debug!("Connection {} closed.", addr);
                        Ok(())
                    }));

                    Ok(())
                })
                // An error occurred during the WebSocket handshake
                .or_else(|err| {
                    debug!("{}", err.description());
                    Ok(())
                })
        });

        // Run the server
        run(init.and_then(|_| server.map_err(|_e| ())));
    }
}
<|endoftext|>


<commit_before>use std::cell::{RefCell};
use std::collections::{HashMap};
use std::net::{SocketAddr};
use std::rc::{Rc};

use super::engine::{Engine};
use super::router::{Router};
use super::middleware::{Middleware, EmptyMiddleware, AuthTokenMiddleware};

use cli::{CliOptions};
use futures::sync::{mpsc};
use futures::{Future, Sink};
use futures::stream::{Stream};
use tokio_core::net::{TcpListener};
use tokio_core::reactor::{Core};
use tokio_tungstenite::{accept_async};
use tungstenite::protocol::{Message};

pub struct Proxy {
    engine: Rc<RefCell<Engine>>,
    connections: Rc<RefCell<HashMap<SocketAddr, mpsc::UnboundedSender<Message>>>>,
    auth_middleware: Rc<RefCell<Box<Middleware>>>,
}

impl Proxy {
    pub fn new(router: Box<Router>, cli: &CliOptions) -> Proxy {
        let auth_middleware: Box<Middleware> = match cli.validate {
            true => Box::new(AuthTokenMiddleware::new(cli)),
            _ => Box::new(EmptyMiddleware::new(cli))
        };

        Proxy {
            engine: Rc::new(RefCell::new(Engine::new(router))),
            connections: Rc::new(RefCell::new(HashMap::new())),
            auth_middleware: Rc::new(RefCell::new(auth_middleware)),
        }
    }

    pub fn run(&self, address: SocketAddr) {
        let mut core = Core::new().unwrap();
        let handle = core.handle();
        let socket = TcpListener::bind(&address, &handle).unwrap();
        println!("Listening on: {}", address);

        let server = socket.incoming().for_each(|(stream, addr)| {
            let engine_inner = self.engine.clone();
            let connections_inner = self.connections.clone();
            let auth_middleware_inner = self.auth_middleware.clone();
            let handle_inner = handle.clone();

            accept_async(stream)
                // Check the Auth header
                .and_then(move |ws_stream| {
                    auth_middleware_inner.borrow().process_request();
                    Ok(ws_stream)
                })
                // Process the messages
                .and_then(move |ws_stream| {
                    // Create a channel for the stream, which other sockets will use to
                    // send us messages. It could be used for broadcasting your data to
                    // another users in the future.
                    let (tx, rx) = mpsc::unbounded();
                    connections_inner.borrow_mut().insert(addr, tx);

                    // Split the WebSocket stream so that it will be possible to work
                    // with the reading and writing halves separately.
                    let (sink, stream) = ws_stream.split();

                    // Read and process each message
                    let connections = connections_inner.clone();
                    let ws_reader = stream.for_each(move |message: Message| {
                        engine_inner.borrow().handle(&message, &addr, &connections);
                        Ok(())
                    });

                    // Write back prepared responses
                    let ws_writer = rx.fold(sink, |mut sink, msg| {
                        sink.start_send(msg).unwrap();
                        Ok(sink)
                    });

                    // Wait for either half to be done to tear down the other
                    let connection = ws_reader.map(|_| ()).map_err(|_| ())
                                              .select(ws_writer.map(|_| ()).map_err(|_| ()));

                    // Close the connection after using
                    handle_inner.spawn(connection.then(move |_| {
                        connections_inner.borrow_mut().remove(&addr);
                        println!("Connection {} closed.", addr);
                        Ok(())
                    }));

                    Ok(())
                }).or_else(|err| {
                    println!("An error occurred during the WebSocket handshake: {}", err);
                    Ok(())
                })
        });

        // Run the server
        core.run(server).unwrap();
    }
}
<commit_msg>Added callback for processing header<commit_after>use std::cell::{RefCell};
use std::collections::{HashMap};
use std::net::{SocketAddr};
use std::rc::{Rc};

use super::engine::{Engine};
use super::router::{Router};
use super::middleware::{Middleware, EmptyMiddleware, AuthTokenMiddleware};

use cli::{CliOptions};
use futures::sync::{mpsc};
use futures::{Future, Sink};
use futures::stream::{Stream};
use tokio_core::net::{TcpListener};
use tokio_core::reactor::{Core};
use tokio_tungstenite::{accept_hdr_async};
use tungstenite::handshake::server::{Request};
use tungstenite::protocol::{Message};

pub struct Proxy {
    engine: Rc<RefCell<Engine>>,
    connections: Rc<RefCell<HashMap<SocketAddr, mpsc::UnboundedSender<Message>>>>,
    auth_middleware: Rc<RefCell<Box<Middleware>>>,
}

impl Proxy {
    pub fn new(router: Box<Router>, cli: &CliOptions) -> Proxy {
        let auth_middleware: Box<Middleware> = match cli.validate {
            true => Box::new(AuthTokenMiddleware::new(cli)),
               _ => Box::new(EmptyMiddleware::new(cli))
        };

        Proxy {
            engine: Rc::new(RefCell::new(Engine::new(router))),
            connections: Rc::new(RefCell::new(HashMap::new())),
            auth_middleware: Rc::new(RefCell::new(auth_middleware)),
        }
    }

    pub fn run(&self, address: SocketAddr) {
        let mut core = Core::new().unwrap();
        let handle = core.handle();
        let socket = TcpListener::bind(&address, &handle).unwrap();
        println!("Listening on: {}", address);

        let server = socket.incoming().for_each(|(stream, addr)| {
            let engine_inner = self.engine.clone();
            let connections_inner = self.connections.clone();
            let handle_inner = handle.clone();

            let auth_middleware_callback = |req: &Request| {
                println!("Received a new ws handshake");
                println!("The request's path is: {}", req.path);
                println!("The request's headers are:");
                for &(ref header, ref value) in req.headers.iter() {
                    use std::str;
                    let v = str::from_utf8(value).unwrap();
                    println!("* {}: {:?}", header, v);
                }

                let auth_middleware_inner = self.auth_middleware.clone();
                let _result = auth_middleware_inner.borrow().process_request();
                Ok(None)
            };

            accept_hdr_async(stream, auth_middleware_callback)
                // Process the messages
                .and_then(move |ws_stream| {
                    // Create a channel for the stream, which other sockets will use to
                    // send us messages. It could be used for broadcasting your data to
                    // another users in the future.
                    let (tx, rx) = mpsc::unbounded();
                    connections_inner.borrow_mut().insert(addr, tx);

                    // Split the WebSocket stream so that it will be possible to work
                    // with the reading and writing halves separately.
                    let (sink, stream) = ws_stream.split();

                    // Read and process each message
                    let connections = connections_inner.clone();
                    let ws_reader = stream.for_each(move |message: Message| {
                        engine_inner.borrow().handle(&message, &addr, &connections);
                        Ok(())
                    });

                    // Write back prepared responses
                    let ws_writer = rx.fold(sink, |mut sink, msg| {
                        sink.start_send(msg).unwrap();
                        Ok(sink)
                    });

                    // Wait for either half to be done to tear down the other
                    let connection = ws_reader.map(|_| ()).map_err(|_| ())
                                              .select(ws_writer.map(|_| ()).map_err(|_| ()));

                    // Close the connection after using
                    handle_inner.spawn(connection.then(move |_| {
                        connections_inner.borrow_mut().remove(&addr);
                        println!("Connection {} closed.", addr);
                        Ok(())
                    }));

                    Ok(())
                }).or_else(|err| {
                    println!("An error occurred during the WebSocket handshake: {}", err);
                    Ok(())
                })
        });

        // Run the server
        core.run(server).unwrap();
    }
}
<|endoftext|>


<commit_before><commit_msg>removed shitty file<commit_after><|endoftext|>


<commit_before>#![allow(dead_code)]
use std::io::Result as IoResult;
use std::io::Error  as IoError;
use std::convert::AsRef;
use super::imagemagick_commands;
use image_tools::image_ops::{ElementaryPageOperations, Pixels, Direction};
use image_tools::image_ops::ImageResolution;
use image_tools::image_ops::RunOperation;
use image_tools::image_ops::OperationResults;

type ImageMagickArg = String;
type FileName = String;
type FilePath = String;

#[derive(Clone, Debug)]
enum ImageMagickOpType {
    Identify,
    IdentifyVerbose,
    Mogrify,
    Convert,
    NoOperation,
}

trait AsShellCommand {
    fn as_shell_command(&self) -> String;
}

impl AsShellCommand for ImageMagickOpType {
    fn as_shell_command(&self) -> String {
        match *self {
            ImageMagickOpType::Identify        => String::from("identify"),
            ImageMagickOpType::IdentifyVerbose => String::from("identify"),
            ImageMagickOpType::Mogrify         => String::from("mogrify"),
            ImageMagickOpType::Convert         => String::from("convert"),
            ImageMagickOpType::NoOperation     => String::from("echo"),
        }
    }
}

#[derive(Clone, Debug)]
struct ImageMagickArgs {
    file_path: FilePath,
    file_name: FileName,
    img_args:  Vec<ImageMagickArg>,
}

impl ImageMagickArgs {
    fn new(file_path: FilePath, file_name: FileName, img_args: &Vec<ImageMagickArg>) -> ImageMagickArgs {
        ImageMagickArgs {
            file_path: file_path,
            file_name: file_name,
            img_args:  img_args.clone(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct ElementaryImageMagickOperation {
    op:   ImageMagickOpType,
    args: ImageMagickArgs,
}

impl ElementaryImageMagickOperation {
    fn new(op: ImageMagickOpType, args: ImageMagickArgs) -> ElementaryImageMagickOperation {
        ElementaryImageMagickOperation {
            op:   op,
            args: args,
        }
    }

    fn arg(&mut self, arg: ImageMagickArg) {
        self.args.img_args.push(arg);
    }

    fn args(&mut self, args: &[ImageMagickArg]) {
        for arg in args {
            self.args.img_args.push(arg.clone());
        }
    }

    fn run_operation(&self) -> IoResult<String> {
        match self.op {
            ImageMagickOpType::Identify        => {
                imagemagick_commands::imagemagick_identify_default(&self.args.file_path)
            }
            ImageMagickOpType::IdentifyVerbose => {
                imagemagick_commands::imagemagick_identify_verbose(&self.args.file_path)
            }
            ImageMagickOpType::Convert         => {
                imagemagick_commands::imagemagick_convert(&self.args.file_path, &self.args.img_args)
            }
            ImageMagickOpType::Mogrify         => {
                imagemagick_commands::imagemagick_mogrify(&self.args.file_path, &self.args.img_args)
            }
            ImageMagickOpType::NoOperation     => {
                imagemagick_commands::imagemagick_no_operation()
            }
        }
    }
}

#[derive(Clone, Debug)]
pub struct ImageMagickOperation {
    ops: Vec<ElementaryImageMagickOperation>,
}

impl ImageMagickOperation {
    fn new() -> ImageMagickOperation {
        ImageMagickOperation {
            ops: Vec::new(),
        }
    }

    fn add_op(&mut self, op: ElementaryImageMagickOperation)  {
        self.ops.push(op);
    }

    fn add_ops(&mut self, ops: &[ElementaryImageMagickOperation]) {
        for op in ops.iter() {
            self.ops.push(op.clone());
        }
    }
}

// This implementation will be the generator for the sequence of
// ImageMagick commands to implement each operation.
impl ElementaryPageOperations for ImageMagickOperation {

    fn identify(file_name: FileName, file_path: FilePath) -> ImageMagickOperation {
        // Identify presesntly generates an identify command without the -verbose flag.
        // This may change in the future where we use verbose as the default.
        let args = ImageMagickArgs::new(file_name, file_path, &Vec::new());
        let elem_op = ElementaryImageMagickOperation::new(ImageMagickOpType::Identify, args);
        let mut op = ImageMagickOperation::new();
        op.add_op(elem_op);

        op

    }

    fn rescale(amount: Pixels, dir: Direction) -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_left_edge(amount: Pixels)        -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_right_edge(amount: Pixels)       -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_top_edge(amount: Pixels)         -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_bottom_edge(amount: Pixels)      -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_left_edge(amount: Pixels)          -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_right_edge(amount: Pixels)         -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_top_edge(amount: Pixels)           -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_bottom_edge(amount: Pixels)        -> ImageMagickOperation {
        unimplemented!();
    }

    fn set_resolution(res: ImageResolution)    -> ImageMagickOperation {
        unimplemented!();
    }
    fn no_operation()                          -> ImageMagickOperation {
        unimplemented!();
    }
}

impl AsRef<[ElementaryImageMagickOperation]> for ImageMagickOperation {
    fn as_ref(&self) -> &[ElementaryImageMagickOperation] {
        self.ops.as_ref()
    }
}

impl RunOperation for ImageMagickOperation {
    fn run_operation(op: ImageMagickOperation) -> OperationResults {
        let mut results = OperationResults::new();

        for action in op.ops {
            let mut result = Vec::new();
            result.push(action.run_operation());
            results.append(&mut OperationResults::from(&mut result));
        }

        results
    }
}
<commit_msg>Implement NoOperation for ImageMagick.<commit_after>#![allow(dead_code)]
use std::io::Result as IoResult;
use std::io::Error  as IoError;
use std::convert::AsRef;
use super::imagemagick_commands;
use image_tools::image_ops::{ElementaryPageOperations, Pixels, Direction};
use image_tools::image_ops::ImageResolution;
use image_tools::image_ops::RunOperation;
use image_tools::image_ops::OperationResults;

type ImageMagickArg = String;
type FileName = String;
type FilePath = String;

#[derive(Clone, Debug)]
enum ImageMagickOpType {
    Identify,
    IdentifyVerbose,
    Mogrify,
    Convert,
    NoOperation,
}

trait AsShellCommand {
    fn as_shell_command(&self) -> String;
}

impl AsShellCommand for ImageMagickOpType {
    fn as_shell_command(&self) -> String {
        match *self {
            ImageMagickOpType::Identify        => String::from("identify"),
            ImageMagickOpType::IdentifyVerbose => String::from("identify"),
            ImageMagickOpType::Mogrify         => String::from("mogrify"),
            ImageMagickOpType::Convert         => String::from("convert"),
            ImageMagickOpType::NoOperation     => String::from("echo"),
        }
    }
}

#[derive(Clone, Debug)]
struct ImageMagickArgs {
    file_path: FilePath,
    file_name: FileName,
    img_args:  Vec<ImageMagickArg>,
}

impl ImageMagickArgs {
    fn new(file_path: FilePath, file_name: FileName, img_args: &Vec<ImageMagickArg>) -> ImageMagickArgs {
        ImageMagickArgs {
            file_path: file_path,
            file_name: file_name,
            img_args:  img_args.clone(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct ElementaryImageMagickOperation {
    op:   ImageMagickOpType,
    args: ImageMagickArgs,
}

impl ElementaryImageMagickOperation {
    fn new(op: ImageMagickOpType, args: ImageMagickArgs) -> ElementaryImageMagickOperation {
        ElementaryImageMagickOperation {
            op:   op,
            args: args,
        }
    }

    fn arg(&mut self, arg: ImageMagickArg) {
        self.args.img_args.push(arg);
    }

    fn args(&mut self, args: &[ImageMagickArg]) {
        for arg in args {
            self.args.img_args.push(arg.clone());
        }
    }

    fn run_operation(&self) -> IoResult<String> {
        match self.op {
            ImageMagickOpType::Identify        => {
                imagemagick_commands::imagemagick_identify_default(&self.args.file_path)
            }
            ImageMagickOpType::IdentifyVerbose => {
                imagemagick_commands::imagemagick_identify_verbose(&self.args.file_path)
            }
            ImageMagickOpType::Convert         => {
                imagemagick_commands::imagemagick_convert(&self.args.file_path, &self.args.img_args)
            }
            ImageMagickOpType::Mogrify         => {
                imagemagick_commands::imagemagick_mogrify(&self.args.file_path, &self.args.img_args)
            }
            ImageMagickOpType::NoOperation     => {
                imagemagick_commands::imagemagick_no_operation()
            }
        }
    }
}

#[derive(Clone, Debug)]
pub struct ImageMagickOperation {
    ops: Vec<ElementaryImageMagickOperation>,
}

impl ImageMagickOperation {
    fn new() -> ImageMagickOperation {
        ImageMagickOperation {
            ops: Vec::new(),
        }
    }

    fn add_op(&mut self, op: ElementaryImageMagickOperation)  {
        self.ops.push(op);
    }

    fn add_ops(&mut self, ops: &[ElementaryImageMagickOperation]) {
        for op in ops.iter() {
            self.ops.push(op.clone());
        }
    }

}

// This implementation will be the generator for the sequence of
// ImageMagick commands to implement each operation.
impl ElementaryPageOperations for ImageMagickOperation {

    fn identify(file_name: FileName, file_path: FilePath) -> ImageMagickOperation {
        // Identify presesntly generates an identify command without the -verbose flag.
        // This may change in the future where we use verbose as the default.
        let args = ImageMagickArgs::new(file_name, file_path, &Vec::new());
        let elem_op = ElementaryImageMagickOperation::new(ImageMagickOpType::Identify, args);
        let mut op = ImageMagickOperation::new();
        op.add_op(elem_op);

        op

    }

    fn rescale(amount: Pixels, dir: Direction) -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_left_edge(amount: Pixels)        -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_right_edge(amount: Pixels)       -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_top_edge(amount: Pixels)         -> ImageMagickOperation {
        unimplemented!();
    }

    fn expand_bottom_edge(amount: Pixels)      -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_left_edge(amount: Pixels)          -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_right_edge(amount: Pixels)         -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_top_edge(amount: Pixels)           -> ImageMagickOperation {
        unimplemented!();
    }

    fn trim_bottom_edge(amount: Pixels)        -> ImageMagickOperation {
        unimplemented!();
    }

    fn set_resolution(res: ImageResolution)    -> ImageMagickOperation {
        unimplemented!();
    }

    fn no_operation()                          -> ImageMagickOperation {
        let mut op = ImageMagickOperation::new();
        let args = ImageMagickArgs::new(String::from(""), String::from(""), &Vec::new());
        let elem_op = ElementaryImageMagickOperation::new(ImageMagickOpType::NoOperation, args);
        op.add_op(elem_op);

        op
    }
}

impl AsRef<[ElementaryImageMagickOperation]> for ImageMagickOperation {
    fn as_ref(&self) -> &[ElementaryImageMagickOperation] {
        self.ops.as_ref()
    }
}

impl RunOperation for ImageMagickOperation {
    fn run_operation(op: ImageMagickOperation) -> OperationResults {
        let mut results = OperationResults::new();

        for action in op.ops {
            let mut result = Vec::new();
            result.push(action.run_operation());
            results.append(&mut OperationResults::from(&mut result));
        }

        results
    }
}
<|endoftext|>