<commit_before>use std::{io::Cursor, path::Path};

use rocksdb::{
    BlockBasedIndexType, BlockBasedOptions, ColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode,
    IteratorMode, MergeOperands, Options, ReadOptions, SliceTransform, WriteBatch, DB,
};

use crate::model::{
    GameId, GameInfo, MasterGame, Month, PersonalEntry, PersonalKey, PersonalKeyPrefix,
    PersonalStatus, UserId,
};

#[derive(Debug)]
pub struct Database {
    inner: DB,
}

impl Database {
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Database, rocksdb::Error> {
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);

        let mut personal_opts = Options::default();
        personal_opts.set_merge_operator_associative("personal merge", personal_merge);
        personal_opts
            .set_prefix_extractor(SliceTransform::create_fixed_prefix(PersonalKeyPrefix::SIZE));
        let mut personal_block_opts = BlockBasedOptions::default();
        personal_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        personal_block_opts.set_block_size(4 * 1024);
        personal_block_opts.set_bloom_filter(8, true);
        personal_opts.set_block_based_table_factory(&personal_block_opts);

        let mut game_opts = Options::default();
        game_opts.set_merge_operator_associative("game merge", game_merge);
        game_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut game_block_opts = BlockBasedOptions::default();
        game_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        game_block_opts.set_block_size(1024);
        game_opts.set_block_based_table_factory(&game_block_opts);

        let mut player_opts = Options::default();
        player_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut player_block_opts = BlockBasedOptions::default();
        player_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        player_block_opts.set_block_size(1024);
        player_opts.set_block_based_table_factory(&player_block_opts);

        let inner = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                ColumnFamilyDescriptor::new("personal", personal_opts),
                ColumnFamilyDescriptor::new("game", game_opts),
                ColumnFamilyDescriptor::new("player", player_opts),
                ColumnFamilyDescriptor::new("master", todo!()),
                ColumnFamilyDescriptor::new("master_game", todo!()),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn queryable(&self) -> QueryableDatabase<'_> {
        QueryableDatabase {
            db: &self.inner,
            cf_personal: self.inner.cf_handle("personal").expect("cf personal"),
            cf_game: self.inner.cf_handle("game").expect("cf game"),
            cf_player: self.inner.cf_handle("player").expect("cf player"),
            cf_master: self.inner.cf_handle("master").expect("cf master"),
            cf_master_game: self.inner.cf_handle("master_game").expect("cf master_game"),
        }
    }
}

pub struct QueryableDatabase<'a> {
    db: &'a DB,
    cf_personal: &'a ColumnFamily,
    cf_game: &'a ColumnFamily,
    cf_player: &'a ColumnFamily,
    cf_master: &'a ColumnFamily,
    cf_master_game: &'a ColumnFamily,
}

impl QueryableDatabase<'_> {
    pub fn db_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value(name)
    }

    pub fn game_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_game, name)
    }

    pub fn personal_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_personal, name)
    }

    pub fn get_game_info(&self, id: GameId) -> Result<Option<GameInfo>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_game, id.to_bytes())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            GameInfo::read(&mut cursor).expect("deserialize game info")
        }))
    }

    pub fn get_personal(
        &self,
        key: &PersonalKeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<PersonalEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_month(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_month(until.add_months_saturating(1)).into_bytes());

        let iterator = self
            .db
            .iterator_cf_opt(self.cf_personal, opt, IteratorMode::Start);

        let mut entry = PersonalEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize personal entry");
        }

        Ok(entry)
    }

    pub fn get_player_status(&self, id: &UserId) -> Result<Option<PersonalStatus>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_player, id.as_str())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            PersonalStatus::read(&mut cursor).expect("deserialize status")
        }))
    }

    pub fn put_player_status(
        &self,
        id: &UserId,
        status: &PersonalStatus,
    ) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalStatus::SIZE_HINT));
        status.write(&mut cursor).expect("serialize status");
        self.db
            .put_cf(self.cf_player, id.as_str(), cursor.into_inner())
    }

    pub fn get_master_game(&self, id: GameId) -> Result<Option<MasterGame>, rocksdb::Error> {
        Ok(self
            .db
            .get_cf(self.cf_master_game, id.to_bytes())?
            .map(|buf| serde_json::from_slice(&buf).expect("deserialize master game")))
    }

    pub fn batch(&self) -> Batch<'_> {
        Batch {
            queryable: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct Batch<'a> {
    queryable: &'a QueryableDatabase<'a>,
    batch: WriteBatch,
}

impl Batch<'_> {
    pub fn merge_personal(&mut self, key: PersonalKey, entry: PersonalEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize personal entry");
        self.batch.merge_cf(
            self.queryable.cf_personal,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

    pub fn merge_game_info(&mut self, id: GameId, info: GameInfo) {
        let mut cursor = Cursor::new(Vec::with_capacity(GameInfo::SIZE_HINT));
        info.write(&mut cursor).expect("serialize game info");
        self.batch
            .merge_cf(self.queryable.cf_game, id.to_bytes(), cursor.into_inner());
    }

    pub fn put_master_game(&mut self, id: GameId, game: &MasterGame) {
        self.batch.put_cf(
            self.queryable.cf_master_game,
            id.to_bytes(),
            serde_json::to_vec(game).expect("serialize master game"),
        );
    }

    pub fn write(self) -> Result<(), rocksdb::Error> {
        self.queryable.db.write(self.batch)
    }
}

fn game_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    // Take latest game info, but merge index status.
    let mut info: Option<GameInfo> = None;
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        let mut new_info = GameInfo::read(&mut cursor).expect("read for game merge");
        if let Some(old_info) = info {
            new_info.indexed.white |= old_info.indexed.white;
            new_info.indexed.black |= old_info.indexed.black;
        }
        info = Some(new_info);
        size_hint = op.len();
    }
    info.map(|info| {
        let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
        info.write(&mut cursor).expect("write game");
        cursor.into_inner()
    })
}

fn personal_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PersonalEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for personal merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write personal entry");
    Some(cursor.into_inner())
}
<commit_msg>prepare cf_master_game opts<commit_after>use std::{io::Cursor, path::Path};

use rocksdb::{
    BlockBasedIndexType, BlockBasedOptions, ColumnFamily, ColumnFamilyDescriptor, DBWithThreadMode,
    IteratorMode, MergeOperands, Options, ReadOptions, SliceTransform, WriteBatch, DB,
};

use crate::model::{
    GameId, GameInfo, MasterGame, Month, PersonalEntry, PersonalKey, PersonalKeyPrefix,
    PersonalStatus, UserId,
};

#[derive(Debug)]
pub struct Database {
    inner: DB,
}

impl Database {
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Database, rocksdb::Error> {
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);

        let mut personal_opts = Options::default();
        personal_opts.set_merge_operator_associative("personal merge", personal_merge);
        personal_opts
            .set_prefix_extractor(SliceTransform::create_fixed_prefix(PersonalKeyPrefix::SIZE));
        let mut personal_block_opts = BlockBasedOptions::default();
        personal_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        personal_block_opts.set_block_size(4 * 1024);
        personal_block_opts.set_bloom_filter(8, true);
        personal_opts.set_block_based_table_factory(&personal_block_opts);

        let mut game_opts = Options::default();
        game_opts.set_merge_operator_associative("game merge", game_merge);
        game_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut game_block_opts = BlockBasedOptions::default();
        game_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        game_block_opts.set_block_size(1024);
        game_opts.set_block_based_table_factory(&game_block_opts);

        let mut player_opts = Options::default();
        player_opts.set_prefix_extractor(SliceTransform::create_noop());
        let mut player_block_opts = BlockBasedOptions::default();
        player_block_opts.set_index_type(BlockBasedIndexType::HashSearch);
        player_block_opts.set_block_size(1024);
        player_opts.set_block_based_table_factory(&player_block_opts);

        let mut master_game_opts = Options::default();
        master_game_opts.set_prefix_extractor(SliceTransform::create_noop());

        let inner = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            path,
            vec![
                ColumnFamilyDescriptor::new("personal", personal_opts),
                ColumnFamilyDescriptor::new("game", game_opts),
                ColumnFamilyDescriptor::new("player", player_opts),
                ColumnFamilyDescriptor::new("master", todo!()),
                ColumnFamilyDescriptor::new("master_game", master_game_opts),
            ],
        )?;

        Ok(Database { inner })
    }

    pub fn queryable(&self) -> QueryableDatabase<'_> {
        QueryableDatabase {
            db: &self.inner,
            cf_personal: self.inner.cf_handle("personal").expect("cf personal"),
            cf_game: self.inner.cf_handle("game").expect("cf game"),
            cf_player: self.inner.cf_handle("player").expect("cf player"),
            cf_master: self.inner.cf_handle("master").expect("cf master"),
            cf_master_game: self.inner.cf_handle("master_game").expect("cf master_game"),
        }
    }
}

pub struct QueryableDatabase<'a> {
    db: &'a DB,
    cf_personal: &'a ColumnFamily,
    cf_game: &'a ColumnFamily,
    cf_player: &'a ColumnFamily,
    cf_master: &'a ColumnFamily,
    cf_master_game: &'a ColumnFamily,
}

impl QueryableDatabase<'_> {
    pub fn db_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value(name)
    }

    pub fn game_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_game, name)
    }

    pub fn personal_property(&self, name: &str) -> Result<Option<String>, rocksdb::Error> {
        self.db.property_value_cf(self.cf_personal, name)
    }

    pub fn get_game_info(&self, id: GameId) -> Result<Option<GameInfo>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_game, id.to_bytes())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            GameInfo::read(&mut cursor).expect("deserialize game info")
        }))
    }

    pub fn get_personal(
        &self,
        key: &PersonalKeyPrefix,
        since: Month,
        until: Month,
    ) -> Result<PersonalEntry, rocksdb::Error> {
        let mut opt = ReadOptions::default();
        opt.set_prefix_same_as_start(true);
        opt.set_iterate_lower_bound(key.with_month(since).into_bytes());
        opt.set_iterate_upper_bound(key.with_month(until.add_months_saturating(1)).into_bytes());

        let iterator = self
            .db
            .iterator_cf_opt(self.cf_personal, opt, IteratorMode::Start);

        let mut entry = PersonalEntry::default();
        for (_key, value) in iterator {
            let mut cursor = Cursor::new(value);
            entry
                .extend_from_reader(&mut cursor)
                .expect("deserialize personal entry");
        }

        Ok(entry)
    }

    pub fn get_player_status(&self, id: &UserId) -> Result<Option<PersonalStatus>, rocksdb::Error> {
        Ok(self.db.get_cf(self.cf_player, id.as_str())?.map(|buf| {
            let mut cursor = Cursor::new(buf);
            PersonalStatus::read(&mut cursor).expect("deserialize status")
        }))
    }

    pub fn put_player_status(
        &self,
        id: &UserId,
        status: &PersonalStatus,
    ) -> Result<(), rocksdb::Error> {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalStatus::SIZE_HINT));
        status.write(&mut cursor).expect("serialize status");
        self.db
            .put_cf(self.cf_player, id.as_str(), cursor.into_inner())
    }

    pub fn get_master_game(&self, id: GameId) -> Result<Option<MasterGame>, rocksdb::Error> {
        Ok(self
            .db
            .get_cf(self.cf_master_game, id.to_bytes())?
            .map(|buf| serde_json::from_slice(&buf).expect("deserialize master game")))
    }

    pub fn batch(&self) -> Batch<'_> {
        Batch {
            queryable: self,
            batch: WriteBatch::default(),
        }
    }
}

pub struct Batch<'a> {
    queryable: &'a QueryableDatabase<'a>,
    batch: WriteBatch,
}

impl Batch<'_> {
    pub fn merge_personal(&mut self, key: PersonalKey, entry: PersonalEntry) {
        let mut cursor = Cursor::new(Vec::with_capacity(PersonalEntry::SIZE_HINT));
        entry.write(&mut cursor).expect("serialize personal entry");
        self.batch.merge_cf(
            self.queryable.cf_personal,
            key.into_bytes(),
            cursor.into_inner(),
        );
    }

    pub fn merge_game_info(&mut self, id: GameId, info: GameInfo) {
        let mut cursor = Cursor::new(Vec::with_capacity(GameInfo::SIZE_HINT));
        info.write(&mut cursor).expect("serialize game info");
        self.batch
            .merge_cf(self.queryable.cf_game, id.to_bytes(), cursor.into_inner());
    }

    pub fn put_master_game(&mut self, id: GameId, game: &MasterGame) {
        self.batch.put_cf(
            self.queryable.cf_master_game,
            id.to_bytes(),
            serde_json::to_vec(game).expect("serialize master game"),
        );
    }

    pub fn write(self) -> Result<(), rocksdb::Error> {
        self.queryable.db.write(self.batch)
    }
}

fn game_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    // Take latest game info, but merge index status.
    let mut info: Option<GameInfo> = None;
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        let mut new_info = GameInfo::read(&mut cursor).expect("read for game merge");
        if let Some(old_info) = info {
            new_info.indexed.white |= old_info.indexed.white;
            new_info.indexed.black |= old_info.indexed.black;
        }
        info = Some(new_info);
        size_hint = op.len();
    }
    info.map(|info| {
        let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
        info.write(&mut cursor).expect("write game");
        cursor.into_inner()
    })
}

fn personal_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PersonalEntry::default();
    let mut size_hint = 0;
    for op in existing.into_iter().chain(operands.into_iter()) {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("deserialize for personal merge");
        size_hint += op.len();
    }
    let mut cursor = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut cursor).expect("write personal entry");
    Some(cursor.into_inner())
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::str;
use serde::de;
use error::{Error, Result};

pub struct BencodeAccess<'a, R: 'a + Read> {
    de: &'a mut Deserializer<R>,
}

impl<'a, R: 'a + Read> BencodeAccess<'a, R> {
    fn new(de: &'a mut Deserializer<R>) -> BencodeAccess<'a, R> {
        BencodeAccess { de: de }
    }
}

impl<'de, 'a, R: 'a + Read> de::VariantAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;

    fn unit_variant(self) -> Result<()> {
        Ok(())
    }

    fn newtype_variant_seed<T: de::DeserializeSeed<'de>>(self, seed: T) -> Result<T::Value> {
        let res = seed.deserialize(&mut *self.de)?;
        if ParseResult::End != self.de.parse()? {
            return Err(Error::InvalidType("expected `d`".to_string()));
        }
        Ok(res)
    }

    fn tuple_variant<V: de::Visitor<'de>>(self, _: usize, visitor: V) -> Result<V::Value> {
        let res = de::Deserializer::deserialize_any(&mut *self.de, visitor)?;
        if ParseResult::End != self.de.parse()? {
            return Err(Error::InvalidType("expected `d`".to_string()));
        }
        Ok(res)
    }

    fn struct_variant<V: de::Visitor<'de>>(self,
                                           _: &'static [&'static str],
                                           visitor: V)
                                           -> Result<V::Value> {
        let res = de::Deserializer::deserialize_any(&mut *self.de, visitor)?;
        if ParseResult::End != self.de.parse()? {
            return Err(Error::InvalidType("expected `d`".to_string()));
        }
        Ok(res)
    }
}

impl<'de, 'a, R: 'a + Read> de::SeqAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;

    fn next_element_seed<T: de::DeserializeSeed<'de>>(&mut self,
                                                      seed: T)
                                                      -> Result<Option<T::Value>> {
        match self.de.parse()? {
            ParseResult::End => Ok(None),
            r @ _ => {
                self.de.next = Some(r);
                Ok(Some(seed.deserialize(&mut *self.de)?))
            }
        }
    }
}

impl<'de, 'a, R: 'a + Read> de::MapAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;
    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
        where K: de::DeserializeSeed<'de>
    {
        match self.de.parse()? {
            ParseResult::End => Ok(None),
            r @ _ => {
                self.de.next = Some(r);
                Ok(Some(seed.deserialize(&mut *self.de)?))
            }
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
        where V: de::DeserializeSeed<'de>
    {
        seed.deserialize(&mut *self.de)
    }
}

impl<'de, 'a, R: 'a + Read> de::EnumAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;
    type Variant = Self;
    fn variant_seed<V: de::DeserializeSeed<'de>>(self, seed: V) -> Result<(V::Value, Self)> {
        match self.de.parse()? {
            t @ ParseResult::Bytes(_) => {
                self.de.next = Some(t);
                Ok((seed.deserialize(&mut *self.de)?, self))
            }
            ParseResult::Map => Ok((seed.deserialize(&mut *self.de)?, self)),
            _ => Err(Error::EndOfStream), // FIXME
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
enum ParseResult {
    Int(i64),
    Bytes(Vec<u8>),
    /// list start
    List,
    /// map start
    Map,
    /// list or map end
    End,
}

#[derive(Debug)]
pub struct Deserializer<R: Read> {
    reader: R,
    next: Option<ParseResult>,
}

impl<'de, R: Read> Deserializer<R> {
    pub fn new(reader: R) -> Deserializer<R> {
        Deserializer {
            reader: reader,
            next: None,
        }
    }

    fn parse_int(&mut self) -> Result<i64> {
        let mut buf = [0; 1];
        let mut result = Vec::new();
        loop {
            if 1 != self.reader.read(&mut buf).map_err(Error::IoError)? {
                return Err(Error::EndOfStream);
            }
            match buf[0] {
                b'e' => {
                    let len_str =
                        String::from_utf8(result)
                            .map_err(|_| {
                                         Error::InvalidValue("Non UTF-8 integer encoding"
                                                                 .to_string())
                                     })?;
                    let len_int = len_str
                        .parse()
                        .map_err(|_| {
                                     Error::InvalidValue(format!("Can't parse `{}` as integer",
                                                                 len_str))
                                 })?;
                    return Ok(len_int);
                }
                n => result.push(n),
            }
        }
    }

    fn parse_bytes_len(&mut self, len_char: u8) -> Result<usize> {
        let mut buf = [0; 1];
        let mut len = Vec::new();
        len.push(len_char);
        loop {
            if 1 != self.reader.read(&mut buf).map_err(Error::IoError)? {
                return Err(Error::EndOfStream);
            }
            match buf[0] {
                b':' => {
                    let len_str =
                        String::from_utf8(len)
                            .map_err(|_| {
                                         Error::InvalidValue("Non UTF-8 integer encoding"
                                                                 .to_string())
                                     })?;
                    let len_int = len_str.parse()
                        .map_err(|_| Error::InvalidValue(format!("Can't parse `{}` as string length", len_str)))?;
                    return Ok(len_int);
                }
                n => len.push(n),
            }
        }
    }

    fn parse_bytes(&mut self, len_char: u8) -> Result<Vec<u8>> {
        let len = self.parse_bytes_len(len_char)?;
        let mut buf = vec![0u8; len];
        let actual_len = self.reader
            .read(buf.as_mut_slice())
            .map_err(Error::IoError)?;
        if len != actual_len {
            return Err(Error::EndOfStream);
        }
        Ok(buf)
    }

    fn parse(&mut self) -> Result<ParseResult> {
        let mut buf = [0; 1];
        if 1 != self.reader.read(&mut buf).map_err(Error::IoError)? {
            return Err(Error::EndOfStream);
        }
        match buf[0] {
            b'i' => Ok(ParseResult::Int(self.parse_int()?)),
            n @ b'0'...b'9' => Ok(ParseResult::Bytes(self.parse_bytes(n)?)),
            b'l' => Ok(ParseResult::List),
            b'd' => Ok(ParseResult::Map),
            b'e' => Ok(ParseResult::End),
            c @ _ => Err(Error::InvalidValue(format!("Invalid charackter `{}`", c as char))),
        }
    }
}

impl<'de, 'a, R: Read> de::Deserializer<'de> for &'a mut Deserializer<R> {
    type Error = Error;

    #[inline]
    fn deserialize_any<V: de::Visitor<'de>>(mut self, visitor: V) -> Result<V::Value> {
        let parser_res = match self.next.take() {
            Some(v) => v,
            None => self.parse()?,
        };
        match parser_res {
            ParseResult::Int(i) => visitor.visit_i64(i),
            ParseResult::Bytes(s) => visitor.visit_bytes(s.as_ref()),
            ParseResult::List => visitor.visit_seq(BencodeAccess::new(&mut self)),
            ParseResult::Map => visitor.visit_map(BencodeAccess::new(&mut self)),
            ParseResult::End => Err(Error::EndOfStream),
        }
    }

    forward_to_deserialize_any! {
        i64 string seq bool i8 i16 i32 u8 u16 u32
        u64 f32 f64 char str unit bytes byte_buf map unit_struct tuple_struct tuple
        newtype_struct ignored_any identifier struct
    }

    #[inline]
    fn deserialize_option<V: de::Visitor<'de>>(self, visitor: V) -> Result<V::Value> {
        visitor.visit_some(self)
    }

    #[inline]
    fn deserialize_enum<V>(self,
                           _name: &str,
                           _variants: &'static [&'static str],
                           visitor: V)
                           -> Result<V::Value>
        where V: de::Visitor<'de>
    {
        visitor.visit_enum(BencodeAccess::new(self))
    }
}

pub fn from_str<'de, T>(s: &'de str) -> Result<T>
    where T: de::Deserialize<'de>
{
    from_bytes(s.as_bytes())
}

pub fn from_bytes<'de, T>(b: &'de [u8]) -> Result<T>
    where T: de::Deserialize<'de>
{
    de::Deserialize::deserialize(&mut Deserializer::new(b))
}
<commit_msg>fix error msg<commit_after>use std::io::Read;
use std::str;
use serde::de;
use error::{Error, Result};

pub struct BencodeAccess<'a, R: 'a + Read> {
    de: &'a mut Deserializer<R>,
}

impl<'a, R: 'a + Read> BencodeAccess<'a, R> {
    fn new(de: &'a mut Deserializer<R>) -> BencodeAccess<'a, R> {
        BencodeAccess { de: de }
    }
}

impl<'de, 'a, R: 'a + Read> de::VariantAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;

    fn unit_variant(self) -> Result<()> {
        Ok(())
    }

    fn newtype_variant_seed<T: de::DeserializeSeed<'de>>(self, seed: T) -> Result<T::Value> {
        let res = seed.deserialize(&mut *self.de)?;
        if ParseResult::End != self.de.parse()? {
            return Err(Error::InvalidType("expected `d`".to_string()));
        }
        Ok(res)
    }

    fn tuple_variant<V: de::Visitor<'de>>(self, _: usize, visitor: V) -> Result<V::Value> {
        let res = de::Deserializer::deserialize_any(&mut *self.de, visitor)?;
        if ParseResult::End != self.de.parse()? {
            return Err(Error::InvalidType("expected `d`".to_string()));
        }
        Ok(res)
    }

    fn struct_variant<V: de::Visitor<'de>>(self,
                                           _: &'static [&'static str],
                                           visitor: V)
                                           -> Result<V::Value> {
        let res = de::Deserializer::deserialize_any(&mut *self.de, visitor)?;
        if ParseResult::End != self.de.parse()? {
            return Err(Error::InvalidType("expected `d`".to_string()));
        }
        Ok(res)
    }
}

impl<'de, 'a, R: 'a + Read> de::SeqAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;

    fn next_element_seed<T: de::DeserializeSeed<'de>>(&mut self,
                                                      seed: T)
                                                      -> Result<Option<T::Value>> {
        match self.de.parse()? {
            ParseResult::End => Ok(None),
            r @ _ => {
                self.de.next = Some(r);
                Ok(Some(seed.deserialize(&mut *self.de)?))
            }
        }
    }
}

impl<'de, 'a, R: 'a + Read> de::MapAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;
    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
        where K: de::DeserializeSeed<'de>
    {
        match self.de.parse()? {
            ParseResult::End => Ok(None),
            r @ _ => {
                self.de.next = Some(r);
                Ok(Some(seed.deserialize(&mut *self.de)?))
            }
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
        where V: de::DeserializeSeed<'de>
    {
        seed.deserialize(&mut *self.de)
    }
}

impl<'de, 'a, R: 'a + Read> de::EnumAccess<'de> for BencodeAccess<'a, R> {
    type Error = Error;
    type Variant = Self;
    fn variant_seed<V: de::DeserializeSeed<'de>>(self, seed: V) -> Result<(V::Value, Self)> {
        match self.de.parse()? {
            t @ ParseResult::Bytes(_) => {
                self.de.next = Some(t);
                Ok((seed.deserialize(&mut *self.de)?, self))
            }
            ParseResult::Map => Ok((seed.deserialize(&mut *self.de)?, self)),
            t @ _ => Err(Error::InvalidValue(format!("Expected bytes or map; got `{:?}`", t))),
        }
    }
}

#[derive(Debug, Eq, PartialEq)]
enum ParseResult {
    Int(i64),
    Bytes(Vec<u8>),
    /// list start
    List,
    /// map start
    Map,
    /// list or map end
    End,
}

#[derive(Debug)]
pub struct Deserializer<R: Read> {
    reader: R,
    next: Option<ParseResult>,
}

impl<'de, R: Read> Deserializer<R> {
    pub fn new(reader: R) -> Deserializer<R> {
        Deserializer {
            reader: reader,
            next: None,
        }
    }

    fn parse_int(&mut self) -> Result<i64> {
        let mut buf = [0; 1];
        let mut result = Vec::new();
        loop {
            if 1 != self.reader.read(&mut buf).map_err(Error::IoError)? {
                return Err(Error::EndOfStream);
            }
            match buf[0] {
                b'e' => {
                    let len_str =
                        String::from_utf8(result)
                            .map_err(|_| {
                                         Error::InvalidValue("Non UTF-8 integer encoding"
                                                                 .to_string())
                                     })?;
                    let len_int = len_str
                        .parse()
                        .map_err(|_| {
                                     Error::InvalidValue(format!("Can't parse `{}` as integer",
                                                                 len_str))
                                 })?;
                    return Ok(len_int);
                }
                n => result.push(n),
            }
        }
    }

    fn parse_bytes_len(&mut self, len_char: u8) -> Result<usize> {
        let mut buf = [0; 1];
        let mut len = Vec::new();
        len.push(len_char);
        loop {
            if 1 != self.reader.read(&mut buf).map_err(Error::IoError)? {
                return Err(Error::EndOfStream);
            }
            match buf[0] {
                b':' => {
                    let len_str =
                        String::from_utf8(len)
                            .map_err(|_| {
                                         Error::InvalidValue("Non UTF-8 integer encoding"
                                                                 .to_string())
                                     })?;
                    let len_int = len_str.parse()
                        .map_err(|_| Error::InvalidValue(format!("Can't parse `{}` as string length", len_str)))?;
                    return Ok(len_int);
                }
                n => len.push(n),
            }
        }
    }

    fn parse_bytes(&mut self, len_char: u8) -> Result<Vec<u8>> {
        let len = self.parse_bytes_len(len_char)?;
        let mut buf = vec![0u8; len];
        let actual_len = self.reader
            .read(buf.as_mut_slice())
            .map_err(Error::IoError)?;
        if len != actual_len {
            return Err(Error::EndOfStream);
        }
        Ok(buf)
    }

    fn parse(&mut self) -> Result<ParseResult> {
        let mut buf = [0; 1];
        if 1 != self.reader.read(&mut buf).map_err(Error::IoError)? {
            return Err(Error::EndOfStream);
        }
        match buf[0] {
            b'i' => Ok(ParseResult::Int(self.parse_int()?)),
            n @ b'0'...b'9' => Ok(ParseResult::Bytes(self.parse_bytes(n)?)),
            b'l' => Ok(ParseResult::List),
            b'd' => Ok(ParseResult::Map),
            b'e' => Ok(ParseResult::End),
            c @ _ => Err(Error::InvalidValue(format!("Invalid charackter `{}`", c as char))),
        }
    }
}

impl<'de, 'a, R: Read> de::Deserializer<'de> for &'a mut Deserializer<R> {
    type Error = Error;

    #[inline]
    fn deserialize_any<V: de::Visitor<'de>>(mut self, visitor: V) -> Result<V::Value> {
        let parser_res = match self.next.take() {
            Some(v) => v,
            None => self.parse()?,
        };
        match parser_res {
            ParseResult::Int(i) => visitor.visit_i64(i),
            ParseResult::Bytes(s) => visitor.visit_bytes(s.as_ref()),
            ParseResult::List => visitor.visit_seq(BencodeAccess::new(&mut self)),
            ParseResult::Map => visitor.visit_map(BencodeAccess::new(&mut self)),
            ParseResult::End => Err(Error::EndOfStream),
        }
    }

    forward_to_deserialize_any! {
        i64 string seq bool i8 i16 i32 u8 u16 u32
        u64 f32 f64 char str unit bytes byte_buf map unit_struct tuple_struct tuple
        newtype_struct ignored_any identifier struct
    }

    #[inline]
    fn deserialize_option<V: de::Visitor<'de>>(self, visitor: V) -> Result<V::Value> {
        visitor.visit_some(self)
    }

    #[inline]
    fn deserialize_enum<V>(self,
                           _name: &str,
                           _variants: &'static [&'static str],
                           visitor: V)
                           -> Result<V::Value>
        where V: de::Visitor<'de>
    {
        visitor.visit_enum(BencodeAccess::new(self))
    }
}

pub fn from_str<'de, T>(s: &'de str) -> Result<T>
    where T: de::Deserialize<'de>
{
    from_bytes(s.as_bytes())
}

pub fn from_bytes<'de, T>(b: &'de [u8]) -> Result<T>
    where T: de::Deserialize<'de>
{
    de::Deserialize::deserialize(&mut Deserializer::new(b))
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(unix)]
mod imp {
    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::{Ok, Err};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    #[cfg(unix)]
    pub struct OsRng {
        inner: ReaderRng<File>
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: reader_rng })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            self.inner.next_u32()
        }
        fn next_u64(&mut self) -> u64 {
            self.inner.next_u64()
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            self.inner.fill_bytes(v)
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use core_collections::Collection;
    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::{Ok, Err};
    use rt::stack;
    use self::libc::{c_ulong, DWORD, BYTE, LPCSTR, BOOL};
    use slice::MutableVector;

    type HCRYPTPROV = c_ulong;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;
    static NTE_BAD_SIGNATURE: DWORD = 0x80090006;

    #[allow(non_snake_case_functions)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let mut ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            // FIXME #13259:
            // It turns out that if we can't acquire a context with the
            // NTE_BAD_SIGNATURE error code, the documentation states:
            //
            //     The provider DLL signature could not be verified. Either the
            //     DLL or the digital signature has been tampered with.
            //
            // Sounds fishy, no? As it turns out, our signature can be bad
            // because our Thread Information Block (TIB) isn't exactly what it
            // expects. As to why, I have no idea. The only data we store in the
            // TIB is the stack limit for each thread, but apparently that's
            // enough to make the signature valid.
            //
            // Furthermore, this error only happens the *first* time we call
            // CryptAcquireContext, so we don't have to worry about future
            // calls.
            //
            // Anyway, the fix employed here is that if we see this error, we
            // pray that we're not close to the end of the stack, temporarily
            // set the stack limit to 0 (what the TIB originally was), acquire a
            // context, and then reset the stack limit.
            //
            // Again, I'm not sure why this is the fix, nor why we're getting
            // this error. All I can say is that this seems to allow libnative
            // to progress where it otherwise would be hindered. Who knew?
            if ret == 0 && os::errno() as DWORD == NTE_BAD_SIGNATURE {
                unsafe {
                    let limit = stack::get_sp_limit();
                    stack::record_sp_limit(0);
                    ret = CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                               PROV_RSA_FULL,
                                               CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
                    stack::record_sp_limit(limit);
                }
            }

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                fail!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(proc() {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<commit_msg>auto merge of #14715 : vhbit/rust/ios-pr2, r=alexcrichton<commit_after>// Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(unix, not(target_os = "ios"))]
mod imp {
    use io::{IoResult, File};
    use path::Path;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use result::{Ok, Err};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    #[cfg(unix)]
    pub struct OsRng {
        inner: ReaderRng<File>
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let reader = try!(File::open(&Path::new("/dev/urandom")));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: reader_rng })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            self.inner.next_u32()
        }
        fn next_u64(&mut self) -> u64 {
            self.inner.next_u64()
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            self.inner.fill_bytes(v)
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use collections::Collection;
    use io::{IoResult};
    use kinds::marker;
    use mem;
    use os;
    use rand::Rng;
    use result::{Ok};
    use self::libc::{c_int, size_t};
    use slice::MutableVector;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed
    /// This does not block.
    pub struct OsRng {
        marker: marker::NoCopy
    }

    struct SecRandom;

    static kSecRandomDefault: *SecRandom = 0 as *SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *SecRandom, count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            Ok(OsRng {marker: marker::NoCopy} )
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate libc;

    use core_collections::Collection;
    use io::{IoResult, IoError};
    use mem;
    use ops::Drop;
    use os;
    use rand::Rng;
    use result::{Ok, Err};
    use rt::stack;
    use self::libc::{c_ulong, DWORD, BYTE, LPCSTR, BOOL};
    use slice::MutableVector;

    type HCRYPTPROV = c_ulong;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    static PROV_RSA_FULL: DWORD = 1;
    static CRYPT_SILENT: DWORD = 64;
    static CRYPT_VERIFYCONTEXT: DWORD = 0xF0000000;
    static NTE_BAD_SIGNATURE: DWORD = 0x80090006;

    #[allow(non_snake_case_functions)]
    extern "system" {
        fn CryptAcquireContextA(phProv: *mut HCRYPTPROV,
                                pszContainer: LPCSTR,
                                pszProvider: LPCSTR,
                                dwProvType: DWORD,
                                dwFlags: DWORD) -> BOOL;
        fn CryptGenRandom(hProv: HCRYPTPROV,
                          dwLen: DWORD,
                          pbBuffer: *mut BYTE) -> BOOL;
        fn CryptReleaseContext(hProv: HCRYPTPROV, dwFlags: DWORD) -> BOOL;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> IoResult<OsRng> {
            let mut hcp = 0;
            let mut ret = unsafe {
                CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            // FIXME #13259:
            // It turns out that if we can't acquire a context with the
            // NTE_BAD_SIGNATURE error code, the documentation states:
            //
            //     The provider DLL signature could not be verified. Either the
            //     DLL or the digital signature has been tampered with.
            //
            // Sounds fishy, no? As it turns out, our signature can be bad
            // because our Thread Information Block (TIB) isn't exactly what it
            // expects. As to why, I have no idea. The only data we store in the
            // TIB is the stack limit for each thread, but apparently that's
            // enough to make the signature valid.
            //
            // Furthermore, this error only happens the *first* time we call
            // CryptAcquireContext, so we don't have to worry about future
            // calls.
            //
            // Anyway, the fix employed here is that if we see this error, we
            // pray that we're not close to the end of the stack, temporarily
            // set the stack limit to 0 (what the TIB originally was), acquire a
            // context, and then reset the stack limit.
            //
            // Again, I'm not sure why this is the fix, nor why we're getting
            // this error. All I can say is that this seems to allow libnative
            // to progress where it otherwise would be hindered. Who knew?
            if ret == 0 && os::errno() as DWORD == NTE_BAD_SIGNATURE {
                unsafe {
                    let limit = stack::get_sp_limit();
                    stack::record_sp_limit(0);
                    ret = CryptAcquireContextA(&mut hcp, 0 as LPCSTR, 0 as LPCSTR,
                                               PROV_RSA_FULL,
                                               CRYPT_VERIFYCONTEXT | CRYPT_SILENT);
                    stack::record_sp_limit(limit);
                }
            }

            if ret == 0 {
                Err(IoError::last_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8, .. 4];
            self.fill_bytes(v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8, .. 8];
            self.fill_bytes(v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                fail!("couldn't generate random bytes: {}", os::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                fail!("couldn't release context: {}", os::last_os_error());
            }
        }
    }
}

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

    use super::OsRng;
    use rand::Rng;
    use task;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8, .. 1000];
        r.fill_bytes(v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in range(0, 20) {
            let (tx, rx) = channel();
            txs.push(tx);
            task::spawn(proc() {
                // wait until all the tasks are ready to go.
                rx.recv();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                task::deschedule();
                let mut v = [0u8, .. 1000];

                for _ in range(0, 100) {
                    r.next_u32();
                    task::deschedule();
                    r.next_u64();
                    task::deschedule();
                    r.fill_bytes(v);
                    task::deschedule();
                }
            })
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(())
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use crate::Error;

extern crate std;
use std::thread_local;

use js_sys::{global, Uint8Array};
use wasm_bindgen::{prelude::wasm_bindgen, JsCast, JsValue};

// Maximum is 65536 bytes see https://developer.mozilla.org/en-US/docs/Web/API/Crypto/getRandomValues
const BROWSER_CRYPTO_BUFFER_SIZE: usize = 256;

enum RngSource {
    Node(NodeCrypto),
    Browser(BrowserCrypto, Uint8Array),
}

// JsValues are always per-thread, so we initialize RngSource for each thread.
//   See: https://github.com/rustwasm/wasm-bindgen/pull/955
thread_local!(
    static RNG_SOURCE: Result<RngSource, Error> = getrandom_init();
);

pub(crate) fn getrandom_inner(dest: &mut [u8]) -> Result<(), Error> {
    RNG_SOURCE.with(|result| {
        let source = result.as_ref().map_err(|&e| e)?;

        match source {
            RngSource::Node(n) => {
                if n.random_fill_sync(dest).is_err() {
                    return Err(Error::NODE_RANDOM_FILL_SYNC);
                }
            }
            RngSource::Browser(crypto, buf) => {
                // getRandomValues does not work with all types of WASM memory,
                // so we initially write to browser memory to avoid exceptions.
                for chunk in dest.chunks_mut(BROWSER_CRYPTO_BUFFER_SIZE) {
                    // The chunk can be smaller than buf's length, so we call to
                    // JS to create a smaller view of buf without allocation.
                    let sub_buf = buf.subarray(0, chunk.len() as u32);

                    if crypto.get_random_values(&sub_buf).is_err() {
                        return Err(Error::WEB_GET_RANDOM_VALUES);
                    }
                    sub_buf.copy_to(chunk);
                }
            }
        };
        Ok(())
    })
}

fn getrandom_init() -> Result<RngSource, Error> {
    let global: Global = global().unchecked_into();
    if is_node(&global) {
        let crypto = NODE_MODULE
            .require("crypto")
            .map_err(|_| Error::NODE_CRYPTO)?;
        return Ok(RngSource::Node(crypto));
    }

    // Assume we are in some Web environment (browser or web worker). We get
    // `self.crypto` (called `msCrypto` on IE), so we can call
    // `crypto.getRandomValues`. If `crypto` isn't defined, we assume that
    // we are in an older web browser and the OS RNG isn't available.
    let crypto = match (global.crypto(), global.ms_crypto()) {
        (c, _) if c.is_object() => c,
        (_, c) if c.is_object() => c,
        _ => return Err(Error::WEB_CRYPTO),
    };

    let buf = Uint8Array::new_with_length(BROWSER_CRYPTO_BUFFER_SIZE as u32);
    Ok(RngSource::Browser(crypto, buf))
}

// Taken from https://www.npmjs.com/package/browser-or-node
fn is_node(global: &Global) -> bool {
    let process = global.process();
    if process.is_object() {
        let versions = process.versions();
        if versions.is_object() {
            return versions.node().is_string();
        }
    }
    false
}

#[wasm_bindgen]
extern "C" {
    type Global; // Return type of js_sys::global()

    // Web Crypto API (https://www.w3.org/TR/WebCryptoAPI/)
    #[wasm_bindgen(method, getter, js_name = "msCrypto")]
    fn ms_crypto(this: &Global) -> BrowserCrypto;
    #[wasm_bindgen(method, getter)]
    fn crypto(this: &Global) -> BrowserCrypto;
    type BrowserCrypto;
    #[wasm_bindgen(method, js_name = getRandomValues, catch)]
    fn get_random_values(this: &BrowserCrypto, buf: &Uint8Array) -> Result<(), JsValue>;

    type NodeModule;
    #[wasm_bindgen(js_name = module)]
    static NODE_MODULE: NodeModule;
    // Node JS crypto module (https://nodejs.org/api/crypto.html)
    #[wasm_bindgen(method, catch)]
    fn require(this: &NodeModule, s: &str) -> Result<NodeCrypto, JsValue>;
    type NodeCrypto;
    #[wasm_bindgen(method, js_name = randomFillSync, catch)]
    fn random_fill_sync(this: &NodeCrypto, buf: &mut [u8]) -> Result<(), JsValue>;

    // Node JS process Object (https://nodejs.org/api/process.html)
    #[wasm_bindgen(method, getter)]
    fn process(this: &Global) -> Process;
    type Process;
    #[wasm_bindgen(method, getter)]
    fn versions(this: &Process) -> Versions;
    type Versions;
    #[wasm_bindgen(method, getter)]
    fn node(this: &Versions) -> JsValue;
}
<commit_msg>js: Add comment explaining why we do this hack<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use crate::Error;

extern crate std;
use std::thread_local;

use js_sys::{global, Uint8Array};
use wasm_bindgen::{prelude::wasm_bindgen, JsCast, JsValue};

// Maximum is 65536 bytes see https://developer.mozilla.org/en-US/docs/Web/API/Crypto/getRandomValues
const BROWSER_CRYPTO_BUFFER_SIZE: usize = 256;

enum RngSource {
    Node(NodeCrypto),
    Browser(BrowserCrypto, Uint8Array),
}

// JsValues are always per-thread, so we initialize RngSource for each thread.
//   See: https://github.com/rustwasm/wasm-bindgen/pull/955
thread_local!(
    static RNG_SOURCE: Result<RngSource, Error> = getrandom_init();
);

pub(crate) fn getrandom_inner(dest: &mut [u8]) -> Result<(), Error> {
    RNG_SOURCE.with(|result| {
        let source = result.as_ref().map_err(|&e| e)?;

        match source {
            RngSource::Node(n) => {
                if n.random_fill_sync(dest).is_err() {
                    return Err(Error::NODE_RANDOM_FILL_SYNC);
                }
            }
            RngSource::Browser(crypto, buf) => {
                // getRandomValues does not work with all types of WASM memory,
                // so we initially write to browser memory to avoid exceptions.
                for chunk in dest.chunks_mut(BROWSER_CRYPTO_BUFFER_SIZE) {
                    // The chunk can be smaller than buf's length, so we call to
                    // JS to create a smaller view of buf without allocation.
                    let sub_buf = buf.subarray(0, chunk.len() as u32);

                    if crypto.get_random_values(&sub_buf).is_err() {
                        return Err(Error::WEB_GET_RANDOM_VALUES);
                    }
                    sub_buf.copy_to(chunk);
                }
            }
        };
        Ok(())
    })
}

fn getrandom_init() -> Result<RngSource, Error> {
    let global: Global = global().unchecked_into();
    if is_node(&global) {
        let crypto = NODE_MODULE
            .require("crypto")
            .map_err(|_| Error::NODE_CRYPTO)?;
        return Ok(RngSource::Node(crypto));
    }

    // Assume we are in some Web environment (browser or web worker). We get
    // `self.crypto` (called `msCrypto` on IE), so we can call
    // `crypto.getRandomValues`. If `crypto` isn't defined, we assume that
    // we are in an older web browser and the OS RNG isn't available.
    let crypto = match (global.crypto(), global.ms_crypto()) {
        (c, _) if c.is_object() => c,
        (_, c) if c.is_object() => c,
        _ => return Err(Error::WEB_CRYPTO),
    };

    let buf = Uint8Array::new_with_length(BROWSER_CRYPTO_BUFFER_SIZE as u32);
    Ok(RngSource::Browser(crypto, buf))
}

// Taken from https://www.npmjs.com/package/browser-or-node
fn is_node(global: &Global) -> bool {
    let process = global.process();
    if process.is_object() {
        let versions = process.versions();
        if versions.is_object() {
            return versions.node().is_string();
        }
    }
    false
}

#[wasm_bindgen]
extern "C" {
    type Global; // Return type of js_sys::global()

    // Web Crypto API (https://www.w3.org/TR/WebCryptoAPI/)
    #[wasm_bindgen(method, getter, js_name = "msCrypto")]
    fn ms_crypto(this: &Global) -> BrowserCrypto;
    #[wasm_bindgen(method, getter)]
    fn crypto(this: &Global) -> BrowserCrypto;
    type BrowserCrypto;
    #[wasm_bindgen(method, js_name = getRandomValues, catch)]
    fn get_random_values(this: &BrowserCrypto, buf: &Uint8Array) -> Result<(), JsValue>;

    // We use a "module" object here instead of just annotating require() with
    // js_name = "module.require", so that Webpack doesn't give a warning. See:
    //   https://github.com/rust-random/getrandom/issues/224
    type NodeModule;
    #[wasm_bindgen(js_name = module)]
    static NODE_MODULE: NodeModule;
    // Node JS crypto module (https://nodejs.org/api/crypto.html)
    #[wasm_bindgen(method, catch)]
    fn require(this: &NodeModule, s: &str) -> Result<NodeCrypto, JsValue>;
    type NodeCrypto;
    #[wasm_bindgen(method, js_name = randomFillSync, catch)]
    fn random_fill_sync(this: &NodeCrypto, buf: &mut [u8]) -> Result<(), JsValue>;

    // Node JS process Object (https://nodejs.org/api/process.html)
    #[wasm_bindgen(method, getter)]
    fn process(this: &Global) -> Process;
    type Process;
    #[wasm_bindgen(method, getter)]
    fn versions(this: &Process) -> Versions;
    type Versions;
    #[wasm_bindgen(method, getter)]
    fn node(this: &Versions) -> JsValue;
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use error::*;
use bytecode;

pub struct VM {
    instruction_ptr: usize,
    stack_ptr: usize,
    globals: Vec<u64>,
    stack: Vec<u64>,
}

impl VM {
    pub fn new() -> VM {
        VM {
            instruction_ptr: 0,
            stack_ptr: 0,
            globals: Vec::new(),
            stack: Vec::new(),
        }
    }

    pub fn exec<P: AsRef<Path>>(&mut self, path: P) -> VMResult<()> {
        let mut image_file = File::open(path).chain_err(|| "unable to open game image")?;
        let mut image_string = String::new();
        image_file.read_to_string(&mut image_string).chain_err(|| "unable to read image")?;

        let byte_iter: Vec<u8> = image_string.bytes().collect();
        while self.instruction_ptr < byte_iter.len() {
            let byte = self.current_byte(&byte_iter)
                .chain_err(|| "unable to read current byte")?;

            match byte {
                bytecode::HALT => break,
                bytecode::ADD => self.add().chain_err(|| "unable to execute 'add' instruction")?,
                bytecode::SUB => self.sub().chain_err(|| "unable to execute 'sub' instruction")?,
                bytecode::MUL => self.mul().chain_err(|| "unable to execute 'mul' instruction")?,
                bytecode::DIV => self.div().chain_err(|| "unable to execute 'div' instruction")?,
                bytecode::PRINT => {
                    self.print().chain_err(|| "unable to execute 'print' instruction")?
                }
                bytecode::PUSH => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    self.push(res).chain_err(|| "unable to push value to the stack")?;
                }
                bytecode::JMP => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    self.jmp(res).chain_err(|| "unable to jump")?;
                    continue;
                }
                bytecode::JZ => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    let top_of_stack = self.peek()
                        .chain_err(|| "unable to get current top of stack")?;

                    if top_of_stack == 0 {
                        self.jmp(res).chain_err(|| "unable to jump")?;
                        continue;
                    }
                }
                bytecode::JNZ => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    let top_of_stack = self.peek()
                        .chain_err(|| "unable to get current top of stack")?;

                    if top_of_stack != 0 {
                        self.jmp(res).chain_err(|| "unable to jump")?;
                        continue;
                    }
                }
                _ => bail!("unexpected opcode: {:08x}", byte),
            }

            self.advance_instruction_ptr().chain_err(|| "unable to advance instruction pointer")?;
        }

        Ok(())
    }

    fn push(&mut self, value: u32) -> VMResult<()> {
        if !self.stack.is_empty() {
            self.stack_ptr += 1;
        }

        self.stack.push(value as u64);

        Ok(())
    }

    fn pop(&mut self) -> VMResult<u64> {
        if self.stack.is_empty() {
            bail!("unable to pop value off an empty Stack");
        }

        let res = self.stack.remove(self.stack_ptr);
        if !self.stack.is_empty() {
            self.stack_ptr -= 1;
        }

        Ok(res)
    }

    fn add(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a + b) as u32).chain_err(|| "unable to push value to the stack")?;
        Ok(())
    }

    fn sub(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a - b) as u32).chain_err(|| "unable to push to stack")?;
        Ok(())
    }

    fn mul(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a * b) as u32).chain_err(|| "unable to push to stack")?;
        Ok(())
    }

    fn div(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a / b) as u32).chain_err(|| "unable to push to stack")?;
        Ok(())
    }

    fn jmp(&mut self, addr: u32) -> VMResult<()> {
        self.instruction_ptr = addr as usize;
        Ok(())
    }

    fn read_word(&mut self, bytes: &Vec<u8>) -> VMResult<u32> {
        // Build a u32 from single bytes
        let mut res: u32 = 0;
        for _ in 0..4 {
            res <<= 8;
            self.advance_instruction_ptr().chain_err(|| "unable to advance instruction pointer")?;
            let current_byte = self.current_byte(&bytes)
                .chain_err(|| "unable to read current byte")?;
            res |= current_byte as u32;
        }

        Ok(res)
    }

    fn advance_instruction_ptr(&mut self) -> VMResult<()> {
        self.instruction_ptr += 1;
        Ok(())
    }

    fn current_byte(&mut self, bytes: &Vec<u8>) -> VMResult<u8> {
        if self.instruction_ptr < bytes.len() {
            Ok(bytes[self.instruction_ptr])
        } else {
            bail!("instruction pointer out of bounds");
        }
    }

    fn peek(&mut self) -> VMResult<u64> {
        if self.stack_ptr < self.stack.len() {
            Ok(self.stack[self.stack_ptr])
        } else {
            bail!("stack pointer out of bounds");
        }
    }

    fn print(&mut self) -> VMResult<()> {
        let top = self.peek().chain_err(|| "unable to get current top of stack")?;
        println!("{:?}", top);
        Ok(())
    }
}
<commit_msg>Pretty opcode printing<commit_after>use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use error::*;
use bytecode;

pub struct VM {
    instruction_ptr: usize,
    stack_ptr: usize,
    globals: Vec<u64>,
    stack: Vec<u64>,
}

impl VM {
    pub fn new() -> VM {
        VM {
            instruction_ptr: 0,
            stack_ptr: 0,
            globals: Vec::new(),
            stack: Vec::new(),
        }
    }

    pub fn exec<P: AsRef<Path>>(&mut self, path: P) -> VMResult<()> {
        let mut image_file = File::open(path).chain_err(|| "unable to open game image")?;
        let mut image_string = String::new();
        image_file.read_to_string(&mut image_string).chain_err(|| "unable to read image")?;

        let byte_iter: Vec<u8> = image_string.bytes().collect();
        while self.instruction_ptr < byte_iter.len() {
            let byte = self.current_byte(&byte_iter)
                .chain_err(|| "unable to read current byte")?;

            match byte {
                bytecode::HALT => break,
                bytecode::ADD => self.add().chain_err(|| "unable to execute 'add' instruction")?,
                bytecode::SUB => self.sub().chain_err(|| "unable to execute 'sub' instruction")?,
                bytecode::MUL => self.mul().chain_err(|| "unable to execute 'mul' instruction")?,
                bytecode::DIV => self.div().chain_err(|| "unable to execute 'div' instruction")?,
                bytecode::PRINT => {
                    self.print().chain_err(|| "unable to execute 'print' instruction")?
                }
                bytecode::PUSH => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    self.push(res).chain_err(|| "unable to push value to the stack")?;
                }
                bytecode::JMP => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    self.jmp(res).chain_err(|| "unable to jump")?;
                    continue;
                }
                bytecode::JZ => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    let top_of_stack = self.peek()
                        .chain_err(|| "unable to get current top of stack")?;

                    if top_of_stack == 0 {
                        self.jmp(res).chain_err(|| "unable to jump")?;
                        continue;
                    }
                }
                bytecode::JNZ => {
                    let res = self.read_word(&byte_iter).chain_err(|| "unable to read word")?;

                    let top_of_stack = self.peek()
                        .chain_err(|| "unable to get current top of stack")?;

                    if top_of_stack != 0 {
                        self.jmp(res).chain_err(|| "unable to jump")?;
                        continue;
                    }
                }
                _ => bail!("unexpected opcode: {:02x}", byte),
            }

            self.advance_instruction_ptr().chain_err(|| "unable to advance instruction pointer")?;
        }

        Ok(())
    }

    fn push(&mut self, value: u32) -> VMResult<()> {
        if !self.stack.is_empty() {
            self.stack_ptr += 1;
        }

        self.stack.push(value as u64);

        Ok(())
    }

    fn pop(&mut self) -> VMResult<u64> {
        if self.stack.is_empty() {
            bail!("unable to pop value off an empty Stack");
        }

        let res = self.stack.remove(self.stack_ptr);
        if !self.stack.is_empty() {
            self.stack_ptr -= 1;
        }

        Ok(res)
    }

    fn add(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a + b) as u32).chain_err(|| "unable to push value to the stack")?;
        Ok(())
    }

    fn sub(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a - b) as u32).chain_err(|| "unable to push to stack")?;
        Ok(())
    }

    fn mul(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a * b) as u32).chain_err(|| "unable to push to stack")?;
        Ok(())
    }

    fn div(&mut self) -> VMResult<()> {
        let b = self.pop().chain_err(|| "unable to pop value off the stack")?;
        let a = self.pop().chain_err(|| "unable to pop value off the stack")?;
        self.push((a / b) as u32).chain_err(|| "unable to push to stack")?;
        Ok(())
    }

    fn jmp(&mut self, addr: u32) -> VMResult<()> {
        self.instruction_ptr = addr as usize;
        Ok(())
    }

    fn read_word(&mut self, bytes: &Vec<u8>) -> VMResult<u32> {
        // Build a u32 from single bytes
        let mut res: u32 = 0;
        for _ in 0..4 {
            res <<= 8;
            self.advance_instruction_ptr().chain_err(|| "unable to advance instruction pointer")?;
            let current_byte = self.current_byte(&bytes)
                .chain_err(|| "unable to read current byte")?;
            res |= current_byte as u32;
        }

        Ok(res)
    }

    fn advance_instruction_ptr(&mut self) -> VMResult<()> {
        self.instruction_ptr += 1;
        Ok(())
    }

    fn current_byte(&mut self, bytes: &Vec<u8>) -> VMResult<u8> {
        if self.instruction_ptr < bytes.len() {
            Ok(bytes[self.instruction_ptr])
        } else {
            bail!("instruction pointer out of bounds");
        }
    }

    fn peek(&mut self) -> VMResult<u64> {
        if self.stack_ptr < self.stack.len() {
            Ok(self.stack[self.stack_ptr])
        } else {
            bail!("stack pointer out of bounds");
        }
    }

    fn print(&mut self) -> VMResult<()> {
        let top = self.peek().chain_err(|| "unable to get current top of stack")?;
        println!("{:?}", top);
        Ok(())
    }
}
<|endoftext|>


<commit_before>use serde::de::{Deserialize, Deserializer, Error};

use types::*;

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Message {
    pub id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward: Option<Forward>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub caption: Option<String>,
    pub kind: MessageKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Forward {
    pub date: Integer,
    pub from: ForwardFrom,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ForwardFrom {
    User {
        user: User,
    },
    Channel {
        channel: Channel,
        message_id: Integer
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageKind {
    Text {
        data: String,
        entities: Vec<MessageEntity>,
    },
    Audio {
        data: Audio,
    },
    Document {
        data: Document,
    },
    Photo {
        data: Vec<PhotoSize>,
    },
    Sticker {
        data: Sticker,
    },
    Video {
        data: Video,
    },
    Voice {
        data: Voice,
    },
    Contact {
        data: Contact,
    },
    Location {
        data: Location,
    },
    Venue {
        data: Venue,
    },
    NewChatMember {
        data: User,
    },
    LeftChatMember {
        data: User,
    },
    NewChatTitle {
        data: String,
    },
    NewChatPhoto {
        data: PhotoSize,
    },
    DeleteChatPhoto,
    GroupChatCreated,
    SupergroupChatCreated,
    ChannelChatCreated,
    MigrateToChatId {
        data: Integer,
    },
    MigrateFromChatId {
        data: Integer,
    },
    PinnedMessage {
        data: Box<Message>,
    },
    #[doc(hidden)]
    Unknown {
        raw: RawMessage
    }
}

impl Deserialize for Message {  // TODO(knsd): Remove .clone()
    fn deserialize<D>(deserializer: D) -> Result<Message, D::Error> where D: Deserializer {
        let raw: RawMessage = Deserialize::deserialize(deserializer)?;

        let id = raw.message_id;
        let from = raw.from.clone();
        let date = raw.date;
        let chat = raw.chat.clone();
        let reply_to_message = raw.reply_to_message.clone();
        let edit_date = raw.edit_date;
        let caption = raw.caption.clone();

        let forward = match (raw.forward_date, &raw.forward_from,
                             &raw.forward_from_chat, raw.forward_from_message_id) {
            (None, &None, &None, None) => None,
            (Some(date), &Some(ref from), &None, None) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::User {
                        user: from.clone(),
                    },
                })
            },
            (Some(date), &None, &Some(Chat::Channel(ref channel)), Some(message_id)) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::Channel {
                        channel: channel.clone(),
                        message_id: message_id,
                    },
                })
            },
            _ => {
                return Err(D::Error::custom("invalid forward fields combination"))
            }
        };

        let make_message = |kind| {
            Ok(Message {
                id: id,
                from: from,
                date: date,
                chat: chat,
                forward: forward,
                reply_to_message: reply_to_message,
                edit_date: edit_date,
                caption: caption,
                kind: kind,
            })
        };

        macro_rules! maybe_field {
            ($name:ident, $variant:ident) => {{
                if let Some(val) = raw.$name {
                    return make_message(MessageKind::$variant {
                        data: val
                    })
                }
            }}
        }

        macro_rules! maybe_true_field {
            ($name:ident, $variant:ident) => {{
                if let Some(True) = raw.$name {
                    return make_message(MessageKind::$variant)
                }
            }}
        }

        if let Some(text) = raw.text {
            let entities = raw.entities.unwrap_or_else(Vec::new);
            return make_message(MessageKind::Text {
                data: text,
                entities: entities,
            })
        }

        maybe_field!(audio, Audio);
        maybe_field!(document, Document);
        maybe_field!(photo, Photo);
        maybe_field!(sticker, Sticker);
        maybe_field!(video, Video);
        maybe_field!(voice, Voice);
        maybe_field!(contact, Contact);
        maybe_field!(location, Location);
        maybe_field!(venue, Venue);
        maybe_field!(new_chat_member, NewChatMember);
        maybe_field!(left_chat_member, LeftChatMember);
        maybe_field!(new_chat_title, NewChatTitle);
        maybe_field!(new_chat_photo, NewChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(group_chat_created, GroupChatCreated);
        maybe_true_field!(supergroup_chat_created, SupergroupChatCreated);
        maybe_true_field!(channel_chat_created, ChannelChatCreated);
        maybe_field!(migrate_to_chat_id, MigrateToChatId);
        maybe_field!(migrate_from_chat_id, MigrateFromChatId);
        maybe_field!(pinned_message, PinnedMessage);

        make_message(MessageKind::Unknown {
            raw: raw,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessage {
    pub message_id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward_from: Option<User>,
    pub forward_from_chat: Option<Chat>,
    pub forward_from_message_id: Option<Integer>,
    pub forward_date: Option<Integer>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub text: Option<String>,
    pub entities: Option<Vec<MessageEntity>>,
    pub audio: Option<Audio>,
    pub document: Option<Document>,
//    pub game: Option<Game>,
    pub photo: Option<Vec<PhotoSize>>,
    pub sticker: Option<Sticker>,
    pub video: Option<Video>,
    pub voice: Option<Voice>,
    pub caption: Option<String>,
    pub contact: Option<Contact>,
    pub location: Option<Location>,
    pub venue: Option<Venue>,
    pub new_chat_member: Option<User>,
    pub left_chat_member: Option<User>,
    pub new_chat_title: Option<String>,
    pub new_chat_photo: Option<PhotoSize>,
    pub delete_chat_photo: Option<True>,
    pub group_chat_created: Option<True>,
    pub supergroup_chat_created: Option<True>,
    pub channel_chat_created: Option<True>,
    pub migrate_to_chat_id: Option<Integer>,
    pub migrate_from_chat_id: Option<Integer>,
    pub pinned_message: Option<Box<Message>>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct MessageEntity {
    offset: Integer,
    length: Integer,
    kind: MessageEntityKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageEntityKind {
    Mention,
    Hashtag,
    BotCommand,
    Url,
    Email,
    Bold,
    Italic,
    Code,
    Pre,
    TextLink(String), // TODO(knsd) URL?
    TextMention(User),
    Unknown(RawMessageEntity),
}

impl Deserialize for MessageEntity {
    fn deserialize<D>(deserializer: D) -> Result<MessageEntity, D::Error> where D: Deserializer {
        use self::MessageEntityKind::*;

        let raw: RawMessageEntity = Deserialize::deserialize(deserializer)?;

        let offset = raw.offset;
        let length = raw.length;

        macro_rules! required_field {
            ($name:ident) => {{
                match raw.$name {
                    Some(val) => val,
                    None => return Err(D::Error::missing_field(stringify!($name)))
                }
            }}
        }

        let kind = match raw.type_.as_str() {
            "mention" => Mention,
            "hashtag" => Hashtag,
            "bot_command" => BotCommand,
            "url" => Url,
            "email" => Email,
            "bold" => Bold,
            "italic" => Italic,
            "code" => Code,
            "pre" => Pre,
            "text_link" => TextLink(required_field!(url)),
            "text_mention" => TextMention(required_field!(user)),
            _ => Unknown(raw),
        };

        Ok(MessageEntity {
            offset: offset,
            length: length,
            kind: kind,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessageEntity {
    #[serde(rename="type")]
    pub type_: String,
    pub offset: Integer,
    pub length: Integer,
    pub url: Option<String>,
    pub user: Option<User>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct PhotoSize {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Audio {
    pub file_id: String,
    pub duration: Integer,
    pub performer: Option<String>,
    pub title: Option<String>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Document {
    pub file_id: String,
    pub thumb: Option<PhotoSize>,
    pub file_name: Option<String>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Sticker {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub thumb: Option<PhotoSize>,
    pub emoji: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Video {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub duration: Integer,
    pub thumb: Option<PhotoSize>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Voice {
    pub file_id: String,
    pub duration: Integer,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Contact {
    pub phone_number: String,
    pub first_name: String,
    pub last_name: Option<String>,
    pub user_id: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Location {
    pub longitude: Float,
    pub latitude: Float,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Venue {
    pub location: Location,
    pub title: String,
    pub address: String,
    pub foursquare_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct UserProfilePhotos {
    pub total_count: Integer,
    pub photos: Vec<Vec<PhotoSize>>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct File {
    pub file_id: String,
    pub file_size: Option<Integer>,
    pub file_path: Option<String>,
}<commit_msg>refactor: hide MessageEntityKind::Unknown<commit_after>use serde::de::{Deserialize, Deserializer, Error};

use types::*;

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Message {
    pub id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward: Option<Forward>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub caption: Option<String>,
    pub kind: MessageKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Forward {
    pub date: Integer,
    pub from: ForwardFrom,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ForwardFrom {
    User {
        user: User,
    },
    Channel {
        channel: Channel,
        message_id: Integer
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageKind {
    Text {
        data: String,
        entities: Vec<MessageEntity>,
    },
    Audio {
        data: Audio,
    },
    Document {
        data: Document,
    },
    Photo {
        data: Vec<PhotoSize>,
    },
    Sticker {
        data: Sticker,
    },
    Video {
        data: Video,
    },
    Voice {
        data: Voice,
    },
    Contact {
        data: Contact,
    },
    Location {
        data: Location,
    },
    Venue {
        data: Venue,
    },
    NewChatMember {
        data: User,
    },
    LeftChatMember {
        data: User,
    },
    NewChatTitle {
        data: String,
    },
    NewChatPhoto {
        data: PhotoSize,
    },
    DeleteChatPhoto,
    GroupChatCreated,
    SupergroupChatCreated,
    ChannelChatCreated,
    MigrateToChatId {
        data: Integer,
    },
    MigrateFromChatId {
        data: Integer,
    },
    PinnedMessage {
        data: Box<Message>,
    },
    #[doc(hidden)]
    Unknown {
        raw: RawMessage
    }
}

impl Deserialize for Message {  // TODO(knsd): Remove .clone()
    fn deserialize<D>(deserializer: D) -> Result<Message, D::Error> where D: Deserializer {
        let raw: RawMessage = Deserialize::deserialize(deserializer)?;

        let id = raw.message_id;
        let from = raw.from.clone();
        let date = raw.date;
        let chat = raw.chat.clone();
        let reply_to_message = raw.reply_to_message.clone();
        let edit_date = raw.edit_date;
        let caption = raw.caption.clone();

        let forward = match (raw.forward_date, &raw.forward_from,
                             &raw.forward_from_chat, raw.forward_from_message_id) {
            (None, &None, &None, None) => None,
            (Some(date), &Some(ref from), &None, None) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::User {
                        user: from.clone(),
                    },
                })
            },
            (Some(date), &None, &Some(Chat::Channel(ref channel)), Some(message_id)) => {
                Some(Forward {
                    date: date,
                    from: ForwardFrom::Channel {
                        channel: channel.clone(),
                        message_id: message_id,
                    },
                })
            },
            _ => {
                return Err(D::Error::custom("invalid forward fields combination"))
            }
        };

        let make_message = |kind| {
            Ok(Message {
                id: id,
                from: from,
                date: date,
                chat: chat,
                forward: forward,
                reply_to_message: reply_to_message,
                edit_date: edit_date,
                caption: caption,
                kind: kind,
            })
        };

        macro_rules! maybe_field {
            ($name:ident, $variant:ident) => {{
                if let Some(val) = raw.$name {
                    return make_message(MessageKind::$variant {
                        data: val
                    })
                }
            }}
        }

        macro_rules! maybe_true_field {
            ($name:ident, $variant:ident) => {{
                if let Some(True) = raw.$name {
                    return make_message(MessageKind::$variant)
                }
            }}
        }

        if let Some(text) = raw.text {
            let entities = raw.entities.unwrap_or_else(Vec::new);
            return make_message(MessageKind::Text {
                data: text,
                entities: entities,
            })
        }

        maybe_field!(audio, Audio);
        maybe_field!(document, Document);
        maybe_field!(photo, Photo);
        maybe_field!(sticker, Sticker);
        maybe_field!(video, Video);
        maybe_field!(voice, Voice);
        maybe_field!(contact, Contact);
        maybe_field!(location, Location);
        maybe_field!(venue, Venue);
        maybe_field!(new_chat_member, NewChatMember);
        maybe_field!(left_chat_member, LeftChatMember);
        maybe_field!(new_chat_title, NewChatTitle);
        maybe_field!(new_chat_photo, NewChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(delete_chat_photo, DeleteChatPhoto);
        maybe_true_field!(group_chat_created, GroupChatCreated);
        maybe_true_field!(supergroup_chat_created, SupergroupChatCreated);
        maybe_true_field!(channel_chat_created, ChannelChatCreated);
        maybe_field!(migrate_to_chat_id, MigrateToChatId);
        maybe_field!(migrate_from_chat_id, MigrateFromChatId);
        maybe_field!(pinned_message, PinnedMessage);

        make_message(MessageKind::Unknown {
            raw: raw,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessage {
    pub message_id: Integer,
    pub from: Option<User>,
    pub date: Integer,
    pub chat: Chat,
    pub forward_from: Option<User>,
    pub forward_from_chat: Option<Chat>,
    pub forward_from_message_id: Option<Integer>,
    pub forward_date: Option<Integer>,
    pub reply_to_message: Option<Box<Message>>,
    pub edit_date: Option<Integer>,
    pub text: Option<String>,
    pub entities: Option<Vec<MessageEntity>>,
    pub audio: Option<Audio>,
    pub document: Option<Document>,
//    pub game: Option<Game>,
    pub photo: Option<Vec<PhotoSize>>,
    pub sticker: Option<Sticker>,
    pub video: Option<Video>,
    pub voice: Option<Voice>,
    pub caption: Option<String>,
    pub contact: Option<Contact>,
    pub location: Option<Location>,
    pub venue: Option<Venue>,
    pub new_chat_member: Option<User>,
    pub left_chat_member: Option<User>,
    pub new_chat_title: Option<String>,
    pub new_chat_photo: Option<PhotoSize>,
    pub delete_chat_photo: Option<True>,
    pub group_chat_created: Option<True>,
    pub supergroup_chat_created: Option<True>,
    pub channel_chat_created: Option<True>,
    pub migrate_to_chat_id: Option<Integer>,
    pub migrate_from_chat_id: Option<Integer>,
    pub pinned_message: Option<Box<Message>>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct MessageEntity {
    offset: Integer,
    length: Integer,
    kind: MessageEntityKind,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum MessageEntityKind {
    Mention,
    Hashtag,
    BotCommand,
    Url,
    Email,
    Bold,
    Italic,
    Code,
    Pre,
    TextLink(String), // TODO(knsd) URL?
    TextMention(User),
    #[doc(hidden)]
    Unknown(RawMessageEntity),
}

impl Deserialize for MessageEntity {
    fn deserialize<D>(deserializer: D) -> Result<MessageEntity, D::Error> where D: Deserializer {
        use self::MessageEntityKind::*;

        let raw: RawMessageEntity = Deserialize::deserialize(deserializer)?;

        let offset = raw.offset;
        let length = raw.length;

        macro_rules! required_field {
            ($name:ident) => {{
                match raw.$name {
                    Some(val) => val,
                    None => return Err(D::Error::missing_field(stringify!($name)))
                }
            }}
        }

        let kind = match raw.type_.as_str() {
            "mention" => Mention,
            "hashtag" => Hashtag,
            "bot_command" => BotCommand,
            "url" => Url,
            "email" => Email,
            "bold" => Bold,
            "italic" => Italic,
            "code" => Code,
            "pre" => Pre,
            "text_link" => TextLink(required_field!(url)),
            "text_mention" => TextMention(required_field!(user)),
            _ => Unknown(raw),
        };

        Ok(MessageEntity {
            offset: offset,
            length: length,
            kind: kind,
        })
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct RawMessageEntity {
    #[serde(rename="type")]
    pub type_: String,
    pub offset: Integer,
    pub length: Integer,
    pub url: Option<String>,
    pub user: Option<User>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct PhotoSize {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Audio {
    pub file_id: String,
    pub duration: Integer,
    pub performer: Option<String>,
    pub title: Option<String>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Document {
    pub file_id: String,
    pub thumb: Option<PhotoSize>,
    pub file_name: Option<String>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Sticker {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub thumb: Option<PhotoSize>,
    pub emoji: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Video {
    pub file_id: String,
    pub width: Integer,
    pub height: Integer,
    pub duration: Integer,
    pub thumb: Option<PhotoSize>,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Voice {
    pub file_id: String,
    pub duration: Integer,
    pub mime_type: Option<String>,
    pub file_size: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Contact {
    pub phone_number: String,
    pub first_name: String,
    pub last_name: Option<String>,
    pub user_id: Option<Integer>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Location {
    pub longitude: Float,
    pub latitude: Float,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct Venue {
    pub location: Location,
    pub title: String,
    pub address: String,
    pub foursquare_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct UserProfilePhotos {
    pub total_count: Integer,
    pub photos: Vec<Vec<PhotoSize>>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Deserialize)]
pub struct File {
    pub file_id: String,
    pub file_size: Option<Integer>,
    pub file_path: Option<String>,
}<|endoftext|>