<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This library implements a schematized DB on top of [RocksDB](https://rocksdb.org/). It makes
//! sure all data passed in and out are structured according to predefined schemas and prevents
//! access to raw keys and values. This library also enforces a set of Libra specific DB options,
//! like custom comparators and schema-to-column-family mapping.
//!
//! It requires that different kinds of key-value pairs be stored in separate column
//! families.  To use this library to store a kind of key-value pairs, the user needs to use the
//! [`define_schema!`] macro to define the schema name, the types of key and value, and name of the
//! column family.

#[macro_use]
pub mod schema;

use crate::schema::{KeyCodec, Schema, SeekKeyCodec, ValueCodec};
use failure::prelude::*;
use lazy_static::lazy_static;
use metrics::OpMetrics;
use rocksdb::{
    rocksdb_options::ColumnFamilyDescriptor, CFHandle, DBOptions, Writable, WriteOptions,
};
use std::{
    collections::{BTreeMap, HashMap},
    iter::Iterator,
    marker::PhantomData,
    path::Path,
};

lazy_static! {
    static ref OP_COUNTER: OpMetrics = OpMetrics::new_and_registered("schemadb");
}

/// Type alias to `rocksdb::ColumnFamilyOptions`. See [`rocksdb doc`](https://github.com/pingcap/rust-rocksdb/blob/master/src/rocksdb_options.rs)
pub type ColumnFamilyOptions = rocksdb::ColumnFamilyOptions;
/// Type alias to `rocksdb::ReadOptions`. See [`rocksdb doc`](https://github.com/pingcap/rust-rocksdb/blob/master/src/rocksdb_options.rs)
pub type ReadOptions = rocksdb::ReadOptions;

/// Type alias to improve readability.
pub type ColumnFamilyName = &'static str;
/// Type alias to improve readability.
pub type ColumnFamilyOptionsMap = HashMap<ColumnFamilyName, ColumnFamilyOptions>;

/// Name for the `default` column family that's always open by RocksDB. We use it to store
/// [`LedgerInfo`](../types/ledger_info/struct.LedgerInfo.html).
pub const DEFAULT_CF_NAME: ColumnFamilyName = "default";

#[derive(Debug)]
enum WriteOp {
    Value(Vec<u8>),
    Deletion,
}

/// `SchemaBatch` holds a collection of updates that can be applied to a DB atomically. The updates
/// will be applied in the order in which they are added to the `SchemaBatch`.
#[derive(Debug, Default)]
pub struct SchemaBatch {
    rows: Vec<(ColumnFamilyName, Vec<u8> /* key */, WriteOp)>,
}

impl SchemaBatch {
    /// Creates an empty batch.
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds an insert/update operation to the batch.
    pub fn put<S: Schema>(&mut self, key: &S::Key, value: &S::Value) -> Result<()> {
        let key = <S::Key as KeyCodec<S>>::encode_key(key)?;
        let value = <S::Value as ValueCodec<S>>::encode_value(value)?;
        self.rows
            .push((S::COLUMN_FAMILY_NAME, key, WriteOp::Value(value)));
        Ok(())
    }

    /// Adds a delete operation to the batch.
    pub fn delete<S: Schema>(&mut self, key: &S::Key) -> Result<()> {
        let key = <S::Key as KeyCodec<S>>::encode_key(key)?;
        self.rows
            .push((S::COLUMN_FAMILY_NAME, key, WriteOp::Deletion));
        Ok(())
    }
}

/// DB Iterator parameterized on [`Schema`] that seeks with [`Schema::Key`] and yields
/// [`Schema::Key`] and [`Schema::Value`]
pub struct SchemaIterator<'a, S> {
    db_iter: rocksdb::DBIterator<&'a rocksdb::DB>,
    phantom: PhantomData<S>,
}

impl<'a, S> SchemaIterator<'a, S>
where
    S: Schema,
{
    fn new(db_iter: rocksdb::DBIterator<&'a rocksdb::DB>) -> Self {
        SchemaIterator {
            db_iter,
            phantom: PhantomData,
        }
    }

    /// Seeks to the first key.
    pub fn seek_to_first(&mut self) -> bool {
        self.db_iter.seek(rocksdb::SeekKey::Start)
    }

    /// Seeks to the last key.
    pub fn seek_to_last(&mut self) -> bool {
        self.db_iter.seek(rocksdb::SeekKey::End)
    }

    /// Seeks to the first key whose binary representation is equal to or greater than that of the
    /// `seek_key`.
    pub fn seek<SK>(&mut self, seek_key: &SK) -> Result<bool>
    where
        SK: SeekKeyCodec<S>,
    {
        let key = <SK as SeekKeyCodec<S>>::encode_seek_key(seek_key)?;
        Ok(self.db_iter.seek(rocksdb::SeekKey::Key(&key)))
    }

    /// Seeks to the last key whose binary representation is less than or equal to that of the
    /// `seek_key`.
    ///
    /// See example in [`RocksDB doc`](https://github.com/facebook/rocksdb/wiki/SeekForPrev).
    pub fn seek_for_prev<SK>(&mut self, seek_key: &SK) -> Result<bool>
    where
        SK: SeekKeyCodec<S>,
    {
        let key = <SK as SeekKeyCodec<S>>::encode_seek_key(seek_key)?;
        Ok(self.db_iter.seek_for_prev(rocksdb::SeekKey::Key(&key)))
    }
}

impl<'a, S> Iterator for SchemaIterator<'a, S>
where
    S: Schema,
{
    type Item = Result<(S::Key, S::Value)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.db_iter.kv().map(|(raw_key, raw_value)| {
            self.db_iter.next();
            Ok((
                <S::Key as KeyCodec<S>>::decode_key(&raw_key)?,
                <S::Value as ValueCodec<S>>::decode_value(&raw_value)?,
            ))
        })
    }
}

/// Checks underlying Rocksdb instance existence by checking `CURRENT` file existence, the same way
/// Rocksdb adopts to detect db existence.
fn db_exists(path: &Path) -> bool {
    let rocksdb_current_file = path.join("CURRENT");
    rocksdb_current_file.is_file()
}

/// All the RocksDB methods return `std::result::Result<T, String>`. Since our methods return
/// `failure::Result<T>`, manual conversion is needed.
fn convert_rocksdb_err(msg: String) -> failure::Error {
    format_err!("RocksDB internal error: {}.", msg)
}

/// This DB is a schematized RocksDB wrapper where all data passed in and out are typed according to
/// [`Schema`]s.
#[derive(Debug)]
pub struct DB {
    inner: rocksdb::DB,
}

impl DB {
    /// Create db with all the column families provided if it doesn't exist at `path`; Otherwise,
    /// try to open it with all the column families.
    pub fn open<P: AsRef<Path>>(path: P, mut cf_opts_map: ColumnFamilyOptionsMap) -> Result<Self> {
        let mut db_opts = DBOptions::new();

        // For now we set the max total WAL size to be 1G. This config can be useful when column
        // families are updated at non-uniform frequencies.
        db_opts.set_max_total_wal_size(1 << 30);

        // If db exists, just open it with all cfs.
        if db_exists(path.as_ref()) {
            return DB::open_cf(db_opts, &path, cf_opts_map.into_iter().collect());
        }

        // If db doesn't exist, create a db first with all column families.
        db_opts.create_if_missing(true);

        let mut db = DB::open_cf(
            db_opts,
            path,
            vec![cf_opts_map
                .remove_entry(&DEFAULT_CF_NAME)
                .ok_or_else(|| format_err!("No \"default\" column family name found"))?],
        )?;
        cf_opts_map
            .into_iter()
            .map(|(cf_name, cf_opts)| db.create_cf((cf_name, cf_opts)))
            .collect::<Result<Vec<_>>>()?;
        Ok(db)
    }

    fn open_cf<'a, P, T>(opts: DBOptions, path: P, cfds: Vec<T>) -> Result<DB>
    where
        P: AsRef<Path>,
        T: Into<ColumnFamilyDescriptor<'a>>,
    {
        let inner = rocksdb::DB::open_cf(
            opts,
            path.as_ref().to_str().ok_or_else(|| {
                format_err!("Path {:?} can not be converted to string.", path.as_ref())
            })?,
            cfds,
        )
        .map_err(convert_rocksdb_err)?;

        Ok(DB { inner })
    }

    fn create_cf<'a, T>(&mut self, cfd: T) -> Result<()>
    where
        T: Into<ColumnFamilyDescriptor<'a>>,
    {
        let _cf_handle = self.inner.create_cf(cfd).map_err(convert_rocksdb_err)?;
        Ok(())
    }

    /// Reads single record by key.
    pub fn get<S: Schema>(&self, schema_key: &S::Key) -> Result<Option<S::Value>> {
        let k = <S::Key as KeyCodec<S>>::encode_key(&schema_key)?;
        let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;

        self.inner
            .get_cf(cf_handle, &k)
            .map_err(convert_rocksdb_err)?
            .map(|raw_value| <S::Value as ValueCodec<S>>::decode_value(&raw_value))
            .transpose()
    }

    /// Writes single record.
    pub fn put<S: Schema>(&self, key: &S::Key, value: &S::Value) -> Result<()> {
        let k = <S::Key as KeyCodec<S>>::encode_key(&key)?;
        let v = <S::Value as ValueCodec<S>>::encode_value(&value)?;
        let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;

        self.inner
            .put_cf_opt(cf_handle, &k, &v, &default_write_options())
            .map_err(convert_rocksdb_err)
    }

    /// Returns a [`SchemaIterator`] on a certain schema.
    pub fn iter<S: Schema>(&self, opts: ReadOptions) -> Result<SchemaIterator<S>> {
        let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;
        Ok(SchemaIterator::new(self.inner.iter_cf_opt(cf_handle, opts)))
    }

    /// Writes a group of records wrapped in a [`SchemaBatch`].
    pub fn write_schemas(&self, batch: SchemaBatch) -> Result<()> {
        let db_batch = rocksdb::WriteBatch::new();
        for (cf_name, key, write_op) in &batch.rows {
            let cf_handle = self.get_cf_handle(cf_name)?;
            match write_op {
                WriteOp::Value(value) => db_batch.put_cf(cf_handle, &key, &value),
                WriteOp::Deletion => db_batch.delete_cf(cf_handle, &key),
            }
            .map_err(convert_rocksdb_err)?;
        }

        self.inner
            .write_opt(&db_batch, &default_write_options())
            .map_err(convert_rocksdb_err)?;

        for (cf_name, key, write_op) in &batch.rows {
            match write_op {
                WriteOp::Value(value) => OP_COUNTER.observe(
                    &format!("db_put_bytes_{}", cf_name),
                    (key.len() + value.len()) as f64,
                ),
                WriteOp::Deletion => OP_COUNTER.inc(&format!("db_delete_{}", cf_name)),
            };
        }

        Ok(())
    }

    fn get_cf_handle(&self, cf_name: &str) -> Result<&CFHandle> {
        self.inner.cf_handle(cf_name).ok_or_else(|| {
            format_err!(
                "DB::cf_handle not found for column family name: {}",
                cf_name
            )
        })
    }

    /// Returns the approximate size of each non-empty column family in bytes.
    pub fn get_approximate_sizes_cf(&self) -> Result<BTreeMap<String, u64>> {
        let mut cf_sizes = BTreeMap::new();

        for cf_name in self.inner.cf_names().into_iter().map(ToString::to_string) {
            let cf_handle = self.get_cf_handle(&cf_name)?;
            let size = self
                .inner
                .get_property_int_cf(cf_handle, "rocksdb.estimate-live-data-size")
                .ok_or_else(|| {
                    format_err!(
                        "Unable to get approximate size of {} column family.",
                        cf_name,
                    )
                })?;
            cf_sizes.insert(cf_name, size);
        }

        Ok(cf_sizes)
    }

    /// Flushes all memtable data. If `sync` is true, the flush will wait until it's done. This is
    /// only used for testing `get_approximate_sizes_cf` in unit tests.
    pub fn flush_all(&self, sync: bool) -> Result<()> {
        for cf_name in self.inner.cf_names() {
            let cf_handle = self.get_cf_handle(cf_name)?;
            self.inner
                .flush_cf(cf_handle, sync)
                .map_err(convert_rocksdb_err)?;
        }
        Ok(())
    }
}

/// For now we always use synchronous writes. This makes sure that once the operation returns
/// `Ok(())` the data is persisted even if the machine crashes. In the future we might consider
/// selectively turning this off for some non-critical writes to improve performance.
fn default_write_options() -> WriteOptions {
    let mut opts = WriteOptions::new();
    opts.set_sync(true);
    opts
}
<commit_msg>[storage] `SchemaBatch` commit content in order<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! This library implements a schematized DB on top of [RocksDB](https://rocksdb.org/). It makes
//! sure all data passed in and out are structured according to predefined schemas and prevents
//! access to raw keys and values. This library also enforces a set of Libra specific DB options,
//! like custom comparators and schema-to-column-family mapping.
//!
//! It requires that different kinds of key-value pairs be stored in separate column
//! families.  To use this library to store a kind of key-value pairs, the user needs to use the
//! [`define_schema!`] macro to define the schema name, the types of key and value, and name of the
//! column family.

#[macro_use]
pub mod schema;

use crate::schema::{KeyCodec, Schema, SeekKeyCodec, ValueCodec};
use failure::prelude::*;
use lazy_static::lazy_static;
use metrics::OpMetrics;
use rocksdb::{
    rocksdb_options::ColumnFamilyDescriptor, CFHandle, DBOptions, Writable, WriteOptions,
};
use std::{
    collections::{BTreeMap, HashMap},
    iter::Iterator,
    marker::PhantomData,
    path::Path,
};

lazy_static! {
    static ref OP_COUNTER: OpMetrics = OpMetrics::new_and_registered("schemadb");
}

/// Type alias to `rocksdb::ColumnFamilyOptions`. See [`rocksdb doc`](https://github.com/pingcap/rust-rocksdb/blob/master/src/rocksdb_options.rs)
pub type ColumnFamilyOptions = rocksdb::ColumnFamilyOptions;
/// Type alias to `rocksdb::ReadOptions`. See [`rocksdb doc`](https://github.com/pingcap/rust-rocksdb/blob/master/src/rocksdb_options.rs)
pub type ReadOptions = rocksdb::ReadOptions;

/// Type alias to improve readability.
pub type ColumnFamilyName = &'static str;
/// Type alias to improve readability.
pub type ColumnFamilyOptionsMap = HashMap<ColumnFamilyName, ColumnFamilyOptions>;

/// Name for the `default` column family that's always open by RocksDB. We use it to store
/// [`LedgerInfo`](../types/ledger_info/struct.LedgerInfo.html).
pub const DEFAULT_CF_NAME: ColumnFamilyName = "default";

#[derive(Debug)]
enum WriteOp {
    Value(Vec<u8>),
    Deletion,
}

/// `SchemaBatch` holds a collection of updates that can be applied to a DB atomically. The updates
/// will be applied in the order in which they are added to the `SchemaBatch`.
#[derive(Debug, Default)]
pub struct SchemaBatch {
    rows: HashMap<ColumnFamilyName, BTreeMap<Vec<u8>, WriteOp>>,
}

impl SchemaBatch {
    /// Creates an empty batch.
    pub fn new() -> Self {
        Self::default()
    }

    /// Adds an insert/update operation to the batch.
    pub fn put<S: Schema>(&mut self, key: &S::Key, value: &S::Value) -> Result<()> {
        let key = <S::Key as KeyCodec<S>>::encode_key(key)?;
        let value = <S::Value as ValueCodec<S>>::encode_value(value)?;
        self.rows
            .entry(S::COLUMN_FAMILY_NAME)
            .or_insert_with(BTreeMap::new)
            .insert(key, WriteOp::Value(value));

        Ok(())
    }

    /// Adds a delete operation to the batch.
    pub fn delete<S: Schema>(&mut self, key: &S::Key) -> Result<()> {
        let key = <S::Key as KeyCodec<S>>::encode_key(key)?;
        self.rows
            .entry(S::COLUMN_FAMILY_NAME)
            .or_insert_with(BTreeMap::new)
            .insert(key, WriteOp::Deletion);

        Ok(())
    }
}

/// DB Iterator parameterized on [`Schema`] that seeks with [`Schema::Key`] and yields
/// [`Schema::Key`] and [`Schema::Value`]
pub struct SchemaIterator<'a, S> {
    db_iter: rocksdb::DBIterator<&'a rocksdb::DB>,
    phantom: PhantomData<S>,
}

impl<'a, S> SchemaIterator<'a, S>
where
    S: Schema,
{
    fn new(db_iter: rocksdb::DBIterator<&'a rocksdb::DB>) -> Self {
        SchemaIterator {
            db_iter,
            phantom: PhantomData,
        }
    }

    /// Seeks to the first key.
    pub fn seek_to_first(&mut self) -> bool {
        self.db_iter.seek(rocksdb::SeekKey::Start)
    }

    /// Seeks to the last key.
    pub fn seek_to_last(&mut self) -> bool {
        self.db_iter.seek(rocksdb::SeekKey::End)
    }

    /// Seeks to the first key whose binary representation is equal to or greater than that of the
    /// `seek_key`.
    pub fn seek<SK>(&mut self, seek_key: &SK) -> Result<bool>
    where
        SK: SeekKeyCodec<S>,
    {
        let key = <SK as SeekKeyCodec<S>>::encode_seek_key(seek_key)?;
        Ok(self.db_iter.seek(rocksdb::SeekKey::Key(&key)))
    }

    /// Seeks to the last key whose binary representation is less than or equal to that of the
    /// `seek_key`.
    ///
    /// See example in [`RocksDB doc`](https://github.com/facebook/rocksdb/wiki/SeekForPrev).
    pub fn seek_for_prev<SK>(&mut self, seek_key: &SK) -> Result<bool>
    where
        SK: SeekKeyCodec<S>,
    {
        let key = <SK as SeekKeyCodec<S>>::encode_seek_key(seek_key)?;
        Ok(self.db_iter.seek_for_prev(rocksdb::SeekKey::Key(&key)))
    }
}

impl<'a, S> Iterator for SchemaIterator<'a, S>
where
    S: Schema,
{
    type Item = Result<(S::Key, S::Value)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.db_iter.kv().map(|(raw_key, raw_value)| {
            self.db_iter.next();
            Ok((
                <S::Key as KeyCodec<S>>::decode_key(&raw_key)?,
                <S::Value as ValueCodec<S>>::decode_value(&raw_value)?,
            ))
        })
    }
}

/// Checks underlying Rocksdb instance existence by checking `CURRENT` file existence, the same way
/// Rocksdb adopts to detect db existence.
fn db_exists(path: &Path) -> bool {
    let rocksdb_current_file = path.join("CURRENT");
    rocksdb_current_file.is_file()
}

/// All the RocksDB methods return `std::result::Result<T, String>`. Since our methods return
/// `failure::Result<T>`, manual conversion is needed.
fn convert_rocksdb_err(msg: String) -> failure::Error {
    format_err!("RocksDB internal error: {}.", msg)
}

/// This DB is a schematized RocksDB wrapper where all data passed in and out are typed according to
/// [`Schema`]s.
#[derive(Debug)]
pub struct DB {
    inner: rocksdb::DB,
}

impl DB {
    /// Create db with all the column families provided if it doesn't exist at `path`; Otherwise,
    /// try to open it with all the column families.
    pub fn open<P: AsRef<Path>>(path: P, mut cf_opts_map: ColumnFamilyOptionsMap) -> Result<Self> {
        let mut db_opts = DBOptions::new();

        // For now we set the max total WAL size to be 1G. This config can be useful when column
        // families are updated at non-uniform frequencies.
        db_opts.set_max_total_wal_size(1 << 30);

        // If db exists, just open it with all cfs.
        if db_exists(path.as_ref()) {
            return DB::open_cf(db_opts, &path, cf_opts_map.into_iter().collect());
        }

        // If db doesn't exist, create a db first with all column families.
        db_opts.create_if_missing(true);

        let mut db = DB::open_cf(
            db_opts,
            path,
            vec![cf_opts_map
                .remove_entry(&DEFAULT_CF_NAME)
                .ok_or_else(|| format_err!("No \"default\" column family name found"))?],
        )?;
        cf_opts_map
            .into_iter()
            .map(|(cf_name, cf_opts)| db.create_cf((cf_name, cf_opts)))
            .collect::<Result<Vec<_>>>()?;
        Ok(db)
    }

    fn open_cf<'a, P, T>(opts: DBOptions, path: P, cfds: Vec<T>) -> Result<DB>
    where
        P: AsRef<Path>,
        T: Into<ColumnFamilyDescriptor<'a>>,
    {
        let inner = rocksdb::DB::open_cf(
            opts,
            path.as_ref().to_str().ok_or_else(|| {
                format_err!("Path {:?} can not be converted to string.", path.as_ref())
            })?,
            cfds,
        )
        .map_err(convert_rocksdb_err)?;

        Ok(DB { inner })
    }

    fn create_cf<'a, T>(&mut self, cfd: T) -> Result<()>
    where
        T: Into<ColumnFamilyDescriptor<'a>>,
    {
        let _cf_handle = self.inner.create_cf(cfd).map_err(convert_rocksdb_err)?;
        Ok(())
    }

    /// Reads single record by key.
    pub fn get<S: Schema>(&self, schema_key: &S::Key) -> Result<Option<S::Value>> {
        let k = <S::Key as KeyCodec<S>>::encode_key(&schema_key)?;
        let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;

        self.inner
            .get_cf(cf_handle, &k)
            .map_err(convert_rocksdb_err)?
            .map(|raw_value| <S::Value as ValueCodec<S>>::decode_value(&raw_value))
            .transpose()
    }

    /// Writes single record.
    pub fn put<S: Schema>(&self, key: &S::Key, value: &S::Value) -> Result<()> {
        let k = <S::Key as KeyCodec<S>>::encode_key(&key)?;
        let v = <S::Value as ValueCodec<S>>::encode_value(&value)?;
        let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;

        self.inner
            .put_cf_opt(cf_handle, &k, &v, &default_write_options())
            .map_err(convert_rocksdb_err)
    }

    /// Returns a [`SchemaIterator`] on a certain schema.
    pub fn iter<S: Schema>(&self, opts: ReadOptions) -> Result<SchemaIterator<S>> {
        let cf_handle = self.get_cf_handle(S::COLUMN_FAMILY_NAME)?;
        Ok(SchemaIterator::new(self.inner.iter_cf_opt(cf_handle, opts)))
    }

    /// Writes a group of records wrapped in a [`SchemaBatch`].
    pub fn write_schemas(&self, batch: SchemaBatch) -> Result<()> {
        let db_batch = rocksdb::WriteBatch::new();
        for (cf_name, rows) in &batch.rows {
            let cf_handle = self.get_cf_handle(cf_name)?;
            for (key, write_op) in rows {
                match write_op {
                    WriteOp::Value(value) => db_batch.put_cf(cf_handle, key, value),
                    WriteOp::Deletion => db_batch.delete_cf(cf_handle, key),
                }
                .map_err(convert_rocksdb_err)?;
            }
        }

        self.inner
            .write_opt(&db_batch, &default_write_options())
            .map_err(convert_rocksdb_err)?;

        // Bump counters only after DB write succeeds.
        for (cf_name, rows) in &batch.rows {
            for (key, write_op) in rows {
                match write_op {
                    WriteOp::Value(value) => OP_COUNTER.observe(
                        &format!("db_put_bytes_{}", cf_name),
                        (key.len() + value.len()) as f64,
                    ),
                    WriteOp::Deletion => OP_COUNTER.inc(&format!("db_delete_{}", cf_name)),
                }
            }
        }

        Ok(())
    }

    fn get_cf_handle(&self, cf_name: &str) -> Result<&CFHandle> {
        self.inner.cf_handle(cf_name).ok_or_else(|| {
            format_err!(
                "DB::cf_handle not found for column family name: {}",
                cf_name
            )
        })
    }

    /// Returns the approximate size of each non-empty column family in bytes.
    pub fn get_approximate_sizes_cf(&self) -> Result<BTreeMap<String, u64>> {
        let mut cf_sizes = BTreeMap::new();

        for cf_name in self.inner.cf_names().into_iter().map(ToString::to_string) {
            let cf_handle = self.get_cf_handle(&cf_name)?;
            let size = self
                .inner
                .get_property_int_cf(cf_handle, "rocksdb.estimate-live-data-size")
                .ok_or_else(|| {
                    format_err!(
                        "Unable to get approximate size of {} column family.",
                        cf_name,
                    )
                })?;
            cf_sizes.insert(cf_name, size);
        }

        Ok(cf_sizes)
    }

    /// Flushes all memtable data. If `sync` is true, the flush will wait until it's done. This is
    /// only used for testing `get_approximate_sizes_cf` in unit tests.
    pub fn flush_all(&self, sync: bool) -> Result<()> {
        for cf_name in self.inner.cf_names() {
            let cf_handle = self.get_cf_handle(cf_name)?;
            self.inner
                .flush_cf(cf_handle, sync)
                .map_err(convert_rocksdb_err)?;
        }
        Ok(())
    }
}

/// For now we always use synchronous writes. This makes sure that once the operation returns
/// `Ok(())` the data is persisted even if the machine crashes. In the future we might consider
/// selectively turning this off for some non-critical writes to improve performance.
fn default_write_options() -> WriteOptions {
    let mut opts = WriteOptions::new();
    opts.set_sync(true);
    opts
}
<|endoftext|>


<commit_before>use slack_api as slack;
use tokio_tungstenite as tungstenite;

use futures_util::{SinkExt, StreamExt};

use tokio::runtime::Runtime;
use tokio::sync::mpsc;

// Alternative if need tokio 0.3
// use tokio_compat_02::FutureExt;

use serde::{Deserialize, Serialize};
use serde_json::json;

use atomic_counter::AtomicCounter;
use atomic_counter::RelaxedCounter;

use std::sync::Arc;
use std::default::Default;
use std::env;

// Type alias for msg_id
type MsgId = Arc<RelaxedCounter>;

fn main() {
    Runtime::new()
        .expect("Failed to create Tokio runtime")
        .block_on(async_main());
}

async fn async_main() -> Result<(), Box<dyn std::error::Error>> {
    // Shared integer counter for message ids
    let msg_id = Arc::new(RelaxedCounter::new(0));

    let token = env::var("SLACK_API_TOKEN").map_err(|_| "SLACK_API_TOKEN env var must be set")?;
    let client = slack::default_client().map_err(|e| format!("Could not get default_client, {:?}", e))?;

    {
        // Post a message
        let msg = slack::chat::PostMessageRequest {
            channel: "karmator-devs",
            text: "Pogger, i'm starting up!",
            as_user: Some(true),
            ..Default::default()
        };
        slack::chat::post_message(&client, &token, &msg).await;

        // We use .compat() from FutureExt specifically because the api client is on reqwest which
        // uses tokio 0.2 and we are going to be using tokio 0.3 (for tungstein for websockets)
        //
        // One possibility could be to:
        // You can provide your own client by implementing the async or sync versions of SlackWebRequestSender.
        //
        // Which would let me get off requwest onto something that is on tokio 0.3 (this may have
        // merit)
        //let request = slack::rtm::StartRequest::default();
        //let response = slack::rtm::start(&client, &token, &request).compat().await;
        let response = slack::rtm::connect(&client, &token).await;

        if let Ok(response) = response {
            println!("Control - Got an ok reply");

            if let Some(ws_url) = response.url {
                println!("Control - \tGot a WS url: {:?}", ws_url);

                let ws_response = tungstenite::connect_async(ws_url).await;

                if let Ok((mut ws_stream, ws_response)) = ws_response {
                    println!("Control - \t\tWS connection established");

                    // Setup the mpsc channel
                    let (mut tx, mut rx) = mpsc::channel(32);

                    // Split the stream
                    let (mut ws_write, mut ws_read) = ws_stream.split();

                    // Spawn the inbound ws stream processor
                    let inbound = tokio::spawn(async move {
                        loop {
                            if let Some(Ok(ws_msg)) = ws_read.next().await {
                                let msg_id = msg_id.clone();
                                let tx2 = tx.clone();

                                tokio::spawn(async move {
                                    process_inbound_message(msg_id, ws_msg, tx2).await;
                                });

                            }
                        }
                    });

                    // Spawn outbound ws processor
                    // TODO: need to make sure to wait on sending till we recieve the hello event
                    // could be a oneshot or some sort of flag which it waits till inbound has
                    // gotten the hello
                    let outbound = tokio::spawn(async move {
                        loop {
                            while let Some(message) = rx.recv().await {
                                println!("Outbound - {:?}", message);

                                // TODO: present some way to do plain vs fancy message, and if
                                // fancy do a webapi post, otherwise dump into the WS
                                //
                                // TODO: look into tracking the sent message with a confirmation
                                // that it was sent (via msg id) and if it fails, resend with
                                // backoff
                                //
                                // TODO: find a way to handle the ping/pong cycle and monitor
                                ws_write.send(message).await;
                            }
                        }
                    });

                    // Wait till either exits (error) then begin recovery
                    let res = tokio::try_join!(inbound, outbound);

                    match res {
                        Ok((first, second)) => println!("Control - \t\t\tBoth exited fine"),
                        Err(err) => println!("Control - \t\t\tSomething failed: {:?}", err),
                    }
                } else {
                    println!("Control - \t\t{:?}", ws_response.err());
                }
            } else {
                println!("Control - \tNo WS url");
            }
        } else {
            println!("Control - {:?}", response);
        }
    }
    Ok(())
}

async fn process_inbound_message(
    msg_id: MsgId,
    msg: tungstenite::tungstenite::Message,
    mut tx: mpsc::Sender<tungstenite::tungstenite::Message>
) -> Result<(), Box<dyn std::error::Error>>
{
    let id = msg_id.inc();

    println!("Inbound - id = {:?}", id);

    // Parse incoming message
    let raw_msg = match msg {
        tungstenite::tungstenite::Message::Text(x) => Some(x),
        _ => {
            println!("Inbound - \tUnsupported ws message type - {:?}", msg);
            None
        },
    };
    println!("Inbound - raw event = {:?}", raw_msg);

    if let Some(s) = raw_msg {

        // Try to parse to an event
        let der_msg: std::result::Result<Event, serde_json::Error> = serde_json::from_str(&s);
        match der_msg {
            Ok(event) => {
                println!("Inbound - \tparsed event = {:?}", event);

                // Check if its a message/certain string, if so, reply
                match event {
                    Event::Message {
                        channel: Some(c),
                        text: t,
                        subtype: _,
                        hidden: _,
                        user: _,
                        ts: _,
                    } if (c == "CAF6S4TRT".to_string()) && (t == "!kappa".to_string()) => {

                        // Send out a message
                        let ws_msg = json!({
                            "id": 1,
                            "type": "message",
                            "channel": "CAF6S4TRT",
                            "text": "mod4 kappa",
                        }).to_string();

                        tx.send(tungstenite::tungstenite::Message::from(ws_msg)).await;

                        println!("Inbound - \t\tKappa was sent");
                    },
                    _ => println!("Inbound - \t\tNo action taken"),
                }
            },
            Err(_) => {
                let der_msg: std::result::Result<MessageSent, serde_json::Error> = serde_json::from_str(&s);
                match der_msg {
                    Ok(event) => println!("Inbound - \tparsed event = {:?}", event),
                    Err(_) => {
                        let der_msg: std::result::Result<MessageError, serde_json::Error> = serde_json::from_str(&s);
                        match der_msg {
                            Ok(event) => println!("Inbound - \tparsed event = {:?}", event),
                            Err(_) => {
                                // TODO: for now print to stderr what didn't get deserialized
                                // later can have config option to log to file the error or not
                                println!("Inbound - \t\tFail parse = {:?}", s);
                            },
                        }
                    },
                }
            },
        }
    }
    Ok(())
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
enum Event {
    // TODO: separate the slack Control messages (hello/bye/ping/pong)
    // from the other messages that will be of interest to plugins
    //
    // First message upon successful connection establishment
    Hello,

    // Client should reconnect after getting such message
    Goodbye,

    // Reply to Ping = { type: ping, id: num }
    Pong { reply_to: usize },

    // These events are what plugins will care for:
    //
    // Slack messages
    // TODO: more involved here for now basics
    Message {
        subtype: Option<String>,
        hidden: Option<bool>,
        channel: Option<String>,
        user: String,
        text: String,
        ts: String,
    },

    // TODO: consider looking at slack_api for types to reuse here
    ReactionAdded {
        user: String,
        reaction: String,
        item_user: Option<String>,
        // item: ListResponseItem
        event_ts: String,
    },
    ReactionRemoved {
        user: String,
        reaction: String,
        item_user: Option<String>,
        // item: ListResponseItem
        event_ts: String,
    },
}

#[derive(Debug, Deserialize)]
struct MessageSent {
    ok: bool,
    reply_to: usize,
    text: String,
    ts: String
}

#[derive(Debug, Deserialize)]
struct MessageError {
    ok: bool,
    reply_to: usize,
    error: ErrorDetail,
}

#[derive(Debug, Deserialize)]
struct ErrorDetail {
    code: usize,
    msg: String,
}
<commit_msg>serde breakup, still need to clean up the parser portion more<commit_after>use slack_api as slack;
use tokio_tungstenite as tungstenite;

use futures_util::{SinkExt, StreamExt};

use tokio::runtime::Runtime;
use tokio::sync::mpsc;

// Alternative if need tokio 0.3
// use tokio_compat_02::FutureExt;

use serde::{Deserialize, Serialize};
use serde_json::json;

use atomic_counter::AtomicCounter;
use atomic_counter::RelaxedCounter;

use std::sync::Arc;
use std::default::Default;
use std::env;
use std::result::Result;

// Type alias for msg_id
type MsgId = Arc<RelaxedCounter>;

fn main() {
    Runtime::new()
        .expect("Failed to create Tokio runtime")
        .block_on(async_main());
}

async fn async_main() -> Result<(), Box<dyn std::error::Error>> {
    // Shared integer counter for message ids
    let msg_id = Arc::new(RelaxedCounter::new(0));

    let token = env::var("SLACK_API_TOKEN").map_err(|_| "SLACK_API_TOKEN env var must be set")?;
    let client = slack::default_client().map_err(|e| format!("Could not get default_client, {:?}", e))?;

    {
        // Post a message
        let msg = slack::chat::PostMessageRequest {
            channel: "karmator-devs",
            text: "Pogger, i'm starting up!",
            as_user: Some(true),
            ..Default::default()
        };
        slack::chat::post_message(&client, &token, &msg).await;

        // We use .compat() from FutureExt specifically because the api client is on reqwest which
        // uses tokio 0.2 and we are going to be using tokio 0.3 (for tungstein for websockets)
        //
        // One possibility could be to:
        // You can provide your own client by implementing the async or sync versions of SlackWebRequestSender.
        //
        // Which would let me get off requwest onto something that is on tokio 0.3 (this may have
        // merit)
        //let request = slack::rtm::StartRequest::default();
        //let response = slack::rtm::start(&client, &token, &request).compat().await;
        let response = slack::rtm::connect(&client, &token).await;

        if let Ok(response) = response {
            println!("Control - Got an ok reply");

            if let Some(ws_url) = response.url {
                println!("Control - \tGot a WS url: {:?}", ws_url);

                let ws_response = tungstenite::connect_async(ws_url).await;

                if let Ok((mut ws_stream, ws_response)) = ws_response {
                    println!("Control - \t\tWS connection established");

                    // Setup the mpsc channel
                    let (mut tx, mut rx) = mpsc::channel(32);

                    // Split the stream
                    let (mut ws_write, mut ws_read) = ws_stream.split();

                    // Spawn the inbound ws stream processor
                    let inbound = tokio::spawn(async move {
                        loop {
                            if let Some(Ok(ws_msg)) = ws_read.next().await {
                                let msg_id = msg_id.clone();
                                let tx2 = tx.clone();

                                tokio::spawn(async move {
                                    process_inbound_message(msg_id, ws_msg, tx2).await;
                                });

                            }
                        }
                    });

                    // Spawn outbound ws processor
                    // TODO: need to make sure to wait on sending till we recieve the hello event
                    // could be a oneshot or some sort of flag which it waits till inbound has
                    // gotten the hello
                    let outbound = tokio::spawn(async move {
                        loop {
                            while let Some(message) = rx.recv().await {
                                println!("Outbound - {:?}", message);

                                // TODO: present some way to do plain vs fancy message, and if
                                // fancy do a webapi post, otherwise dump into the WS
                                //
                                // TODO: look into tracking the sent message with a confirmation
                                // that it was sent (via msg id) and if it fails, resend with
                                // backoff
                                //
                                // TODO: find a way to handle the ping/pong cycle and monitor
                                ws_write.send(message).await;
                            }
                        }
                    });

                    // Wait till either exits (error) then begin recovery
                    let res = tokio::try_join!(inbound, outbound);

                    match res {
                        Ok((first, second)) => println!("Control - \t\t\tBoth exited fine"),
                        Err(err) => println!("Control - \t\t\tSomething failed: {:?}", err),
                    }
                } else {
                    println!("Control - \t\t{:?}", ws_response.err());
                }
            } else {
                println!("Control - \tNo WS url");
            }
        } else {
            println!("Control - {:?}", response);
        }
    }
    Ok(())
}

async fn process_inbound_message(
    msg_id: MsgId,
    msg: tungstenite::tungstenite::Message,
    mut tx: mpsc::Sender<tungstenite::tungstenite::Message>
) -> Result<(), Box<dyn std::error::Error>>
{
    let id = msg_id.inc();

    println!("Inbound - id = {:?}", id);

    // Parse incoming message
    let raw_msg = match msg {
        tungstenite::tungstenite::Message::Text(x) => Some(x),
        _ => {
            println!("Inbound - \tUnsupported ws message type - {:?}", msg);
            None
        },
    };
    println!("Inbound - raw event = {:?}", raw_msg);

    if let Some(s) = raw_msg {
        if let Some(e) = parse_event(s) {
            match e {
                Event::UserEvent(event) => {
                    // Check if its a message/certain string, if so, reply
                    match event {
                        UserEvent::Message {
                            channel: Some(c),
                            text: t,
                            subtype: _,
                            hidden: _,
                            user: _,
                            ts: _,
                        } if (c == "CAF6S4TRT".to_string()) && (t == "!kappa".to_string()) => {
                            // Send out a message
                            let ws_msg = json!({
                                "id": 1,
                                "type": "message",
                                "channel": "CAF6S4TRT",
                                "text": "mod4 kappa",
                            }).to_string();

                            tx.send(tungstenite::tungstenite::Message::from(ws_msg)).await;

                            println!("Inbound - \t\tKappa was sent");
                        },
                        _ => println!("Inbound - \t\tNo action taken"),
                    }
                },

                Event::SystemControl(sc) => (),
                Event::MessageControl(mc) => (),
            }
        }
    }
    Ok(())
}

fn parse_event(s: String) -> Option<Event> {
    // Try to parse an System Control or Message Control, then an Event (to pass on downstream)
    let sc: Result<SystemControl, serde_json::Error> = serde_json::from_str(&s);
    match sc {
        Ok(sc) => {
            println!("Inbound - \tSystem Control = {:?}", sc);
            Some(Event::SystemControl(sc))
        },
        Err(_) => {
            // Try to parse a Message Control
            let mc: Result<MessageControl, serde_json::Error> = serde_json::from_str(&s);
            match mc {
                Ok(mc) => {
                    println!("Inbound - \tMessage Control = {:?}", mc);
                    Some(Event::MessageControl(mc))
                },
                Err(_) => {
                    // Try to parse an event
                    let event: Result<UserEvent, serde_json::Error> = serde_json::from_str(&s);
                    match event {
                        Ok(event) => {
                            println!("Inbound - \tEvent = {:?}", event);
                            Some(Event::UserEvent(event))
                        },
                        Err(_) => {
                            // TODO: for now print to stderr what didn't get deserialized
                            // later can have config option to log to file the error or not
                            println!("Inbound - \t\tFail parse = {:?}", s);
                            None
                        },
                    }
                },
            }
        },
    }
}

// Containing enum to ease the result chains
#[derive(Debug)]
enum Event {
    UserEvent(UserEvent),
    SystemControl(SystemControl),
    MessageControl(MessageControl),
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
enum UserEvent {
    // These events are what plugins will care for:
    //
    // Slack messages
    // TODO: more involved here for now basics
    Message {
        subtype: Option<String>,
        hidden: Option<bool>,
        channel: Option<String>,
        user: String,
        text: String,
        ts: String,
    },

    // TODO: consider looking at slack_api for types to reuse here
    ReactionAdded {
        user: String,
        reaction: String,
        item_user: Option<String>,
        // item: ListResponseItem
        event_ts: String,
    },
    ReactionRemoved {
        user: String,
        reaction: String,
        item_user: Option<String>,
        // item: ListResponseItem
        event_ts: String,
    },
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
enum SystemControl {
    // TODO: separate the slack Control messages (hello/bye/ping/pong)
    // from the other messages that will be of interest to plugins
    //
    // First message upon successful connection establishment
    Hello,

    // Client should reconnect after getting such message
    Goodbye,

    // Reply to Ping = { type: ping, id: num }
    Pong { reply_to: usize },
}

#[derive(Debug, Deserialize)]
#[serde(untagged)]
enum MessageControl {
    MessageSent {
        ok: bool,
        reply_to: usize,
        text: String,
        ts: String
    },

    MessageError {
        ok: bool,
        reply_to: usize,
        error: ErrorDetail,
    },
}

#[derive(Debug, Deserialize)]
struct ErrorDetail {
    code: usize,
    msg: String,
}
<|endoftext|>


<commit_before>use bindings::*;
use compile::Compile;
use function::ABI;
use libc::c_uint;
use std::kinds::marker::{
    ContravariantLifetime,
    Managed
};
use std::str::raw::from_c_str;
use std::c_str::ToCStr;
use util::NativeRef;
/// The types that a value can be
bitflags!(
    flags TypeKind: i32 {
        static Invalid      = -1,
        static Void         = 0,
        static SByte        = 1,
        static UByte        = 2,
        static Short        = 3,
        static UShort       = 4,
        static Int          = 5,
        static UInt         = 6,
        static NInt         = 7,
        static NUInt        = 8,
        static Long         = 9,
        static ULong        = 10,
        static Float32      = 11,
        static Float64      = 12,
        static NFloat       = 13,
        static MaxPrimitive = 13,
        static Struct       = 14,
        static Union        = 15,
        static Signature    = 16,
        static Pointer      = 17,
        static FirstTagged  = 32,
        static SysBool      = 10009,
        static SysChar      = 10010
    }
)
/// A single field of a struct
pub struct Field<'a> {
    /// The index of the field
    pub index: c_uint,
    _type: jit_type_t,
    marker: ContravariantLifetime<'a>
}
impl<'a> Field<'a> {
    #[inline]
    /// Get the field's name or none if it lacks one
    pub fn get_name(&self) -> Option<String> {
        unsafe {
            let c_name = jit_type_get_name(self._type, self.index);
            if c_name.is_null() {
                None
            } else {
                Some(from_c_str(c_name))
            }
        }
    }
    #[inline]
    /// Get the type of the field
    pub fn get_type(&self) -> Type {
        unsafe {
            NativeRef::from_ptr(jit_type_get_field(self._type, self.index))
        }
    }
    #[inline]
    /// Get the offset of the field
    pub fn get_offset(&self) -> uint {
        unsafe {
            jit_type_get_offset(self._type, self.index) as uint
        }
    }
}
/// A type field iterator
pub struct Fields<'a> {
    _type: jit_type_t,
    index: c_uint,
    length: c_uint,
    marker: ContravariantLifetime<'a>
}
impl<'a> Fields<'a> {
    #[inline(always)]
    fn new(ty:Type) -> Fields<'a> {
        unsafe {
            Fields {
                _type: ty.as_ptr(),
                index: 0 as c_uint,
                length: jit_type_num_fields(ty.as_ptr()),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Field<'a>> for Fields<'a> {
    fn next(&mut self) -> Option<Field<'a>> {
        let index = self.index;
        self.index += 1;
        if index < self.length {
            Some(Field {
                _type: self._type.clone(),
                index: index,
                marker: ContravariantLifetime::<'a>
            })
        } else {
            None
        }
    }
    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        ((self.length - self.index) as uint, None)
    }
    #[inline]
    fn count(&mut self) -> uint {
        let count = self.length - self.index;
        self.index = self.length;
        count as uint
    }
    #[inline]
    fn nth(&mut self, n:uint) -> Option<Field<'a>> {
        self.index += n as u32;
        self.next()
    }
}
pub struct Params<'a> {
    _type: jit_type_t,
    index: c_uint,
    length: c_uint,
    marker: ContravariantLifetime<'a>
}
impl<'a> Params<'a> {
    #[inline(always)]
    fn new(ty:Type) -> Params<'a> {
        unsafe {
            Params {
                _type: ty.as_ptr(),
                index: 0 as c_uint,
                length: jit_type_num_params(ty.as_ptr()),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Type> for Params<'a> {
    fn next(&mut self) -> Option<Type> {
        let index = self.index;
        self.index += 1;
        if index < self.length {
            Some(unsafe {
                NativeRef::from_ptr(jit_type_get_param(self._type, index))
            })
        } else {
            None
        }
    }
    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        ((self.length - self.index) as uint, None)
    }
    #[inline]
    fn count(&mut self) -> uint {
        let count = self.length - self.index;
        self.index = self.length;
        count as uint
    }
    #[inline]
    fn nth(&mut self, n:uint) -> Option<Type> {
        self.index += n as u32;
        self.next()
    }
}
/**
 * An object that represents a native system type.
 *
 * Each `Type` represents a basic system type,
 * be it a primitive, a struct, a union, a pointer, or a function signature.
 * The library uses this information to lay out values in memory.
*/
#[deriving(PartialEq)]
pub struct Type {
    _type: jit_type_t,
    marker: Managed
}
impl NativeRef for Type {
    #[inline]
    unsafe fn as_ptr(&self) -> jit_type_t {
        self._type
    }
    #[inline]
    unsafe fn from_ptr(ptr:jit_type_t) -> Type {
        Type {
            _type: ptr,
            marker: Managed
        }
    }
}
impl Clone for Type {
    #[inline]
    /// Make a copy of the type descriptor by increasing its reference count.
    fn clone(&self) -> Type {
        unsafe {
            NativeRef::from_ptr(jit_type_copy(self.as_ptr()))
        }
    }
}
#[unsafe_destructor]
impl Drop for Type {
    #[inline]
    /**
     * Free a type descriptor by decreasing its reference count.
     *
     * This function is save to use on pre-defined types, which
     * are never actually freed.
     */
    fn drop(&mut self) {
        unsafe {
            jit_type_free(self.as_ptr());
        }
    }
}
impl Type {
    /// Create a type descriptor for a function signature.
    pub fn create_signature(abi: ABI, return_type: Type, params: &mut [Type]) -> Type {
        unsafe {
            let mut native_params:Vec<jit_type_t> = params.iter().map(|param| param.as_ptr()).collect();
            let signature = jit_type_create_signature(abi as jit_abi_t, return_type.as_ptr(), native_params.as_mut_ptr(), params.len() as c_uint, 1);
            NativeRef::from_ptr(signature)
        }
    }

    fn create_complex(fields: &mut [Type], union: bool) -> Type {
        unsafe {
            let mut native_fields:Vec<jit_type_t> = fields.iter().map(|field| field.as_ptr()).collect();
            let f = if union { jit_type_create_union } else { jit_type_create_struct };
            let ty:jit_type_t = f(native_fields.as_mut_ptr(), fields.len() as c_uint, 1);
            NativeRef::from_ptr(ty)
        }
    }
    #[inline(always)]
    /// Create a type descriptor for a structure.
    pub fn create_struct(fields: &mut [Type]) -> Type {
        Type::create_complex(fields, false)
    }
    #[inline(always)]
    /// Create a type descriptor for a union.
    pub fn create_union(fields: &mut [Type]) -> Type {
        let inner = Type::create_complex(fields, true);
        Type::create_struct(&mut [get::<int>(), inner])
    }
    #[inline(always)]
    /// Create a type descriptor for a pointer to another type.
    pub fn create_pointer(pointee: Type) -> Type {
        unsafe {
            let ptr = jit_type_create_pointer(pointee.as_ptr(), 1);
            NativeRef::from_ptr(ptr)
        }
    }
    #[inline]
    /// Get the size of this type in bytes.
    pub fn get_size(&self) -> jit_nuint {
        unsafe {
            jit_type_get_size(self.as_ptr())
        }
    }
    /**
     * Get a value that indicates the kind of this type. This allows
     * callers to quickly classify a type to determine how it should
     * be handled further.
     */
    #[inline]
    pub fn get_kind(&self) -> TypeKind {
        unsafe {
            TypeKind::from_bits(jit_type_get_kind(self.as_ptr())).unwrap()
        }
    }
    #[inline]
    /// Get the type that is referred to by this pointer type.
    pub fn get_ref(&self) -> Type {
        unsafe {
            NativeRef::from_ptr(jit_type_get_ref(self.as_ptr()))
        }
    }
    /// Set the field or parameter names of this type.
    pub fn set_names<T:ToCStr>(&self, names:&[T]) -> bool {
        unsafe {
            let native_names : Vec<*i8> = names.iter().map(|name| name.to_c_str().unwrap()).collect();
            jit_type_set_names(self.as_ptr(), native_names.as_ptr() as *mut *mut i8, names.len() as u32) != 0
        }
    }
    #[inline]
    /// Iterator over the type's fields
    pub fn iter_fields<'a>(&'a self) -> Fields<'a> {
        Fields::new(self.clone())
    }
    #[inline]
    /// Iterator over the function signature's parameters
    pub fn iter_params<'a>(&'a self) -> Params<'a> {
        Params::new(self.clone())
    }
    #[inline]
    /// Find the field/parameter index for a particular name.
    pub fn find_name<T:ToCStr>(&self, name:T) -> uint {
        name.with_c_str(|c_name| unsafe {
            jit_type_find_name(self.as_ptr(), c_name) as uint
        })
    }
}
#[inline]
/// Get the Rust type given as a type descriptor
pub fn get<T:Compile>() -> Type {
    Compile::jit_type(None::<T>)
}<commit_msg>Move `create_struct` and `create_union` implementation to their respective functions<commit_after>use bindings::*;
use compile::Compile;
use function::ABI;
use libc::c_uint;
use std::kinds::marker::{
    ContravariantLifetime,
    Managed
};
use std::str::raw::from_c_str;
use std::c_str::ToCStr;
use util::NativeRef;
/// The types that a value can be
bitflags!(
    flags TypeKind: i32 {
        static Invalid      = -1,
        static Void         = 0,
        static SByte        = 1,
        static UByte        = 2,
        static Short        = 3,
        static UShort       = 4,
        static Int          = 5,
        static UInt         = 6,
        static NInt         = 7,
        static NUInt        = 8,
        static Long         = 9,
        static ULong        = 10,
        static Float32      = 11,
        static Float64      = 12,
        static NFloat       = 13,
        static MaxPrimitive = 13,
        static Struct       = 14,
        static Union        = 15,
        static Signature    = 16,
        static Pointer      = 17,
        static FirstTagged  = 32,
        static SysBool      = 10009,
        static SysChar      = 10010
    }
)
/// A single field of a struct
pub struct Field<'a> {
    /// The index of the field
    pub index: c_uint,
    _type: jit_type_t,
    marker: ContravariantLifetime<'a>
}
impl<'a> Field<'a> {
    #[inline]
    /// Get the field's name or none if it lacks one
    pub fn get_name(&self) -> Option<String> {
        unsafe {
            let c_name = jit_type_get_name(self._type, self.index);
            if c_name.is_null() {
                None
            } else {
                Some(from_c_str(c_name))
            }
        }
    }
    #[inline]
    /// Get the type of the field
    pub fn get_type(&self) -> Type {
        unsafe {
            NativeRef::from_ptr(jit_type_get_field(self._type, self.index))
        }
    }
    #[inline]
    /// Get the offset of the field
    pub fn get_offset(&self) -> uint {
        unsafe {
            jit_type_get_offset(self._type, self.index) as uint
        }
    }
}
/// A type field iterator
pub struct Fields<'a> {
    _type: jit_type_t,
    index: c_uint,
    length: c_uint,
    marker: ContravariantLifetime<'a>
}
impl<'a> Fields<'a> {
    #[inline(always)]
    fn new(ty:Type) -> Fields<'a> {
        unsafe {
            Fields {
                _type: ty.as_ptr(),
                index: 0 as c_uint,
                length: jit_type_num_fields(ty.as_ptr()),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Field<'a>> for Fields<'a> {
    fn next(&mut self) -> Option<Field<'a>> {
        let index = self.index;
        self.index += 1;
        if index < self.length {
            Some(Field {
                _type: self._type.clone(),
                index: index,
                marker: ContravariantLifetime::<'a>
            })
        } else {
            None
        }
    }
    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        ((self.length - self.index) as uint, None)
    }
    #[inline]
    fn count(&mut self) -> uint {
        let count = self.length - self.index;
        self.index = self.length;
        count as uint
    }
    #[inline]
    fn nth(&mut self, n:uint) -> Option<Field<'a>> {
        self.index += n as u32;
        self.next()
    }
}
pub struct Params<'a> {
    _type: jit_type_t,
    index: c_uint,
    length: c_uint,
    marker: ContravariantLifetime<'a>
}
impl<'a> Params<'a> {
    #[inline(always)]
    fn new(ty:Type) -> Params<'a> {
        unsafe {
            Params {
                _type: ty.as_ptr(),
                index: 0 as c_uint,
                length: jit_type_num_params(ty.as_ptr()),
                marker: ContravariantLifetime::<'a>
            }
        }
    }
}
impl<'a> Iterator<Type> for Params<'a> {
    fn next(&mut self) -> Option<Type> {
        let index = self.index;
        self.index += 1;
        if index < self.length {
            Some(unsafe {
                NativeRef::from_ptr(jit_type_get_param(self._type, index))
            })
        } else {
            None
        }
    }
    #[inline]
    fn size_hint(&self) -> (uint, Option<uint>) {
        ((self.length - self.index) as uint, None)
    }
    #[inline]
    fn count(&mut self) -> uint {
        let count = self.length - self.index;
        self.index = self.length;
        count as uint
    }
    #[inline]
    fn nth(&mut self, n:uint) -> Option<Type> {
        self.index += n as u32;
        self.next()
    }
}
/**
 * An object that represents a native system type.
 *
 * Each `Type` represents a basic system type,
 * be it a primitive, a struct, a union, a pointer, or a function signature.
 * The library uses this information to lay out values in memory.
*/
#[deriving(PartialEq)]
pub struct Type {
    _type: jit_type_t,
    marker: Managed
}
impl NativeRef for Type {
    #[inline]
    unsafe fn as_ptr(&self) -> jit_type_t {
        self._type
    }
    #[inline]
    unsafe fn from_ptr(ptr:jit_type_t) -> Type {
        Type {
            _type: ptr,
            marker: Managed
        }
    }
}
impl Clone for Type {
    #[inline]
    /// Make a copy of the type descriptor by increasing its reference count.
    fn clone(&self) -> Type {
        unsafe {
            NativeRef::from_ptr(jit_type_copy(self.as_ptr()))
        }
    }
}
#[unsafe_destructor]
impl Drop for Type {
    #[inline]
    /**
     * Free a type descriptor by decreasing its reference count.
     *
     * This function is save to use on pre-defined types, which
     * are never actually freed.
     */
    fn drop(&mut self) {
        unsafe {
            jit_type_free(self.as_ptr());
        }
    }
}
impl Type {
    /// Create a type descriptor for a function signature.
    pub fn create_signature(abi: ABI, return_type: Type, params: &mut [Type]) -> Type {
        unsafe {
            let mut native_params:Vec<jit_type_t> = params.iter().map(|param| param.as_ptr()).collect();
            let signature = jit_type_create_signature(abi as jit_abi_t, return_type.as_ptr(), native_params.as_mut_ptr(), params.len() as c_uint, 1);
            NativeRef::from_ptr(signature)
        }
    }
    #[inline(always)]
    /// Create a type descriptor for a structure.
    pub fn create_struct(fields: &mut [Type]) -> Type {
        unsafe {
            let mut native_fields:Vec<_> = fields.iter().map(|field| field.as_ptr()).collect();
            NativeRef::from_ptr(jit_type_create_struct(native_fields.as_mut_ptr(), fields.len() as c_uint, 1))
        }
    }
    #[inline(always)]
    /// Create a type descriptor for a union.
    pub fn create_union(fields: &mut [Type]) -> Type {
        unsafe {
            let mut native_fields:Vec<_> = fields.iter().map(|field| field.as_ptr()).collect();
            NativeRef::from_ptr(jit_type_create_union(native_fields.as_mut_ptr(), fields.len() as c_uint, 1))
        }
    }
    #[inline(always)]
    /// Create a type descriptor for a pointer to another type.
    pub fn create_pointer(pointee: Type) -> Type {
        unsafe {
            let ptr = jit_type_create_pointer(pointee.as_ptr(), 1);
            NativeRef::from_ptr(ptr)
        }
    }
    #[inline]
    /// Get the size of this type in bytes.
    pub fn get_size(&self) -> jit_nuint {
        unsafe {
            jit_type_get_size(self.as_ptr())
        }
    }
    /**
     * Get a value that indicates the kind of this type. This allows
     * callers to quickly classify a type to determine how it should
     * be handled further.
     */
    #[inline]
    pub fn get_kind(&self) -> TypeKind {
        unsafe {
            TypeKind::from_bits(jit_type_get_kind(self.as_ptr())).unwrap()
        }
    }
    #[inline]
    /// Get the type that is referred to by this pointer type.
    pub fn get_ref(&self) -> Type {
        unsafe {
            NativeRef::from_ptr(jit_type_get_ref(self.as_ptr()))
        }
    }
    /// Set the field or parameter names of this type.
    pub fn set_names<T:ToCStr>(&self, names:&[T]) -> bool {
        unsafe {
            let native_names : Vec<*i8> = names.iter().map(|name| name.to_c_str().unwrap()).collect();
            jit_type_set_names(self.as_ptr(), native_names.as_ptr() as *mut *mut i8, names.len() as u32) != 0
        }
    }
    #[inline]
    /// Iterator over the type's fields
    pub fn iter_fields<'a>(&'a self) -> Fields<'a> {
        Fields::new(self.clone())
    }
    #[inline]
    /// Iterator over the function signature's parameters
    pub fn iter_params<'a>(&'a self) -> Params<'a> {
        Params::new(self.clone())
    }
    #[inline]
    /// Find the field/parameter index for a particular name.
    pub fn find_name<T:ToCStr>(&self, name:T) -> uint {
        name.with_c_str(|c_name| unsafe {
            jit_type_find_name(self.as_ptr(), c_name) as uint
        })
    }
}
#[inline]
/// Get the Rust type given as a type descriptor
pub fn get<T:Compile>() -> Type {
    Compile::jit_type(None::<T>)
}<|endoftext|>


<commit_before>//! Routing entity for handling endpoints
//!
//! This module is intended for matching and converting a passed URLs by a client
//! in request into certain queue/topic names.
//!

use std::clone::Clone;
use std::collections::HashMap;

use super::endpoint::ReadOnlyEndpoint;
use error::{PathfinderError, Result};

/// A struct which is stores a mapping of resources that can be
/// represented in requests and transformed into certain API endpoints of
/// microservices.
///
/// # Examples
///
/// Attempt to get an endpoint with correct URL will return all expected data:
///
/// ```
/// use config::{get_config};
/// use engine::router::{Router, Endpoint, extract_endpoints};
///
/// let url = "/api/matchmaking/search";
/// let config = get_config(&"../../../tests/files/config_with_valid_endpoints.yaml");
/// let endpoints = extract_endpoints(config);
/// let router = Box::new(Router::new(endpoints))
///
/// let endpoint = route.match_url(url).unwrap();
/// assert_eq!(endpoint.get_url(), "/api/matchmaking/search");
/// assert_eq!(endpoint.get_microservice(), "microservice.search");
/// ```
///
/// For not matched URL will be returned an error:
///
/// ```
/// use config::{get_config};
/// use engine::router::{Router, Endpoint, extract_endpoints};
///
/// let url = "/api/matchmaking/search";
/// let config = get_config(&"../../../tests/files/config_with_invalid_endpoints.yaml");
/// let endpoints = extract_endpoints(config);
/// let router = Box::new(Router::new(endpoints))
///
/// assert_eq!(route.match_url(url).is_err(), true);
/// ```
///
pub struct Router {
    endpoints: HashMap<String, ReadOnlyEndpoint>
}

impl Router {
    /// Returns a new instance of `Router` that contains a mapping for resources.
    pub fn new(endpoints: HashMap<String, ReadOnlyEndpoint>) -> Router {
        Router {
            endpoints: endpoints
        }
    }

    /// Returns an endpoint that was found for a passed URL.
    pub fn match_url(&self, url: &str) -> Result<ReadOnlyEndpoint> {
        match self.endpoints.contains_key(url) {
            true => {
                let endpoint = self.endpoints[url].clone();
                Ok(endpoint)
            }
            false => Err(PathfinderError::EndpointNotFound(url.to_string()))
        }
    }
}

#[cfg(test)]
mod tests {
    use config::get_config;
    use engine::router::{extract_endpoints, Router, REQUEST_EXCHANGE, RESPONSE_EXCHANGE};

    fn get_router(file_path: &str) -> Box<Router> {
        let config = get_config(file_path);
        let endpoints = extract_endpoints(config);
        Box::new(Router::new(endpoints))
    }

    #[test]
    fn test_router_match_url_returns_an_endpoint_for_a_matched_url() {
        let router = get_router(&"./tests/files/config_with_valid_endpoints.yaml");
        let result_match = router.match_url(&"/api/matchmaking/search");

        assert_eq!(result_match.is_ok(), true);
        let endpoint = result_match.unwrap();
        assert_eq!(endpoint.get_url(), "/api/matchmaking/search");
        assert_eq!(endpoint.get_microservice(), "microservice.search");
    }

    #[test]
    fn test_routers_match_url_returns_an_error_for_an_unknown_url() {
        let router = get_router(&"./tests/files/config_with_invalid_endpoints.yaml");
        let result_match = router.match_url(&"/api/matchmaking/search");

        assert_eq!(result_match.is_err(), true);
    }

    #[test]
    fn test_router_returns_a_custom_endpont_for_a_match() {
        let router = get_router(&"./tests/files/config_with_custom_url_mapping.yaml");
        let endpoint = router.match_url_or_default(&"/api/matchmaking/leaderboard");

        assert_eq!(endpoint.get_url(), "/api/matchmaking/leaderboard");
        assert_eq!(endpoint.get_microservice(), "microservice.leaderboard");
        assert_eq!(endpoint.get_request_exchange(), "amqp.direct");
        assert_eq!(endpoint.get_response_exchange(), "open-matchmaking.default");
    }
}
<commit_msg>Removed excessive tests and imports for the router.rs module<commit_after>//! Routing entity for handling endpoints
//!
//! This module is intended for matching and converting a passed URLs by a client
//! in request into certain queue/topic names.
//!

use std::clone::Clone;
use std::collections::HashMap;

use super::endpoint::ReadOnlyEndpoint;
use error::{PathfinderError, Result};

/// A struct which is stores a mapping of resources that can be
/// represented in requests and transformed into certain API endpoints of
/// microservices.
///
/// # Examples
///
/// Attempt to get an endpoint with correct URL will return all expected data:
///
/// ```
/// use config::{get_config};
/// use engine::router::{Router, Endpoint, extract_endpoints};
///
/// let url = "/api/matchmaking/search";
/// let config = get_config(&"../../../tests/files/config_with_valid_endpoints.yaml");
/// let endpoints = extract_endpoints(config);
/// let router = Box::new(Router::new(endpoints))
///
/// let endpoint = route.match_url(url).unwrap();
/// assert_eq!(endpoint.get_url(), "/api/matchmaking/search");
/// assert_eq!(endpoint.get_microservice(), "microservice.search");
/// ```
///
/// For not matched URL will be returned an error:
///
/// ```
/// use config::{get_config};
/// use engine::router::{Router, Endpoint, extract_endpoints};
///
/// let url = "/api/matchmaking/search";
/// let config = get_config(&"../../../tests/files/config_with_invalid_endpoints.yaml");
/// let endpoints = extract_endpoints(config);
/// let router = Box::new(Router::new(endpoints))
///
/// assert_eq!(route.match_url(url).is_err(), true);
/// ```
///
pub struct Router {
    endpoints: HashMap<String, ReadOnlyEndpoint>
}

impl Router {
    /// Returns a new instance of `Router` that contains a mapping for resources.
    pub fn new(endpoints: HashMap<String, ReadOnlyEndpoint>) -> Router {
        Router {
            endpoints: endpoints
        }
    }

    /// Returns an endpoint that was found for a passed URL.
    pub fn match_url(&self, url: &str) -> Result<ReadOnlyEndpoint> {
        match self.endpoints.contains_key(url) {
            true => {
                let endpoint = self.endpoints[url].clone();
                Ok(endpoint)
            }
            false => Err(PathfinderError::EndpointNotFound(url.to_string()))
        }
    }
}

#[cfg(test)]
mod tests {
    use config::get_config;
    use engine::router::{extract_endpoints, Router};

    fn get_router(file_path: &str) -> Box<Router> {
        let config = get_config(file_path);
        let endpoints = extract_endpoints(config);
        Box::new(Router::new(endpoints))
    }

    #[test]
    fn test_router_match_url_returns_an_endpoint_for_a_matched_url() {
        let router = get_router(&"./tests/files/config_with_valid_endpoints.yaml");
        let result_match = router.match_url(&"/api/matchmaking/search");

        assert_eq!(result_match.is_ok(), true);
        let endpoint = result_match.unwrap();
        assert_eq!(endpoint.get_url(), "/api/matchmaking/search");
        assert_eq!(endpoint.get_microservice(), "microservice.search");
    }

    #[test]
    fn test_routers_match_url_returns_an_error_for_an_unknown_url() {
        let router = get_router(&"./tests/files/config_with_invalid_endpoints.yaml");
        let result_match = router.match_url(&"/api/matchmaking/search");

        assert_eq!(result_match.is_err(), true);
    }
}
<|endoftext|>


<commit_before>use std::io;

#[cfg(test)] mod tests;
mod read_ext;
mod progress;

pub use progress::{Progress, ProgressFn};

use read_ext::ReadExt;

const SOH: u8 = 0x01;
const EOT: u8 = 0x04;
const ACK: u8 = 0x06;
const NAK: u8 = 0x15;
const CAN: u8 = 0x18;

/// Implementation of the XMODEM protocol.
pub struct Xmodem<R> {
    packet: u8,
    inner: R,
    started: bool,
    progress: ProgressFn
}

impl Xmodem<()> {
    /// Transmits `data` to the receiver `to` using the XMODEM protocol. If the
    /// length of the total data yielded by `data` is not a multiple of 128
    /// bytes, the data is padded with zeroes and sent to the receiver.
    ///
    /// Returns the number of bytes written to `to`, excluding padding zeroes.
    #[inline]
    pub fn transmit<R, W>(data: R, to: W) -> io::Result<usize>
        where W: io::Read + io::Write, R: io::Read
    {
        Self::transmit_with_progress(data, to, progress::noop)
    }

    /// Transmits `data` to the receiver `to` using the XMODEM protocol. If the
    /// length of the total data yielded by `data` is not a multiple of 128
    /// bytes, the data is padded with zeroes and sent to the receiver.
    ///
    /// The function `f` is used as a callback to indicate progress throughout
    /// the transmission. See the [`Progress`] enum for more information.
    ///
    /// Returns the number of bytes written to `to`, excluding padding zeroes.
    pub fn transmit_with_progress<R, W>(mut data: R, to: W, f: ProgressFn) -> io::Result<usize>
        where W: io::Read + io::Write, R: io::Read
    {
        let mut transmitter = Xmodem::new_with_progress(to, f);
        let mut packet = [0u8; 128];
        let mut written = 0;
        'next_packet: loop {
            let n = data.read_max(&mut packet)?;
            packet[n..].iter_mut().for_each(|b| *b = 0);

            if n == 0 {
                transmitter.write_packet(&[])?;
                return Ok(written);
            }

            for _ in 0..10 {
                match transmitter.write_packet(&packet) {
                    Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                    Err(e) => return Err(e),
                    Ok(_) => {
                        written += n;
                        continue 'next_packet;
                    }
                }
            }

            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "bad transmit"));
        }
    }

    /// Receives `data` from `from` using the XMODEM protocol and writes it into
    /// `into`. Returns the number of bytes read from `from`, a multiple of 128.
    #[inline]
    pub fn receive<R, W>(from: R, into: W) -> io::Result<usize>
       where R: io::Read + io::Write, W: io::Write
    {
        Self::receive_with_progress(from, into, progress::noop)
    }

    /// Receives `data` from `from` using the XMODEM protocol and writes it into
    /// `into`. Returns the number of bytes read from `from`, a multiple of 128.
    ///
    /// The function `f` is used as a callback to indicate progress throughout
    /// the reception. See the [`Progress`] enum for more information.
    pub fn receive_with_progress<R, W>(from: R, mut into: W, f: ProgressFn) -> io::Result<usize>
       where R: io::Read + io::Write, W: io::Write
    {
        let mut receiver = Xmodem::new_with_progress(from, f);
        let mut packet = [0u8; 128];
        let mut received = 0;
        'next_packet: loop {
            for _ in 0..10 {
                match receiver.read_packet(&mut packet) {
                    Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                    Err(e) => return Err(e),
                    Ok(0) => break 'next_packet,
                    Ok(n) => {
                        received += n;
                        into.write_all(&packet)?;
                        continue 'next_packet;
                    }
                }
            }

            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "bad receive"));
        }

        Ok(received)
    }
}

impl<T: io::Read + io::Write> Xmodem<T> {
    /// Returns a new `Xmodem` instance with the internal reader/writer set to
    /// `inner`. The returned instance can be used for both receiving
    /// (downloading) and sending (uploading).
    pub fn new(inner: T) -> Self {
        Self { packet: 1, started: false, inner, progress: progress::noop}
    }

    /// Returns a new `Xmodem` instance with the internal reader/writer set to
    /// `inner`. The returned instance can be used for both receiving
    /// (downloading) and sending (uploading). The function `f` is used as a
    /// callback to indicate progress throughout the transfer. See the
    /// [`Progress`] enum for more information.
    pub fn new_with_progress(inner: T, f: ProgressFn) -> Self {
        Self { packet: 1, started: false, inner, progress: f }
    }

    /// Reads a single byte from the inner I/O stream. If `abort_on_can` is
    /// `true`, an error of `ConnectionAborted` is returned if the read byte is
    /// `CAN`.
    ///
    /// # Errors
    ///
    /// Returns an error if reading from the inner stream fails or if
    /// `abort_on_can` is `true` and the read byte is `CAN`.
    fn read_byte(&mut self, abort_on_can: bool) -> io::Result<u8> {
        let mut buf = [0u8; 1];
        self.inner.read_exact(&mut buf)?;

        let byte = buf[0];
        if abort_on_can && byte == CAN {
            return Err(io::Error::new(io::ErrorKind::ConnectionAborted, "received CAN"));
        }

        Ok(byte)
    }

    /// Writes a single byte to the inner I/O stream.
    ///
    /// # Errors
    ///
    /// Returns an error if writing to the inner stream fails.
    fn write_byte(&mut self, byte: u8) -> io::Result<()> {
        self.inner.write_all(&[byte])
    }

    /// Reads a single byte from the inner I/O stream and compares it to `byte`.
    /// If the bytes match, the byte is returned as an `Ok`. If they differ and
    /// the read byte is not `CAN`, an error of `InvalidData` with the message
    /// `expected` is returned. If they differ and the read byte is `CAN`, an
    /// error of `ConnectionAborted` is returned. In either case, if they bytes
    /// differ, a `CAN` byte is written out to the inner stream.
    ///
    /// # Errors
    ///
    /// Returns an error if reading from the inner stream fails, if the read
    /// byte was not `byte`, if the read byte was `CAN` and `byte` is not `CAN`,
    /// or if writing the `CAN` byte failed on byte mismatch.
    fn expect_byte_or_cancel(&mut self, byte: u8, msg: &'static str) -> io::Result<u8> {
        let b = self.read_byte(false)?;
        if b == byte {
            Ok(byte)
        } else {
            self.write_byte(CAN)?;
            if b == CAN {
                Err(io::Error::new(io::ErrorKind::ConnectionAborted, msg))
            } else {
                Err(io::Error::new(io::ErrorKind::InvalidData, msg))
            }
        }
    }

    /// Reads a single byte from the inner I/O stream and compares it to `byte`.
    /// If they differ, an error of `InvalidData` with the message `expected` is
    /// returned. Otherwise the byte is returned. If `byte` is not `CAN` and the
    /// read byte is `CAN`, a `ConnectionAborted` error is returned.
    ///
    /// # Errors
    ///
    /// Returns an error if reading from the inner stream fails, or if the read
    /// byte was not `byte`. If the read byte differed and was `CAN`, an error
    /// of `ConnectionAborted` is returned. Otherwise, the error kind is
    /// `InvalidData`.
    fn expect_byte(&mut self, byte: u8, expected: &'static str) -> io::Result<u8> {
        let b = self.read_byte(byte != CAN)?;
        if b == byte {
            Ok(b)
        } else {
            Err(io::Error::new(io::ErrorKind::InvalidData, expected))
        }
    }

    /// Reads (downloads) a single packet from the inner stream using the XMODEM
    /// protocol. On success, returns the number of bytes read (always 128).
    ///
    /// The progress callback is called with `Progress::Start` when reception
    /// for the first packet has started and subsequently with
    /// `Progress::Packet` when a packet is received successfully.
    ///
    /// # Errors
    ///
    /// Returns an error if reading or writing to the inner stream fails at any
    /// point. Also returns an error if the XMODEM protocol indicates an error.
    /// In particular, an `InvalidData` error is returned when:
    ///
    ///   * The sender's first byte for a packet isn't `EOT` or `SOH`.
    ///   * The sender doesn't send a second `EOT` after the first.
    ///   * The received packet numbers don't match the expected values.
    ///
    /// An error of kind `Interrupted` is returned if a packet checksum fails.
    ///
    /// An error of kind `ConnectionAborted` is returned if a `CAN` byte is
    /// received when not expected.
    ///
    /// An error of kind `UnexpectedEof` is returned if `buf.len() < 128`.
    pub fn read_packet(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.len() < 128 {
            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buf.len() < 128"))
        }

        if !self.started {
            self.write_byte(NAK)?;
            (self.progress)(Progress::Started);
            self.started = true;
        }

        // Waits for an SOH or EOT byte from the sender.
        match self.read_byte(true)? {
            SOH => {
                let p = self.packet;
                let c = 0xFFu8.wrapping_sub(p);
                // Reads the next byte and compares it to the current packet number.
                self.expect_byte_or_cancel(p, "packet")?;
                // Reads the next byte and compares it to the 1s complement of the packet number.
                self.expect_byte_or_cancel(c, "complement")?;

                // Reads a packet (128 bytes) from the sender.
                let mut sum = 0u8;
                for b in buf.iter_mut() {
                    *b = self.read_byte(false)?;
                    sum = sum.wrapping_add(*b);
                }

                // Reads the next byte and compares it to the computed checksum.
                if self.read_byte(false)? != sum {
                    self.write_byte(NAK)?;
                    Err(io::Error::new(io::ErrorKind::Interrupted, "fail checksum"))
                } else {
                    self.write_byte(ACK)?;
                    (self.progress)(Progress::Packet(self.packet));
                    self.packet = self.packet.wrapping_add(1);
                    Ok(128)
                }
            },
            EOT => {
                self.started = false;
                self.write_byte(NAK)?;
                self.expect_byte_or_cancel(EOT, "end of transmission")?;
                self.write_byte(ACK)?;
                Ok(0)
            },

            _ => {
                self.write_byte(CAN)?;
                Err(io::Error::new(io::ErrorKind::InvalidData, "wtf?"))
            },
        }
    }

    /// Sends (uploads) a single packet to the inner stream using the XMODEM
    /// protocol. If `buf` is empty, end of transmissions is sent. Users of this
    /// interface should ensure that `write_packet(&[])` is called when data
    /// transmission is complete. On success, returns the number of bytes
    /// written.
    ///
    /// The progress callback is called with `Progress::Waiting` before waiting
    /// for the receiver's `NAK`, `Progress::Start` when transmission of the
    /// first packet has started and subsequently with `Progress::Packet` when a
    /// packet is sent successfully.
    ///
    /// # Errors
    ///
    /// Returns an error if reading or writing to the inner stream fails at any
    /// point. Also returns an error if the XMODEM protocol indicates an error.
    /// In particular, an `InvalidData` error is returned when:
    ///
    ///   * The receiver's first byte isn't a `NAK`.
    ///   * The receiver doesn't respond with a `NAK` to the first `EOT`.
    ///   * The receiver doesn't respond with an `ACK` to the second `EOT`.
    ///   * The receiver responds to a complete packet with something besides
    ///     `ACK` or `NAK`.
    ///
    /// An error of kind `UnexpectedEof` is returned if `buf.len() < 128 &&
    /// buf.len() != 0`.
    ///
    /// An error of kind `ConnectionAborted` is returned if a `CAN` byte is
    /// received when not expected.
    ///
    /// An error of kind `Interrupted` is returned if a packet checksum fails.
    pub fn write_packet(&mut self, buf: &[u8]) -> io::Result<usize> {
        (self.progress)(Progress::Waiting);
        if !self.started {
            self.expect_byte(NAK, "wait NAK")?;
            self.started = true;
            (self.progress)(Progress::Started);
        }

        if buf.len() == 0 {
            self.started = false;
            self.write_byte(EOT)?;
            self.expect_byte(NAK, "wait NAK")?;
            self.write_byte(EOT)?;
            self.expect_byte(ACK, "wait ACK")?;
            Ok((0))
        } else if buf.len() < 128 {
            Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buf.len() < 128"))
        } else {
            self.write_byte(SOH)?;

            let num = self.packet;
            self.write_byte(num)?;
            let num = 0xFFu8.wrapping_sub(num);
            self.write_byte(num)?;

            let mut sum = 0u8;
            for &b in buf {
                self.write_byte(b)?;
                sum = sum.wrapping_add(b);
            }
            self.write_byte(sum)?;

            match self.read_byte(false)? {
                ACK => {
                    (self.progress)(Progress::Packet(self.packet));
                    self.packet = self.packet.wrapping_add(1);
                    Ok(128)
                }
                NAK => Err(io::Error::new(io::ErrorKind::Interrupted, "interrupted")),
                _ =>   Err(io::Error::new(io::ErrorKind::InvalidData, "end")),
            }

        }
    }

    /// Flush this output stream, ensuring that all intermediately buffered
    /// contents reach their destination.
    ///
    /// # Errors
    ///
    /// It is considered an error if not all bytes could be written due to I/O
    /// errors or EOF being reached.
    pub fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}
<commit_msg>fix xmodem<commit_after>use std::io;

#[cfg(test)] mod tests;
mod read_ext;
mod progress;

pub use progress::{Progress, ProgressFn};

use read_ext::ReadExt;

const SOH: u8 = 0x01;
const EOT: u8 = 0x04;
const ACK: u8 = 0x06;
const NAK: u8 = 0x15;
const CAN: u8 = 0x18;

/// Implementation of the XMODEM protocol.
pub struct Xmodem<R> {
    packet: u8,
    inner: R,
    started: bool,
    progress: ProgressFn
}

impl Xmodem<()> {
    /// Transmits `data` to the receiver `to` using the XMODEM protocol. If the
    /// length of the total data yielded by `data` is not a multiple of 128
    /// bytes, the data is padded with zeroes and sent to the receiver.
    ///
    /// Returns the number of bytes written to `to`, excluding padding zeroes.
    #[inline]
    pub fn transmit<R, W>(data: R, to: W) -> io::Result<usize>
        where W: io::Read + io::Write, R: io::Read
    {
        Self::transmit_with_progress(data, to, progress::noop)
    }

    /// Transmits `data` to the receiver `to` using the XMODEM protocol. If the
    /// length of the total data yielded by `data` is not a multiple of 128
    /// bytes, the data is padded with zeroes and sent to the receiver.
    ///
    /// The function `f` is used as a callback to indicate progress throughout
    /// the transmission. See the [`Progress`] enum for more information.
    ///
    /// Returns the number of bytes written to `to`, excluding padding zeroes.
    pub fn transmit_with_progress<R, W>(mut data: R, to: W, f: ProgressFn) -> io::Result<usize>
        where W: io::Read + io::Write, R: io::Read
    {
        let mut transmitter = Xmodem::new_with_progress(to, f);
        let mut packet = [0u8; 128];
        let mut written = 0;
        'next_packet: loop {
            let n = data.read_max(&mut packet)?;
            packet[n..].iter_mut().for_each(|b| *b = 0);

            if n == 0 {
                transmitter.write_packet(&[])?;
                return Ok(written);
            }

            for _ in 0..10 {
                match transmitter.write_packet(&packet) {
                    Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                    Err(e) => return Err(e),
                    Ok(_) => {
                        written += n;
                        continue 'next_packet;
                    }
                }
            }

            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "bad transmit"));
        }
    }

    /// Receives `data` from `from` using the XMODEM protocol and writes it into
    /// `into`. Returns the number of bytes read from `from`, a multiple of 128.
    #[inline]
    pub fn receive<R, W>(from: R, into: W) -> io::Result<usize>
       where R: io::Read + io::Write, W: io::Write
    {
        Self::receive_with_progress(from, into, progress::noop)
    }

    /// Receives `data` from `from` using the XMODEM protocol and writes it into
    /// `into`. Returns the number of bytes read from `from`, a multiple of 128.
    ///
    /// The function `f` is used as a callback to indicate progress throughout
    /// the reception. See the [`Progress`] enum for more information.
    pub fn receive_with_progress<R, W>(from: R, mut into: W, f: ProgressFn) -> io::Result<usize>
       where R: io::Read + io::Write, W: io::Write
    {
        let mut receiver = Xmodem::new_with_progress(from, f);
        let mut packet = [0u8; 128];
        let mut received = 0;
        'next_packet: loop {
            for _ in 0..10 {
                match receiver.read_packet(&mut packet) {
                    Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                    Err(e) => return Err(e),
                    Ok(0) => break 'next_packet,
                    Ok(n) => {
                        received += n;
                        into.write_all(&packet)?;
                        continue 'next_packet;
                    }
                }
            }

            return Err(io::Error::new(io::ErrorKind::BrokenPipe, "bad receive"));
        }

        Ok(received)
    }
}

impl<T: io::Read + io::Write> Xmodem<T> {
    /// Returns a new `Xmodem` instance with the internal reader/writer set to
    /// `inner`. The returned instance can be used for both receiving
    /// (downloading) and sending (uploading).
    pub fn new(inner: T) -> Self {
        Self { packet: 1, started: false, inner, progress: progress::noop}
    }

    /// Returns a new `Xmodem` instance with the internal reader/writer set to
    /// `inner`. The returned instance can be used for both receiving
    /// (downloading) and sending (uploading). The function `f` is used as a
    /// callback to indicate progress throughout the transfer. See the
    /// [`Progress`] enum for more information.
    pub fn new_with_progress(inner: T, f: ProgressFn) -> Self {
        Self { packet: 1, started: false, inner, progress: f }
    }

    /// Reads a single byte from the inner I/O stream. If `abort_on_can` is
    /// `true`, an error of `ConnectionAborted` is returned if the read byte is
    /// `CAN`.
    ///
    /// # Errors
    ///
    /// Returns an error if reading from the inner stream fails or if
    /// `abort_on_can` is `true` and the read byte is `CAN`.
    fn read_byte(&mut self, abort_on_can: bool) -> io::Result<u8> {
        let mut buf = [0u8; 1];
        self.inner.read_exact(&mut buf)?;

        let byte = buf[0];
        if abort_on_can && byte == CAN {
            return Err(io::Error::new(io::ErrorKind::ConnectionAborted, "received CAN"));
        }

        Ok(byte)
    }

    /// Writes a single byte to the inner I/O stream.
    ///
    /// # Errors
    ///
    /// Returns an error if writing to the inner stream fails.
    fn write_byte(&mut self, byte: u8) -> io::Result<()> {
        self.inner.write_all(&[byte])
    }

    /// Reads a single byte from the inner I/O stream and compares it to `byte`.
    /// If the bytes match, the byte is returned as an `Ok`. If they differ and
    /// the read byte is not `CAN`, an error of `InvalidData` with the message
    /// `expected` is returned. If they differ and the read byte is `CAN`, an
    /// error of `ConnectionAborted` is returned. In either case, if they bytes
    /// differ, a `CAN` byte is written out to the inner stream.
    ///
    /// # Errors
    ///
    /// Returns an error if reading from the inner stream fails, if the read
    /// byte was not `byte`, if the read byte was `CAN` and `byte` is not `CAN`,
    /// or if writing the `CAN` byte failed on byte mismatch.
    fn expect_byte_or_cancel(&mut self, byte: u8, msg: &'static str) -> io::Result<u8> {
        let b = self.read_byte(false)?;
        if b == byte {
            Ok(byte)
        } else {
            self.write_byte(CAN)?;
            if b == CAN {
                Err(io::Error::new(io::ErrorKind::ConnectionAborted, msg))
            } else {
                Err(io::Error::new(io::ErrorKind::InvalidData, msg))
            }
        }
    }

    /// Reads a single byte from the inner I/O stream and compares it to `byte`.
    /// If they differ, an error of `InvalidData` with the message `expected` is
    /// returned. Otherwise the byte is returned. If `byte` is not `CAN` and the
    /// read byte is `CAN`, a `ConnectionAborted` error is returned.
    ///
    /// # Errors
    ///
    /// Returns an error if reading from the inner stream fails, or if the read
    /// byte was not `byte`. If the read byte differed and was `CAN`, an error
    /// of `ConnectionAborted` is returned. Otherwise, the error kind is
    /// `InvalidData`.
    fn expect_byte(&mut self, byte: u8, expected: &'static str) -> io::Result<u8> {
        let b = self.read_byte(byte != CAN)?;
        if b == byte {
            Ok(b)
        } else {
            Err(io::Error::new(io::ErrorKind::InvalidData, expected))
        }
    }

    /// Reads (downloads) a single packet from the inner stream using the XMODEM
    /// protocol. On success, returns the number of bytes read (always 128).
    ///
    /// The progress callback is called with `Progress::Start` when reception
    /// for the first packet has started and subsequently with
    /// `Progress::Packet` when a packet is received successfully.
    ///
    /// # Errors
    ///
    /// Returns an error if reading or writing to the inner stream fails at any
    /// point. Also returns an error if the XMODEM protocol indicates an error.
    /// In particular, an `InvalidData` error is returned when:
    ///
    ///   * The sender's first byte for a packet isn't `EOT` or `SOH`.
    ///   * The sender doesn't send a second `EOT` after the first.
    ///   * The received packet numbers don't match the expected values.
    ///
    /// An error of kind `Interrupted` is returned if a packet checksum fails.
    ///
    /// An error of kind `ConnectionAborted` is returned if a `CAN` byte is
    /// received when not expected.
    ///
    /// An error of kind `UnexpectedEof` is returned if `buf.len() < 128`.
    pub fn read_packet(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        if buf.len() < 128 {
            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buf.len() < 128"))
        }

        if !self.started {
            self.write_byte(NAK)?;
            (self.progress)(Progress::Started);
            self.started = true;
        }

        // Waits for an SOH or EOT byte from the sender.
        match self.read_byte(true)? {
            SOH => {
                let p = self.packet;
                let c = 0xFFu8.wrapping_sub(p);
                // Reads the next byte and compares it to the current packet number.
                self.expect_byte_or_cancel(p, "packet")?;
                // Reads the next byte and compares it to the 1s complement of the packet number.
                self.expect_byte_or_cancel(c, "complement")?;

                // Reads a packet (128 bytes) from the sender.
                let mut sum = 0u8;
                for b in buf.iter_mut() {
                    *b = self.read_byte(false)?;
                    sum = sum.wrapping_add(*b);
                }

                // Reads the next byte and compares it to the computed checksum.
                if self.read_byte(false)? != sum {
                    self.write_byte(NAK)?;
                    Err(io::Error::new(io::ErrorKind::Interrupted, "fail checksum"))
                } else {
                    self.write_byte(ACK)?;
                    (self.progress)(Progress::Packet(self.packet));
                    self.packet = self.packet.wrapping_add(1);
                    Ok(128)
                }
            },
            EOT => {
                self.started = false;
                self.write_byte(NAK)?;
                self.expect_byte_or_cancel(EOT, "end of transmission")?;
                self.write_byte(ACK)?;
                Ok(0)
            },

            _ => {
                self.write_byte(CAN)?;
                Err(io::Error::new(io::ErrorKind::InvalidData, "wtf?"))
            },
        }
    }

    /// Sends (uploads) a single packet to the inner stream using the XMODEM
    /// protocol. If `buf` is empty, end of transmissions is sent. Users of this
    /// interface should ensure that `write_packet(&[])` is called when data
    /// transmission is complete. On success, returns the number of bytes
    /// written.
    ///
    /// The progress callback is called with `Progress::Waiting` before waiting
    /// for the receiver's `NAK`, `Progress::Start` when transmission of the
    /// first packet has started and subsequently with `Progress::Packet` when a
    /// packet is sent successfully.
    ///
    /// # Errors
    ///
    /// Returns an error if reading or writing to the inner stream fails at any
    /// point. Also returns an error if the XMODEM protocol indicates an error.
    /// In particular, an `InvalidData` error is returned when:
    ///
    ///   * The receiver's first byte isn't a `NAK`.
    ///   * The receiver doesn't respond with a `NAK` to the first `EOT`.
    ///   * The receiver doesn't respond with an `ACK` to the second `EOT`.
    ///   * The receiver responds to a complete packet with something besides
    ///     `ACK` or `NAK`.
    ///
    /// An error of kind `UnexpectedEof` is returned if `buf.len() < 128 &&
    /// buf.len() != 0`.
    ///
    /// An error of kind `ConnectionAborted` is returned if a `CAN` byte is
    /// received when not expected.
    ///
    /// An error of kind `Interrupted` is returned if a packet checksum fails.
    pub fn write_packet(&mut self, buf: &[u8]) -> io::Result<usize> {
        (self.progress)(Progress::Waiting);
        if !self.started {
            self.expect_byte(NAK, "wait NAK")?;
            self.started = true;
            (self.progress)(Progress::Started);
        }

        if buf.len() == 0 {
            self.started = false;
            self.write_byte(EOT)?;
            self.expect_byte(NAK, "wait NAK")?;
            self.write_byte(EOT)?;
            self.expect_byte(ACK, "wait ACK")?;
            Ok(0)
        } else if buf.len() < 128 {
            Err(io::Error::new(io::ErrorKind::UnexpectedEof, "buf.len() < 128"))
        } else {
            self.write_byte(SOH)?;

            let num = self.packet;
            self.write_byte(num)?;
            let num = 0xFFu8.wrapping_sub(num);
            self.write_byte(num)?;

            let mut sum = 0u8;
            for &b in buf {
                self.write_byte(b)?;
                sum = sum.wrapping_add(b);
            }
            self.write_byte(sum)?;

            match self.read_byte(false)? {
                ACK => {
                    (self.progress)(Progress::Packet(self.packet));
                    self.packet = self.packet.wrapping_add(1);
                    Ok(128)
                }
                NAK => Err(io::Error::new(io::ErrorKind::Interrupted, "interrupted")),
                _ =>   Err(io::Error::new(io::ErrorKind::InvalidData, "end")),
            }

        }
    }

    /// Flush this output stream, ensuring that all intermediately buffered
    /// contents reach their destination.
    ///
    /// # Errors
    ///
    /// It is considered an error if not all bytes could be written due to I/O
    /// errors or EOF being reached.
    pub fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}
<|endoftext|>