<commit_before>use wpilib::wpilib_hal::*;
use wpilib::hal_call::*;
use wpilib::sensor;
use wpilib::usage::*;

/// A digital output used to control lights, etc from the RoboRIO.
pub struct DigitalOutput {
    channel: i32,
    handle: HAL_DigitalHandle,
}

impl DigitalOutput {
    /// Create a new digital output on the specificed channel, returning an error if initialization
    /// fails.
    pub fn new(channel: i32) -> HalResult<DigitalOutput> {
        if !sensor::check_digital_channel(channel) {
            return Err(0);
        }

        let handle = hal_call!(HAL_InitializeDIOPort(HAL_GetPort(channel), false as i32))?;

        report_usage(ResourceType::DigitalOutput, channel);

        Ok(DigitalOutput {
            channel: channel,
            handle: handle,
        })
    }

    /// Set the value to output.
    pub fn set(&mut self, value: bool) -> HalResult<()> {
        hal_call!(HAL_SetDIO(self.handle, value as i32))
    }

    /// Get the previously-written output.
    pub fn get(&self) -> HalResult<bool> {
        Ok(hal_call!(HAL_GetDIO(self.handle))? != 0)
    }

    /// Get the channel for this DIO.
    pub fn get_channel(&self) -> i32 {
        self.channel
    }

    /// Get a handle to this DIO.
    pub fn get_handle(&self) -> HAL_DigitalHandle {
        self.handle
    }

    /// Write a pulse to this output.
    pub fn pulse(&mut self, length: f64) -> HalResult<()> {
        hal_call!(HAL_Pulse(self.handle, length))
    }

    /// Is this output currently in the middle of a pulse?
    pub fn is_pulsing(&self) -> HalResult<bool> {
        Ok(hal_call!(HAL_IsPulsing(self.handle))? != 0)
    }
}

impl Drop for DigitalOutput {
    fn drop(&mut self) {
        unsafe {
            HAL_FreeDIOPort(self.handle);
        }
    }
}
<commit_msg>Add back in DO PWM code<commit_after>use wpilib::wpilib_hal::*;
use wpilib::hal_call::*;
use wpilib::sensor;
use wpilib::usage::*;

/// A digital output used to control lights, etc from the RoboRIO.
pub struct DigitalOutput {
    channel: i32,
    handle: HAL_DigitalHandle,
    pwm: Option<HAL_DigitalPWMHandle>,
}

impl DigitalOutput {
    /// Create a new digital output on the specificed channel, returning an error if initialization
    /// fails.
    pub fn new(channel: i32) -> HalResult<DigitalOutput> {
        if !sensor::check_digital_channel(channel) {
            return Err(0);
        }

        let handle = hal_call!(HAL_InitializeDIOPort(HAL_GetPort(channel), false as i32))?;

        report_usage(ResourceType::DigitalOutput, channel);

        Ok(DigitalOutput {
            channel: channel,
            handle: handle,
            pwm: None,
        })
    }

    /// Set the PWM rate for this output, from 0.6Hz to 19kHz. Will return an error if PWM has not
    /// been enabled. All digital channels will use the same PWM rate.
    pub fn set_pwm_rate(rate: f64) -> HalResult<()> {
        hal_call!(HAL_SetDigitalPWMRate(rate))
    }

    /// Set the value to output.
    pub fn set(&mut self, value: bool) -> HalResult<()> {
        hal_call!(HAL_SetDIO(self.handle, value as i32))
    }

    /// Get the previously-written output.
    pub fn get(&self) -> HalResult<bool> {
        Ok(hal_call!(HAL_GetDIO(self.handle))? != 0)
    }

    /// Get the channel for this DIO.
    pub fn get_channel(&self) -> i32 {
        self.channel
    }

    /// Get a handle to this DIO.
    pub fn get_handle(&self) -> HAL_DigitalHandle {
        self.handle
    }

    /// Write a pulse to this output.
    pub fn pulse(&mut self, length: f64) -> HalResult<()> {
        hal_call!(HAL_Pulse(self.handle, length))
    }

    /// Is this output currently in the middle of a pulse?
    pub fn is_pulsing(&self) -> HalResult<bool> {
        Ok(hal_call!(HAL_IsPulsing(self.handle))? != 0)
    }

    /// Enable PWM for this output.
    pub fn enable_pwm(&mut self, initial_duty_cycle: f64) -> HalResult<()> {
        let pwm = hal_call!(HAL_AllocateDigitalPWM())?;
        hal_call!(HAL_SetDigitalPWMDutyCycle(pwm, initial_duty_cycle))?;
        hal_call!(HAL_SetDigitalPWMOutputChannel(pwm, self.channel))?;
        self.pwm = Some(pwm);
        Ok(())
    }

    /// Turn off PWM for this output.
    pub fn disable_pwm(&mut self) -> HalResult<()> {
        if let Some(pwm) = self.pwm {
            hal_call!(HAL_SetDigitalPWMOutputChannel(pwm, sensor::num_digital_channels()))?;
            hal_call!(HAL_FreeDigitalPWM(pwm))?;
            self.pwm = None;
        }
        Ok(())
    }

    /// Set a new duty cycle to use in PWM on this output.
    pub fn update_duty_cycle(&mut self, duty_cycle: f64) -> HalResult<()> {
        if let Some(pwm) = self.pwm {
            hal_call!(HAL_SetDigitalPWMDutyCycle(pwm, duty_cycle))
        } else {
            Ok(())
        }
    }
}

impl Drop for DigitalOutput {
    fn drop(&mut self) {
        let _ = self.disable_pwm();
        unsafe {
            HAL_FreeDIOPort(self.handle);
        }
    }
}
<|endoftext|>


<commit_before>use std::iter::IntoIterator;
use std::str::FromStr;

use crate::error::Result;
use crate::map::Map;
use crate::source::AsyncSource;
use crate::{config::Config, path::Expression, source::Source, value::Value};

/// A configuration builder
///
/// It registers ordered sources of configuration to later build consistent [`Config`] from them.
/// Configuration sources it defines are defaults, [`Source`]s and overrides.
///
/// Defaults are alaways loaded first and can be overwritten by any of two other sources.
/// Overrides are always loaded last, thus cannot be overridden.
/// Both can be only set explicitly key by key in code
/// using [`set_default`](Self::set_default) or [`set_override`](Self::set_override).
///
/// An intermediate category, [`Source`], set groups of keys at once implicitly using data coming from external sources
/// like files, environment variables or others that one implements. Defining a [`Source`] is as simple as implementing
/// a trait for a struct.
///
/// Adding sources, setting defaults and overrides does not invoke any I/O nor builds a config.
/// It happens on demand when [`build`](Self::build) (or its alternative) is called.
/// Therefore all errors, related to any of the [`Source`] will only show up then.
///
/// # Sync and async builder
///
/// [`ConfigBuilder`] uses type parameter to keep track of builder state.
///
/// In [`DefaultState`] builder only supports [`Source`]s
///
/// In [`AsyncState`] it supports both [`Source`]s and [`AsyncSource`]s at the price of building using `async fn`.
///
/// # Examples
///
/// ```rust
/// # use config::*;
/// # use std::error::Error;
/// # fn main() -> Result<(), Box<dyn Error>> {
/// let mut builder = Config::builder()
///     .set_default("default", "1")?
///     .add_source(File::new("config/settings", FileFormat::Json))
/// //  .add_async_source(...)
///     .set_override("override", "1")?;
///
/// match builder.build() {
///     Ok(config) => {
///         // use your config
///     },
///     Err(e) => {
///         // something went wrong
///     }
/// }
/// # Ok(())
/// # }
/// ```
///
/// If any [`AsyncSource`] is used, the builder will transition to [`AsyncState`].
/// In such case, it is required to _await_ calls to [`build`](Self::build) and its non-consuming sibling.
///
/// Calls can be not chained as well
/// ```rust
/// # use std::error::Error;
/// # use config::*;
/// # fn main() -> Result<(), Box<dyn Error>> {
/// let mut builder = Config::builder();
/// builder = builder.set_default("default", "1")?;
/// builder = builder.add_source(File::new("config/settings", FileFormat::Json));
/// builder = builder.add_source(File::new("config/settings.prod", FileFormat::Json));
/// builder = builder.set_override("override", "1")?;
/// # Ok(())
/// # }
/// ```
///
/// Calling [`Config::builder`](Config::builder) yields builder in the default state.
/// If having an asynchronous state as the initial state is desired, _turbofish_ notation needs to be used.
/// ```rust
/// # use config::{*, builder::AsyncState};
/// let mut builder = ConfigBuilder::<AsyncState>::default();
/// ```
///
/// If for some reason acquiring builder in default state is required without calling [`Config::builder`](Config::builder)
/// it can also be achieved.
/// ```rust
/// # use config::{*, builder::DefaultState};
/// let mut builder = ConfigBuilder::<DefaultState>::default();
/// ```
#[derive(Debug, Clone, Default)]
pub struct ConfigBuilder<St: BuilderState> {
    defaults: Map<Expression, Value>,
    overrides: Map<Expression, Value>,
    state: St,
}

/// Represents [`ConfigBuilder`] state.
pub trait BuilderState {}

/// Represents data specific to builder in default, sychronous state, without support for async.
#[derive(Debug, Default)]
pub struct DefaultState {
    sources: Vec<Box<dyn Source + Send + Sync>>,
}

/// The asynchronous configuration builder.
///
/// Similar to a [`ConfigBuilder`] it maintains a set of defaults, a set of sources, and overrides.
///
/// Defaults do not override anything, sources override defaults, and overrides override anything else.
/// Within those three groups order of adding them at call site matters - entities added later take precedence.
///
/// For more detailed description and examples see [`ConfigBuilder`].
/// [`AsyncConfigBuilder`] is just an extension of it that takes async functions into account.
///
/// To obtain a [`Config`] call [`build`](AsyncConfigBuilder::build) or [`build_cloned`](AsyncConfigBuilder::build_cloned)
///
/// # Example
/// Since this library does not implement any [`AsyncSource`] an example in rustdocs cannot be given.
/// Detailed explanation about why such a source is not implemented is in [`AsyncSource`]'s documentation.
///
/// Refer to [`ConfigBuilder`] for similar API sample usage or to the examples folder of the crate, where such a source is implemented.
#[derive(Debug, Clone, Default)]
pub struct AsyncConfigBuilder {
    defaults: Map<Expression, Value>,
    overrides: Map<Expression, Value>,
}

/// Represents data specific to builder in asychronous state, with support for async.
#[derive(Debug, Default)]
pub struct AsyncState {
    sources: Vec<SourceType>,
}

#[derive(Debug, Clone)]
enum SourceType {
    Sync(Box<dyn Source + Send + Sync>),
    Async(Box<dyn AsyncSource + Send + Sync>),
}

impl BuilderState for DefaultState {}
impl BuilderState for AsyncState {}

impl<St: BuilderState> ConfigBuilder<St> {
    // operations allowed in any state

    /// Set a default `value` at `key`
    ///
    /// This value can be overwritten by any [`Source`], [`AsyncSource`] or override.
    ///
    /// # Errors
    ///
    /// Fails if `Expression::from_str(key)` fails.
    pub fn set_default<S, T>(mut self, key: S, value: T) -> Result<ConfigBuilder<St>>
    where
        S: AsRef<str>,
        T: Into<Value>,
    {
        self.defaults
            .insert(Expression::from_str(key.as_ref())?, value.into());
        Ok(self)
    }

    /// Set an override
    ///
    /// This function sets an overwrite value. It will not be altered by any default, [`Source`] nor [`AsyncSource`]
    ///
    /// # Errors
    ///
    /// Fails if `Expression::from_str(key)` fails.
    pub fn set_override<S, T>(mut self, key: S, value: T) -> Result<ConfigBuilder<St>>
    where
        S: AsRef<str>,
        T: Into<Value>,
    {
        self.overrides
            .insert(Expression::from_str(key.as_ref())?, value.into());
        Ok(self)
    }
}

impl ConfigBuilder<DefaultState> {
    // operations allowed in sync state

    /// Registers new [`Source`] in this builder.
    ///
    /// Calling this method does not invoke any I/O. [`Source`] is only saved in internal register for later use.
    pub fn add_source<T>(mut self, source: T) -> Self
    where
        T: Source + Send + Sync + 'static,
    {
        self.state.sources.push(Box::new(source));
        self
    }

    /// Registers new [`AsyncSource`] in this builder and forces transition to [`AsyncState`].
    ///
    /// Calling this method does not invoke any I/O. [`AsyncSource`] is only saved in internal register for later use.
    pub fn add_async_source<T>(self, source: T) -> ConfigBuilder<AsyncState>
    where
        T: AsyncSource + Send + Sync + 'static,
    {
        let async_state = ConfigBuilder {
            state: AsyncState {
                sources: self
                    .state
                    .sources
                    .into_iter()
                    .map(|s| SourceType::Sync(s))
                    .collect(),
            },
            defaults: self.defaults,
            overrides: self.overrides,
        };

        async_state.add_async_source(source)
    }

    /// Reads all registered [`Source`]s.
    ///
    /// This is the method that invokes all I/O operations.
    /// For a non consuming alternative see [`build_cloned`](Self::build_cloned)
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub fn build(self) -> Result<Config> {
        Self::build_internal(self.defaults, self.overrides, &self.state.sources)
    }

    /// Reads all registered [`Source`]s.
    ///
    /// Similar to [`build`](Self::build), but it does not take ownership of `ConfigBuilder` to allow later reuse.
    /// Internally it clones data to achieve it.
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub fn build_cloned(&self) -> Result<Config> {
        Self::build_internal(
            self.defaults.clone(),
            self.overrides.clone(),
            &self.state.sources,
        )
    }

    fn build_internal(
        defaults: Map<Expression, Value>,
        overrides: Map<Expression, Value>,
        sources: &[Box<dyn Source + Send + Sync>],
    ) -> Result<Config> {
        let mut cache: Value = Map::<String, Value>::new().into();

        // Add defaults
        for (key, val) in defaults.into_iter() {
            key.set(&mut cache, val);
        }

        // Add sources
        sources.collect_to(&mut cache)?;

        // Add overrides
        for (key, val) in overrides.into_iter() {
            key.set(&mut cache, val);
        }

        Ok(Config::new(cache))
    }
}

impl ConfigBuilder<AsyncState> {
    // operations allowed in async state

    /// Registers new [`Source`] in this builder.
    ///
    /// Calling this method does not invoke any I/O. [`Source`] is only saved in internal register for later use.
    pub fn add_source<T>(mut self, source: T) -> ConfigBuilder<AsyncState>
    where
        T: Source + Send + Sync + 'static,
    {
        self.state.sources.push(SourceType::Sync(Box::new(source)));
        self
    }

    /// Registers new [`AsyncSource`] in this builder.
    ///
    /// Calling this method does not invoke any I/O. [`AsyncSource`] is only saved in internal register for later use.
    pub fn add_async_source<T>(mut self, source: T) -> ConfigBuilder<AsyncState>
    where
        T: AsyncSource + Send + Sync + 'static,
    {
        self.state.sources.push(SourceType::Async(Box::new(source)));
        self
    }

    /// Reads all registered defaults, [`Source`]s, [`AsyncSource`]s and overrides.
    ///
    /// This is the method that invokes all I/O operations.
    /// For a non consuming alternative see [`build_cloned`](Self::build_cloned)
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub async fn build(self) -> Result<Config> {
        Self::build_internal(self.defaults, self.overrides, &self.state.sources).await
    }

    /// Reads all registered defaults, [`Source`]s, [`AsyncSource`]s and overrides.
    ///
    /// Similar to [`build`](Self::build), but it does not take ownership of `ConfigBuilder` to allow later reuse.
    /// Internally it clones data to achieve it.
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub async fn build_cloned(&self) -> Result<Config> {
        Self::build_internal(
            self.defaults.clone(),
            self.overrides.clone(),
            &self.state.sources,
        )
        .await
    }

    async fn build_internal(
        defaults: Map<Expression, Value>,
        overrides: Map<Expression, Value>,
        sources: &[SourceType],
    ) -> Result<Config> {
        let mut cache: Value = Map::<String, Value>::new().into();

        // Add defaults
        for (key, val) in defaults.into_iter() {
            key.set(&mut cache, val);
        }

        for source in sources.iter() {
            match source {
                SourceType::Sync(source) => source.collect_to(&mut cache)?,
                SourceType::Async(source) => source.collect_to(&mut cache).await?,
            }
        }

        // Add overrides
        for (key, val) in overrides.into_iter() {
            key.set(&mut cache, val);
        }

        Ok(Config::new(cache))
    }
}
<commit_msg>Fix clippy: Remove unused field: AsyncConfigBuilder::overrides<commit_after>use std::iter::IntoIterator;
use std::str::FromStr;

use crate::error::Result;
use crate::map::Map;
use crate::source::AsyncSource;
use crate::{config::Config, path::Expression, source::Source, value::Value};

/// A configuration builder
///
/// It registers ordered sources of configuration to later build consistent [`Config`] from them.
/// Configuration sources it defines are defaults, [`Source`]s and overrides.
///
/// Defaults are alaways loaded first and can be overwritten by any of two other sources.
/// Overrides are always loaded last, thus cannot be overridden.
/// Both can be only set explicitly key by key in code
/// using [`set_default`](Self::set_default) or [`set_override`](Self::set_override).
///
/// An intermediate category, [`Source`], set groups of keys at once implicitly using data coming from external sources
/// like files, environment variables or others that one implements. Defining a [`Source`] is as simple as implementing
/// a trait for a struct.
///
/// Adding sources, setting defaults and overrides does not invoke any I/O nor builds a config.
/// It happens on demand when [`build`](Self::build) (or its alternative) is called.
/// Therefore all errors, related to any of the [`Source`] will only show up then.
///
/// # Sync and async builder
///
/// [`ConfigBuilder`] uses type parameter to keep track of builder state.
///
/// In [`DefaultState`] builder only supports [`Source`]s
///
/// In [`AsyncState`] it supports both [`Source`]s and [`AsyncSource`]s at the price of building using `async fn`.
///
/// # Examples
///
/// ```rust
/// # use config::*;
/// # use std::error::Error;
/// # fn main() -> Result<(), Box<dyn Error>> {
/// let mut builder = Config::builder()
///     .set_default("default", "1")?
///     .add_source(File::new("config/settings", FileFormat::Json))
/// //  .add_async_source(...)
///     .set_override("override", "1")?;
///
/// match builder.build() {
///     Ok(config) => {
///         // use your config
///     },
///     Err(e) => {
///         // something went wrong
///     }
/// }
/// # Ok(())
/// # }
/// ```
///
/// If any [`AsyncSource`] is used, the builder will transition to [`AsyncState`].
/// In such case, it is required to _await_ calls to [`build`](Self::build) and its non-consuming sibling.
///
/// Calls can be not chained as well
/// ```rust
/// # use std::error::Error;
/// # use config::*;
/// # fn main() -> Result<(), Box<dyn Error>> {
/// let mut builder = Config::builder();
/// builder = builder.set_default("default", "1")?;
/// builder = builder.add_source(File::new("config/settings", FileFormat::Json));
/// builder = builder.add_source(File::new("config/settings.prod", FileFormat::Json));
/// builder = builder.set_override("override", "1")?;
/// # Ok(())
/// # }
/// ```
///
/// Calling [`Config::builder`](Config::builder) yields builder in the default state.
/// If having an asynchronous state as the initial state is desired, _turbofish_ notation needs to be used.
/// ```rust
/// # use config::{*, builder::AsyncState};
/// let mut builder = ConfigBuilder::<AsyncState>::default();
/// ```
///
/// If for some reason acquiring builder in default state is required without calling [`Config::builder`](Config::builder)
/// it can also be achieved.
/// ```rust
/// # use config::{*, builder::DefaultState};
/// let mut builder = ConfigBuilder::<DefaultState>::default();
/// ```
#[derive(Debug, Clone, Default)]
pub struct ConfigBuilder<St: BuilderState> {
    defaults: Map<Expression, Value>,
    overrides: Map<Expression, Value>,
    state: St,
}

/// Represents [`ConfigBuilder`] state.
pub trait BuilderState {}

/// Represents data specific to builder in default, sychronous state, without support for async.
#[derive(Debug, Default)]
pub struct DefaultState {
    sources: Vec<Box<dyn Source + Send + Sync>>,
}

/// The asynchronous configuration builder.
///
/// Similar to a [`ConfigBuilder`] it maintains a set of defaults, a set of sources, and overrides.
///
/// Defaults do not override anything, sources override defaults, and overrides override anything else.
/// Within those three groups order of adding them at call site matters - entities added later take precedence.
///
/// For more detailed description and examples see [`ConfigBuilder`].
/// [`AsyncConfigBuilder`] is just an extension of it that takes async functions into account.
///
/// To obtain a [`Config`] call [`build`](AsyncConfigBuilder::build) or [`build_cloned`](AsyncConfigBuilder::build_cloned)
///
/// # Example
/// Since this library does not implement any [`AsyncSource`] an example in rustdocs cannot be given.
/// Detailed explanation about why such a source is not implemented is in [`AsyncSource`]'s documentation.
///
/// Refer to [`ConfigBuilder`] for similar API sample usage or to the examples folder of the crate, where such a source is implemented.
#[derive(Debug, Clone, Default)]
pub struct AsyncConfigBuilder {
    defaults: Map<Expression, Value>,
}

/// Represents data specific to builder in asychronous state, with support for async.
#[derive(Debug, Default)]
pub struct AsyncState {
    sources: Vec<SourceType>,
}

#[derive(Debug, Clone)]
enum SourceType {
    Sync(Box<dyn Source + Send + Sync>),
    Async(Box<dyn AsyncSource + Send + Sync>),
}

impl BuilderState for DefaultState {}
impl BuilderState for AsyncState {}

impl<St: BuilderState> ConfigBuilder<St> {
    // operations allowed in any state

    /// Set a default `value` at `key`
    ///
    /// This value can be overwritten by any [`Source`], [`AsyncSource`] or override.
    ///
    /// # Errors
    ///
    /// Fails if `Expression::from_str(key)` fails.
    pub fn set_default<S, T>(mut self, key: S, value: T) -> Result<ConfigBuilder<St>>
    where
        S: AsRef<str>,
        T: Into<Value>,
    {
        self.defaults
            .insert(Expression::from_str(key.as_ref())?, value.into());
        Ok(self)
    }

    /// Set an override
    ///
    /// This function sets an overwrite value. It will not be altered by any default, [`Source`] nor [`AsyncSource`]
    ///
    /// # Errors
    ///
    /// Fails if `Expression::from_str(key)` fails.
    pub fn set_override<S, T>(mut self, key: S, value: T) -> Result<ConfigBuilder<St>>
    where
        S: AsRef<str>,
        T: Into<Value>,
    {
        self.overrides
            .insert(Expression::from_str(key.as_ref())?, value.into());
        Ok(self)
    }
}

impl ConfigBuilder<DefaultState> {
    // operations allowed in sync state

    /// Registers new [`Source`] in this builder.
    ///
    /// Calling this method does not invoke any I/O. [`Source`] is only saved in internal register for later use.
    pub fn add_source<T>(mut self, source: T) -> Self
    where
        T: Source + Send + Sync + 'static,
    {
        self.state.sources.push(Box::new(source));
        self
    }

    /// Registers new [`AsyncSource`] in this builder and forces transition to [`AsyncState`].
    ///
    /// Calling this method does not invoke any I/O. [`AsyncSource`] is only saved in internal register for later use.
    pub fn add_async_source<T>(self, source: T) -> ConfigBuilder<AsyncState>
    where
        T: AsyncSource + Send + Sync + 'static,
    {
        let async_state = ConfigBuilder {
            state: AsyncState {
                sources: self
                    .state
                    .sources
                    .into_iter()
                    .map(|s| SourceType::Sync(s))
                    .collect(),
            },
            defaults: self.defaults,
            overrides: self.overrides,
        };

        async_state.add_async_source(source)
    }

    /// Reads all registered [`Source`]s.
    ///
    /// This is the method that invokes all I/O operations.
    /// For a non consuming alternative see [`build_cloned`](Self::build_cloned)
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub fn build(self) -> Result<Config> {
        Self::build_internal(self.defaults, self.overrides, &self.state.sources)
    }

    /// Reads all registered [`Source`]s.
    ///
    /// Similar to [`build`](Self::build), but it does not take ownership of `ConfigBuilder` to allow later reuse.
    /// Internally it clones data to achieve it.
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub fn build_cloned(&self) -> Result<Config> {
        Self::build_internal(
            self.defaults.clone(),
            self.overrides.clone(),
            &self.state.sources,
        )
    }

    fn build_internal(
        defaults: Map<Expression, Value>,
        overrides: Map<Expression, Value>,
        sources: &[Box<dyn Source + Send + Sync>],
    ) -> Result<Config> {
        let mut cache: Value = Map::<String, Value>::new().into();

        // Add defaults
        for (key, val) in defaults.into_iter() {
            key.set(&mut cache, val);
        }

        // Add sources
        sources.collect_to(&mut cache)?;

        // Add overrides
        for (key, val) in overrides.into_iter() {
            key.set(&mut cache, val);
        }

        Ok(Config::new(cache))
    }
}

impl ConfigBuilder<AsyncState> {
    // operations allowed in async state

    /// Registers new [`Source`] in this builder.
    ///
    /// Calling this method does not invoke any I/O. [`Source`] is only saved in internal register for later use.
    pub fn add_source<T>(mut self, source: T) -> ConfigBuilder<AsyncState>
    where
        T: Source + Send + Sync + 'static,
    {
        self.state.sources.push(SourceType::Sync(Box::new(source)));
        self
    }

    /// Registers new [`AsyncSource`] in this builder.
    ///
    /// Calling this method does not invoke any I/O. [`AsyncSource`] is only saved in internal register for later use.
    pub fn add_async_source<T>(mut self, source: T) -> ConfigBuilder<AsyncState>
    where
        T: AsyncSource + Send + Sync + 'static,
    {
        self.state.sources.push(SourceType::Async(Box::new(source)));
        self
    }

    /// Reads all registered defaults, [`Source`]s, [`AsyncSource`]s and overrides.
    ///
    /// This is the method that invokes all I/O operations.
    /// For a non consuming alternative see [`build_cloned`](Self::build_cloned)
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub async fn build(self) -> Result<Config> {
        Self::build_internal(self.defaults, self.overrides, &self.state.sources).await
    }

    /// Reads all registered defaults, [`Source`]s, [`AsyncSource`]s and overrides.
    ///
    /// Similar to [`build`](Self::build), but it does not take ownership of `ConfigBuilder` to allow later reuse.
    /// Internally it clones data to achieve it.
    ///
    /// # Errors
    /// If source collection fails, be it technical reasons or related to inability to read data as `Config` for different reasons,
    /// this method returns error.
    pub async fn build_cloned(&self) -> Result<Config> {
        Self::build_internal(
            self.defaults.clone(),
            self.overrides.clone(),
            &self.state.sources,
        )
        .await
    }

    async fn build_internal(
        defaults: Map<Expression, Value>,
        overrides: Map<Expression, Value>,
        sources: &[SourceType],
    ) -> Result<Config> {
        let mut cache: Value = Map::<String, Value>::new().into();

        // Add defaults
        for (key, val) in defaults.into_iter() {
            key.set(&mut cache, val);
        }

        for source in sources.iter() {
            match source {
                SourceType::Sync(source) => source.collect_to(&mut cache)?,
                SourceType::Async(source) => source.collect_to(&mut cache).await?,
            }
        }

        // Add overrides
        for (key, val) in overrides.into_iter() {
            key.set(&mut cache, val);
        }

        Ok(Config::new(cache))
    }
}
<|endoftext|>


<commit_before>#![deny(
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_mut,
    unused_qualifications,
    while_true,
)]

extern crate clap;
extern crate log;
extern crate interactor;

extern crate libimagstore;
#[macro_use]
extern crate libimagerror;

pub mod ui;

<commit_msg>Add modules: error, result<commit_after>#![deny(
    non_camel_case_types,
    non_snake_case,
    path_statements,
    trivial_numeric_casts,
    unstable_features,
    unused_allocation,
    unused_import_braces,
    unused_imports,
    unused_mut,
    unused_qualifications,
    while_true,
)]

extern crate clap;
extern crate log;
extern crate interactor;

extern crate libimagstore;
#[macro_use]
extern crate libimagerror;

pub mod error;
pub mod result;
pub mod ui;

<|endoftext|>


<commit_before>//! Shadowsocks Server Context

use std::{
    io,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

#[cfg(feature = "local-dns-relay")]
use std::net::IpAddr;

use bloomfilter::Bloom;
#[cfg(feature = "local-dns-relay")]
use lru_time_cache::LruCache;
use spin::Mutex;
use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;

#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
#[cfg(feature = "local-flow-stat")]
use crate::relay::flow::ServerFlowStatistic;
use crate::{
    config::{Config, ConfigType, ServerConfig},
    relay::{dns_resolver::resolve, socks5::Address},
};

// Entries for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_SERVER: usize = 1_000_000;

// Entries for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_CLIENT: usize = 10_000;

// Error rate for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_SERVER: f64 = 1e-6;

// Error rate for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_CLIENT: f64 = 1e-15;

// A bloom filter borrowed from shadowsocks-libev's `ppbloom`
//
// It contains 2 bloom filters and each one holds 1/2 entries.
// Use them as a ring buffer.
struct PingPongBloom {
    blooms: [Bloom<[u8]>; 2],
    bloom_count: [usize; 2],
    item_count: usize,
    current: usize,
}

impl PingPongBloom {
    fn new(ty: ConfigType) -> PingPongBloom {
        let (mut item_count, fp_p) = if ty.is_local() {
            (BF_NUM_ENTRIES_FOR_CLIENT, BF_ERROR_RATE_FOR_CLIENT)
        } else {
            (BF_NUM_ENTRIES_FOR_SERVER, BF_ERROR_RATE_FOR_SERVER)
        };

        item_count /= 2;

        PingPongBloom {
            blooms: [
                Bloom::new_for_fp_rate(item_count, fp_p),
                Bloom::new_for_fp_rate(item_count, fp_p),
            ],
            bloom_count: [0, 0],
            item_count,
            current: 0,
        }
    }

    // Check if data in `buf` exist.
    //
    // Set into the current bloom filter if not exist.
    //
    // Return `true` if data exist in bloom filter.
    fn check_and_set(&mut self, buf: &[u8]) -> bool {
        for bloom in &self.blooms {
            if bloom.check(buf) {
                return true;
            }
        }

        if self.bloom_count[self.current] >= self.item_count {
            // Current bloom filter is full,
            // Create a new one and use that one as current.

            self.current = (self.current + 1) % 2;

            self.bloom_count[self.current] = 0;
            self.blooms[self.current].clear();
        }

        // Cannot be optimized by `check_and_set`
        // Because we have to check every filters in `blooms` before `set`
        self.blooms[self.current].set(buf);
        self.bloom_count[self.current] += 1;

        false
    }
}

/// Server's global running status
///
/// Shared between UDP and TCP servers
pub struct ServerState {
    #[cfg(feature = "trust-dns")]
    dns_resolver: Option<TokioAsyncResolver>,
}

#[cfg(feature = "trust-dns")]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(config: &Config, rt: Handle) -> SharedServerState {
        let state = ServerState {
            dns_resolver: match create_resolver(config.get_dns_config(), config.timeout, config.ipv6_first, rt).await {
                Ok(resolver) => Some(resolver),
                Err(..) => None,
            },
        };

        Arc::new(state)
    }

    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.dns_resolver.as_ref()
    }
}

#[cfg(not(feature = "trust-dns"))]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(_config: &Config, _rt: Handle) -> SharedServerState {
        Arc::new(ServerState {})
    }
}

/// `ServerState` wrapped in `Arc`
pub type SharedServerState = Arc<ServerState>;

/// Shared basic configuration for the whole server
pub struct Context {
    config: Config,

    // Shared variables for all servers
    server_state: SharedServerState,

    // Server's running indicator
    // For killing all background jobs
    server_running: AtomicBool,

    // Check for duplicated IV/Nonce, for prevent replay attack
    // https://github.com/shadowsocks/shadowsocks-org/issues/44
    nonce_ppbloom: Mutex<PingPongBloom>,

    // For Android's flow stat report
    #[cfg(feature = "local-flow-stat")]
    local_flow_statistic: ServerFlowStatistic,

    // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded
    #[cfg(feature = "local-dns-relay")]
    reverse_lookup_cache: Mutex<LruCache<IpAddr, bool>>,
}

/// Unique context thw whole server
pub type SharedContext = Arc<Context>;

impl Context {
    /// Create a non-shared Context
    fn new(config: Config, server_state: SharedServerState) -> Context {
        let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type));
        #[cfg(feature = "local-dns-relay")]
        let reverse_lookup_cache = Mutex::new(LruCache::<IpAddr, bool>::with_capacity(8192));

        Context {
            config,
            server_state,
            server_running: AtomicBool::new(true),
            nonce_ppbloom,
            #[cfg(feature = "local-flow-stat")]
            local_flow_statistic: ServerFlowStatistic::new(),
            #[cfg(feature = "local-dns-relay")]
            reverse_lookup_cache,
        }
    }

    /// Create a shared Context, wrapped in `Arc`
    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
        SharedContext::new(Context::new(config, server_state))
    }

    /// Config for TCP server
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Clone the internal ServerState
    pub fn clone_server_state(&self) -> SharedServerState {
        self.server_state.clone()
    }

    /// Mutable Config for TCP server
    ///
    /// NOTE: Only for launching plugins
    pub fn config_mut(&mut self) -> &mut Config {
        &mut self.config
    }

    /// Get ServerConfig by index
    pub fn server_config(&self, idx: usize) -> &ServerConfig {
        &self.config.server[idx]
    }

    /// Get mutable ServerConfig by index
    pub fn server_config_mut(&mut self, idx: usize) -> &mut ServerConfig {
        &mut self.config.server[idx]
    }

    #[cfg(feature = "trust-dns")]
    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.server_state.dns_resolver()
    }

    /// Perform a DNS resolution
    pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
        resolve(self, host, port).await
    }

    /// Check if the server is still in running state
    pub fn server_running(&self) -> bool {
        self.server_running.load(Ordering::Acquire)
    }

    /// Stops the server, kills all detached running tasks
    pub fn set_server_stopped(&self) {
        self.server_running.store(false, Ordering::Release)
    }

    /// Check if nonce exist or not
    ///
    /// If not, set into the current bloom filter
    pub fn check_nonce_and_set(&self, nonce: &[u8]) -> bool {
        // Plain cipher doesn't have a nonce
        // Always treated as non-duplicated
        if nonce.is_empty() {
            return false;
        }

        let mut ppbloom = self.nonce_ppbloom.lock();
        ppbloom.check_and_set(nonce)
    }

    /// Check client ACL (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_client_blocked(addr),
        }
    }

    /// Check outbound address ACL (for server)
    pub fn check_outbound_blocked(&self, addr: &Address) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_outbound_blocked(addr),
        }
    }

    /// Check resolved outbound address ACL (for server)
    pub fn check_resolved_outbound_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_resolved_outbound_blocked(addr),
        }
    }

    /// Add a record to the reverse lookup cache
    #[cfg(feature = "local-dns-relay")]
    pub fn add_to_reverse_lookup_cache(&self, addr: &IpAddr, forward: bool) {
        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
        reverse_lookup_cache.insert(addr.clone(), forward);
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_qname_in_proxy_list(&self, qname: &Address) -> Option<bool> {
        match self.config.acl {
            // Proxy everything by default
            None => None,
            Some(ref a) => a.check_qname_in_proxy_list(qname),
        }
    }

    #[cfg(feature = "local-dns-relay")]
    pub fn check_ip_in_proxy_list(&self, ip: &IpAddr) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => true,
            Some(ref a) => {
                // first check the IP list
                if a.check_ip_in_proxy_list(ip) {
                    true
                } else {
                    // do the reverse lookup in our local cache
                    let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                    // if a qname is found
                    if let Some(forward) = reverse_lookup_cache.get(ip) {
                        // if qname is resolved by remote,
                        // we should proxy it as well
                        *forward
                    } else {
                        // by default not proxied
                        false
                    }
                }
            }
        }
    }

    /// Check target address ACL (for client)
    pub async fn check_target_bypassed(&self, target: &Address) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => false,
            Some(ref a) => {
                #[cfg(feature = "local-dns-relay")]
                {
                    if let Address::SocketAddress(ref saddr) = target {
                        // do the reverse lookup in our local cache
                        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                        // if a qname is found
                        if let Some(forward) = reverse_lookup_cache.get(&saddr.ip()) {
                            return !*forward
                        }
                    }
                }
                a.check_target_bypassed(self, target).await
            }
        }
    }

    /// Get client flow statistics
    #[cfg(feature = "local-flow-stat")]
    pub fn local_flow_statistic(&self) -> &ServerFlowStatistic {
        &self.local_flow_statistic
    }
}
<commit_msg>Revert the compare order in check_ip_in_proxy_list<commit_after>//! Shadowsocks Server Context

use std::{
    io,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

#[cfg(feature = "local-dns-relay")]
use std::net::IpAddr;

use bloomfilter::Bloom;
#[cfg(feature = "local-dns-relay")]
use lru_time_cache::LruCache;
use spin::Mutex;
use tokio::runtime::Handle;
#[cfg(feature = "trust-dns")]
use trust_dns_resolver::TokioAsyncResolver;

#[cfg(feature = "trust-dns")]
use crate::relay::dns_resolver::create_resolver;
#[cfg(feature = "local-flow-stat")]
use crate::relay::flow::ServerFlowStatistic;
use crate::{
    config::{Config, ConfigType, ServerConfig},
    relay::{dns_resolver::resolve, socks5::Address},
};

// Entries for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_SERVER: usize = 1_000_000;

// Entries for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_NUM_ENTRIES_FOR_CLIENT: usize = 10_000;

// Error rate for server's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_SERVER: f64 = 1e-6;

// Error rate for client's bloom filter
//
// Borrowed from shadowsocks-libev's default value
const BF_ERROR_RATE_FOR_CLIENT: f64 = 1e-15;

// A bloom filter borrowed from shadowsocks-libev's `ppbloom`
//
// It contains 2 bloom filters and each one holds 1/2 entries.
// Use them as a ring buffer.
struct PingPongBloom {
    blooms: [Bloom<[u8]>; 2],
    bloom_count: [usize; 2],
    item_count: usize,
    current: usize,
}

impl PingPongBloom {
    fn new(ty: ConfigType) -> PingPongBloom {
        let (mut item_count, fp_p) = if ty.is_local() {
            (BF_NUM_ENTRIES_FOR_CLIENT, BF_ERROR_RATE_FOR_CLIENT)
        } else {
            (BF_NUM_ENTRIES_FOR_SERVER, BF_ERROR_RATE_FOR_SERVER)
        };

        item_count /= 2;

        PingPongBloom {
            blooms: [
                Bloom::new_for_fp_rate(item_count, fp_p),
                Bloom::new_for_fp_rate(item_count, fp_p),
            ],
            bloom_count: [0, 0],
            item_count,
            current: 0,
        }
    }

    // Check if data in `buf` exist.
    //
    // Set into the current bloom filter if not exist.
    //
    // Return `true` if data exist in bloom filter.
    fn check_and_set(&mut self, buf: &[u8]) -> bool {
        for bloom in &self.blooms {
            if bloom.check(buf) {
                return true;
            }
        }

        if self.bloom_count[self.current] >= self.item_count {
            // Current bloom filter is full,
            // Create a new one and use that one as current.

            self.current = (self.current + 1) % 2;

            self.bloom_count[self.current] = 0;
            self.blooms[self.current].clear();
        }

        // Cannot be optimized by `check_and_set`
        // Because we have to check every filters in `blooms` before `set`
        self.blooms[self.current].set(buf);
        self.bloom_count[self.current] += 1;

        false
    }
}

/// Server's global running status
///
/// Shared between UDP and TCP servers
pub struct ServerState {
    #[cfg(feature = "trust-dns")]
    dns_resolver: Option<TokioAsyncResolver>,
}

#[cfg(feature = "trust-dns")]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(config: &Config, rt: Handle) -> SharedServerState {
        let state = ServerState {
            dns_resolver: match create_resolver(config.get_dns_config(), config.timeout, config.ipv6_first, rt).await {
                Ok(resolver) => Some(resolver),
                Err(..) => None,
            },
        };

        Arc::new(state)
    }

    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.dns_resolver.as_ref()
    }
}

#[cfg(not(feature = "trust-dns"))]
impl ServerState {
    /// Create a global shared server state
    pub async fn new_shared(_config: &Config, _rt: Handle) -> SharedServerState {
        Arc::new(ServerState {})
    }
}

/// `ServerState` wrapped in `Arc`
pub type SharedServerState = Arc<ServerState>;

/// Shared basic configuration for the whole server
pub struct Context {
    config: Config,

    // Shared variables for all servers
    server_state: SharedServerState,

    // Server's running indicator
    // For killing all background jobs
    server_running: AtomicBool,

    // Check for duplicated IV/Nonce, for prevent replay attack
    // https://github.com/shadowsocks/shadowsocks-org/issues/44
    nonce_ppbloom: Mutex<PingPongBloom>,

    // For Android's flow stat report
    #[cfg(feature = "local-flow-stat")]
    local_flow_statistic: ServerFlowStatistic,

    // For DNS relay's ACL domain name reverse lookup -- whether the IP shall be forwarded
    #[cfg(feature = "local-dns-relay")]
    reverse_lookup_cache: Mutex<LruCache<IpAddr, bool>>,
}

/// Unique context thw whole server
pub type SharedContext = Arc<Context>;

impl Context {
    /// Create a non-shared Context
    fn new(config: Config, server_state: SharedServerState) -> Context {
        let nonce_ppbloom = Mutex::new(PingPongBloom::new(config.config_type));
        #[cfg(feature = "local-dns-relay")]
        let reverse_lookup_cache = Mutex::new(LruCache::<IpAddr, bool>::with_capacity(8192));

        Context {
            config,
            server_state,
            server_running: AtomicBool::new(true),
            nonce_ppbloom,
            #[cfg(feature = "local-flow-stat")]
            local_flow_statistic: ServerFlowStatistic::new(),
            #[cfg(feature = "local-dns-relay")]
            reverse_lookup_cache,
        }
    }

    /// Create a shared Context, wrapped in `Arc`
    pub fn new_shared(config: Config, server_state: SharedServerState) -> SharedContext {
        SharedContext::new(Context::new(config, server_state))
    }

    /// Config for TCP server
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Clone the internal ServerState
    pub fn clone_server_state(&self) -> SharedServerState {
        self.server_state.clone()
    }

    /// Mutable Config for TCP server
    ///
    /// NOTE: Only for launching plugins
    pub fn config_mut(&mut self) -> &mut Config {
        &mut self.config
    }

    /// Get ServerConfig by index
    pub fn server_config(&self, idx: usize) -> &ServerConfig {
        &self.config.server[idx]
    }

    /// Get mutable ServerConfig by index
    pub fn server_config_mut(&mut self, idx: usize) -> &mut ServerConfig {
        &mut self.config.server[idx]
    }

    #[cfg(feature = "trust-dns")]
    /// Get the global shared resolver
    pub fn dns_resolver(&self) -> Option<&TokioAsyncResolver> {
        self.server_state.dns_resolver()
    }

    /// Perform a DNS resolution
    pub async fn dns_resolve(&self, host: &str, port: u16) -> io::Result<Vec<SocketAddr>> {
        resolve(self, host, port).await
    }

    /// Check if the server is still in running state
    pub fn server_running(&self) -> bool {
        self.server_running.load(Ordering::Acquire)
    }

    /// Stops the server, kills all detached running tasks
    pub fn set_server_stopped(&self) {
        self.server_running.store(false, Ordering::Release)
    }

    /// Check if nonce exist or not
    ///
    /// If not, set into the current bloom filter
    pub fn check_nonce_and_set(&self, nonce: &[u8]) -> bool {
        // Plain cipher doesn't have a nonce
        // Always treated as non-duplicated
        if nonce.is_empty() {
            return false;
        }

        let mut ppbloom = self.nonce_ppbloom.lock();
        ppbloom.check_and_set(nonce)
    }

    /// Check client ACL (for server)
    pub fn check_client_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_client_blocked(addr),
        }
    }

    /// Check outbound address ACL (for server)
    pub fn check_outbound_blocked(&self, addr: &Address) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_outbound_blocked(addr),
        }
    }

    /// Check resolved outbound address ACL (for server)
    pub fn check_resolved_outbound_blocked(&self, addr: &SocketAddr) -> bool {
        match self.config.acl {
            None => false,
            Some(ref a) => a.check_resolved_outbound_blocked(addr),
        }
    }

    /// Add a record to the reverse lookup cache
    #[cfg(feature = "local-dns-relay")]
    pub fn add_to_reverse_lookup_cache(&self, addr: &IpAddr, forward: bool) {
        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
        reverse_lookup_cache.insert(addr.clone(), forward);
    }

    /// Check if domain name is in proxy_list.
    /// If so, it should be resolved from remote (for Android's DNS relay)
    pub fn check_qname_in_proxy_list(&self, qname: &Address) -> Option<bool> {
        match self.config.acl {
            // Proxy everything by default
            None => None,
            Some(ref a) => a.check_qname_in_proxy_list(qname),
        }
    }

    #[cfg(feature = "local-dns-relay")]
    pub fn check_ip_in_proxy_list(&self, ip: &IpAddr) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => true,
            Some(ref a) => {
                // do the reverse lookup in our local cache
                let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                // if a qname is found
                if let Some(forward) = reverse_lookup_cache.get(ip) {
                    // if qname is resolved by remote,
                    // we should proxy it as well
                    *forward
                } else {
                    a.check_ip_in_proxy_list(ip)
                }
            }
        }
    }

    /// Check target address ACL (for client)
    pub async fn check_target_bypassed(&self, target: &Address) -> bool {
        match self.config.acl {
            // Proxy everything by default
            None => false,
            Some(ref a) => {
                #[cfg(feature = "local-dns-relay")]
                {
                    if let Address::SocketAddress(ref saddr) = target {
                        // do the reverse lookup in our local cache
                        let mut reverse_lookup_cache = self.reverse_lookup_cache.lock();
                        // if a qname is found
                        if let Some(forward) = reverse_lookup_cache.get(&saddr.ip()) {
                            return !*forward
                        }
                    }
                }
                a.check_target_bypassed(self, target).await
            }
        }
    }

    /// Get client flow statistics
    #[cfg(feature = "local-flow-stat")]
    pub fn local_flow_statistic(&self) -> &ServerFlowStatistic {
        &self.local_flow_statistic
    }
}
<|endoftext|>


<commit_before>use ffi;
use types::{Result, PortMidiDeviceId, Error};
use io::{InputPort, OutputPort};
use device::DeviceInfo;

/// The elements are static after initializing
pub struct PortMidi {
    device_cnt: i32,
    buffer_size: i32,
}
impl PortMidi {
    /// Initializes the underlying PortMidi C library.
    /// It does not support *hot plugging*, this means
    /// that devices that are connect after calling `new`
    /// are not picked up by PortMidi.
    pub fn new() -> Result<Self> {
        try!(Result::from(unsafe { ffi::Pm_Initialize() }));

        Ok(PortMidi {
            device_cnt: unsafe { ffi::Pm_CountDevices() },
            buffer_size: 1024, // TODO: argument
        })
    }

    /// Return the number of devices. This number will not change during the lifetime
    /// of the program.
    pub fn device_cnt(&self) -> PortMidiDeviceId {
        self.device_cnt
    }

    /// Gets the `PortMidiDeviceId` for the default input, or `None` if
    /// there isn't one
    ///
    /// See the PortMidi documentation for details of how to set the default device
    pub fn default_input_device_id(&self) -> Result<PortMidiDeviceId> {
        match unsafe { ffi::Pm_GetDefaultInputDeviceID() } {
            ffi::PM_NO_DEVICE => Err(Error::NoDefaultDevice),
            id @ _ => Ok(id),
        }
    }

    /// Gets the `PortMidiDeviceId` for the default output, or `None` if
    /// there isn't one
    ///
    /// See the PortMidi documentation for details of how to set the default device
    pub fn default_output_device_id(&self) -> Result<PortMidiDeviceId> {
        match unsafe { ffi::Pm_GetDefaultOutputDeviceID() } {
            ffi::PM_NO_DEVICE => Err(Error::NoDefaultDevice),
            id @ _ => Ok(id),
        }
    }

    pub fn device(&self, id: PortMidiDeviceId) -> Result<DeviceInfo> {
        DeviceInfo::new(id)
    }

    pub fn devices(&self) -> Result<Vec<DeviceInfo>> {
        let mut devices = Vec::with_capacity(self.device_cnt() as usize);
        for res in (0..self.device_cnt()).map(|id| self.device(id)) {
            match res {
                Ok(device) => devices.push(device),
                Err(err) => return Err(err),
            }
        }
        Ok(devices)
    }

    pub fn default_input_port(&self) -> Result<InputPort> {
        let info = try!(self.default_input_device_id().and_then(|id| self.device(id)));
        InputPort::new(info, self.buffer_size)
    }

    pub fn input_port(&self, device: DeviceInfo) -> Result<InputPort> {
        if device.is_input() {
            InputPort::new(device, self.buffer_size)
        } else {
            Err(Error::NotAnInputDevice)
        }
    }

    pub fn default_output_port(&self) -> Result<OutputPort> {
        let info = try!(self.default_output_device_id().and_then(|id| self.device(id)));
        OutputPort::new(info, self.buffer_size)
    }

    pub fn output_port(&self, device: DeviceInfo) -> Result<InputPort> {
        if device.is_output() {
            InputPort::new(device, self.buffer_size)
        } else {
            Err(Error::NotAnOutputDevice)
        }
    }
}
impl Drop for PortMidi {
    fn drop(&mut self) {
        Result::from(unsafe { ffi::Pm_Terminate() })
            .map_err(|err| println!("Could not terminate library: {}", err));
    }
}
<commit_msg>-library<commit_after>use ffi;
use types::{Result, PortMidiDeviceId, Error};
use io::{InputPort, OutputPort};
use device::DeviceInfo;

/// The elements are static after initializing
pub struct PortMidi {
    device_cnt: i32,
    buffer_size: i32,
}
impl PortMidi {
    /// Initializes the underlying PortMidi C library.
    /// It does not support *hot plugging*, this means
    /// that devices that are connect after calling `new`
    /// are not picked up by PortMidi.
    pub fn new() -> Result<Self> {
        try!(Result::from(unsafe { ffi::Pm_Initialize() }));

        Ok(PortMidi {
            device_cnt: unsafe { ffi::Pm_CountDevices() },
            buffer_size: 1024, // TODO: argument
        })
    }

    /// Return the number of devices. This number will not change during the lifetime
    /// of the program.
    pub fn device_cnt(&self) -> PortMidiDeviceId {
        self.device_cnt
    }

    /// Gets the `PortMidiDeviceId` for the default input, or `None` if
    /// there isn't one
    ///
    /// See the PortMidi documentation for details of how to set the default device
    pub fn default_input_device_id(&self) -> Result<PortMidiDeviceId> {
        match unsafe { ffi::Pm_GetDefaultInputDeviceID() } {
            ffi::PM_NO_DEVICE => Err(Error::NoDefaultDevice),
            id @ _ => Ok(id),
        }
    }

    /// Gets the `PortMidiDeviceId` for the default output, or `None` if
    /// there isn't one
    ///
    /// See the PortMidi documentation for details of how to set the default device
    pub fn default_output_device_id(&self) -> Result<PortMidiDeviceId> {
        match unsafe { ffi::Pm_GetDefaultOutputDeviceID() } {
            ffi::PM_NO_DEVICE => Err(Error::NoDefaultDevice),
            id @ _ => Ok(id),
        }
    }

    pub fn device(&self, id: PortMidiDeviceId) -> Result<DeviceInfo> {
        DeviceInfo::new(id)
    }

    pub fn devices(&self) -> Result<Vec<DeviceInfo>> {
        let mut devices = Vec::with_capacity(self.device_cnt() as usize);
        for res in (0..self.device_cnt()).map(|id| self.device(id)) {
            match res {
                Ok(device) => devices.push(device),
                Err(err) => return Err(err),
            }
        }
        Ok(devices)
    }

    pub fn default_input_port(&self) -> Result<InputPort> {
        let info = try!(self.default_input_device_id().and_then(|id| self.device(id)));
        InputPort::new(info, self.buffer_size)
    }

    pub fn input_port(&self, device: DeviceInfo) -> Result<InputPort> {
        if device.is_input() {
            InputPort::new(device, self.buffer_size)
        } else {
            Err(Error::NotAnInputDevice)
        }
    }

    pub fn default_output_port(&self) -> Result<OutputPort> {
        let info = try!(self.default_output_device_id().and_then(|id| self.device(id)));
        OutputPort::new(info, self.buffer_size)
    }

    pub fn output_port(&self, device: DeviceInfo) -> Result<InputPort> {
        if device.is_output() {
            InputPort::new(device, self.buffer_size)
        } else {
            Err(Error::NotAnOutputDevice)
        }
    }
}
impl Drop for PortMidi {
    fn drop(&mut self) {
        Result::from(unsafe { ffi::Pm_Terminate() })
            .map_err(|err| println!("Could not terminate: {}", err));
    }
}
<|endoftext|>


<commit_before>// $ cargo bench --features full --bench file

#![feature(rustc_private, test)]
#![recursion_limit = "1024"]
#![allow(clippy::missing_panics_doc, clippy::must_use_candidate)]

extern crate test;

#[macro_use]
#[path = "../tests/macros/mod.rs"]
mod macros;

#[path = "../tests/common/mod.rs"]
mod common;
#[path = "../tests/repo/mod.rs"]
pub mod repo;

use proc_macro2::TokenStream;
use std::fs;
use std::str::FromStr;
use test::Bencher;

const FILE: &str = "tests/rust/library/core/src/str/mod.rs";

fn get_tokens() -> TokenStream {
    repo::clone_rust();
    let content = fs::read_to_string(FILE).unwrap();
    TokenStream::from_str(&content).unwrap()
}

#[bench]
fn baseline(b: &mut Bencher) {
    let tokens = get_tokens();
    b.iter(|| drop(tokens.clone()));
}

#[bench]
fn parse_file(b: &mut Bencher) {
    let tokens = get_tokens();
    b.iter(|| syn::parse2::<syn::File>(tokens.clone()));
}
<commit_msg>Add a benchmark of TokenBuffer::inner_new<commit_after>// $ cargo bench --features full --bench file

#![feature(rustc_private, test)]
#![recursion_limit = "1024"]
#![allow(
    clippy::items_after_statements,
    clippy::missing_panics_doc,
    clippy::must_use_candidate
)]

extern crate test;

#[macro_use]
#[path = "../tests/macros/mod.rs"]
mod macros;

#[path = "../tests/common/mod.rs"]
mod common;
#[path = "../tests/repo/mod.rs"]
pub mod repo;

use proc_macro2::{Span, TokenStream};
use std::fs;
use std::str::FromStr;
use syn::parse::{ParseStream, Parser};
use test::Bencher;

const FILE: &str = "tests/rust/library/core/src/str/mod.rs";

fn get_tokens() -> TokenStream {
    repo::clone_rust();
    let content = fs::read_to_string(FILE).unwrap();
    TokenStream::from_str(&content).unwrap()
}

#[bench]
fn baseline(b: &mut Bencher) {
    let tokens = get_tokens();
    b.iter(|| drop(tokens.clone()));
}

#[bench]
fn create_token_buffer(b: &mut Bencher) {
    let tokens = get_tokens();
    fn immediate_fail(_input: ParseStream) -> syn::Result<()> {
        Err(syn::Error::new(Span::call_site(), ""))
    }
    b.iter(|| immediate_fail.parse2(tokens.clone()));
}

#[bench]
fn parse_file(b: &mut Bencher) {
    let tokens = get_tokens();
    b.iter(|| syn::parse2::<syn::File>(tokens.clone()));
}
<|endoftext|>


<commit_before>// $ cargo bench --features full --bench rust
//
// Syn only, useful for profiling:
// $ RUSTFLAGS='--cfg syn_only' cargo build --release --features full --bench rust

#![cfg_attr(not(syn_only), feature(rustc_private))]
#![recursion_limit = "1024"]
#![allow(clippy::cast_lossless, clippy::unnecessary_wraps)]

#[macro_use]
#[path = "../tests/macros/mod.rs"]
mod macros;

#[path = "../tests/common/mod.rs"]
mod common;
#[path = "../tests/repo/mod.rs"]
mod repo;

use std::fs;
use std::time::{Duration, Instant};

#[cfg(not(syn_only))]
mod tokenstream_parse {
    use proc_macro2::TokenStream;
    use std::str::FromStr;

    pub fn bench(content: &str) -> Result<(), ()> {
        TokenStream::from_str(content).map(drop).map_err(drop)
    }
}

mod syn_parse {
    pub fn bench(content: &str) -> Result<(), ()> {
        syn::parse_file(content).map(drop).map_err(drop)
    }
}

#[cfg(not(syn_only))]
mod librustc_parse {
    extern crate rustc_data_structures;
    extern crate rustc_error_messages;
    extern crate rustc_errors;
    extern crate rustc_parse;
    extern crate rustc_session;
    extern crate rustc_span;

    use rustc_data_structures::sync::Lrc;
    use rustc_error_messages::FluentBundle;
    use rustc_errors::{emitter::Emitter, Diagnostic, Handler};
    use rustc_session::parse::ParseSess;
    use rustc_span::source_map::{FilePathMapping, SourceMap};
    use rustc_span::{edition::Edition, FileName};

    pub fn bench(content: &str) -> Result<(), ()> {
        struct SilentEmitter;

        impl Emitter for SilentEmitter {
            fn emit_diagnostic(&mut self, _diag: &Diagnostic) {}
            fn source_map(&self) -> Option<&Lrc<SourceMap>> {
                None
            }
            fn fluent_bundle(&self) -> Option<&Lrc<FluentBundle>> {
                None
            }
            fn fallback_fluent_bundle(&self) -> &FluentBundle {
                panic!("silent emitter attempted to translate a diagnostic");
            }
        }

        rustc_span::create_session_if_not_set_then(Edition::Edition2018, |_| {
            let cm = Lrc::new(SourceMap::new(FilePathMapping::empty()));
            let emitter = Box::new(SilentEmitter);
            let handler = Handler::with_emitter(false, None, emitter);
            let sess = ParseSess::with_span_handler(handler, cm);
            if let Err(diagnostic) = rustc_parse::parse_crate_from_source_str(
                FileName::Custom("bench".to_owned()),
                content.to_owned(),
                &sess,
            ) {
                diagnostic.cancel();
                return Err(());
            };
            Ok(())
        })
    }
}

#[cfg(not(syn_only))]
mod read_from_disk {
    pub fn bench(content: &str) -> Result<(), ()> {
        let _ = content;
        Ok(())
    }
}

fn exec(mut codepath: impl FnMut(&str) -> Result<(), ()>) -> Duration {
    let begin = Instant::now();
    let mut success = 0;
    let mut total = 0;

    walkdir::WalkDir::new("tests/rust/src")
        .into_iter()
        .filter_entry(repo::base_dir_filter)
        .for_each(|entry| {
            let entry = entry.unwrap();
            let path = entry.path();
            if path.is_dir() {
                return;
            }
            let content = fs::read_to_string(path).unwrap();
            let ok = codepath(&content).is_ok();
            success += ok as usize;
            total += 1;
            if !ok {
                eprintln!("FAIL {}", path.display());
            }
        });

    assert_eq!(success, total);
    begin.elapsed()
}

fn main() {
    repo::clone_rust();

    macro_rules! testcases {
        ($($(#[$cfg:meta])* $name:ident,)*) => {
            [
                $(
                    $(#[$cfg])*
                    (stringify!($name), $name::bench as fn(&str) -> Result<(), ()>),
                )*
            ]
        };
    }

    #[cfg(not(syn_only))]
    {
        let mut lines = 0;
        let mut files = 0;
        exec(|content| {
            lines += content.lines().count();
            files += 1;
            Ok(())
        });
        eprintln!("\n{} lines in {} files", lines, files);
    }

    for (name, f) in testcases!(
        #[cfg(not(syn_only))]
        read_from_disk,
        #[cfg(not(syn_only))]
        tokenstream_parse,
        syn_parse,
        #[cfg(not(syn_only))]
        librustc_parse,
    ) {
        eprint!("{:20}", format!("{}:", name));
        let elapsed = exec(f);
        eprintln!(
            "elapsed={}.{:03}s",
            elapsed.as_secs(),
            elapsed.subsec_millis(),
        );
    }
    eprintln!();
}
<commit_msg>Update benches to nightly-2022-08-18<commit_after>// $ cargo bench --features full --bench rust
//
// Syn only, useful for profiling:
// $ RUSTFLAGS='--cfg syn_only' cargo build --release --features full --bench rust

#![cfg_attr(not(syn_only), feature(rustc_private))]
#![recursion_limit = "1024"]
#![allow(clippy::cast_lossless, clippy::unnecessary_wraps)]

#[macro_use]
#[path = "../tests/macros/mod.rs"]
mod macros;

#[path = "../tests/common/mod.rs"]
mod common;
#[path = "../tests/repo/mod.rs"]
mod repo;

use std::fs;
use std::time::{Duration, Instant};

#[cfg(not(syn_only))]
mod tokenstream_parse {
    use proc_macro2::TokenStream;
    use std::str::FromStr;

    pub fn bench(content: &str) -> Result<(), ()> {
        TokenStream::from_str(content).map(drop).map_err(drop)
    }
}

mod syn_parse {
    pub fn bench(content: &str) -> Result<(), ()> {
        syn::parse_file(content).map(drop).map_err(drop)
    }
}

#[cfg(not(syn_only))]
mod librustc_parse {
    extern crate rustc_data_structures;
    extern crate rustc_error_messages;
    extern crate rustc_errors;
    extern crate rustc_parse;
    extern crate rustc_session;
    extern crate rustc_span;

    use rustc_data_structures::sync::Lrc;
    use rustc_error_messages::FluentBundle;
    use rustc_errors::{emitter::Emitter, translation::Translate, Diagnostic, Handler};
    use rustc_session::parse::ParseSess;
    use rustc_span::source_map::{FilePathMapping, SourceMap};
    use rustc_span::{edition::Edition, FileName};

    pub fn bench(content: &str) -> Result<(), ()> {
        struct SilentEmitter;

        impl Emitter for SilentEmitter {
            fn emit_diagnostic(&mut self, _diag: &Diagnostic) {}
            fn source_map(&self) -> Option<&Lrc<SourceMap>> {
                None
            }
        }

        impl Translate for SilentEmitter {
            fn fluent_bundle(&self) -> Option<&Lrc<FluentBundle>> {
                None
            }
            fn fallback_fluent_bundle(&self) -> &FluentBundle {
                panic!("silent emitter attempted to translate a diagnostic");
            }
        }

        rustc_span::create_session_if_not_set_then(Edition::Edition2018, |_| {
            let cm = Lrc::new(SourceMap::new(FilePathMapping::empty()));
            let emitter = Box::new(SilentEmitter);
            let handler = Handler::with_emitter(false, None, emitter);
            let sess = ParseSess::with_span_handler(handler, cm);
            if let Err(diagnostic) = rustc_parse::parse_crate_from_source_str(
                FileName::Custom("bench".to_owned()),
                content.to_owned(),
                &sess,
            ) {
                diagnostic.cancel();
                return Err(());
            };
            Ok(())
        })
    }
}

#[cfg(not(syn_only))]
mod read_from_disk {
    pub fn bench(content: &str) -> Result<(), ()> {
        let _ = content;
        Ok(())
    }
}

fn exec(mut codepath: impl FnMut(&str) -> Result<(), ()>) -> Duration {
    let begin = Instant::now();
    let mut success = 0;
    let mut total = 0;

    walkdir::WalkDir::new("tests/rust/src")
        .into_iter()
        .filter_entry(repo::base_dir_filter)
        .for_each(|entry| {
            let entry = entry.unwrap();
            let path = entry.path();
            if path.is_dir() {
                return;
            }
            let content = fs::read_to_string(path).unwrap();
            let ok = codepath(&content).is_ok();
            success += ok as usize;
            total += 1;
            if !ok {
                eprintln!("FAIL {}", path.display());
            }
        });

    assert_eq!(success, total);
    begin.elapsed()
}

fn main() {
    repo::clone_rust();

    macro_rules! testcases {
        ($($(#[$cfg:meta])* $name:ident,)*) => {
            [
                $(
                    $(#[$cfg])*
                    (stringify!($name), $name::bench as fn(&str) -> Result<(), ()>),
                )*
            ]
        };
    }

    #[cfg(not(syn_only))]
    {
        let mut lines = 0;
        let mut files = 0;
        exec(|content| {
            lines += content.lines().count();
            files += 1;
            Ok(())
        });
        eprintln!("\n{} lines in {} files", lines, files);
    }

    for (name, f) in testcases!(
        #[cfg(not(syn_only))]
        read_from_disk,
        #[cfg(not(syn_only))]
        tokenstream_parse,
        syn_parse,
        #[cfg(not(syn_only))]
        librustc_parse,
    ) {
        eprint!("{:20}", format!("{}:", name));
        let elapsed = exec(f);
        eprintln!(
            "elapsed={}.{:03}s",
            elapsed.as_secs(),
            elapsed.subsec_millis(),
        );
    }
    eprintln!();
}
<|endoftext|>


<commit_before>use av_format::error::*;
use std::io::SeekFrom;
use av_data::packet::Packet;
use av_data::timeinfo::TimeInfo;
use av_data::params::*;
use av_format::stream::Stream;
use av_format::buffer::Buffered;
use av_format::demuxer::{Demuxer, Event};
use av_format::demuxer::{Descr, Descriptor};
use av_format::common::GlobalInfo;
use std::collections::VecDeque;
use rational::Rational64;

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

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

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

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

        self.header = Some(header);

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

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

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

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

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

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

            input = i3;
        }
    }
}

use nom::Needed;

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

    fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            // println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            // println!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets();
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        _el => {
                            // println!("got element: {:#?}", el);
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                },
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(size))
                },
                e => {
                    println!("parsing issue: {:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS"   => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1"    => Some("av1".to_owned()),
        "V_VP8"    => Some("vp8".to_owned()),
        "V_VP9"    => Some("vp9".to_owned()),
        _ => None,
    }
}

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

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    use av_data::audiosample::*;
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let map = if audio.channel_positions.is_none() {
            Some(ChannelMap::default_map(audio.channels as usize))
        } else {
            unimplemented!("Convert matroska map to rust-av map")
        };
        let a = AudioInfo {
            rate: rate as usize,
            map: map,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

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

// TODO: make sure the timecode_scale isn't 0
pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = if let Some(ts) = t.track_timecode_scale  {
        (ts * info.timecode_scale as f64) as i64
    } else {
        info.timecode_scale as i64
    };

    Stream {
        id: t.track_uid as isize,
        index: t.track_number as usize,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
    }
}

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

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

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

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe<'a>(&'a self) -> &'a Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    }
};

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::Offset;
    use av_format::demuxer::Context;
    use av_format::buffer::*;

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

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

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

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

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

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => println!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => println!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

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

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

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

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

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

        self.header = Some(header);

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

        self.seek_head = None;
        self.info = None;
        self.tracks = None;

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

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

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

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

            input = i3;
        }
    }
}

use nom::Needed;

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

    fn read_event(&mut self, buf: &Box<Buffered>) -> Result<(SeekFrom, Event)> {
        if let Some(event) = self.queue.pop_front() {
            Ok((SeekFrom::Current(0), event))
        } else {
            // println!("no more stored packet, parsing a new one");
            match segment_element(buf.data()) {
                Ok((i, element)) => {
                    let seek = SeekFrom::Current(buf.data().offset(i) as i64);
                    match element {
                        SegmentElement::Cluster(c) => {
                            //self.clusters.push(c);
                            // println!("got cluster element at timecode: {}", c.timecode);
                            let mut packets = c.generate_packets();
                            self.queue.extend(packets.drain(..));
                            if let Some(event) = self.queue.pop_front() {
                                return Ok((seek, event));
                            }
                        }
                        _el => {
                            // println!("got element: {:#?}", el);
                        }
                    }
                    Ok((seek, Event::MoreDataNeeded(0)))
                },
                Err(Err::Incomplete(Needed::Size(size))) => {
                    Err(Error::MoreDataNeeded(size))
                },
                e => {
                    println!("parsing issue: {:?}", e);
                    Err(Error::InvalidData)
                }
            }
        }
    }
}

fn track_entry_codec_id(t: &TrackEntry) -> Option<String> {
    // TODO: Support V_QUICKTIME and V_MS/VFW/FOURCC
    match t.codec_id.as_ref() {
        "A_OPUS"   => Some("opus".to_owned()),
        "A_VORBIS" => Some("vorbis".to_owned()),
        "V_AV1"    => Some("av1".to_owned()),
        "V_VP8"    => Some("vp8".to_owned()),
        "V_VP9"    => Some("vp9".to_owned()),
        _ => None,
    }
}

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

fn track_entry_audio_kind(t: &TrackEntry) -> Option<MediaKind> {
    use av_data::audiosample::*;
    // TODO: Validate that a track::video exists for track::type video before.
    if let Some(ref audio) = t.audio {
        let rate = if let Some(r) = audio.output_sampling_frequency {
            r
        } else {
            audio.sampling_frequency
        };
        // TODO: complete it
        let map = if audio.channel_positions.is_none() {
            Some(ChannelMap::default_map(audio.channels as usize))
        } else {
            unimplemented!("Convert matroska map to rust-av map")
        };
        let a = AudioInfo {
            rate: rate as usize,
            map: map,
            format: None,
        };
        Some(MediaKind::Audio(a))
    } else {
        None
    }
}

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

// TODO: make sure the timecode_scale isn't 0
pub fn track_to_stream(info: &Info, t: &TrackEntry) -> Stream {
    let num = if let Some(ts) = t.track_timecode_scale  {
        (ts * info.timecode_scale as f64) as i64
    } else {
        info.timecode_scale as i64
    };

    Stream {
        id: t.track_uid as isize,
        index: t.track_number as usize,
        start: None,
        duration: t.default_duration,
        timebase: Rational64::new(num, 1000 * 1000 * 1000),
        // TODO: Extend CodecParams and fill it with the remaining information
        params: CodecParams {
            extradata: t.codec_private.clone(),
            bit_rate: 0,
            delay: t.codec_delay.unwrap_or(0) as usize,
            convergence_window: t.seek_pre_roll.unwrap_or(0) as usize,
            codec_id: track_entry_codec_id(t),
            kind: track_entry_media_kind(t),
        },
        user_private: None
    }
}

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

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

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

        v
    }
}

struct Des {
    d: Descr,
}

impl Descriptor for Des {
    fn create(&self) -> Box<Demuxer> {
        Box::new(MkvDemuxer::new())
    }
    fn describe<'a>(&'a self) -> &'a Descr {
        &self.d
    }
    fn probe(&self, data: &[u8]) -> u8 {
        match ebml_header(&data[..100]) {
            Ok(_) => 100,
            _ => 0,
        }
    }
}

pub const MKV_DESC: &Descriptor = &Des {
    d: Descr {
        name: "matroska",
        demuxer: "mkv",
        description: "Nom-based Matroska demuxer",
        extensions: &["mkv", "webm", "mka"],
        mime: &["video/x-matroska", "audio/x-matroska"],
    }
};

#[cfg(test)]
#[allow(non_upper_case_globals)]
mod tests {
    use super::*;
    use std::io::Cursor;
    use nom::Offset;
    use av_format::demuxer::Context;
    use av_format::buffer::*;

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

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

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

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

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

        for n in 100..2000 {
            let res = demuxer.parse_until_tracks(&webm[0..n]);
            match res {
                Ok(_) => println!("Size {} ok", n),
                Err(Err::Incomplete(needed)) => println!("Incomplete {} needs {:?}", n, needed),
                Err(e) => {
                    panic!("Error at size {}: {:?}", n, e);
                }
            }
        }
    }

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