<commit_before>//! Interpretation of DICOM data sets as streams of tokens.
use dicom_core::header::{DataElementHeader, Length, VR};
use dicom_core::value::{DicomValueType, PrimitiveValue};
use dicom_core::{value::Value, DataElement, Tag};
use std::fmt;

pub mod read;
pub mod write;

pub use self::read::DataSetReader;
pub use self::write::DataSetWriter;

/// A token of a DICOM data set stream. This is part of the interpretation of a
/// data set as a stream of symbols, which may either represent data headers or
/// actual value data.
#[derive(Debug, Clone)]
pub enum DataToken {
    /// A data header of a primitive value.
    ElementHeader(DataElementHeader),
    /// The beginning of a sequence element.
    SequenceStart { tag: Tag, len: Length },
    /// The beginning of an encapsulated pixel data element.
    PixelSequenceStart,
    /// The ending delimiter of a sequence or encapsulated pixel data.
    SequenceEnd,
    /// The beginning of a new item in the sequence.
    ItemStart { len: Length },
    /// The ending delimiter of an item.
    ItemEnd,
    /// A primitive data element value.
    PrimitiveValue(PrimitiveValue),
    /// An owned piece of raw data representing an item's value.
    ///
    /// This variant is used to represent the value of an offset table or a
    /// compressed fragment. It should not be used to represent nested data
    /// sets.
    ItemValue(Vec<u8>),
}

impl fmt::Display for DataToken {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            DataToken::PrimitiveValue(ref v) => write!(f, "PrimitiveValue({:?})", v.value_type()),
            other => write!(f, "{:?}", other),
        }
    }
}

/// This implementation treats undefined lengths as equal.
impl PartialEq<Self> for DataToken {
    fn eq(&self, other: &Self) -> bool {
        use DataToken::*;
        match (self, other) {
            (
                ElementHeader(DataElementHeader {
                    tag: tag1,
                    vr: vr1,
                    len: len1,
                }),
                ElementHeader(DataElementHeader {
                    tag: tag2,
                    vr: vr2,
                    len: len2,
                }),
            ) => tag1 == tag2 && vr1 == vr2 && len1.inner_eq(*len2),
            (
                SequenceStart {
                    tag: tag1,
                    len: len1,
                },
                SequenceStart {
                    tag: tag2,
                    len: len2,
                },
            ) => tag1 == tag2 && len1.inner_eq(*len2),
            (ItemStart { len: len1 }, ItemStart { len: len2 }) => len1.inner_eq(*len2),
            (PrimitiveValue(v1), PrimitiveValue(v2)) => v1 == v2,
            (ItemValue(v1), ItemValue(v2)) => v1 == v2,
            (ItemEnd, ItemEnd)
            | (SequenceEnd, SequenceEnd)
            | (PixelSequenceStart, PixelSequenceStart) => true,
            _ => false,
        }
    }
}

impl From<DataElementHeader> for DataToken {
    fn from(header: DataElementHeader) -> Self {
        match (header.vr(), header.tag) {
            (VR::OB, Tag(0x7fe0, 0x0010)) if header.len.is_undefined() => {
                DataToken::PixelSequenceStart
            }
            (VR::SQ, _) => DataToken::SequenceStart {
                tag: header.tag,
                len: header.len,
            },
            _ => DataToken::ElementHeader(header),
        }
    }
}

impl DataToken {
    /// Check whether this token represents the start of a sequence
    /// of nested data sets.
    pub fn is_sequence_start(&self) -> bool {
        match self {
            DataToken::SequenceStart { .. } => true,
            _ => false,
        }
    }

    /// Check whether this token represents the end of a sequence
    /// or the end of an encapsulated element.
    pub fn is_sequence_end(&self) -> bool {
        match self {
            DataToken::SequenceEnd => true,
            _ => false,
        }
    }
}

/// The type of delimiter: sequence or item.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum SeqTokenType {
    Sequence,
    Item,
}

/// A trait for converting structured DICOM data into a stream of data tokens.
pub trait IntoTokens {
    /// The iterator type through which tokens are obtained.
    type Iter: Iterator<Item = DataToken>;

    /// Convert the value into tokens.
    fn into_tokens(self) -> Self::Iter;
}

impl IntoTokens for dicom_core::header::EmptyObject {
    type Iter = std::iter::Empty<DataToken>;

    fn into_tokens(self) -> Self::Iter {
        unreachable!()
    }
}

/// Token generator from a DICOM data element.
pub enum DataElementTokens
where
    I: IntoTokens,
{
    /// initial state, at the beginning of the element
    Start(
        // Option is used for easy taking from a &mut,
        // should always be Some in practice
        Option<DataElement>,
    ),
    /// the header of a plain primitive element was read
    Header(
        // Option is used for easy taking from a &mut,
        // should always be Some in practice
        Option<DataElement>,
    ),
    /// reading tokens from items
    Items(
        FlattenTokens<
            <dicom_core::value::C<AsItem> as IntoIterator>::IntoIter,
            ItemTokens,
        >,
    ),
    /// the header of encapsulated pixel data was read, will read
    /// the offset table next
    PixelData(
        // Pixel fragments; Option is used for easy taking from a &mut,
        // should always be Some in practice
        Option<dicom_core::value::C
>,
        ItemValueTokens<dicom_core::value::C<u8>>,
    ),
    /// the header and offset of encapsulated pixel data was read,
    /// fragments come next
    PixelDataFragments(
        FlattenTokens<
            <dicom_core::value::C<ItemValue
> as IntoIterator>::IntoIter,
            ItemValueTokens
,
        >,
    ),
    /// no more elements
    End,
}

impl Iterator for DataElementTokens
where
    I: IntoTokens,
    P: AsRef<[u8]>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        let (out, next_state) = match self {
            DataElementTokens::Start(elem) => {
                let elem = elem.take().unwrap();
                // data element header token
                let header = *elem.header();

                let token = dbg!(DataToken::from(header));
                match token {
                    DataToken::SequenceStart { .. } => {
                        // retrieve sequence value, begin item sequence
                        match elem.into_value() {
                            Value::Primitive(_) | Value::PixelSequence { .. } => unreachable!(),
                            Value::Sequence { items, size: _ } => {
                                let items: dicom_core::value::C<_> = items
                                    .into_iter()
                                    .map(|o| AsItem(Length::UNDEFINED, o))
                                    .collect();
                                (Some(token), DataElementTokens::Items(items.into_tokens()))
                            }
                        }
                    }
                    DataToken::PixelSequenceStart => {
                        match elem.into_value() {
                            Value::PixelSequence {
                                fragments,
                                offset_table,
                            } => {
                                (
                                    // begin pixel sequence
                                    Some(DataToken::PixelSequenceStart),
                                    DataElementTokens::PixelData(
                                        Some(fragments),
                                        ItemValue(offset_table).into_tokens(),
                                    ),
                                )
                            }
                            Value::Primitive(_) | Value::Sequence { .. } => unreachable!(),
                        }
                    }
                    _ => (
                        Some(DataToken::ElementHeader(*elem.header())),
                        DataElementTokens::Header(Some(elem)),
                    ),
                }
            }
            DataElementTokens::Header(elem) => {
                let elem = elem.take().unwrap();
                match elem.into_value() {
                    Value::Sequence { .. } | Value::PixelSequence { .. } => unreachable!(),
                    Value::Primitive(value) => {
                        // return primitive value, done
                        let token = DataToken::PrimitiveValue(value);
                        (Some(token), DataElementTokens::End)
                    }
                }
            }
            DataElementTokens::Items(tokens) => {
                if let Some(token) = tokens.next() {
                    // bypass manual state transition
                    return Some(token);
                } else {
                    // sequence end token, end
                    (Some(DataToken::SequenceEnd), DataElementTokens::End)
                }
            }
            DataElementTokens::PixelData(fragments, tokens) => {
                if let Some(token) = tokens.next() {
                    // bypass manual state transition
                    return Some(token);
                }
                // pixel data fragments next
                let fragments = fragments.take().unwrap();
                let tokens: dicom_core::value::C<_> =
                    fragments.into_iter().map(|o| ItemValue(o)).collect();
                *self = DataElementTokens::PixelDataFragments(tokens.into_tokens());
                // recursive call to ensure the retrieval of a data token
                return self.next();
            }
            DataElementTokens::PixelDataFragments(tokens) => {
                if let Some(token) = tokens.next() {
                    // bypass manual state transition
                    return Some(token);
                } else {
                    // sequence end token, end
                    (Some(DataToken::SequenceEnd), DataElementTokens::End)
                }
            }
            DataElementTokens::End => return None,
        };
        *self = next_state;

        out
    }
}

impl IntoTokens for DataElement
where
    I: IntoTokens,
    P: AsRef<[u8]>,
{
    type Iter = DataElementTokens;

    fn into_tokens(self) -> Self::Iter {
        DataElementTokens::Start(Some(self))
    }
}

/// Flatten a sequence of elements into their respective
/// token sequence in order.
#[derive(Debug, PartialEq)]
pub struct FlattenTokens<O, K> {
    seq: O,
    tokens: Option<K>,
}

impl<O, K> Iterator for FlattenTokens<O, K>
where
    O: Iterator,
    O::Item: IntoTokens<Iter = K>,
    K: Iterator<Item = DataToken>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        // ensure a token sequence
        if self.tokens.is_none() {
            match self.seq.next() {
                Some(entries) => {
                    self.tokens = Some(entries.into_tokens());
                }
                None => return None,
            }
        }

        // retrieve the next token
        match self.tokens.as_mut().map(|s| s.next()) {
            Some(Some(token)) => Some(token),
            Some(None) => {
                self.tokens = None;
                self.next()
            }
            None => unreachable!(),
        }
    }
}

impl<T> IntoTokens for Vec<T>
where
    T: IntoTokens,
{
    type Iter = FlattenTokens<<Vec<T> as IntoIterator>::IntoIter, <T as IntoTokens>::Iter>;

    fn into_tokens(self) -> Self::Iter {
        FlattenTokens {
            seq: self.into_iter(),
            tokens: None,
        }
    }
}

impl<T> IntoTokens for dicom_core::value::C<T>
where
    T: IntoTokens,
{
    type Iter =
        FlattenTokens<<dicom_core::value::C<T> as IntoIterator>::IntoIter, <T as IntoTokens>::Iter>;

    fn into_tokens(self) -> Self::Iter {
        FlattenTokens {
            seq: self.into_iter(),
            tokens: None,
        }
    }
}

// A stream of tokens from a DICOM item.
#[derive(Debug)]
pub enum ItemTokens<T> {
    /// Just started, an item header token will come next
    Start {
        len: Length,
        object_tokens: Option<T>,
    },
    /// Will return tokens from the inner object, then an end of item token
    /// when it ends
    Object { object_tokens: T },
    /// Just ended, no more tokens
    End,
}

impl<T> ItemTokens<T>
where
    T: Iterator<Item = DataToken>,
{
    pub fn new<O>(len: Length, object: O) -> Self
    where
        O: IntoTokens<Iter = T>,
    {
        ItemTokens::Start {
            len,
            object_tokens: Some(object.into_tokens()),
        }
    }
}

impl<T> Iterator for ItemTokens<T>
where
    T: Iterator<Item = DataToken>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        let (next_state, out) = match self {
            ItemTokens::Start { len, object_tokens } => (
                ItemTokens::Object {
                    object_tokens: object_tokens.take().unwrap(),
                },
                Some(DataToken::ItemStart { len: *len }),
            ),
            ItemTokens::Object { object_tokens } => {
                if let Some(token) = object_tokens.next() {
                    return Some(token);
                } else {
                    (ItemTokens::End, Some(DataToken::ItemEnd))
                }
            }
            ItemTokens::End => {
                return None;
            }
        };

        *self = next_state;
        out
    }
}

/// A newtype for interpreting the given data as an item.
/// When converting a value of this type into tokens, the inner value's tokens
/// will be surrounded by an item start and an item delimiter.
#[derive(Debug, Clone, PartialEq)]
pub struct AsItem(Length, I);

impl IntoTokens for AsItem
where
    I: IntoTokens,
{
    type Iter = ItemTokens;

    fn into_tokens(self) -> Self::Iter {
        ItemTokens::new(self.0, self.1)
    }
}

/// A newtype for wrapping a piece of raw data into an item.
/// When converting a value of this type into tokens, the algorithm
/// will create an item start with an explicit length, followed by
/// an item value token, then an item delimiter.
#[derive(Debug, Clone, PartialEq)]
pub struct ItemValue
(P);

impl
 IntoTokens for ItemValue

where
    P: AsRef<[u8]>,
{
    type Iter = ItemValueTokens
;

    fn into_tokens(self) -> Self::Iter {
        ItemValueTokens::new(self.0)
    }
}

#[derive(Debug)]
pub enum ItemValueTokens
 {
    /// Just started, an item header token will come next
    Start(Option
),
    /// Will return a token of the value
    Value(P),
    /// Will return an end of item token
    Done,
    /// Just ended, no more tokens
    End,
}

impl
 ItemValueTokens
 {
    pub fn new(value: P) -> Self {
        ItemValueTokens::Start(Some(value))
    }
}

impl
 Iterator for ItemValueTokens

where
    P: AsRef<[u8]>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        let (out, next_state) = match self {
            ItemValueTokens::Start(value) => {
                let value = value.take().unwrap();
                let len = Length(value.as_ref().len() as u32);

                (
                    Some(DataToken::ItemStart { len }),
                    if len == Length(0) {
                        ItemValueTokens::Done
                    } else {
                        ItemValueTokens::Value(value)
                    },
                )
            }
            ItemValueTokens::Value(value) => (
                Some(DataToken::ItemValue(value.as_ref().to_owned())),
                ItemValueTokens::Done,
            ),
            ItemValueTokens::Done => (Some(DataToken::ItemEnd), ItemValueTokens::End),
            ItemValueTokens::End => return None,
        };

        *self = next_state;
        out
    }
}
<commit_msg>Remove dbg line<commit_after>//! Interpretation of DICOM data sets as streams of tokens.
use dicom_core::header::{DataElementHeader, Length, VR};
use dicom_core::value::{DicomValueType, PrimitiveValue};
use dicom_core::{value::Value, DataElement, Tag};
use std::fmt;

pub mod read;
pub mod write;

pub use self::read::DataSetReader;
pub use self::write::DataSetWriter;

/// A token of a DICOM data set stream. This is part of the interpretation of a
/// data set as a stream of symbols, which may either represent data headers or
/// actual value data.
#[derive(Debug, Clone)]
pub enum DataToken {
    /// A data header of a primitive value.
    ElementHeader(DataElementHeader),
    /// The beginning of a sequence element.
    SequenceStart { tag: Tag, len: Length },
    /// The beginning of an encapsulated pixel data element.
    PixelSequenceStart,
    /// The ending delimiter of a sequence or encapsulated pixel data.
    SequenceEnd,
    /// The beginning of a new item in the sequence.
    ItemStart { len: Length },
    /// The ending delimiter of an item.
    ItemEnd,
    /// A primitive data element value.
    PrimitiveValue(PrimitiveValue),
    /// An owned piece of raw data representing an item's value.
    ///
    /// This variant is used to represent the value of an offset table or a
    /// compressed fragment. It should not be used to represent nested data
    /// sets.
    ItemValue(Vec<u8>),
}

impl fmt::Display for DataToken {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            DataToken::PrimitiveValue(ref v) => write!(f, "PrimitiveValue({:?})", v.value_type()),
            other => write!(f, "{:?}", other),
        }
    }
}

/// This implementation treats undefined lengths as equal.
impl PartialEq<Self> for DataToken {
    fn eq(&self, other: &Self) -> bool {
        use DataToken::*;
        match (self, other) {
            (
                ElementHeader(DataElementHeader {
                    tag: tag1,
                    vr: vr1,
                    len: len1,
                }),
                ElementHeader(DataElementHeader {
                    tag: tag2,
                    vr: vr2,
                    len: len2,
                }),
            ) => tag1 == tag2 && vr1 == vr2 && len1.inner_eq(*len2),
            (
                SequenceStart {
                    tag: tag1,
                    len: len1,
                },
                SequenceStart {
                    tag: tag2,
                    len: len2,
                },
            ) => tag1 == tag2 && len1.inner_eq(*len2),
            (ItemStart { len: len1 }, ItemStart { len: len2 }) => len1.inner_eq(*len2),
            (PrimitiveValue(v1), PrimitiveValue(v2)) => v1 == v2,
            (ItemValue(v1), ItemValue(v2)) => v1 == v2,
            (ItemEnd, ItemEnd)
            | (SequenceEnd, SequenceEnd)
            | (PixelSequenceStart, PixelSequenceStart) => true,
            _ => false,
        }
    }
}

impl From<DataElementHeader> for DataToken {
    fn from(header: DataElementHeader) -> Self {
        match (header.vr(), header.tag) {
            (VR::OB, Tag(0x7fe0, 0x0010)) if header.len.is_undefined() => {
                DataToken::PixelSequenceStart
            }
            (VR::SQ, _) => DataToken::SequenceStart {
                tag: header.tag,
                len: header.len,
            },
            _ => DataToken::ElementHeader(header),
        }
    }
}

impl DataToken {
    /// Check whether this token represents the start of a sequence
    /// of nested data sets.
    pub fn is_sequence_start(&self) -> bool {
        match self {
            DataToken::SequenceStart { .. } => true,
            _ => false,
        }
    }

    /// Check whether this token represents the end of a sequence
    /// or the end of an encapsulated element.
    pub fn is_sequence_end(&self) -> bool {
        match self {
            DataToken::SequenceEnd => true,
            _ => false,
        }
    }
}

/// The type of delimiter: sequence or item.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum SeqTokenType {
    Sequence,
    Item,
}

/// A trait for converting structured DICOM data into a stream of data tokens.
pub trait IntoTokens {
    /// The iterator type through which tokens are obtained.
    type Iter: Iterator<Item = DataToken>;

    /// Convert the value into tokens.
    fn into_tokens(self) -> Self::Iter;
}

impl IntoTokens for dicom_core::header::EmptyObject {
    type Iter = std::iter::Empty<DataToken>;

    fn into_tokens(self) -> Self::Iter {
        unreachable!()
    }
}

/// Token generator from a DICOM data element.
pub enum DataElementTokens
where
    I: IntoTokens,
{
    /// initial state, at the beginning of the element
    Start(
        // Option is used for easy taking from a &mut,
        // should always be Some in practice
        Option<DataElement>,
    ),
    /// the header of a plain primitive element was read
    Header(
        // Option is used for easy taking from a &mut,
        // should always be Some in practice
        Option<DataElement>,
    ),
    /// reading tokens from items
    Items(
        FlattenTokens<
            <dicom_core::value::C<AsItem> as IntoIterator>::IntoIter,
            ItemTokens,
        >,
    ),
    /// the header of encapsulated pixel data was read, will read
    /// the offset table next
    PixelData(
        // Pixel fragments; Option is used for easy taking from a &mut,
        // should always be Some in practice
        Option<dicom_core::value::C
>,
        ItemValueTokens<dicom_core::value::C<u8>>,
    ),
    /// the header and offset of encapsulated pixel data was read,
    /// fragments come next
    PixelDataFragments(
        FlattenTokens<
            <dicom_core::value::C<ItemValue
> as IntoIterator>::IntoIter,
            ItemValueTokens
,
        >,
    ),
    /// no more elements
    End,
}

impl Iterator for DataElementTokens
where
    I: IntoTokens,
    P: AsRef<[u8]>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        let (out, next_state) = match self {
            DataElementTokens::Start(elem) => {
                let elem = elem.take().unwrap();
                // data element header token
                let header = *elem.header();

                let token = DataToken::from(header);
                match token {
                    DataToken::SequenceStart { .. } => {
                        // retrieve sequence value, begin item sequence
                        match elem.into_value() {
                            Value::Primitive(_) | Value::PixelSequence { .. } => unreachable!(),
                            Value::Sequence { items, size: _ } => {
                                let items: dicom_core::value::C<_> = items
                                    .into_iter()
                                    .map(|o| AsItem(Length::UNDEFINED, o))
                                    .collect();
                                (Some(token), DataElementTokens::Items(items.into_tokens()))
                            }
                        }
                    }
                    DataToken::PixelSequenceStart => {
                        match elem.into_value() {
                            Value::PixelSequence {
                                fragments,
                                offset_table,
                            } => {
                                (
                                    // begin pixel sequence
                                    Some(DataToken::PixelSequenceStart),
                                    DataElementTokens::PixelData(
                                        Some(fragments),
                                        ItemValue(offset_table).into_tokens(),
                                    ),
                                )
                            }
                            Value::Primitive(_) | Value::Sequence { .. } => unreachable!(),
                        }
                    }
                    _ => (
                        Some(DataToken::ElementHeader(*elem.header())),
                        DataElementTokens::Header(Some(elem)),
                    ),
                }
            }
            DataElementTokens::Header(elem) => {
                let elem = elem.take().unwrap();
                match elem.into_value() {
                    Value::Sequence { .. } | Value::PixelSequence { .. } => unreachable!(),
                    Value::Primitive(value) => {
                        // return primitive value, done
                        let token = DataToken::PrimitiveValue(value);
                        (Some(token), DataElementTokens::End)
                    }
                }
            }
            DataElementTokens::Items(tokens) => {
                if let Some(token) = tokens.next() {
                    // bypass manual state transition
                    return Some(token);
                } else {
                    // sequence end token, end
                    (Some(DataToken::SequenceEnd), DataElementTokens::End)
                }
            }
            DataElementTokens::PixelData(fragments, tokens) => {
                if let Some(token) = tokens.next() {
                    // bypass manual state transition
                    return Some(token);
                }
                // pixel data fragments next
                let fragments = fragments.take().unwrap();
                let tokens: dicom_core::value::C<_> =
                    fragments.into_iter().map(|o| ItemValue(o)).collect();
                *self = DataElementTokens::PixelDataFragments(tokens.into_tokens());
                // recursive call to ensure the retrieval of a data token
                return self.next();
            }
            DataElementTokens::PixelDataFragments(tokens) => {
                if let Some(token) = tokens.next() {
                    // bypass manual state transition
                    return Some(token);
                } else {
                    // sequence end token, end
                    (Some(DataToken::SequenceEnd), DataElementTokens::End)
                }
            }
            DataElementTokens::End => return None,
        };
        *self = next_state;

        out
    }
}

impl IntoTokens for DataElement
where
    I: IntoTokens,
    P: AsRef<[u8]>,
{
    type Iter = DataElementTokens;

    fn into_tokens(self) -> Self::Iter {
        DataElementTokens::Start(Some(self))
    }
}

/// Flatten a sequence of elements into their respective
/// token sequence in order.
#[derive(Debug, PartialEq)]
pub struct FlattenTokens<O, K> {
    seq: O,
    tokens: Option<K>,
}

impl<O, K> Iterator for FlattenTokens<O, K>
where
    O: Iterator,
    O::Item: IntoTokens<Iter = K>,
    K: Iterator<Item = DataToken>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        // ensure a token sequence
        if self.tokens.is_none() {
            match self.seq.next() {
                Some(entries) => {
                    self.tokens = Some(entries.into_tokens());
                }
                None => return None,
            }
        }

        // retrieve the next token
        match self.tokens.as_mut().map(|s| s.next()) {
            Some(Some(token)) => Some(token),
            Some(None) => {
                self.tokens = None;
                self.next()
            }
            None => unreachable!(),
        }
    }
}

impl<T> IntoTokens for Vec<T>
where
    T: IntoTokens,
{
    type Iter = FlattenTokens<<Vec<T> as IntoIterator>::IntoIter, <T as IntoTokens>::Iter>;

    fn into_tokens(self) -> Self::Iter {
        FlattenTokens {
            seq: self.into_iter(),
            tokens: None,
        }
    }
}

impl<T> IntoTokens for dicom_core::value::C<T>
where
    T: IntoTokens,
{
    type Iter =
        FlattenTokens<<dicom_core::value::C<T> as IntoIterator>::IntoIter, <T as IntoTokens>::Iter>;

    fn into_tokens(self) -> Self::Iter {
        FlattenTokens {
            seq: self.into_iter(),
            tokens: None,
        }
    }
}

// A stream of tokens from a DICOM item.
#[derive(Debug)]
pub enum ItemTokens<T> {
    /// Just started, an item header token will come next
    Start {
        len: Length,
        object_tokens: Option<T>,
    },
    /// Will return tokens from the inner object, then an end of item token
    /// when it ends
    Object { object_tokens: T },
    /// Just ended, no more tokens
    End,
}

impl<T> ItemTokens<T>
where
    T: Iterator<Item = DataToken>,
{
    pub fn new<O>(len: Length, object: O) -> Self
    where
        O: IntoTokens<Iter = T>,
    {
        ItemTokens::Start {
            len,
            object_tokens: Some(object.into_tokens()),
        }
    }
}

impl<T> Iterator for ItemTokens<T>
where
    T: Iterator<Item = DataToken>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        let (next_state, out) = match self {
            ItemTokens::Start { len, object_tokens } => (
                ItemTokens::Object {
                    object_tokens: object_tokens.take().unwrap(),
                },
                Some(DataToken::ItemStart { len: *len }),
            ),
            ItemTokens::Object { object_tokens } => {
                if let Some(token) = object_tokens.next() {
                    return Some(token);
                } else {
                    (ItemTokens::End, Some(DataToken::ItemEnd))
                }
            }
            ItemTokens::End => {
                return None;
            }
        };

        *self = next_state;
        out
    }
}

/// A newtype for interpreting the given data as an item.
/// When converting a value of this type into tokens, the inner value's tokens
/// will be surrounded by an item start and an item delimiter.
#[derive(Debug, Clone, PartialEq)]
pub struct AsItem(Length, I);

impl IntoTokens for AsItem
where
    I: IntoTokens,
{
    type Iter = ItemTokens;

    fn into_tokens(self) -> Self::Iter {
        ItemTokens::new(self.0, self.1)
    }
}

/// A newtype for wrapping a piece of raw data into an item.
/// When converting a value of this type into tokens, the algorithm
/// will create an item start with an explicit length, followed by
/// an item value token, then an item delimiter.
#[derive(Debug, Clone, PartialEq)]
pub struct ItemValue
(P);

impl
 IntoTokens for ItemValue

where
    P: AsRef<[u8]>,
{
    type Iter = ItemValueTokens
;

    fn into_tokens(self) -> Self::Iter {
        ItemValueTokens::new(self.0)
    }
}

#[derive(Debug)]
pub enum ItemValueTokens
 {
    /// Just started, an item header token will come next
    Start(Option
),
    /// Will return a token of the value
    Value(P),
    /// Will return an end of item token
    Done,
    /// Just ended, no more tokens
    End,
}

impl
 ItemValueTokens
 {
    pub fn new(value: P) -> Self {
        ItemValueTokens::Start(Some(value))
    }
}

impl
 Iterator for ItemValueTokens

where
    P: AsRef<[u8]>,
{
    type Item = DataToken;

    fn next(&mut self) -> Option<Self::Item> {
        let (out, next_state) = match self {
            ItemValueTokens::Start(value) => {
                let value = value.take().unwrap();
                let len = Length(value.as_ref().len() as u32);

                (
                    Some(DataToken::ItemStart { len }),
                    if len == Length(0) {
                        ItemValueTokens::Done
                    } else {
                        ItemValueTokens::Value(value)
                    },
                )
            }
            ItemValueTokens::Value(value) => (
                Some(DataToken::ItemValue(value.as_ref().to_owned())),
                ItemValueTokens::Done,
            ),
            ItemValueTokens::Done => (Some(DataToken::ItemEnd), ItemValueTokens::End),
            ItemValueTokens::End => return None,
        };

        *self = next_state;
        out
    }
}
<|endoftext|>


<commit_before>use dungeon::{Dungeon};
use celloption::{Occupant};
use genotype::{Genotype};
use phenotype::{Seed};
use config::{Config};
use util::{odds};

use std::iter::{range_step};
use std::collections::{HashMap, HashSet};
use rand::{Rng};

#[derive(Clone, Debug)]
pub struct DesirableProperties {
    seed: Seed,
    room_number: u32,
    room_size: u32,
    doors: u32,
    monsters: u32,
    path_length: u32,
    branching: f64,
    occupants: Vec<(Occupant, (u32, u32))>,
    rooms: Vec<Room>,
    mazes: Vec<Maze>,
}

#[derive(Clone, Debug)]
struct Room {
    x: u32,
    y: u32,
    w: u32,
    h: u32
}

#[derive(Clone, Debug)]
struct Maze {
    region: u32,
    path: HashSet<(u32, u32)>
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum Direction {
    North,
    East,
    South,
    West
}

pub type CollisionFn<'a> = Box<Fn(&(u32, u32)) -> bool + 'a>;

impl Maze {
    pub fn new<T: Rng>(rng: &mut T, positions: &HashSet<(u32, u32)>, collides: CollisionFn, branching: f64, region: u32, x: u32, y: u32) -> Maze {
        let path = Maze::grow(rng, positions, collides, branching, x, y);
        Maze {
            region: region,
            path: path
        }
    }

    fn grow<T: Rng>(rng: &mut T, positions: &HashSet<(u32, u32)>, collides: CollisionFn, branching: f64, x: u32, y: u32) -> HashSet<(u32, u32)> {
        let branching_factor = (branching * 100.0) as u64;
        let all = vec!(Direction::North, Direction::East, Direction::South, Direction::West);
        let mut path: HashSet<(u32, u32)> = HashSet::new();
        let mut all_paths: HashSet<(u32, u32)> = positions.clone();
        let mut direction = rng.choose(all.as_slice()).unwrap().clone();
        path.insert((x, y));
        all_paths.insert((x, y));
        while !path.is_empty() {
            let &(x, y) = path.iter().next().unwrap();
            let sx = x as i32;
            let sy = y as i32;
            let sc = vec!((Direction::North, (0, 1)),
                          (Direction::East, (1, 0)),
                          (Direction::South, (0, -1)),
                          (Direction::West, (-1, 0)));
            let uncarved: HashMap<Direction, (i32, i32)> = sc.iter().filter(|&&(_, (dir_x, dir_y))| {
                let coord1 = ((sx + dir_x*1) as u32, (sy + dir_y*1) as u32);
                let coord2 = ((sx + dir_x*2) as u32, (sy + dir_y*2) as u32);
                let coord3 = ((sx + dir_x*3) as u32, (sy + dir_y*3) as u32);
                !all_paths.contains(&coord1) &&
                    !all_paths.contains(&coord2) &&
                    !all_paths.contains(&coord3) &&
                    !collides(&coord2) &&
                    !collides(&coord3)
            }).cloned().collect();
            if !uncarved.is_empty() {
                let same_direction = odds(rng, branching_factor, 100);
                let (rel_x, rel_y) = if same_direction && uncarved.contains_key(&direction) {
                    uncarved[direction]
                } else {
                    let (dir, coords) = uncarved.into_iter().next().unwrap();
                    direction = dir;
                    coords
                };
                let coord1 = ((sx + rel_x*1) as u32, (sy + rel_y*1) as u32);
                let coord2 = ((sx + rel_x*2) as u32, (sy + rel_y*2) as u32);
                path.insert(coord2);
                all_paths.insert(coord1);
                all_paths.insert(coord2);
            } else {
                path.remove(&(x, y));
            }
        }
        all_paths
    }
}

impl Room {
    pub fn random<T: Rng>(rng: &mut T, width: u32, height: u32, room_size: u32) -> Room {
        let w: u32 = rng.gen_range(3, room_size);
        let h: u32 = rng.gen_range(3, room_size);
        assert!(width >= w);
        assert!(height >= h);
        let x: u32 = rng.gen_range(1, width - w);
        let y: u32 = rng.gen_range(1, height - h);
        Room {
            x: Room::make_odd(x),
            y: Room::make_odd(y),
            w: Room::make_odd(w),
            h: Room::make_odd(h),
        }
    }

    fn make_odd(n: u32) -> u32 {
        if n % 2 == 0 { n - 1 } else { n }
    }

    pub fn intersects(&self, rooms: &Vec<Room>) -> bool {
        for r in rooms {
            let rw = r.x + r.w;
            let rh = r.y + r.h;
            let roomw = self.x + self.w;
            let roomh = self.y + self.h;
            if roomw >= r.x && roomh >= r.y && self.x <= rw && self.y <= rh {
                return true
            }
        }
        false
    }

    pub fn contains(&self, x: u32, y: u32) -> bool {
        let roomw = self.x + self.w;
        let roomh = self.y + self.h;
        x >= self.x && x <= roomw && y >= self.y && y <= roomh
    }
}

impl DesirableProperties {
    pub fn new(config: &Config, seed: &Seed) -> DesirableProperties {
        // read in the configurations
        let desirables = config.get_table(None, "desirable_patterns");
        let room_number = config.get_integer(desirables, "room_number") as u32;
        let room_size = config.get_integer(desirables, "room_size") as u32;
        let doors = config.get_integer(desirables, "doors") as u32;
        let monsters = config.get_integer(desirables, "monsters") as u32;
        let path_length = config.get_integer(desirables, "path_length") as u32;
        let branching = config.get_float(desirables, "branching");
        DesirableProperties {
            seed: seed.clone(),
            room_size: room_size,
            room_number: room_number,
            doors: doors,
            monsters: monsters,
            path_length: path_length,
            branching: branching,
            occupants: vec![],
            rooms: vec![],
            mazes: vec![]
        }
    }

}

impl Genotype for DesirableProperties {
    fn initialize<T: Rng>(&self, rng: &mut T) -> DesirableProperties {
        let w = self.seed.width;
        let h = self.seed.height;
        let occupants = self.seed.random_occupants(rng).iter().take(self.monsters as usize).cloned().collect();
        // randomly generate rooms
        let mut rooms = vec![];
        for _ in range(0, self.room_number) {
            for _ in range(0, 10) {
                let room = Room::random(rng, w, h, self.room_size);
                if !room.intersects(&rooms) {
                    rooms.push(room);
                    break;
                }
            }
        }
        // fill in mazes
        let mut mazes = vec![];
        let mut positions: HashSet<(u32, u32)> = HashSet::new();
        let mut region = 0;
        let is_occupied = |x, y| rooms.clone().iter().fold(false, |accum, ref m| accum || m.contains(x, y));
        for x in range_step(1, w, 2) {
            for y in range_step(1, h, 2) {
                if !is_occupied(x, y) && !positions.contains(&(x, y)) {
                    let collides: CollisionFn = box |&(cx, cy)| cx >= w || cy >= h || is_occupied(cx, cy);
                    let maze = Maze::new(rng, &positions, collides, self.branching, region, x, y);
                    positions = positions.union(&maze.path).cloned().collect();
                    mazes.push(maze);
                    region += 1;
                }
            }
        }
        // find connectors
        // remove dead ends
        DesirableProperties {
            seed: self.seed.clone(),
            room_number: self.room_number,
            room_size: self.room_size,
            doors: self.doors,
            monsters: self.monsters,
            path_length: self.path_length,
            branching: self.branching,
            occupants: occupants,
            rooms: rooms.clone(),
            mazes: mazes,
        }
    }

    fn mutate<T: Rng>(&mut self, _: &mut T, _: f64) {
        // mutate a certain % of the rooms and start again
    }

    fn generate(&self) -> Dungeon {
        let w = self.seed.width;
        let h = self.seed.height;
        let mut dungeon = Dungeon::new(w, h, None);
        let wall = self.seed.tiles.get("wall").unwrap();
        let floor = self.seed.tiles.get("floor").unwrap();
        for room in self.rooms.iter() {
            for i in range(room.x, room.x + room.w) {
                for j in range(room.y, room.y + room.h) {
                    dungeon.cells[i as usize][j as usize].tile = Some(wall.clone())
                }
            }
        }
        for maze in self.mazes.iter() {
            for path in maze.path.iter() {
                let (x, y): (u32, u32) = *path;
                dungeon.cells[x as usize][y as usize].tile = Some(floor.clone());
            }
        }
        dungeon.clone()
    }
}
<commit_msg>Cleaning up maze path code.<commit_after>use dungeon::{Dungeon};
use celloption::{Occupant};
use genotype::{Genotype};
use phenotype::{Seed};
use config::{Config};
use util::{odds};

use std::iter::{range_step};
use std::collections::{HashMap, HashSet};
use rand::{Rng};

#[derive(Clone, Debug)]
pub struct DesirableProperties {
    seed: Seed,
    room_number: u32,
    room_size: u32,
    doors: u32,
    monsters: u32,
    path_length: u32,
    branching: f64,
    occupants: Vec<(Occupant, (u32, u32))>,
    rooms: Vec<Room>,
    mazes: Vec<Maze>,
}

#[derive(Clone, Debug)]
struct Room {
    x: u32,
    y: u32,
    w: u32,
    h: u32
}

#[derive(Clone, Debug)]
struct Maze {
    region: u32,
    path: HashSet<(u32, u32)>
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum Direction {
    North,
    East,
    South,
    West
}

pub type CollisionFn<'a> = Box<Fn(&(u32, u32)) -> bool + 'a>;

impl Maze {
    pub fn new<T: Rng>(rng: &mut T,
                       previous: &HashSet<(u32, u32)>,
                       collides: CollisionFn,
                       branching: f64,
                       region: u32,
                       x: u32,
                       y: u32) -> Maze {
        let path = Maze::grow(rng, previous, collides, branching, x, y);
        Maze {
            region: region,
            path: path
        }
    }

    fn grow<T: Rng>(rng: &mut T,
                    previous: &HashSet<(u32, u32)>,
                    collides: CollisionFn,
                    branching: f64,
                    x: u32,
                    y: u32) -> HashSet<(u32, u32)> {
        let branching_factor = (branching * 100.0) as u64;
        let mut path: Vec<(u32, u32)> = Vec::new();
        let mut all_paths: HashSet<(u32, u32)> = previous.clone();
        let all = vec!(Direction::North, Direction::East, Direction::South, Direction::West);
        let mut direction = rng.choose(all.as_slice()).unwrap().clone();
        let sc = vec!((Direction::North, (0, 1)),
                      (Direction::East, (1, 0)),
                      (Direction::South, (0, -1)),
                      (Direction::West, (-1, 0)));
        path.push((x, y));
        all_paths.insert((x, y));
        while !path.is_empty() {
            let (x, y) = path[path.len()-1];
            let coord_at = |(vx, vy): (i32, i32), n| ((x as i32 + vx*n) as u32, (y as i32 + vy*n) as u32);
            let uncarved: HashMap<Direction, (i32, i32)> = sc.iter().filter(|&&(_, new_dir)| {
                let coord1 = coord_at(new_dir, 1);
                let coord2 = coord_at(new_dir, 2);
                let coord3 = coord_at(new_dir, 3);
                !all_paths.contains(&coord1) &&
                    !all_paths.contains(&coord2) &&
                    !all_paths.contains(&coord3) &&
                    !collides(&coord2) &&
                    !collides(&coord3)
            }).cloned().collect();
            if !uncarved.is_empty() {
                let same_direction = odds(rng, branching_factor, 100);
                let rel = if same_direction && uncarved.contains_key(&direction) {
                    uncarved[direction]
                } else {
                    let (dir, coords) = uncarved.into_iter().next().unwrap();
                    direction = dir;
                    coords
                };
                let coord1 = coord_at(rel, 1);
                let coord2 = coord_at(rel, 2);
                all_paths.insert(coord1);
                all_paths.insert(coord2);
                path.push(coord2);
            } else {
                path.pop();
            }
        }
        all_paths.difference(previous).cloned().collect()
    }
}

impl Room {
    pub fn random<T: Rng>(rng: &mut T, width: u32, height: u32, room_size: u32) -> Room {
        let w: u32 = rng.gen_range(3, room_size);
        let h: u32 = rng.gen_range(3, room_size);
        assert!(width >= w);
        assert!(height >= h);
        let x: u32 = rng.gen_range(1, width - w);
        let y: u32 = rng.gen_range(1, height - h);
        let make_odd = |n| if n % 2 == 0 { n - 1 } else { n };
        Room {
            x: make_odd(x),
            y: make_odd(y),
            w: make_odd(w),
            h: make_odd(h),
        }
    }

    pub fn intersects(&self, rooms: &Vec<Room>) -> bool {
        for r in rooms {
            let rw = r.x + r.w;
            let rh = r.y + r.h;
            let roomw = self.x + self.w;
            let roomh = self.y + self.h;
            if roomw > r.x && roomh > r.y && self.x < rw && self.y < rh {
                return true
            }
        }
        false
    }

    pub fn contains(&self, x: u32, y: u32) -> bool {
        let roomw = self.x + self.w;
        let roomh = self.y + self.h;
        x >= self.x && x < roomw && y >= self.y && y < roomh
    }
}

impl DesirableProperties {
    pub fn new(config: &Config, seed: &Seed) -> DesirableProperties {
        // read in the configurations
        let desirables = config.get_table(None, "desirable_patterns");
        let room_number = config.get_integer(desirables, "room_number") as u32;
        let room_size = config.get_integer(desirables, "room_size") as u32;
        let doors = config.get_integer(desirables, "doors") as u32;
        let monsters = config.get_integer(desirables, "monsters") as u32;
        let path_length = config.get_integer(desirables, "path_length") as u32;
        let branching = config.get_float(desirables, "branching");
        DesirableProperties {
            seed: seed.clone(),
            room_size: room_size,
            room_number: room_number,
            doors: doors,
            monsters: monsters,
            path_length: path_length,
            branching: branching,
            occupants: vec![],
            rooms: vec![],
            mazes: vec![]
        }
    }

}

impl Genotype for DesirableProperties {
    fn initialize<T: Rng>(&self, rng: &mut T) -> DesirableProperties {
        let w = self.seed.width;
        let h = self.seed.height;
        let occupants = self.seed.random_occupants(rng).iter().take(self.monsters as usize).cloned().collect();
        // randomly generate rooms
        let mut rooms = vec![];
        for _ in range(0, self.room_number) {
            for _ in range(0, 10) {
                let room = Room::random(rng, w, h, self.room_size);
                if !room.intersects(&rooms) {
                    rooms.push(room);
                    break;
                }
            }
        }
        // fill in mazes
        let mut mazes = vec![];
        let mut positions: HashSet<(u32, u32)> = HashSet::new();
        let mut region = 0;
        let is_occupied = |x, y| rooms.clone().iter().fold(false, |accum, ref m| accum || m.contains(x, y));
        for x in range_step(1, w, 2) {
            for y in range_step(1, h, 2) {
                if !is_occupied(x, y) && !positions.contains(&(x, y)) {
                    let collides: CollisionFn = box |&(cx, cy)| cx >= w || cy >= h || is_occupied(cx, cy);
                    let maze = Maze::new(rng, &positions, collides, self.branching, region, x, y);
                    positions = positions.union(&maze.path).cloned().collect();
                    mazes.push(maze);
                    region += 1;
                }
            }
        }
        // find connectors
        // remove dead ends
        DesirableProperties {
            seed: self.seed.clone(),
            room_number: self.room_number,
            room_size: self.room_size,
            doors: self.doors,
            monsters: self.monsters,
            path_length: self.path_length,
            branching: self.branching,
            occupants: occupants,
            rooms: rooms.clone(),
            mazes: mazes,
        }
    }

    fn mutate<T: Rng>(&mut self, _: &mut T, _: f64) {
        // mutate a certain % of the rooms and start again
    }

    fn generate(&self) -> Dungeon {
        let w = self.seed.width;
        let h = self.seed.height;
        let mut dungeon = Dungeon::new(w, h, None);
        let wall = self.seed.tiles.get("wall").unwrap();
        let floor = self.seed.tiles.get("floor").unwrap();
        for room in self.rooms.iter() {
            for i in range(room.x, room.x + room.w) {
                for j in range(room.y, room.y + room.h) {
                    dungeon.cells[i as usize][j as usize].tile = Some(wall.clone())
                }
            }
        }
        for maze in self.mazes.iter() {
            for path in maze.path.iter() {
                let (x, y): (u32, u32) = *path;
                dungeon.cells[x as usize][y as usize].tile = Some(floor.clone());
            }
        }
        dungeon.clone()
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io::stderr;
use std::path::Path;
use std::process;

const PREVIEW_LEN: usize = 50;

#[derive(Debug, PartialEq)]
struct Match {
    offset: usize,
    key: u8,
    preview: String,
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();
    if args.len() < 2 {
        writeln!(stderr(), "usage: roxor <file> <crib>").unwrap();
        process::exit(1);
    }

    let path = Path::new(&args[0]);
    let crib = args[1].as_bytes();

    if crib.len() < 2 {
        writeln!(stderr(), "error: crib too short").unwrap();
        process::exit(1);
    }

    let mut file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", path.display(), why),
        Ok(file) => file,
    };

    let mut ciphertext: Vec<u8> = Vec::new();
    file.read_to_end(&mut ciphertext).unwrap();

    if ciphertext.len() < crib.len() {
        writeln!(stderr(), "error: ciphertext too short").unwrap();
        process::exit(1);
    }

    for m in attack_cipher(&ciphertext[..], crib) {
        println!("Found text at 0x{:x} (XOR key 0x{:02x})", m.offset, m.key);
        println!("  preview: {}", m.preview);
    }
}

fn attack_cipher(ciphertext: &[u8], crib: &[u8]) -> Vec<Match> {
    let mut matches: Vec<Match> = Vec::new();
    for (i, x) in ciphertext.iter().enumerate() {
        let mut old_key = x ^ crib[0];
        let mut j = 1;
        while j < crib.len() && (i+j) < ciphertext.len() {
            let key = ciphertext[i+j] ^ crib[j];
            if key != old_key { break; }
            old_key = key;
            j += 1;
            if j == crib.len() {
                let preview: String = ciphertext[i..]
                    .iter()
                    .take(PREVIEW_LEN)
                    .map(|x| x ^ key)
                    .map(|x| if x >= 32 && x <= 126 {x as char} else {'.'})
                    .collect();
                matches.push(Match { offset: i, key: key, preview: preview });
            }
        }
    }
    return matches;
}

#[cfg(test)]
mod tests {
    use super::{Match, attack_cipher};

    struct Test {
        ciphertext: &'static [u8],
        crib: &'static [u8],
        matches: Vec<Match>,
    }

    #[test]
    fn test_attack_cipher() {
        let tests = vec![
            Test {
                ciphertext: &[],
                crib: &[],
                matches: vec![],
            },
            Test {
                ciphertext: b"haystack",
                crib: b"needle",
                matches: vec![],
            },
            Test {
                ciphertext: b"needle in haystack",
                crib: b"needle",
                matches: vec![
                    Match { offset: 0, key: 0, preview: "needle in haystack".to_string() },
                ],
            },
            Test {
                ciphertext: b"a needle, another needle",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0, preview: "needle".to_string() },
                ],
            },
            Test {
                ciphertext: b"\x23\x62\x2c\x27\x27\x26\x2e\x27\x6e\x62\x23\x2c\x2d\x36\x2a\x27\x30\x62\x2c\x27\x27\x26\x2e\x27",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0x42, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0x42, preview: "needle".to_string() },
                ],
            },
        ];

        for t in tests.iter() {
            assert_eq!(t.matches, attack_cipher(t.ciphertext, t.crib));
        }
    }
}
<commit_msg>Add standard copyright note<commit_after>// Rust implementation of roxor.c
//
// Copyright (C) 2015 Mathias Lafeldt <mathias.lafeldt@gmail.com>
//
// This file is part of roxor.
//
// roxor is licensed under the terms of the MIT License. See LICENSE file.

use std::env;
use std::fs::File;
use std::io::prelude::*;
use std::io::stderr;
use std::path::Path;
use std::process;

const PREVIEW_LEN: usize = 50;

#[derive(Debug, PartialEq)]
struct Match {
    offset: usize,
    key: u8,
    preview: String,
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();
    if args.len() < 2 {
        writeln!(stderr(), "usage: roxor <file> <crib>").unwrap();
        process::exit(1);
    }

    let path = Path::new(&args[0]);
    let crib = args[1].as_bytes();

    if crib.len() < 2 {
        writeln!(stderr(), "error: crib too short").unwrap();
        process::exit(1);
    }

    let mut file = match File::open(&path) {
        Err(why) => panic!("couldn't open {}: {}", path.display(), why),
        Ok(file) => file,
    };

    let mut ciphertext: Vec<u8> = Vec::new();
    file.read_to_end(&mut ciphertext).unwrap();

    if ciphertext.len() < crib.len() {
        writeln!(stderr(), "error: ciphertext too short").unwrap();
        process::exit(1);
    }

    for m in attack_cipher(&ciphertext[..], crib) {
        println!("Found text at 0x{:x} (XOR key 0x{:02x})", m.offset, m.key);
        println!("  preview: {}", m.preview);
    }
}

fn attack_cipher(ciphertext: &[u8], crib: &[u8]) -> Vec<Match> {
    let mut matches: Vec<Match> = Vec::new();
    for (i, x) in ciphertext.iter().enumerate() {
        let mut old_key = x ^ crib[0];
        let mut j = 1;
        while j < crib.len() && (i+j) < ciphertext.len() {
            let key = ciphertext[i+j] ^ crib[j];
            if key != old_key { break; }
            old_key = key;
            j += 1;
            if j == crib.len() {
                let preview: String = ciphertext[i..]
                    .iter()
                    .take(PREVIEW_LEN)
                    .map(|x| x ^ key)
                    .map(|x| if x >= 32 && x <= 126 {x as char} else {'.'})
                    .collect();
                matches.push(Match { offset: i, key: key, preview: preview });
            }
        }
    }
    return matches;
}

#[cfg(test)]
mod tests {
    use super::{Match, attack_cipher};

    struct Test {
        ciphertext: &'static [u8],
        crib: &'static [u8],
        matches: Vec<Match>,
    }

    #[test]
    fn test_attack_cipher() {
        let tests = vec![
            Test {
                ciphertext: &[],
                crib: &[],
                matches: vec![],
            },
            Test {
                ciphertext: b"haystack",
                crib: b"needle",
                matches: vec![],
            },
            Test {
                ciphertext: b"needle in haystack",
                crib: b"needle",
                matches: vec![
                    Match { offset: 0, key: 0, preview: "needle in haystack".to_string() },
                ],
            },
            Test {
                ciphertext: b"a needle, another needle",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0, preview: "needle".to_string() },
                ],
            },
            Test {
                ciphertext: b"\x23\x62\x2c\x27\x27\x26\x2e\x27\x6e\x62\x23\x2c\x2d\x36\x2a\x27\x30\x62\x2c\x27\x27\x26\x2e\x27",
                crib: b"needle",
                matches: vec![
                    Match { offset: 2, key: 0x42, preview: "needle, another needle".to_string() },
                    Match { offset: 18, key: 0x42, preview: "needle".to_string() },
                ],
            },
        ];

        for t in tests.iter() {
            assert_eq!(t.matches, attack_cipher(t.ciphertext, t.crib));
        }
    }
}
<|endoftext|>


<commit_before>use std::ops::Range;

use byteorder::{ByteOrder, LittleEndian};

use gameboy::{CPU_FREQUENCY, MAX_CPU_CYCLES, MAX_DIV_REG_CYCLES, Gpu, Irq, Memory, Timer};
use gameboy::cartridge::{Cartridge, CartridgeDetails};
use gameboy::frame::Frame;
use super::memory_map::{self, Address};

const MAIN_MEM_SIZE: usize = 0x2000;
const ZRAM_SIZE: usize = 0x80;
const MMAP_SIZE: usize = 0x80;

pub struct Interconnect {
    pub gpu: Gpu,
    pub ram: Memory,
    pub zram: Memory,
    pub mmap_io: Memory,
    pub timer: Timer,
    pub irq: Irq,
    pub cart: Option<Cartridge>,
    pub interrupt: u8,
}

impl Interconnect {
    pub fn new() -> Interconnect {
        Interconnect {
            gpu: Gpu::new(),
            ram: Memory::new(MAIN_MEM_SIZE),
            zram: Memory::new(ZRAM_SIZE),
            mmap_io: Memory::new(MMAP_SIZE),
            timer: Timer::new(),
            irq: Irq::new(),
            cart: None,
            interrupt: 0x00,
        }
    }

    pub fn with_cart(cart: Cartridge) -> Interconnect {
        Interconnect {
            gpu: Gpu::new(),
            ram: Memory::new(MAIN_MEM_SIZE),
            zram: Memory::new(ZRAM_SIZE),
            mmap_io: Memory::new(MMAP_SIZE),
            timer: Timer::new(),
            irq: Irq::new(),
            cart: Some(cart),
            interrupt: 0x00,
        }
    }

    pub fn step(&mut self, cycles: usize) -> Result<(), String> {
        self.gpu.step(&mut self.irq, cycles)?;
        self.timer.step(&mut self.irq, cycles)?;

        Ok(())
    }

    pub fn reset(&mut self) {
        self.gpu.reset();
        self.timer.reset();
        self.irq.reset();
    }

    pub fn cart_details(&self) -> CartridgeDetails {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        cart.details(&self)
    }

    pub fn request_frame(&self) -> &Frame {
        &self.gpu.frame
    }

    pub fn write_u8(&mut self, addr: u16, byte: u8) {
        // Special case - DMA transfer
        if addr == 0xFF46 {
            self.dma_transfer(byte);
            return;
        }
        let cart = self.cart.as_mut().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(a) => self.ram.write_u8(a, byte),
            Address::RamShadow(a) => self.ram.write_u8(a, byte),
            Address::Gfx(a) => self.gpu.ram.write_u8(a, byte),
            Address::CartRam(a) => cart.ram.write_u8(a, byte),
            Address::CartRom(a) => cart.rom.write_u8(a, byte),
            Address::ZRam(a) => {
                // if a != 0x00 {
                self.zram.write_u8(a, byte);
                // }
            }
            Address::Oam(a) => self.gpu.sprite_data.write_u8(a, byte),
            Address::Unused(_) => (),
            Address::Io(a) => {
                match a {
                    0x00 => println!("err: write to joypad not supported"),
                    0x01...0x02 => println!("err: write to serial driver not supported"),
                    0x04...0x07 => self.timer.write_u8(a, byte),
                    0x0F => self.irq.request_flag = byte,
                    0x10...0x26 => println!("err: write to sound driver not supported"),
                    0x40...0x45 => self.gpu.write_u8(a, byte),
                    0x47...0x49 => self.gpu.write_u8(a, byte),
                    0x4A...0x4B => self.gpu.write_u8(a, byte),
                    0x7F => self.mmap_io.write_u8(a, byte),
                    _ => panic!("write memory mapped I/O in unsupported range: {:04X}", a),
                }
            }
            Address::InterruptEnableRegister(a) => self.irq.enable_flag = byte,
            _ => {
                panic!("Unable to write byte to: {:#X}, invalid memory region.",
                       addr)
            }
        }
    }

    pub fn read_u8(&self, addr: u16) -> u8 {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(addr) => self.ram.read_u8(addr),
            Address::RamShadow(addr) => self.ram.read_u8(addr),
            Address::CartRom(addr) |
            Address::CartRomOtherBank(addr) => cart.rom.read_u8(addr),
            Address::Gfx(value) => self.gpu.ram.read_u8(value),
            Address::CartRam(a) => cart.ram.read_u8(a),
            Address::ZRam(a) => self.zram.read_u8(a),
            Address::Oam(a) => self.gpu.sprite_data.read_u8(a),
            Address::Unused(_) => 0xFF, // Always return high
            Address::Io(a) => {
                match a {
                    0x00 => {
                        println!("err: read from joypad not supported");
                        0xFF
                    }
                    0x01...0x02 => {
                        println!("err: read from serial driver not supported");
                        0
                    }
                    0x04...0x07 => self.timer.read_u8(a),
                    0x0F => self.irq.request_flag,
                    0x10...0x26 => {
                        println!("err: read from sound driver not supported");
                        0
                    }
                    0x40...0x45 => self.gpu.read_u8(a),
                    0x47...0x49 => self.gpu.read_u8(a),
                    0x4A...0x4B => self.gpu.read_u8(a),
                    0x7F => self.mmap_io.read_u8(a),
                    _ => panic!("read memory mapped I/O in unsupported range"),
                }
            }
            Address::InterruptEnableRegister(a) => self.irq.enable_flag,
            _ => panic!("Unable to read address: {:#X}", addr),
        }
    }

    pub fn read_bytes(&self, r: Range<u16>) -> &[u8] {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(r.start) {
            Address::Ram(_) |
            Address::RamShadow(_) => self.ram.read_bytes(r),
            Address::CartRom(_) |
            Address::CartRomOtherBank(_) => cart.rom.read_bytes(r),
            Address::Gfx(_) => self.gpu.ram.read_bytes(r),
            Address::CartRam(_) => cart.ram.read_bytes(r),
            Address::ZRam(a) => self.zram.read_bytes(r),
            Address::Io(a) => self.mmap_io.read_bytes(r),
            _ => panic!("Unable to read address range: {:?}", r),
        }
    }

    pub fn write_u16(&mut self, addr: u16, val: u16) {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(a) => self.ram.write_u16(a, val),
            Address::RamShadow(a) => self.ram.write_u16(a, val),
            Address::ZRam(a) => self.zram.write_u16(a, val),
            _ => panic!("Unable to write address: {:#X}", addr),
        }
    }

    pub fn read_u16(&self, addr: u16) -> u16 {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(addr) => self.ram.read_u16(addr),
            Address::RamShadow(addr) => self.ram.read_u16(addr),
            Address::CartRom(addr) |
            Address::CartRomOtherBank(addr) => cart.rom.read_u16(addr),
            Address::Gfx(value) => self.gpu.ram.read_u16(value),
            Address::CartRam(a) => cart.ram.read_u16(a),
            Address::ZRam(a) => self.zram.read_u16(a),
            Address::Io(a) => self.mmap_io.read_u16(a),
            _ => panic!("Unable to read address: {:#X}", addr),
        }
    }

    fn dma_transfer(&mut self, byte: u8) {
        let addr = ((byte as u16) << 0x08); // "The written value specifies the transfer source address divided by 0x100"
        for x in 0..0xA0 {
            let val = self.read_u8(addr + x);
            self.write_u8(0xFE00 + x, val);
        }
    }
}
<commit_msg>Add sink for wave pattern RAM write/read<commit_after>use std::ops::Range;

use byteorder::{ByteOrder, LittleEndian};

use gameboy::{CPU_FREQUENCY, MAX_CPU_CYCLES, MAX_DIV_REG_CYCLES, Gpu, Irq, Memory, Timer};
use gameboy::cartridge::{Cartridge, CartridgeDetails};
use gameboy::frame::Frame;
use super::memory_map::{self, Address};

const MAIN_MEM_SIZE: usize = 0x2000;
const ZRAM_SIZE: usize = 0x80;
const MMAP_SIZE: usize = 0x80;

pub struct Interconnect {
    pub gpu: Gpu,
    pub ram: Memory,
    pub zram: Memory,
    pub mmap_io: Memory,
    pub timer: Timer,
    pub irq: Irq,
    pub cart: Option<Cartridge>,
    pub interrupt: u8,
}

impl Interconnect {
    pub fn new() -> Interconnect {
        Interconnect {
            gpu: Gpu::new(),
            ram: Memory::new(MAIN_MEM_SIZE),
            zram: Memory::new(ZRAM_SIZE),
            mmap_io: Memory::new(MMAP_SIZE),
            timer: Timer::new(),
            irq: Irq::new(),
            cart: None,
            interrupt: 0x00,
        }
    }

    pub fn with_cart(cart: Cartridge) -> Interconnect {
        Interconnect {
            gpu: Gpu::new(),
            ram: Memory::new(MAIN_MEM_SIZE),
            zram: Memory::new(ZRAM_SIZE),
            mmap_io: Memory::new(MMAP_SIZE),
            timer: Timer::new(),
            irq: Irq::new(),
            cart: Some(cart),
            interrupt: 0x00,
        }
    }

    pub fn step(&mut self, cycles: usize) -> Result<(), String> {
        self.gpu.step(&mut self.irq, cycles)?;
        self.timer.step(&mut self.irq, cycles)?;

        Ok(())
    }

    pub fn reset(&mut self) {
        self.gpu.reset();
        self.timer.reset();
        self.irq.reset();
    }

    pub fn cart_details(&self) -> CartridgeDetails {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        cart.details(&self)
    }

    pub fn request_frame(&self) -> &Frame {
        &self.gpu.frame
    }

    pub fn write_u8(&mut self, addr: u16, byte: u8) {
        // Special case - DMA transfer
        if addr == 0xFF46 {
            self.dma_transfer(byte);
            return;
        }
        let cart = self.cart.as_mut().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(a) => self.ram.write_u8(a, byte),
            Address::RamShadow(a) => self.ram.write_u8(a, byte),
            Address::Gfx(a) => self.gpu.ram.write_u8(a, byte),
            Address::CartRam(a) => cart.ram.write_u8(a, byte),
            Address::CartRom(a) => cart.rom.write_u8(a, byte),
            Address::ZRam(a) => {
                // if a != 0x00 {
                self.zram.write_u8(a, byte);
                // }
            }
            Address::Oam(a) => self.gpu.sprite_data.write_u8(a, byte),
            Address::Unused(_) => (),
            Address::Io(a) => {
                match a {
                    0x00 => println!("err: write to joypad not supported"),
                    0x01...0x02 => println!("err: write to serial driver not supported"),
                    0x04...0x07 => self.timer.write_u8(a, byte),
                    0x0F => self.irq.request_flag = byte,
                    0x10...0x26 => println!("err: write to sound driver not supported"),
                    0x30...0x3F => println!("err: write to wave pattern RAM not supported"),
                    0x40...0x45 => self.gpu.write_u8(a, byte),
                    0x47...0x49 => self.gpu.write_u8(a, byte),
                    0x4A...0x4B => self.gpu.write_u8(a, byte),
                    0x7F => self.mmap_io.write_u8(a, byte),
                    _ => panic!("write memory mapped I/O in unsupported range: {:04X}", a),
                }
            }
            Address::InterruptEnableRegister(a) => self.irq.enable_flag = byte,
            _ => {
                panic!("Unable to write byte to: {:#X}, invalid memory region.",
                       addr)
            }
        }
    }

    pub fn read_u8(&self, addr: u16) -> u8 {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(addr) => self.ram.read_u8(addr),
            Address::RamShadow(addr) => self.ram.read_u8(addr),
            Address::CartRom(addr) |
            Address::CartRomOtherBank(addr) => cart.rom.read_u8(addr),
            Address::Gfx(value) => self.gpu.ram.read_u8(value),
            Address::CartRam(a) => cart.ram.read_u8(a),
            Address::ZRam(a) => self.zram.read_u8(a),
            Address::Oam(a) => self.gpu.sprite_data.read_u8(a),
            Address::Unused(_) => 0xFF, // Always return high
            Address::Io(a) => {
                match a {
                    0x00 => {
                        println!("err: read from joypad not supported");
                        0xFF
                    }
                    0x01...0x02 => {
                        println!("err: read from serial driver not supported");
                        0
                    }
                    0x04...0x07 => self.timer.read_u8(a),
                    0x0F => self.irq.request_flag,
                    0x10...0x26 => {
                        println!("err: read from sound driver not supported");
                        0
                    }
                    0x30...0x3F => {
                        println!("err: write to wave pattern RAM not supported");
                        0
                    }
                    0x40...0x45 => self.gpu.read_u8(a),
                    0x47...0x49 => self.gpu.read_u8(a),
                    0x4A...0x4B => self.gpu.read_u8(a),
                    0x7F => self.mmap_io.read_u8(a),
                    _ => panic!("read memory mapped I/O in unsupported range"),
                }
            }
            Address::InterruptEnableRegister(a) => self.irq.enable_flag,
            _ => panic!("Unable to read address: {:#X}", addr),
        }
    }

    pub fn read_bytes(&self, r: Range<u16>) -> &[u8] {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(r.start) {
            Address::Ram(_) |
            Address::RamShadow(_) => self.ram.read_bytes(r),
            Address::CartRom(_) |
            Address::CartRomOtherBank(_) => cart.rom.read_bytes(r),
            Address::Gfx(_) => self.gpu.ram.read_bytes(r),
            Address::CartRam(_) => cart.ram.read_bytes(r),
            Address::ZRam(a) => self.zram.read_bytes(r),
            Address::Io(a) => self.mmap_io.read_bytes(r),
            _ => panic!("Unable to read address range: {:?}", r),
        }
    }

    pub fn write_u16(&mut self, addr: u16, val: u16) {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(a) => self.ram.write_u16(a, val),
            Address::RamShadow(a) => self.ram.write_u16(a, val),
            Address::ZRam(a) => self.zram.write_u16(a, val),
            _ => panic!("Unable to write address: {:#X}", addr),
        }
    }

    pub fn read_u16(&self, addr: u16) -> u16 {
        let cart = self.cart.as_ref().expect("Cartridge is empty");

        match memory_map::map_address(addr) {
            Address::Ram(addr) => self.ram.read_u16(addr),
            Address::RamShadow(addr) => self.ram.read_u16(addr),
            Address::CartRom(addr) |
            Address::CartRomOtherBank(addr) => cart.rom.read_u16(addr),
            Address::Gfx(value) => self.gpu.ram.read_u16(value),
            Address::CartRam(a) => cart.ram.read_u16(a),
            Address::ZRam(a) => self.zram.read_u16(a),
            Address::Io(a) => self.mmap_io.read_u16(a),
            _ => panic!("Unable to read address: {:#X}", addr),
        }
    }

    fn dma_transfer(&mut self, byte: u8) {
        let addr = ((byte as u16) << 0x08); // "The written value specifies the transfer source address divided by 0x100"
        for x in 0..0xA0 {
            let val = self.read_u8(addr + x);
            self.write_u8(0xFE00 + x, val);
        }
    }
}
<|endoftext|>


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

#![feature(box_syntax)]
#![feature(custom_attribute)]
#![feature(custom_derive)]
#![feature(plugin)]
#![feature(slice_patterns)]
#![feature(step_by)]
#![feature(vec_push_all)]
#![feature(custom_attribute)]
#![plugin(serde_macros, plugins)]
#![plugin(regex_macros)]

#[macro_use]
extern crate log;
extern crate euclid;
extern crate hyper;
extern crate ipc_channel;
extern crate msg;
extern crate png;
extern crate regex;
extern crate serde;
extern crate stb_image;
extern crate url;
extern crate util;

use hyper::header::{ContentType, Headers};
use hyper::http::RawStatus;
use hyper::method::Method;
use hyper::mime::{Attr, Mime};
use hyper::status::StatusCode;
use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
use msg::constellation_msg::{PipelineId};
use regex::Regex;
use serde::{Deserializer, Serializer};
use std::sync::mpsc::Receiver;
use std::thread;
use url::Url;
use util::mem::HeapSizeOf;

pub mod hosts;
pub mod image_cache_task;
pub mod net_error_list;
pub mod storage_task;

pub static IPV4_REGEX: Regex = regex!(
    r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
pub static IPV6_REGEX: Regex = regex!(r"^([a-fA-F0-9]{0,4}[:]?){1,8}(/\d{1,3})?$");

/// [Response type](https://fetch.spec.whatwg.org/#concept-response-type)
#[derive(Clone, PartialEq, Copy)]
pub enum ResponseType {
    Basic,
    CORS,
    Default,
    Error,
    Opaque
}

/// [Response termination reason](https://fetch.spec.whatwg.org/#concept-response-termination-reason)
#[derive(Clone, Copy)]
pub enum TerminationReason {
    EndUserAbort,
    Fatal,
    Timeout
}

/// The response body can still be pushed to after fetch
/// This provides a way to store unfinished response bodies
#[derive(Clone)]
pub enum ResponseBody {
    Empty, // XXXManishearth is this necessary, or is Done(vec![]) enough?
    Receiving(Vec<u8>),
    Done(Vec<u8>),
}

pub enum ResponseMsg {
    Chunk(Vec<u8>),
    Finished,
    Errored
}

/// A [Response](https://fetch.spec.whatwg.org/#concept-response) as defined by the Fetch spec
#[derive(Clone)]
pub struct Response {
    pub response_type: ResponseType,
    pub termination_reason: Option<TerminationReason>,
    pub url: Option<Url>,
    /// `None` can be considered a StatusCode of `0`.
    pub status: Option<StatusCode>,
    pub headers: Headers,
    pub body: ResponseBody,
    /// [Internal response](https://fetch.spec.whatwg.org/#concept-internal-response), only used if the Response
    /// is a filtered response
    pub internal_response: Option<Box<Response>>,
}

impl Response {
    pub fn network_error() -> Response {
        Response {
            response_type: ResponseType::Error,
            termination_reason: None,
            url: None,
            status: None,
            headers: Headers::new(),
            body: ResponseBody::Empty,
            internal_response: None
        }
    }

    pub fn is_network_error(&self) -> bool {
        match self.response_type {
            ResponseType::Error => true,
            _ => false
        }
    }
}

/// Image handling.
///
/// It may be surprising that this goes in the network crate as opposed to the graphics crate.
/// However, image handling is generally very integrated with the network stack (especially where
/// caching is involved) and as a result it must live in here.
pub mod image {
    pub mod base;
}

#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct LoadData {
    pub url: Url,
    pub method: Method,
    #[ignore_heap_size_of = "Defined in hyper"]
    /// Headers that will apply to the initial request only
    pub headers: Headers,
    #[ignore_heap_size_of = "Defined in hyper"]
    /// Headers that will apply to the initial request and any redirects
    pub preserved_headers: Headers,
    pub data: Option<Vec<u8>>,
    pub cors: Option<ResourceCORSData>,
    pub pipeline_id: Option<PipelineId>,
}

impl LoadData {
    pub fn new(url: Url, id: Option<PipelineId>) -> LoadData {
        LoadData {
            url: url,
            method: Method::Get,
            headers: Headers::new(),
            preserved_headers: Headers::new(),
            data: None,
            cors: None,
            pipeline_id: id,
        }
    }
}

/// Interface for observing the final response for an asynchronous fetch operation.
pub trait AsyncFetchListener {
    fn response_available(&self, response: Response);
}

/// A listener for asynchronous network events. Cancelling the underlying request is unsupported.
pub trait AsyncResponseListener {
    /// The response headers for a request have been received.
    fn headers_available(&self, metadata: Metadata);
    /// A portion of the response body has been received. This data is unavailable after
    /// this method returned, and must be stored accordingly.
    fn data_available(&self, payload: Vec<u8>);
    /// The response is complete. If the provided status is an Err value, there is no guarantee
    /// that the response body was completely read.
    fn response_complete(&self, status: Result<(), String>);
}

/// Data for passing between threads/processes to indicate a particular action to
/// take on a provided network listener.
#[derive(Deserialize, Serialize)]
pub enum ResponseAction {
    /// Invoke headers_available
    HeadersAvailable(Metadata),
    /// Invoke data_available
    DataAvailable(Vec<u8>),
    /// Invoke response_complete
    ResponseComplete(Result<(), String>)
}

impl ResponseAction {
    /// Execute the default action on a provided listener.
    pub fn process(self, listener: &AsyncResponseListener) {
        match self {
            ResponseAction::HeadersAvailable(m) => listener.headers_available(m),
            ResponseAction::DataAvailable(d) => listener.data_available(d),
            ResponseAction::ResponseComplete(r) => listener.response_complete(r),
        }
    }
}

/// A target for async networking events. Commonly used to dispatch a runnable event to another
/// thread storing the wrapped closure for later execution.
#[derive(Deserialize, Serialize)]
pub struct AsyncResponseTarget {
    pub sender: IpcSender<ResponseAction>,
}

impl AsyncResponseTarget {
    pub fn invoke_with_listener(&self, action: ResponseAction) {
        self.sender.send(action).unwrap()
    }
}

/// A wrapper for a network load that can either be channel or event-based.
#[derive(Deserialize, Serialize)]
pub enum LoadConsumer {
    Channel(IpcSender<LoadResponse>),
    Listener(AsyncResponseTarget),
}

/// Handle to a resource task
pub type ResourceTask = IpcSender<ControlMsg>;

#[derive(PartialEq, Copy, Clone, Deserialize, Serialize)]
pub enum IncludeSubdomains {
    Included,
    NotIncluded
}

#[derive(Deserialize, Serialize)]
pub enum ControlMsg {
    /// Request the data associated with a particular URL
    Load(LoadData, LoadConsumer),
    /// Store a set of cookies for a given originating URL
    SetCookiesForUrl(Url, String, CookieSource),
    /// Retrieve the stored cookies for a given URL
    GetCookiesForUrl(Url, IpcSender<Option<String>>, CookieSource),
    Exit
}

/// Initialized but unsent request. Encapsulates everything necessary to instruct
/// the resource task to make a new request. The `load` method *must* be called before
/// destruction or the task will panic.
pub struct PendingAsyncLoad {
    resource_task: ResourceTask,
    url: Url,
    pipeline: Option<PipelineId>,
    guard: PendingLoadGuard,
}

struct PendingLoadGuard {
    loaded: bool,
}

impl PendingLoadGuard {
    fn neuter(&mut self) {
        self.loaded = true;
    }
}

impl Drop for PendingLoadGuard {
    fn drop(&mut self) {
        if !thread::panicking() {
            assert!(self.loaded)
        }
    }
}

impl PendingAsyncLoad {
    pub fn new(resource_task: ResourceTask, url: Url, pipeline: Option<PipelineId>)
               -> PendingAsyncLoad {
        PendingAsyncLoad {
            resource_task: resource_task,
            url: url,
            pipeline: pipeline,
            guard: PendingLoadGuard { loaded: false, },
        }
    }

    /// Initiate the network request associated with this pending load.
    pub fn load(mut self) -> IpcReceiver<LoadResponse> {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let (sender, receiver) = ipc::channel().unwrap();
        let consumer = LoadConsumer::Channel(sender);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
        receiver
    }

    /// Initiate the network request associated with this pending load, using the provided target.
    pub fn load_async(mut self, listener: AsyncResponseTarget) {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let consumer = LoadConsumer::Listener(listener);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
    }
}

/// Message sent in response to `Load`.  Contains metadata, and a port
/// for receiving the data.
///
/// Even if loading fails immediately, we send one of these and the
/// progress_port will provide the error.
#[derive(Serialize, Deserialize)]
pub struct LoadResponse {
    /// Metadata, such as from HTTP headers.
    pub metadata: Metadata,
    /// Port for reading data.
    pub progress_port: IpcReceiver<ProgressMsg>,
}

#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct ResourceCORSData {
    /// CORS Preflight flag
    pub preflight: bool,
    /// Origin of CORS Request
    pub origin: Url,
}

/// Metadata about a loaded resource, such as is obtained from HTTP headers.
#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct Metadata {
    /// Final URL after redirects.
    pub final_url: Url,

    /// MIME type / subtype.
    pub content_type: Option<(ContentType)>,

    /// Character set.
    pub charset: Option<String>,

    #[ignore_heap_size_of = "Defined in hyper"]
    /// Headers
    pub headers: Option<Headers>,

    /// HTTP Status
    pub status: Option<RawStatus>,
}

impl Metadata {
    /// Metadata with defaults for everything optional.
    pub fn default(url: Url) -> Self {
        Metadata {
            final_url:    url,
            content_type: None,
            charset:      None,
            headers: None,
            // https://fetch.spec.whatwg.org/#concept-response-status-message
            status: Some(RawStatus(200, "OK".into())),
        }
    }

    /// Extract the parts of a Mime that we care about.
    pub fn set_content_type(&mut self, content_type: Option<&Mime>) {
        match content_type {
            None => (),
            Some(mime) => {
                self.content_type = Some(ContentType(mime.clone()));
                let &Mime(_, _, ref parameters) = mime;
                for &(ref k, ref v) in parameters {
                    if &Attr::Charset == k {
                        self.charset = Some(v.to_string());
                    }
                }
            }
        }
    }
}

/// The creator of a given cookie
#[derive(PartialEq, Copy, Clone, Deserialize, Serialize)]
pub enum CookieSource {
    /// An HTTP API
    HTTP,
    /// A non-HTTP API
    NonHTTP,
}

/// Messages sent in response to a `Load` message
#[derive(PartialEq, Debug, Deserialize, Serialize)]
pub enum ProgressMsg {
    /// Binary data - there may be multiple of these
    Payload(Vec<u8>),
    /// Indicates loading is complete, either successfully or not
    Done(Result<(), String>)
}

/// Convenience function for synchronously loading a whole resource.
pub fn load_whole_resource(resource_task: &ResourceTask, url: Url)
        -> Result<(Metadata, Vec<u8>), String> {
    let (start_chan, start_port) = ipc::channel().unwrap();
    resource_task.send(ControlMsg::Load(LoadData::new(url, None),
                       LoadConsumer::Channel(start_chan))).unwrap();
    let response = start_port.recv().unwrap();

    let mut buf = vec!();
    loop {
        match response.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => buf.push_all(&data),
            ProgressMsg::Done(Ok(())) => return Ok((response.metadata, buf)),
            ProgressMsg::Done(Err(e)) => return Err(e)
        }
    }
}

/// Load a URL asynchronously and iterate over chunks of bytes from the response.
pub fn load_bytes_iter(pending: PendingAsyncLoad) -> (Metadata, ProgressMsgPortIterator) {
    let input_port = pending.load();
    let response = input_port.recv().unwrap();
    let iter = ProgressMsgPortIterator {
        progress_port: response.progress_port
    };
    (response.metadata, iter)
}

/// Iterator that reads chunks of bytes from a ProgressMsg port
pub struct ProgressMsgPortIterator {
    progress_port: IpcReceiver<ProgressMsg>,
}

impl Iterator for ProgressMsgPortIterator {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => Some(data),
            ProgressMsg::Done(Ok(()))  => None,
            ProgressMsg::Done(Err(e))  => {
                error!("error receiving bytes: {}", e);
                None
            }
        }
    }
}
<commit_msg>Remove unused import in net_traits<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(box_syntax)]
#![feature(custom_attribute)]
#![feature(custom_derive)]
#![feature(plugin)]
#![feature(slice_patterns)]
#![feature(step_by)]
#![feature(vec_push_all)]
#![feature(custom_attribute)]
#![plugin(serde_macros, plugins)]
#![plugin(regex_macros)]

#[macro_use]
extern crate log;
extern crate euclid;
extern crate hyper;
extern crate ipc_channel;
extern crate msg;
extern crate png;
extern crate regex;
extern crate serde;
extern crate stb_image;
extern crate url;
extern crate util;

use hyper::header::{ContentType, Headers};
use hyper::http::RawStatus;
use hyper::method::Method;
use hyper::mime::{Attr, Mime};
use hyper::status::StatusCode;
use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
use msg::constellation_msg::{PipelineId};
use regex::Regex;
use serde::{Deserializer, Serializer};
use std::thread;
use url::Url;
use util::mem::HeapSizeOf;

pub mod hosts;
pub mod image_cache_task;
pub mod net_error_list;
pub mod storage_task;

pub static IPV4_REGEX: Regex = regex!(
    r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
pub static IPV6_REGEX: Regex = regex!(r"^([a-fA-F0-9]{0,4}[:]?){1,8}(/\d{1,3})?$");

/// [Response type](https://fetch.spec.whatwg.org/#concept-response-type)
#[derive(Clone, PartialEq, Copy)]
pub enum ResponseType {
    Basic,
    CORS,
    Default,
    Error,
    Opaque
}

/// [Response termination reason](https://fetch.spec.whatwg.org/#concept-response-termination-reason)
#[derive(Clone, Copy)]
pub enum TerminationReason {
    EndUserAbort,
    Fatal,
    Timeout
}

/// The response body can still be pushed to after fetch
/// This provides a way to store unfinished response bodies
#[derive(Clone)]
pub enum ResponseBody {
    Empty, // XXXManishearth is this necessary, or is Done(vec![]) enough?
    Receiving(Vec<u8>),
    Done(Vec<u8>),
}

pub enum ResponseMsg {
    Chunk(Vec<u8>),
    Finished,
    Errored
}

/// A [Response](https://fetch.spec.whatwg.org/#concept-response) as defined by the Fetch spec
#[derive(Clone)]
pub struct Response {
    pub response_type: ResponseType,
    pub termination_reason: Option<TerminationReason>,
    pub url: Option<Url>,
    /// `None` can be considered a StatusCode of `0`.
    pub status: Option<StatusCode>,
    pub headers: Headers,
    pub body: ResponseBody,
    /// [Internal response](https://fetch.spec.whatwg.org/#concept-internal-response), only used if the Response
    /// is a filtered response
    pub internal_response: Option<Box<Response>>,
}

impl Response {
    pub fn network_error() -> Response {
        Response {
            response_type: ResponseType::Error,
            termination_reason: None,
            url: None,
            status: None,
            headers: Headers::new(),
            body: ResponseBody::Empty,
            internal_response: None
        }
    }

    pub fn is_network_error(&self) -> bool {
        match self.response_type {
            ResponseType::Error => true,
            _ => false
        }
    }
}

/// Image handling.
///
/// It may be surprising that this goes in the network crate as opposed to the graphics crate.
/// However, image handling is generally very integrated with the network stack (especially where
/// caching is involved) and as a result it must live in here.
pub mod image {
    pub mod base;
}

#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct LoadData {
    pub url: Url,
    pub method: Method,
    #[ignore_heap_size_of = "Defined in hyper"]
    /// Headers that will apply to the initial request only
    pub headers: Headers,
    #[ignore_heap_size_of = "Defined in hyper"]
    /// Headers that will apply to the initial request and any redirects
    pub preserved_headers: Headers,
    pub data: Option<Vec<u8>>,
    pub cors: Option<ResourceCORSData>,
    pub pipeline_id: Option<PipelineId>,
}

impl LoadData {
    pub fn new(url: Url, id: Option<PipelineId>) -> LoadData {
        LoadData {
            url: url,
            method: Method::Get,
            headers: Headers::new(),
            preserved_headers: Headers::new(),
            data: None,
            cors: None,
            pipeline_id: id,
        }
    }
}

/// Interface for observing the final response for an asynchronous fetch operation.
pub trait AsyncFetchListener {
    fn response_available(&self, response: Response);
}

/// A listener for asynchronous network events. Cancelling the underlying request is unsupported.
pub trait AsyncResponseListener {
    /// The response headers for a request have been received.
    fn headers_available(&self, metadata: Metadata);
    /// A portion of the response body has been received. This data is unavailable after
    /// this method returned, and must be stored accordingly.
    fn data_available(&self, payload: Vec<u8>);
    /// The response is complete. If the provided status is an Err value, there is no guarantee
    /// that the response body was completely read.
    fn response_complete(&self, status: Result<(), String>);
}

/// Data for passing between threads/processes to indicate a particular action to
/// take on a provided network listener.
#[derive(Deserialize, Serialize)]
pub enum ResponseAction {
    /// Invoke headers_available
    HeadersAvailable(Metadata),
    /// Invoke data_available
    DataAvailable(Vec<u8>),
    /// Invoke response_complete
    ResponseComplete(Result<(), String>)
}

impl ResponseAction {
    /// Execute the default action on a provided listener.
    pub fn process(self, listener: &AsyncResponseListener) {
        match self {
            ResponseAction::HeadersAvailable(m) => listener.headers_available(m),
            ResponseAction::DataAvailable(d) => listener.data_available(d),
            ResponseAction::ResponseComplete(r) => listener.response_complete(r),
        }
    }
}

/// A target for async networking events. Commonly used to dispatch a runnable event to another
/// thread storing the wrapped closure for later execution.
#[derive(Deserialize, Serialize)]
pub struct AsyncResponseTarget {
    pub sender: IpcSender<ResponseAction>,
}

impl AsyncResponseTarget {
    pub fn invoke_with_listener(&self, action: ResponseAction) {
        self.sender.send(action).unwrap()
    }
}

/// A wrapper for a network load that can either be channel or event-based.
#[derive(Deserialize, Serialize)]
pub enum LoadConsumer {
    Channel(IpcSender<LoadResponse>),
    Listener(AsyncResponseTarget),
}

/// Handle to a resource task
pub type ResourceTask = IpcSender<ControlMsg>;

#[derive(PartialEq, Copy, Clone, Deserialize, Serialize)]
pub enum IncludeSubdomains {
    Included,
    NotIncluded
}

#[derive(Deserialize, Serialize)]
pub enum ControlMsg {
    /// Request the data associated with a particular URL
    Load(LoadData, LoadConsumer),
    /// Store a set of cookies for a given originating URL
    SetCookiesForUrl(Url, String, CookieSource),
    /// Retrieve the stored cookies for a given URL
    GetCookiesForUrl(Url, IpcSender<Option<String>>, CookieSource),
    Exit
}

/// Initialized but unsent request. Encapsulates everything necessary to instruct
/// the resource task to make a new request. The `load` method *must* be called before
/// destruction or the task will panic.
pub struct PendingAsyncLoad {
    resource_task: ResourceTask,
    url: Url,
    pipeline: Option<PipelineId>,
    guard: PendingLoadGuard,
}

struct PendingLoadGuard {
    loaded: bool,
}

impl PendingLoadGuard {
    fn neuter(&mut self) {
        self.loaded = true;
    }
}

impl Drop for PendingLoadGuard {
    fn drop(&mut self) {
        if !thread::panicking() {
            assert!(self.loaded)
        }
    }
}

impl PendingAsyncLoad {
    pub fn new(resource_task: ResourceTask, url: Url, pipeline: Option<PipelineId>)
               -> PendingAsyncLoad {
        PendingAsyncLoad {
            resource_task: resource_task,
            url: url,
            pipeline: pipeline,
            guard: PendingLoadGuard { loaded: false, },
        }
    }

    /// Initiate the network request associated with this pending load.
    pub fn load(mut self) -> IpcReceiver<LoadResponse> {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let (sender, receiver) = ipc::channel().unwrap();
        let consumer = LoadConsumer::Channel(sender);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
        receiver
    }

    /// Initiate the network request associated with this pending load, using the provided target.
    pub fn load_async(mut self, listener: AsyncResponseTarget) {
        self.guard.neuter();
        let load_data = LoadData::new(self.url, self.pipeline);
        let consumer = LoadConsumer::Listener(listener);
        self.resource_task.send(ControlMsg::Load(load_data, consumer)).unwrap();
    }
}

/// Message sent in response to `Load`.  Contains metadata, and a port
/// for receiving the data.
///
/// Even if loading fails immediately, we send one of these and the
/// progress_port will provide the error.
#[derive(Serialize, Deserialize)]
pub struct LoadResponse {
    /// Metadata, such as from HTTP headers.
    pub metadata: Metadata,
    /// Port for reading data.
    pub progress_port: IpcReceiver<ProgressMsg>,
}

#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct ResourceCORSData {
    /// CORS Preflight flag
    pub preflight: bool,
    /// Origin of CORS Request
    pub origin: Url,
}

/// Metadata about a loaded resource, such as is obtained from HTTP headers.
#[derive(Clone, Deserialize, Serialize, HeapSizeOf)]
pub struct Metadata {
    /// Final URL after redirects.
    pub final_url: Url,

    /// MIME type / subtype.
    pub content_type: Option<(ContentType)>,

    /// Character set.
    pub charset: Option<String>,

    #[ignore_heap_size_of = "Defined in hyper"]
    /// Headers
    pub headers: Option<Headers>,

    /// HTTP Status
    pub status: Option<RawStatus>,
}

impl Metadata {
    /// Metadata with defaults for everything optional.
    pub fn default(url: Url) -> Self {
        Metadata {
            final_url:    url,
            content_type: None,
            charset:      None,
            headers: None,
            // https://fetch.spec.whatwg.org/#concept-response-status-message
            status: Some(RawStatus(200, "OK".into())),
        }
    }

    /// Extract the parts of a Mime that we care about.
    pub fn set_content_type(&mut self, content_type: Option<&Mime>) {
        match content_type {
            None => (),
            Some(mime) => {
                self.content_type = Some(ContentType(mime.clone()));
                let &Mime(_, _, ref parameters) = mime;
                for &(ref k, ref v) in parameters {
                    if &Attr::Charset == k {
                        self.charset = Some(v.to_string());
                    }
                }
            }
        }
    }
}

/// The creator of a given cookie
#[derive(PartialEq, Copy, Clone, Deserialize, Serialize)]
pub enum CookieSource {
    /// An HTTP API
    HTTP,
    /// A non-HTTP API
    NonHTTP,
}

/// Messages sent in response to a `Load` message
#[derive(PartialEq, Debug, Deserialize, Serialize)]
pub enum ProgressMsg {
    /// Binary data - there may be multiple of these
    Payload(Vec<u8>),
    /// Indicates loading is complete, either successfully or not
    Done(Result<(), String>)
}

/// Convenience function for synchronously loading a whole resource.
pub fn load_whole_resource(resource_task: &ResourceTask, url: Url)
        -> Result<(Metadata, Vec<u8>), String> {
    let (start_chan, start_port) = ipc::channel().unwrap();
    resource_task.send(ControlMsg::Load(LoadData::new(url, None),
                       LoadConsumer::Channel(start_chan))).unwrap();
    let response = start_port.recv().unwrap();

    let mut buf = vec!();
    loop {
        match response.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => buf.push_all(&data),
            ProgressMsg::Done(Ok(())) => return Ok((response.metadata, buf)),
            ProgressMsg::Done(Err(e)) => return Err(e)
        }
    }
}

/// Load a URL asynchronously and iterate over chunks of bytes from the response.
pub fn load_bytes_iter(pending: PendingAsyncLoad) -> (Metadata, ProgressMsgPortIterator) {
    let input_port = pending.load();
    let response = input_port.recv().unwrap();
    let iter = ProgressMsgPortIterator {
        progress_port: response.progress_port
    };
    (response.metadata, iter)
}

/// Iterator that reads chunks of bytes from a ProgressMsg port
pub struct ProgressMsgPortIterator {
    progress_port: IpcReceiver<ProgressMsg>,
}

impl Iterator for ProgressMsgPortIterator {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        match self.progress_port.recv().unwrap() {
            ProgressMsg::Payload(data) => Some(data),
            ProgressMsg::Done(Ok(()))  => None,
            ProgressMsg::Done(Err(e))  => {
                error!("error receiving bytes: {}", e);
                None
            }
        }
    }
}
<|endoftext|>