<commit_before>use std::vec::Vec;

use hyper::Client;
use hyper::header::Connection;

use std::io::Read;
use std::default::Default;

use tendril::StrTendril;

use html5ever::tokenizer::{Tokenizer, TokenizerOpts};

pub fn get_rows(url: &str) -> Vec<String> {
    let mut tok = Tokenizer::new(collector::RowCollector::new(), TokenizerOpts {
        .. Default::default()
    });
    tok.feed(get_html(url));
    tok.end();
    tok.unwrap().rows
}

fn get_html(url: &str) -> StrTendril {
    let mut res = Client::new().get(url)
        .header(Connection::close()).send().unwrap();
    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();
    let input = StrTendril::try_from_byte_slice(body.as_bytes()).unwrap();
    input.try_reinterpret().unwrap()
}

mod collector {
    use std::vec::Vec;
    use std::str::FromStr;

    use html5ever::tokenizer::{TokenSink, Token, Tag};
    use html5ever::tokenizer::{TagToken, StartTag, EndTag};

    pub struct RowCollector {
        pub rows: Vec<String>
    }

    impl RowCollector {
        pub fn new() -> Self {
            RowCollector {
                rows: Vec::new(),
            }
        }

        fn get_url(&self, tag: Tag) -> Option<String> {
            if match tag.kind {
                StartTag => true,
                EndTag => false
            } && tag.name.as_ref() == "p" {
                for attr in tag.attrs.iter() {
                    if attr.name.local.as_ref() == "data-pid" {
                        return Some(String::from_str(attr.value.as_ref()).unwrap())
                    }
                }
            }

            return None
        }

        fn is_tag_token(&self, token: &Token) -> bool {
            match *token {
                TagToken(_) => true,
                _ => false
            }
        }

        fn get_tag_token(&self, token: Token) -> Tag {
            match token {
                TagToken(tag) => tag,
                _ => panic!()
            }
        }
    }

    impl TokenSink for RowCollector {
        fn process_token(&mut self, token: Token) {
            if self.is_tag_token(&token) {
                let url = self.get_url(self.get_tag_token(token));
                if url.is_some() {
                    self.rows.push(url.unwrap());
                }
            }
        }
    }
}
<commit_msg>static strings<commit_after>use std::vec::Vec;

use hyper::Client;
use hyper::header::Connection;

use std::io::Read;
use std::default::Default;

use tendril::StrTendril;

use html5ever::tokenizer::{Tokenizer, TokenizerOpts};

pub fn get_rows(url: &str) -> Vec<String> {
    let mut tok = Tokenizer::new(collector::RowCollector::new(), TokenizerOpts {
        .. Default::default()
    });
    tok.feed(get_html(url));
    tok.end();
    tok.unwrap().rows
}

fn get_html(url: &str) -> StrTendril {
    let mut res = Client::new().get(url)
        .header(Connection::close()).send().unwrap();
    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();
    let input = StrTendril::try_from_byte_slice(body.as_bytes()).unwrap();
    input.try_reinterpret().unwrap()
}

mod collector {
    use std::vec::Vec;
    use std::str::FromStr;

    use html5ever::tokenizer::{TokenSink, Token, Tag};
    use html5ever::tokenizer::{TagToken, StartTag, EndTag};

    static TARGET_TAG: &'static str = "p";
    static TARGET_ATTR: &'static str = "data-pid";

    pub struct RowCollector {
        pub rows: Vec<String>
    }

    impl RowCollector {
        pub fn new() -> Self {
            RowCollector {
                rows: Vec::new(),
            }
        }

        fn get_url(&self, tag: Tag) -> Option<String> {
            if match tag.kind {
                StartTag => true,
                EndTag => false
            } && tag.name.as_ref() == TARGET_TAG {
                for attr in tag.attrs.iter() {
                    if attr.name.local.as_ref() == TARGET_ATTR {
                        return Some(String::from_str(attr.value.as_ref()).unwrap())
                    }
                }
            }

            return None
        }

        fn is_tag_token(&self, token: &Token) -> bool {
            match *token {
                TagToken(_) => true,
                _ => false
            }
        }

        fn get_tag_token(&self, token: Token) -> Tag {
            match token {
                TagToken(tag) => tag,
                _ => panic!()
            }
        }
    }

    impl TokenSink for RowCollector {
        fn process_token(&mut self, token: Token) {
            if self.is_tag_token(&token) {
                let url = self.get_url(self.get_tag_token(token));
                if url.is_some() {
                    self.rows.push(url.unwrap());
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use events::publish;
use eyre::Result;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::Path;
use strum::{Display, EnumString, EnumVariantNames};
use tracing::Event;
use validator::Validate;

/// Logging level
#[derive(
    Debug,
    Display,
    Clone,
    Copy,
    PartialEq,
    PartialOrd,
    JsonSchema,
    Deserialize,
    Serialize,
    EnumString,
    EnumVariantNames,
)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum LoggingLevel {
    Trace,
    Debug,
    Info,
    Warn,
    Error,
    Never,
}

/// Create a `LoggingLevel` from a `tracing::Level`
impl From<&tracing::Level> for LoggingLevel {
    fn from(level: &tracing::Level) -> Self {
        match *level {
            tracing::Level::TRACE => Self::Trace,
            tracing::Level::DEBUG => Self::Debug,
            tracing::Level::INFO => Self::Info,
            tracing::Level::WARN => Self::Warn,
            tracing::Level::ERROR => Self::Error,
        }
    }
}

/// Logging format
#[derive(
    Debug, PartialEq, Clone, Copy, JsonSchema, Deserialize, Serialize, EnumString, EnumVariantNames,
)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum LoggingFormat {
    Simple,
    Detail,
    Json,
}

pub mod config {
    use super::*;
    use defaults::Defaults;
    use std::fs;
    use std::{env, path::PathBuf};

    /// Get the directory where logs are stored
    pub fn dir(ensure: bool) -> Result<PathBuf> {
        let config = crate::config::dir(false)?;
        let dir = match env::consts::OS {
            "macos" | "windows" => config.join("Logs"),
            _ => config.join("logs"),
        };
        if ensure {
            fs::create_dir_all(&dir)?;
        }
        Ok(dir)
    }

    /// Logging to standard error stream
    ///
    /// Configuration settings for log entries printed to stderr when using the CLI
    #[derive(
        Debug, Defaults, PartialEq, Clone, Copy, JsonSchema, Deserialize, Serialize, Validate,
    )]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingStdErrConfig {
        /// The maximum log level to emit
        #[def = "LoggingLevel::Info"]
        pub level: LoggingLevel,

        /// The format for the logs entries
        #[def = "LoggingFormat::Simple"]
        pub format: LoggingFormat,
    }

    /// Logging to desktop notifications
    ///
    /// Configuration settings for log entries shown to the user in the desktop
    #[derive(
        Debug, Defaults, PartialEq, Clone, Copy, JsonSchema, Deserialize, Serialize, Validate,
    )]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingDesktopConfig {
        /// The maximum log level to emit
        #[def = "LoggingLevel::Info"]
        pub level: LoggingLevel,
    }

    /// Logging to file
    ///
    /// Configuration settings for logs entries written to file
    #[derive(Debug, Defaults, PartialEq, Clone, JsonSchema, Deserialize, Serialize, Validate)]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingFileConfig {
        /// The path of the log file
        #[def = "default_file_path()"]
        pub path: String,

        /// The maximum log level to emit
        #[def = "LoggingLevel::Info"]
        pub level: LoggingLevel,
    }

    /// Get the default value for `logging.file.path`
    pub fn default_file_path() -> String {
        dir(true)
            .expect("Unable to get logs directory")
            .join("log.json")
            .into_os_string()
            .into_string()
            .expect("Unable to convert path to string")
    }

    /// Logging
    ///
    /// Configuration settings for logging
    #[derive(Debug, Default, PartialEq, Clone, JsonSchema, Deserialize, Serialize, Validate)]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingConfig {
        pub stderr: LoggingStdErrConfig,
        pub desktop: LoggingDesktopConfig,
        pub file: LoggingFileConfig,
    }
}

/// Create a preliminary logging subscriber.
///
/// This can be necessary to ensure that any log events that get emitted during
/// initialization are displayed to the user.
pub fn prelim() -> tracing::subscriber::DefaultGuard {
    let subscriber = tracing_subscriber::fmt()
        .pretty()
        .with_max_level(tracing::Level::INFO)
        .with_writer(std::io::stderr)
        .finish();
    tracing::subscriber::set_default(subscriber)
}

/// Custom tracing_subscriber layer that prints events to stderr filtered
/// by level for the "plain" format. Other formats are handled by `tracing_subscriber`
/// formatters (see below).
struct StderrPlainLayer {
    level: LoggingLevel,
}

#[derive(Default)]
struct StderrVisitor {
    message: String,
}

impl tracing::field::Visit for StderrVisitor {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        if field.name() == "message" {
            self.message = format!("{:?}", value);
        }
    }
}

impl<S: tracing::subscriber::Subscriber> tracing_subscriber::layer::Layer<S> for StderrPlainLayer {
    fn on_event(&self, event: &Event<'_>, _ctx: tracing_subscriber::layer::Context<'_, S>) {
        let level = LoggingLevel::from(event.metadata().level());
        if level >= self.level {
            let level_name = level.to_string().to_uppercase();

            #[cfg(feature = "ansi_term")]
            let level_name = {
                use ansi_term::Color::{Blue, Green, Purple, Red, White, Yellow};
                match level {
                    LoggingLevel::Trace => Purple,
                    LoggingLevel::Debug => Blue,
                    LoggingLevel::Info => Green,
                    LoggingLevel::Warn => Yellow,
                    LoggingLevel::Error => Red,
                    _ => White,
                }
                .bold()
                .paint(format!("{:5}", level_name))
            };

            let mut visitor = StderrVisitor::default();
            event.record(&mut visitor);
            eprintln!("{} {}", level_name, visitor.message)
        }
    }
}

/// Custom tracing_subscriber layer that publishes events
/// under the pubsub "logging" topic as a JSON value.
struct PubSubLayer {
    level: LoggingLevel,
}

impl<S: tracing::subscriber::Subscriber> tracing_subscriber::layer::Layer<S> for PubSubLayer {
    fn on_event(&self, event: &Event<'_>, _ctx: tracing_subscriber::layer::Context<'_, S>) {
        use tracing_serde::AsSerde;
        if LoggingLevel::from(event.metadata().level()) >= self.level {
            publish("logging", &event.as_serde())
        }
    }
}

/// Initialize logging
///
/// This initializes a logging subscriber based on configuration and
/// context (e.g. stderr should not be written to if the context
/// is the desktop application).
///
/// # Arguments
///
/// - `stderr`: should stderr logging be enabled
/// - `pubsub`: should pubsub logging be enabled (for desktop notifications)
/// - `file`: should file logging be enabled
/// - `config`: the logging configuration
pub fn init(
    stderr: bool,
    pubsub: bool,
    file: bool,
    config: &config::LoggingConfig,
) -> Result<tracing_appender::non_blocking::WorkerGuard> {
    use tracing_error::ErrorLayer;
    use tracing_subscriber::prelude::*;
    use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter};

    // Stderr logging layer
    let stderr_level = if stderr {
        config.stderr.level
    } else {
        LoggingLevel::Never
    };

    // Pubsub logging layer (used for desktop notifications)
    let pubsub_level = if pubsub {
        config.desktop.level
    } else {
        LoggingLevel::Never
    };
    let pubsub_layer = PubSubLayer {
        level: pubsub_level,
    };

    // File logging layer
    let file_level = if file {
        config.file.level
    } else {
        LoggingLevel::Never
    };
    let (file_writer, file_guard) = if file_level != LoggingLevel::Never {
        let path = Path::new(&config.file.path);
        let file_appender =
            tracing_appender::rolling::daily(&path.parent().unwrap(), &path.file_name().unwrap());
        tracing_appender::non_blocking(file_appender)
    } else {
        tracing_appender::non_blocking(std::io::sink())
    };
    let file_layer = fmt::Layer::new().json().with_writer(file_writer);

    // Error reporting layer (necessary for using `eyre` crate)
    let error_layer = ErrorLayer::default();

    // tracing_subscriber does not currently allow for different layers to have different
    // levels so work out the minimum level and filter by that in the root subscriber.
    let mut min_level = LoggingLevel::Never;
    if stderr_level < min_level {
        min_level = stderr_level
    }
    if pubsub_level < min_level {
        min_level = pubsub_level
    }
    if file_level < min_level {
        min_level = file_level
    }

    // Filter out debug log entries from some crates to avoid excessive noise.
    // We may want to show entries from other crates during development
    // so we may add another flag for this in the future.
    // e.g. `--log-scope=stencila` vs `--log-scope=all`.
    let directives = format!(
        "{},reqwest=info,hyper=info,warp=info",
        min_level.to_string()
    );

    let registry = tracing_subscriber::registry()
        .with(EnvFilter::new(directives))
        .with(pubsub_layer)
        .with(file_layer)
        .with(error_layer);

    if config.stderr.format == LoggingFormat::Detail {
        registry.with(fmt::Layer::new().pretty()).init();
    } else if config.stderr.format == LoggingFormat::Json {
        registry.with(fmt::Layer::new().json()).init();
    } else {
        registry
            .with(StderrPlainLayer {
                level: stderr_level,
            })
            .init();
    }

    Ok(file_guard)
}

/// Generate some test tracing events.
///
/// Can be used for testing that events are propagated
/// to subscribers as expected.
#[tracing::instrument]
pub fn test_events() {
    tracing::trace!("A trace event");
    tracing::debug!("A debug event");
    tracing::info!("An info event");
    tracing::warn!("A warn event");
    tracing::error!("An error event");
}

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

    #[test]
    #[allow(clippy::eq_op)]
    fn test_logging_level_ordering() -> Result<()> {
        assert!(LoggingLevel::Debug > LoggingLevel::Trace);
        assert!(LoggingLevel::Debug >= LoggingLevel::Debug);
        assert!(LoggingLevel::Debug == LoggingLevel::Debug);
        assert!(LoggingLevel::Trace <= LoggingLevel::Debug);
        assert!(LoggingLevel::Trace < LoggingLevel::Debug);

        assert!(LoggingLevel::Info > LoggingLevel::Debug);
        assert!(LoggingLevel::Warn > LoggingLevel::Info);
        assert!(LoggingLevel::Error > LoggingLevel::Warn);
        Ok(())
    }
}
<commit_msg>chore(Logging): Filter html5ever debug logs<commit_after>use events::publish;
use eyre::Result;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::Path;
use strum::{Display, EnumString, EnumVariantNames};
use tracing::Event;
use validator::Validate;

/// Logging level
#[derive(
    Debug,
    Display,
    Clone,
    Copy,
    PartialEq,
    PartialOrd,
    JsonSchema,
    Deserialize,
    Serialize,
    EnumString,
    EnumVariantNames,
)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum LoggingLevel {
    Trace,
    Debug,
    Info,
    Warn,
    Error,
    Never,
}

/// Create a `LoggingLevel` from a `tracing::Level`
impl From<&tracing::Level> for LoggingLevel {
    fn from(level: &tracing::Level) -> Self {
        match *level {
            tracing::Level::TRACE => Self::Trace,
            tracing::Level::DEBUG => Self::Debug,
            tracing::Level::INFO => Self::Info,
            tracing::Level::WARN => Self::Warn,
            tracing::Level::ERROR => Self::Error,
        }
    }
}

/// Logging format
#[derive(
    Debug, PartialEq, Clone, Copy, JsonSchema, Deserialize, Serialize, EnumString, EnumVariantNames,
)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
pub enum LoggingFormat {
    Simple,
    Detail,
    Json,
}

pub mod config {
    use super::*;
    use defaults::Defaults;
    use std::fs;
    use std::{env, path::PathBuf};

    /// Get the directory where logs are stored
    pub fn dir(ensure: bool) -> Result<PathBuf> {
        let config = crate::config::dir(false)?;
        let dir = match env::consts::OS {
            "macos" | "windows" => config.join("Logs"),
            _ => config.join("logs"),
        };
        if ensure {
            fs::create_dir_all(&dir)?;
        }
        Ok(dir)
    }

    /// Logging to standard error stream
    ///
    /// Configuration settings for log entries printed to stderr when using the CLI
    #[derive(
        Debug, Defaults, PartialEq, Clone, Copy, JsonSchema, Deserialize, Serialize, Validate,
    )]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingStdErrConfig {
        /// The maximum log level to emit
        #[def = "LoggingLevel::Info"]
        pub level: LoggingLevel,

        /// The format for the logs entries
        #[def = "LoggingFormat::Simple"]
        pub format: LoggingFormat,
    }

    /// Logging to desktop notifications
    ///
    /// Configuration settings for log entries shown to the user in the desktop
    #[derive(
        Debug, Defaults, PartialEq, Clone, Copy, JsonSchema, Deserialize, Serialize, Validate,
    )]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingDesktopConfig {
        /// The maximum log level to emit
        #[def = "LoggingLevel::Info"]
        pub level: LoggingLevel,
    }

    /// Logging to file
    ///
    /// Configuration settings for logs entries written to file
    #[derive(Debug, Defaults, PartialEq, Clone, JsonSchema, Deserialize, Serialize, Validate)]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingFileConfig {
        /// The path of the log file
        #[def = "default_file_path()"]
        pub path: String,

        /// The maximum log level to emit
        #[def = "LoggingLevel::Info"]
        pub level: LoggingLevel,
    }

    /// Get the default value for `logging.file.path`
    pub fn default_file_path() -> String {
        dir(true)
            .expect("Unable to get logs directory")
            .join("log.json")
            .into_os_string()
            .into_string()
            .expect("Unable to convert path to string")
    }

    /// Logging
    ///
    /// Configuration settings for logging
    #[derive(Debug, Default, PartialEq, Clone, JsonSchema, Deserialize, Serialize, Validate)]
    #[serde(default)]
    #[schemars(deny_unknown_fields)]
    pub struct LoggingConfig {
        pub stderr: LoggingStdErrConfig,
        pub desktop: LoggingDesktopConfig,
        pub file: LoggingFileConfig,
    }
}

/// Create a preliminary logging subscriber.
///
/// This can be necessary to ensure that any log events that get emitted during
/// initialization are displayed to the user.
pub fn prelim() -> tracing::subscriber::DefaultGuard {
    let subscriber = tracing_subscriber::fmt()
        .pretty()
        .with_max_level(tracing::Level::INFO)
        .with_writer(std::io::stderr)
        .finish();
    tracing::subscriber::set_default(subscriber)
}

/// Custom tracing_subscriber layer that prints events to stderr filtered
/// by level for the "plain" format. Other formats are handled by `tracing_subscriber`
/// formatters (see below).
struct StderrPlainLayer {
    level: LoggingLevel,
}

#[derive(Default)]
struct StderrVisitor {
    message: String,
}

impl tracing::field::Visit for StderrVisitor {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        if field.name() == "message" {
            self.message = format!("{:?}", value);
        }
    }
}

impl<S: tracing::subscriber::Subscriber> tracing_subscriber::layer::Layer<S> for StderrPlainLayer {
    fn on_event(&self, event: &Event<'_>, _ctx: tracing_subscriber::layer::Context<'_, S>) {
        let level = LoggingLevel::from(event.metadata().level());
        if level >= self.level {
            let level_name = level.to_string().to_uppercase();

            #[cfg(feature = "ansi_term")]
            let level_name = {
                use ansi_term::Color::{Blue, Green, Purple, Red, White, Yellow};
                match level {
                    LoggingLevel::Trace => Purple,
                    LoggingLevel::Debug => Blue,
                    LoggingLevel::Info => Green,
                    LoggingLevel::Warn => Yellow,
                    LoggingLevel::Error => Red,
                    _ => White,
                }
                .bold()
                .paint(format!("{:5}", level_name))
            };

            let mut visitor = StderrVisitor::default();
            event.record(&mut visitor);
            eprintln!("{} {}", level_name, visitor.message)
        }
    }
}

/// Custom tracing_subscriber layer that publishes events
/// under the pubsub "logging" topic as a JSON value.
struct PubSubLayer {
    level: LoggingLevel,
}

impl<S: tracing::subscriber::Subscriber> tracing_subscriber::layer::Layer<S> for PubSubLayer {
    fn on_event(&self, event: &Event<'_>, _ctx: tracing_subscriber::layer::Context<'_, S>) {
        use tracing_serde::AsSerde;
        if LoggingLevel::from(event.metadata().level()) >= self.level {
            publish("logging", &event.as_serde())
        }
    }
}

/// Initialize logging
///
/// This initializes a logging subscriber based on configuration and
/// context (e.g. stderr should not be written to if the context
/// is the desktop application).
///
/// # Arguments
///
/// - `stderr`: should stderr logging be enabled
/// - `pubsub`: should pubsub logging be enabled (for desktop notifications)
/// - `file`: should file logging be enabled
/// - `config`: the logging configuration
pub fn init(
    stderr: bool,
    pubsub: bool,
    file: bool,
    config: &config::LoggingConfig,
) -> Result<tracing_appender::non_blocking::WorkerGuard> {
    use tracing_error::ErrorLayer;
    use tracing_subscriber::prelude::*;
    use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter};

    // Stderr logging layer
    let stderr_level = if stderr {
        config.stderr.level
    } else {
        LoggingLevel::Never
    };

    // Pubsub logging layer (used for desktop notifications)
    let pubsub_level = if pubsub {
        config.desktop.level
    } else {
        LoggingLevel::Never
    };
    let pubsub_layer = PubSubLayer {
        level: pubsub_level,
    };

    // File logging layer
    let file_level = if file {
        config.file.level
    } else {
        LoggingLevel::Never
    };
    let (file_writer, file_guard) = if file_level != LoggingLevel::Never {
        let path = Path::new(&config.file.path);
        let file_appender =
            tracing_appender::rolling::daily(&path.parent().unwrap(), &path.file_name().unwrap());
        tracing_appender::non_blocking(file_appender)
    } else {
        tracing_appender::non_blocking(std::io::sink())
    };
    let file_layer = fmt::Layer::new().json().with_writer(file_writer);

    // Error reporting layer (necessary for using `eyre` crate)
    let error_layer = ErrorLayer::default();

    // tracing_subscriber does not currently allow for different layers to have different
    // levels so work out the minimum level and filter by that in the root subscriber.
    let mut min_level = LoggingLevel::Never;
    if stderr_level < min_level {
        min_level = stderr_level
    }
    if pubsub_level < min_level {
        min_level = pubsub_level
    }
    if file_level < min_level {
        min_level = file_level
    }

    // Filter out debug log entries from some crates to avoid excessive noise.
    // We may want to show entries from other crates during development
    // so we may add another flag for this in the future.
    // e.g. `--log-scope=stencila` vs `--log-scope=all`.
    let directives = format!(
        "{},reqwest=info,html5ever=info,hyper=info,warp=info",
        min_level.to_string()
    );

    let registry = tracing_subscriber::registry()
        .with(EnvFilter::new(directives))
        .with(pubsub_layer)
        .with(file_layer)
        .with(error_layer);

    if config.stderr.format == LoggingFormat::Detail {
        registry.with(fmt::Layer::new().pretty()).init();
    } else if config.stderr.format == LoggingFormat::Json {
        registry.with(fmt::Layer::new().json()).init();
    } else {
        registry
            .with(StderrPlainLayer {
                level: stderr_level,
            })
            .init();
    }

    Ok(file_guard)
}

/// Generate some test tracing events.
///
/// Can be used for testing that events are propagated
/// to subscribers as expected.
#[tracing::instrument]
pub fn test_events() {
    tracing::trace!("A trace event");
    tracing::debug!("A debug event");
    tracing::info!("An info event");
    tracing::warn!("A warn event");
    tracing::error!("An error event");
}

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

    #[test]
    #[allow(clippy::eq_op)]
    fn test_logging_level_ordering() -> Result<()> {
        assert!(LoggingLevel::Debug > LoggingLevel::Trace);
        assert!(LoggingLevel::Debug >= LoggingLevel::Debug);
        assert!(LoggingLevel::Debug == LoggingLevel::Debug);
        assert!(LoggingLevel::Trace <= LoggingLevel::Debug);
        assert!(LoggingLevel::Trace < LoggingLevel::Debug);

        assert!(LoggingLevel::Info > LoggingLevel::Debug);
        assert!(LoggingLevel::Warn > LoggingLevel::Info);
        assert!(LoggingLevel::Error > LoggingLevel::Warn);
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Provides a Quaternion type for properly interpolating rotations

use std::f32;
use std::ops::{Add, Sub, Mul, Div};

use linalg::{self, Vector, Transform, Matrix4};

/// Quaternions describe a rotation in 3d space but can be
/// properly interpolated unlike rotation matrices. The quaternion
/// also provides methods to compute the corresponding rotation
/// transform so that it can be composed with scaling and translation
#[derive(Debug, Copy, Clone)]
pub struct Quaternion {
    pub v: Vector,
    pub w: f32,
}

impl Quaternion {
    /// Construct an identity quaternion
    pub fn identity() -> Quaternion {
        Quaternion { v: Vector::broadcast(0.0), w: 1.0 }
    }
    /// Construct a quaternion from the rotation matrix
    /// Based on Shoemake 1991
    pub fn from_matrix(m: &Matrix4) -> Quaternion {
        let trace = *m.at(0, 0) + *m.at(1, 1) + *m.at(2, 2);
        if trace > 0.0 {
            // Compute w from matrix trace, then the vector
            let mut s = f32::sqrt(trace + 1.0);
            let w = s / 2.0;
            s = 0.5 / s;
            Quaternion { v: Vector::new(s * (*m.at(2, 1) - *m.at(1, 2)), s * (*m.at(0, 2) - *m.at(2, 0)),
                                        s * (*m.at(1, 0) - *m.at(0, 1))),
                         w: w
            }
        } else {
            // Compute largest of x, y or z then the remaining components
            let next = [1, 2, 0];
            let mut q = Vector::broadcast(0.0);
            let mut i = 0;
            if *m.at(1, 1) > *m.at(0, 0) {
                i = 1;
            }
            if *m.at(2, 2) > *m.at(i, i) {
                i = 2;
            }
            let j = next[i];
            let k = next[j];
            let mut s = f32::sqrt((*m.at(i, i) - (*m.at(j, j) + *m.at(k, k))) + 1.0);
            q[i] = s * 0.5;
            if s != 0.0 {
                s = 0.5 / s;
            }
            let w = (*m.at(k, j) - *m.at(j, k)) * s;
            q[j] = (*m.at(j, i) + *m.at(i, j)) * s;
            q[k] = (*m.at(k, i) + *m.at(i, k)) * s;
            Quaternion { v: q, w: w }
        }
    }
    /// Construct the quaternion from the transform
    pub fn from_transform(t: &Transform) -> Quaternion {
        Quaternion::from_matrix(&t.mat)
    }
    /// Get the rotation transform described by this quaternion
    pub fn to_transform(&self) -> Transform {
        Transform::from_mat(&Matrix4::new(
            [1.0 - 2.0 * (f32::powf(self.v.y, 2.0) + f32::powf(self.v.x, 2.0)),
             2.0 * (self.v.x * self.v.y + self.v.z * self.w),
             2.0 * (self.v.x * self.v.z - self.v.y * self.w),
             0.0,

             2.0 * (self.v.x * self.v.y - self.v.z * self.w),
             1.0 - 2.0 * (f32::powf(self.v.x, 2.0) + f32::powf(self.v.z, 2.0)),
             2.0 * (self.v.y * self.v.z + self.v.x * self.w),
             0.0,

             2.0 * (self.v.x * self.v.z + self.v.y * self.w),
             2.0 * (self.v.y * self.v.z - self.v.x * self.w),
             1.0 - 2.0 * (f32::powf(self.v.x, 2.0) + f32::powf(self.v.y, 2.0)),
             0.0,

             0.0, 0.0, 0.0, 1.0
            ]))
    }
    /// Get the normalized quaternion for this rotation
    pub fn normalized(&self) -> Quaternion {
        *self / f32::sqrt(dot(self, self))
    }
}

/// Compute the dot product of the two quaternions
pub fn dot(a: &Quaternion, b: &Quaternion) -> f32 {
    linalg::dot(&a.v, &b.v) + a.w * b.w
}

/// Use spherical linear interpolation to interpolate between the two quaternions
pub fn slerp(t: f32, a: &Quaternion, b: &Quaternion) -> Quaternion {
    // Check if a and b are nearly parallel. To avoid numerical instability we do
    // regular linear interpolation in this case
    let cos_theta = dot(a, b);
    if cos_theta > 0.9995 {
        ((1.0 - t) * *a + t * *b).normalized()
    } else {
        let theta = f32::acos(linalg::clamp(cos_theta, -1.0, 1.0));
        let theta_t = theta * t;
        let q_perp = (*b - *a * cos_theta).normalized();
        *a * f32::cos(theta_t) + q_perp * f32::sin(theta_t)
    }
}

impl Add for Quaternion {
    type Output = Quaternion;
    /// Add two quaternions
    fn add(self, rhs: Quaternion) -> Quaternion {
        Quaternion { v: self.v + rhs.v, w: self.w + rhs.w }
    }
}

impl Sub for Quaternion {
    type Output = Quaternion;
    /// Subtract two quaternions
    fn sub(self, rhs: Quaternion) -> Quaternion {
        Quaternion { v: self.v - rhs.v, w: self.w - rhs.w }
    }
}

impl Mul<f32> for Quaternion {
    type Output = Quaternion;
    /// Multiply the quaternion by a scalar
    fn mul(self, rhs: f32) -> Quaternion {
        Quaternion { v: self.v * rhs, w: self.w * rhs }
    }
}

impl Mul<Quaternion> for f32 {
    type Output = Quaternion;
    /// Multiply the quaternion by a scalar
    fn mul(self, rhs: Quaternion) -> Quaternion {
        Quaternion { v: self * rhs.v, w: self * rhs.w }
    }
}

impl Div<f32> for Quaternion {
    type Output = Quaternion;
    /// Divide the quaternion by a scalar
    fn div(self, rhs: f32) -> Quaternion {
        Quaternion { v: self.v * rhs, w: self. w * rhs }
    }
}

<commit_msg>Add method to go from quat to matrix<commit_after>//! Provides a Quaternion type for properly interpolating rotations

use std::f32;
use std::ops::{Add, Sub, Mul, Div};

use linalg::{self, Vector, Transform, Matrix4};

/// Quaternions describe a rotation in 3d space but can be
/// properly interpolated unlike rotation matrices. The quaternion
/// also provides methods to compute the corresponding rotation
/// transform so that it can be composed with scaling and translation
#[derive(Debug, Copy, Clone)]
pub struct Quaternion {
    pub v: Vector,
    pub w: f32,
}

impl Quaternion {
    /// Construct an identity quaternion
    pub fn identity() -> Quaternion {
        Quaternion { v: Vector::broadcast(0.0), w: 1.0 }
    }
    /// Construct a quaternion from the rotation matrix
    /// Based on Shoemake 1991
    pub fn from_matrix(m: &Matrix4) -> Quaternion {
        let trace = *m.at(0, 0) + *m.at(1, 1) + *m.at(2, 2);
        if trace > 0.0 {
            // Compute w from matrix trace, then the vector
            let mut s = f32::sqrt(trace + 1.0);
            let w = s / 2.0;
            s = 0.5 / s;
            Quaternion { v: Vector::new(s * (*m.at(2, 1) - *m.at(1, 2)), s * (*m.at(0, 2) - *m.at(2, 0)),
                                        s * (*m.at(1, 0) - *m.at(0, 1))),
                         w: w
            }
        } else {
            // Compute largest of x, y or z then the remaining components
            let next = [1, 2, 0];
            let mut q = Vector::broadcast(0.0);
            let mut i = 0;
            if *m.at(1, 1) > *m.at(0, 0) {
                i = 1;
            }
            if *m.at(2, 2) > *m.at(i, i) {
                i = 2;
            }
            let j = next[i];
            let k = next[j];
            let mut s = f32::sqrt((*m.at(i, i) - (*m.at(j, j) + *m.at(k, k))) + 1.0);
            q[i] = s * 0.5;
            if s != 0.0 {
                s = 0.5 / s;
            }
            let w = (*m.at(k, j) - *m.at(j, k)) * s;
            q[j] = (*m.at(j, i) + *m.at(i, j)) * s;
            q[k] = (*m.at(k, i) + *m.at(i, k)) * s;
            Quaternion { v: q, w: w }
        }
    }
    /// Construct the quaternion from the transform
    pub fn from_transform(t: &Transform) -> Quaternion {
        Quaternion::from_matrix(&t.mat)
    }
    /// Get the rotation transform described by this quaternion
    pub fn to_matrix(&self) -> Matrix4 {
        Matrix4::new(
            [1.0 - 2.0 * (f32::powf(self.v.y, 2.0) + f32::powf(self.v.x, 2.0)),
             2.0 * (self.v.x * self.v.y + self.v.z * self.w),
             2.0 * (self.v.x * self.v.z - self.v.y * self.w),
             0.0,

             2.0 * (self.v.x * self.v.y - self.v.z * self.w),
             1.0 - 2.0 * (f32::powf(self.v.x, 2.0) + f32::powf(self.v.z, 2.0)),
             2.0 * (self.v.y * self.v.z + self.v.x * self.w),
             0.0,

             2.0 * (self.v.x * self.v.z + self.v.y * self.w),
             2.0 * (self.v.y * self.v.z - self.v.x * self.w),
             1.0 - 2.0 * (f32::powf(self.v.x, 2.0) + f32::powf(self.v.y, 2.0)),
             0.0,

             0.0, 0.0, 0.0, 1.0
            ])
    }
    /// Get the rotation transform described by this quaternion
    pub fn to_transform(&self) -> Transform {
        Transform::from_mat(&self.to_matrix())
    }
    /// Get the normalized quaternion for this rotation
    pub fn normalized(&self) -> Quaternion {
        *self / f32::sqrt(dot(self, self))
    }
}

/// Compute the dot product of the two quaternions
pub fn dot(a: &Quaternion, b: &Quaternion) -> f32 {
    linalg::dot(&a.v, &b.v) + a.w * b.w
}

/// Use spherical linear interpolation to interpolate between the two quaternions
pub fn slerp(t: f32, a: &Quaternion, b: &Quaternion) -> Quaternion {
    // Check if a and b are nearly parallel. To avoid numerical instability we do
    // regular linear interpolation in this case
    let cos_theta = dot(a, b);
    if cos_theta > 0.9995 {
        ((1.0 - t) * *a + t * *b).normalized()
    } else {
        let theta = f32::acos(linalg::clamp(cos_theta, -1.0, 1.0));
        let theta_t = theta * t;
        let q_perp = (*b - *a * cos_theta).normalized();
        *a * f32::cos(theta_t) + q_perp * f32::sin(theta_t)
    }
}

impl Add for Quaternion {
    type Output = Quaternion;
    /// Add two quaternions
    fn add(self, rhs: Quaternion) -> Quaternion {
        Quaternion { v: self.v + rhs.v, w: self.w + rhs.w }
    }
}

impl Sub for Quaternion {
    type Output = Quaternion;
    /// Subtract two quaternions
    fn sub(self, rhs: Quaternion) -> Quaternion {
        Quaternion { v: self.v - rhs.v, w: self.w - rhs.w }
    }
}

impl Mul<f32> for Quaternion {
    type Output = Quaternion;
    /// Multiply the quaternion by a scalar
    fn mul(self, rhs: f32) -> Quaternion {
        Quaternion { v: self.v * rhs, w: self.w * rhs }
    }
}

impl Mul<Quaternion> for f32 {
    type Output = Quaternion;
    /// Multiply the quaternion by a scalar
    fn mul(self, rhs: Quaternion) -> Quaternion {
        Quaternion { v: self * rhs.v, w: self * rhs.w }
    }
}

impl Div<f32> for Quaternion {
    type Output = Quaternion;
    /// Divide the quaternion by a scalar
    fn div(self, rhs: f32) -> Quaternion {
        Quaternion { v: self.v * rhs, w: self. w * rhs }
    }
}

<|endoftext|>


<commit_before>use nom::{Consumer, FileProducer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use metadata::{
  Block, BlockData,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
};
use metadata::types::MetaDataConsumer;

// Will return true when the unwrapped value of `$option` and `$compare`
// match or `$option` is `Option::None`, otherwise false.
macro_rules! optional_eq (
  ($compare: expr, $option: expr) => (
    $option.map_or(true, |compare| $compare == compare);
  );
);

// With the given filename, return all metadata blocks available.
//
// This function expects a flac file, but will return a proper `Result::Err`
// when things go wrong.
//
// # Failures
//
// * `ErrorKind::NotFound` is returned when the given filename isn't found.
// * `ErrorKind::InvalidData` is returned when the data within the file
//   isn't valid FLAC data.
pub fn get_metadata(filename: &str) -> Result<Vec<Block>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer = MetaDataConsumer::new();

    consumer.run(&mut producer);

    if !consumer.data.is_empty() {
      Ok(consumer.data)
    } else {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_stream_info("path/to/file.flac") {
///   Ok(stream_info) => {
///     // Use the stream_info variable...
///   }
///   Err(error)      => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let stream_info = metadata::get_stream_info("path/to/file.flac").unwrap();
/// ```
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_vorbis_comment("path/to/file.flac") {
///   Ok(vorbis_comment) => {
///     // Use the vorbis_comment variable...
///   }
///   Err(error)         => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let vorbis_comment =
///   metadata::get_vorbis_comment("path/to/file.flac").unwrap();
/// ```
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file and this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// Putting `None` into any of the optional constraints conveys that you
/// want any of that parameter. Otherwise it will try to look for the image
/// that matches within the given constraints.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let BlockData::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq!(picture.picture_type, picture_type) &&
           optional_eq!(picture.mime_type, mime_type) &&
           optional_eq!(picture.description, description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

#[cfg(test)]
mod tests {
  use super::*;
  use std::io::ErrorKind;

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq!(0, Some(1)));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq!(0, None), "Should always return true when None");
    assert!(optional_eq!(0, Some(0)), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<commit_msg>Add examples for `get_cue_sheet`<commit_after>use nom::{Consumer, FileProducer};
use std::io::{Error, ErrorKind, Result};
use std::u32;

use metadata::{
  Block, BlockData,
  StreamInfo, CueSheet, VorbisComment, Picture,
  PictureType,
};
use metadata::types::MetaDataConsumer;

// Will return true when the unwrapped value of `$option` and `$compare`
// match or `$option` is `Option::None`, otherwise false.
macro_rules! optional_eq (
  ($compare: expr, $option: expr) => (
    $option.map_or(true, |compare| $compare == compare);
  );
);

// With the given filename, return all metadata blocks available.
//
// This function expects a flac file, but will return a proper `Result::Err`
// when things go wrong.
//
// # Failures
//
// * `ErrorKind::NotFound` is returned when the given filename isn't found.
// * `ErrorKind::InvalidData` is returned when the data within the file
//   isn't valid FLAC data.
pub fn get_metadata(filename: &str) -> Result<Vec<Block>> {
  FileProducer::new(filename, 1024).and_then(|mut producer| {
    let mut consumer = MetaDataConsumer::new();

    consumer.run(&mut producer);

    if !consumer.data.is_empty() {
      Ok(consumer.data)
    } else {
      let error_str = "parser: couldn't find any metadata";

      Err(Error::new(ErrorKind::InvalidData, error_str))
    }
  })
}

/// Reads and returns the `StreamInfo` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `StreamInfo` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_stream_info("path/to/file.flac") {
///   Ok(stream_info) => {
///     // Use the stream_info variable...
///   }
///   Err(error)      => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let stream_info = metadata::get_stream_info("path/to/file.flac").unwrap();
/// ```
pub fn get_stream_info(filename: &str) -> Result<StreamInfo> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find StreamInfo";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::StreamInfo(stream_info) = block.data {
        result = Ok(stream_info);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `VorbisComment` metadata block of the given FLAC
/// file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `VorbisComment` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_vorbis_comment("path/to/file.flac") {
///   Ok(vorbis_comment) => {
///     // Use the vorbis_comment variable...
///   }
///   Err(error)         => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors:
///
/// ```no_run
/// use flac::metadata;
///
/// let vorbis_comment =
///   metadata::get_vorbis_comment("path/to/file.flac").unwrap();
/// ```
pub fn get_vorbis_comment(filename: &str) -> Result<VorbisComment> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find VorbisComment";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::VorbisComment(vorbis_comment) = block.data {
        result = Ok(vorbis_comment);
        break;
      }
    }

    result
  })
}

/// Reads and returns the `CueSheet` metadata block of the given FLAC file.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found
///   or there is no `CueSheet` within the file.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
///
/// # Examples
///
/// Handling errors might look something like this:
///
/// ```
/// use flac::metadata;
///
/// match metadata::get_cue_sheet("path/to/file.flac") {
///   Ok(cue_sheet) => {
///     // Use the cue_sheet variable...
///   }
///   Err(error)    => println!("{}", error),
/// }
/// ```
///
/// Or just ignore the errors
///
/// ```no_run
/// use flac::metadata;
///
/// let cue_sheet = metadata::get_cue_sheet("path/to/file.flac").unwrap();
/// ```
pub fn get_cue_sheet(filename: &str) -> Result<CueSheet> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find CueSheet";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    for block in blocks {
      if let BlockData::CueSheet(cue_sheet) = block.data {
        result = Ok(cue_sheet);
        break;
      }
    }

    result
  })
}

/// Reads and returns a `Picture` metadata block of the given FLAC file.
///
/// There can be more than one `Picture` block in a file and this function
/// takes optional, that being `Option<T>`, parameters that act as
/// constraints to search within. The `Picture` with the largest area
/// matching all constraints will be returned.
///
/// Putting `None` into any of the optional constraints conveys that you
/// want any of that parameter. Otherwise it will try to look for the image
/// that matches within the given constraints.
///
/// # Failures
///
/// * `ErrorKind::NotFound` is returned when the given filename isn't found,
///   there is no `Picture` within the file, or no `Picture` that fits the
///   given constraints.
/// * `ErrorKind::InvalidData` is returned when the data within the file
///   isn't valid FLAC data.
pub fn get_picture(filename: &str,
                   picture_type: Option<PictureType>,
                   mime_type: Option<&str>,
                   description: Option<&str>,
                   max_width: Option<u32>,
                   max_height: Option<u32>,
                   max_depth: Option<u32>,
                   max_colors: Option<u32>)
                   -> Result<Picture> {
  get_metadata(filename).and_then(|blocks| {
    let error_str  = "metadata: couldn't find any Picture";
    let mut result = Err(Error::new(ErrorKind::NotFound, error_str));

    let mut max_area_seen  = 0;
    let mut max_depth_seen = 0;

    let max_value      = u32::max_value();
    let max_width_num  = max_width.unwrap_or(max_value);
    let max_height_num = max_height.unwrap_or(max_value);
    let max_depth_num  = max_depth.unwrap_or(max_value);
    let max_colors_num = max_colors.unwrap_or(max_value);

    for block in blocks {
      if let BlockData::Picture(picture) = block.data {
        let area = (picture.width as u64) * (picture.height as u64);

        if optional_eq!(picture.picture_type, picture_type) &&
           optional_eq!(picture.mime_type, mime_type) &&
           optional_eq!(picture.description, description) &&
           picture.width <= max_width_num &&
           picture.height <= max_height_num &&
           picture.depth <= max_depth_num &&
           picture.colors <= max_colors_num &&
           (area > max_area_seen || (area == max_area_seen &&
                                     picture.depth > max_depth_seen)) {
          max_area_seen  = area;
          max_depth_seen = picture.depth;
          result         = Ok(picture);
        }
      }
    }

    result
  })
}

#[cfg(test)]
mod tests {
  use super::*;
  use std::io::ErrorKind;

  #[test]
  #[should_panic]
  fn test_panic_optional_eq() {
    assert!(optional_eq!(0, Some(1)));
  }

  #[test]
  fn test_optional_eq() {
    assert!(optional_eq!(0, None), "Should always return true when None");
    assert!(optional_eq!(0, Some(0)), "Should return true (Some(0) == 0)");
  }

  #[test]
  fn test_get_metadata() {
    let not_found    = get_metadata("non-existent/file.txt");
    let invalid_data = get_metadata("README.md");
    let result       = get_metadata("tests/assets/input-SVAUP.flac");

    assert_eq!(not_found.unwrap_err().kind(), ErrorKind::NotFound);
    assert_eq!(invalid_data.unwrap_err().kind(), ErrorKind::InvalidData);
    assert!(result.is_ok());
  }
}
<|endoftext|>


<commit_before>use std::io;

use super::datagram::{DatagramReceiver, ReadAdapter};
use super::ipv4_header::{IPv4Header, IPv4HeaderData, IPv4HeaderMut};
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::transport_header::{TransportHeader, TransportHeaderData, TransportHeaderMut};

/// Convert from level 5 to level 3 by appending correct IP and transport headers.
pub struct Packetizer {
    buffer: Box<[u8; MAX_PACKET_LENGTH]>,
    transport_index: usize,
    payload_index: usize,
    ipv4_header_data: IPv4HeaderData,
    transport_header_data: TransportHeaderData,
    last_packet_length: Option<usize>,
}

impl Packetizer {
    pub fn new(reference_ipv4_header: &IPv4Header, reference_transport_header: &TransportHeader) -> Self {
        let mut buffer = Box::new([0; MAX_PACKET_LENGTH]);

        let transport_index = reference_ipv4_header.header_length() as usize;
        let payload_index = transport_index + reference_transport_header.header_length() as usize;

        let mut ipv4_header_data = reference_ipv4_header.data().clone();
        let mut transport_header_data = reference_transport_header.data_clone();

        {
            let ipv4_header_raw = &mut buffer[..transport_index];
            ipv4_header_raw.copy_from_slice(reference_ipv4_header.raw());
            let mut ipv4_header = ipv4_header_data.bind_mut(ipv4_header_raw);
            ipv4_header.swap_source_and_destination();
        }

        {
            let transport_header_raw = &mut buffer[transport_index..payload_index];
            transport_header_raw.copy_from_slice(reference_transport_header.raw());
            let mut transport_header = transport_header_data.bind_mut(transport_header_raw);
            transport_header.swap_source_and_destination();
        }

        Self {
            buffer: buffer,
            transport_index: transport_index,
            payload_index: payload_index,
            ipv4_header_data: ipv4_header_data,
            transport_header_data: transport_header_data,
            last_packet_length: None,
        }
    }

    pub fn packetize_empty_payload(&mut self) -> IPv4Packet {
        self.inflate(0)
    }

    pub fn packetize<R: DatagramReceiver>(&mut self, source: &mut R) -> io::Result<IPv4Packet> {
        self.last_packet_length = None;
        let r = source.recv(&mut self.buffer[self.payload_index..])?;
        let ipv4_packet = self.inflate(r as u16);
        Ok(ipv4_packet)
    }

    pub fn packetize_read<R: io::Read>(&mut self, source: &mut R, max_chunk_size: Option<usize>) -> io::Result<IPv4Packet> {
        // FIXME handle EOF (read returning 0) for TCP stream
        let mut adapter = ReadAdapter::new(source, max_chunk_size);
        self.packetize(&mut adapter)
    }

    fn ipv4_header_mut(&mut self) -> IPv4HeaderMut {
        let raw = &mut self.buffer[..self.transport_index];
        self.ipv4_header_data.bind_mut(raw)
    }

    fn transport_header_mut(&mut self) -> TransportHeaderMut {
        let raw = &mut self.buffer[self.transport_index..self.payload_index];
        self.transport_header_data.bind_mut(raw)
    }

    fn inflate(&mut self, payload_length: u16) -> IPv4Packet {
        let total_length = self.payload_index as u16 + payload_length;

        self.ipv4_header_mut().set_total_length(total_length);
        self.transport_header_mut().set_payload_length(payload_length);

        self.last_packet_length = Some(total_length as usize);

        let mut ipv4_packet = IPv4Packet::new(&mut self.buffer[..total_length as usize], self.ipv4_header_data.clone(), self.transport_header_data.clone());
        ipv4_packet.compute_checksums();
        ipv4_packet
    }

    pub fn last_packet(&mut self) -> Option<IPv4Packet> {
        if let Some(packet_length) = self.last_packet_length {
            Some(IPv4Packet::new(&mut self.buffer[..packet_length as usize], self.ipv4_header_data.clone(), self.transport_header_data.clone()))
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;
    use byteorder::{BigEndian, WriteBytesExt};
    use relay::datagram::tests::MockDatagramSocket;

    fn create_packet() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.write_u8(4u8 << 4 | 5).unwrap();
        raw.write_u8(0).unwrap(); // ToS
        raw.write_u16::<BigEndian>(32).unwrap(); // total length 20 + 8 + 4
        raw.write_u32::<BigEndian>(0).unwrap(); // id_flags_fragment_offset
        raw.write_u8(0).unwrap(); // TTL
        raw.write_u8(17).unwrap(); // protocol (UDP)
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum
        raw.write_u32::<BigEndian>(0x12345678).unwrap(); // source address
        raw.write_u32::<BigEndian>(0x42424242).unwrap(); // destination address

        raw.write_u16::<BigEndian>(1234).unwrap(); // source port
        raw.write_u16::<BigEndian>(5678).unwrap(); // destination port
        raw.write_u16::<BigEndian>(12).unwrap(); // length
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum

        raw.write_u32::<BigEndian>(0x11223344).unwrap(); // payload
        raw
    }

    #[test]
    fn merge_headers_and_payload() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut mock = MockDatagramSocket::from_data(&data);

        let ipv4_header = reference_packet.ipv4_header();
        let transport_header = reference_packet.transport_header().unwrap();
        let mut packetizer = Packetizer::new(&ipv4_header, &transport_header);

        let packet = packetizer.packetize(&mut mock).unwrap();
        assert_eq!(36, packet.ipv4_header().total_length());
        assert_eq!(data, &packet.raw()[28..36]);
    }

    #[test]
    fn last_packet() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut mock = MockDatagramSocket::from_data(&data);

        let ipv4_header = reference_packet.ipv4_header();
        let transport_header = reference_packet.transport_header().unwrap();
        let mut packetizer = Packetizer::new(&ipv4_header, &transport_header);

        let _ = packetizer.packetize(&mut mock).unwrap();
        let packet = packetizer.last_packet().unwrap();
        assert_eq!(36, packet.ipv4_header().total_length());
        assert_eq!(data, &packet.raw()[28..36]);
    }

    #[test]
    fn packetize_chunks() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut cursor = io::Cursor::new(&data);

        let ipv4_header = reference_packet.ipv4_header();
        let transport_header = reference_packet.transport_header().unwrap();
        let mut packetizer = Packetizer::new(&ipv4_header, &transport_header);

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(2)).unwrap();
            assert_eq!(30, packet.ipv4_header().total_length());
            assert_eq!([0x11, 0x22], packet.payload().unwrap());
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(3)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!([0x33, 0x44, 0x55], packet.payload().unwrap());
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(1024)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!([0x66, 0x77, 0x88], packet.payload().unwrap());
        }
    }
}
<commit_msg>Expose Packetizer manual inflate()<commit_after>use std::io;

use super::datagram::{DatagramReceiver, ReadAdapter};
use super::ipv4_header::{IPv4Header, IPv4HeaderData, IPv4HeaderMut};
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::transport_header::{TransportHeader, TransportHeaderData, TransportHeaderMut};

/// Convert from level 5 to level 3 by appending correct IP and transport headers.
pub struct Packetizer {
    buffer: Box<[u8; MAX_PACKET_LENGTH]>,
    transport_index: usize,
    payload_index: usize,
    ipv4_header_data: IPv4HeaderData,
    transport_header_data: TransportHeaderData,
}

impl Packetizer {
    pub fn new(reference_ipv4_header: &IPv4Header, reference_transport_header: &TransportHeader) -> Self {
        let mut buffer = Box::new([0; MAX_PACKET_LENGTH]);

        let transport_index = reference_ipv4_header.header_length() as usize;
        let payload_index = transport_index + reference_transport_header.header_length() as usize;

        let mut ipv4_header_data = reference_ipv4_header.data().clone();
        let mut transport_header_data = reference_transport_header.data_clone();

        {
            let ipv4_header_raw = &mut buffer[..transport_index];
            ipv4_header_raw.copy_from_slice(reference_ipv4_header.raw());
            let mut ipv4_header = ipv4_header_data.bind_mut(ipv4_header_raw);
            ipv4_header.swap_source_and_destination();
        }

        {
            let transport_header_raw = &mut buffer[transport_index..payload_index];
            transport_header_raw.copy_from_slice(reference_transport_header.raw());
            let mut transport_header = transport_header_data.bind_mut(transport_header_raw);
            transport_header.swap_source_and_destination();
        }

        Self {
            buffer: buffer,
            transport_index: transport_index,
            payload_index: payload_index,
            ipv4_header_data: ipv4_header_data,
            transport_header_data: transport_header_data,
        }
    }

    pub fn packetize_empty_payload(&mut self) -> IPv4Packet {
        self.build(0)
    }

    pub fn packetize<R: DatagramReceiver>(&mut self, source: &mut R) -> io::Result<IPv4Packet> {
        let r = source.recv(&mut self.buffer[self.payload_index..])?;
        let ipv4_packet = self.build(r as u16);
        Ok(ipv4_packet)
    }

    pub fn packetize_read<R: io::Read>(&mut self, source: &mut R, max_chunk_size: Option<usize>) -> io::Result<IPv4Packet> {
        // FIXME handle EOF (read returning 0) for TCP stream
        let mut adapter = ReadAdapter::new(source, max_chunk_size);
        self.packetize(&mut adapter)
    }

    fn ipv4_header_mut(&mut self) -> IPv4HeaderMut {
        let raw = &mut self.buffer[..self.transport_index];
        self.ipv4_header_data.bind_mut(raw)
    }

    fn transport_header_mut(&mut self) -> TransportHeaderMut {
        let raw = &mut self.buffer[self.transport_index..self.payload_index];
        self.transport_header_data.bind_mut(raw)
    }

    fn build(&mut self, payload_length: u16) -> IPv4Packet {
        let total_length = self.payload_index as u16 + payload_length;

        self.ipv4_header_mut().set_total_length(total_length);
        self.transport_header_mut().set_payload_length(payload_length);

        let mut ipv4_packet = IPv4Packet::new(&mut self.buffer[..total_length as usize], self.ipv4_header_data.clone(), self.transport_header_data.clone());
        ipv4_packet.compute_checksums();
        ipv4_packet
    }

    pub fn inflate(&mut self, packet_length: u16) -> IPv4Packet {
        IPv4Packet::new(&mut self.buffer[..packet_length as usize], self.ipv4_header_data.clone(), self.transport_header_data.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io;
    use byteorder::{BigEndian, WriteBytesExt};
    use relay::datagram::tests::MockDatagramSocket;

    fn create_packet() -> Vec<u8> {
        let mut raw = Vec::new();
        raw.write_u8(4u8 << 4 | 5).unwrap();
        raw.write_u8(0).unwrap(); // ToS
        raw.write_u16::<BigEndian>(32).unwrap(); // total length 20 + 8 + 4
        raw.write_u32::<BigEndian>(0).unwrap(); // id_flags_fragment_offset
        raw.write_u8(0).unwrap(); // TTL
        raw.write_u8(17).unwrap(); // protocol (UDP)
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum
        raw.write_u32::<BigEndian>(0x12345678).unwrap(); // source address
        raw.write_u32::<BigEndian>(0x42424242).unwrap(); // destination address

        raw.write_u16::<BigEndian>(1234).unwrap(); // source port
        raw.write_u16::<BigEndian>(5678).unwrap(); // destination port
        raw.write_u16::<BigEndian>(12).unwrap(); // length
        raw.write_u16::<BigEndian>(0).unwrap(); // checksum

        raw.write_u32::<BigEndian>(0x11223344).unwrap(); // payload
        raw
    }

    #[test]
    fn merge_headers_and_payload() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut mock = MockDatagramSocket::from_data(&data);

        let ipv4_header = reference_packet.ipv4_header();
        let transport_header = reference_packet.transport_header().unwrap();
        let mut packetizer = Packetizer::new(&ipv4_header, &transport_header);

        let packet = packetizer.packetize(&mut mock).unwrap();
        assert_eq!(36, packet.ipv4_header().total_length());
        assert_eq!(data, &packet.raw()[28..36]);
    }

    #[test]
    fn last_packet() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut mock = MockDatagramSocket::from_data(&data);

        let ipv4_header = reference_packet.ipv4_header();
        let transport_header = reference_packet.transport_header().unwrap();
        let mut packetizer = Packetizer::new(&ipv4_header, &transport_header);

        let packet_length = packetizer.packetize(&mut mock).unwrap().length();
        let packet = packetizer.inflate(packet_length);
        assert_eq!(36, packet.ipv4_header().total_length());
        assert_eq!(data, &packet.raw()[28..36]);
    }

    #[test]
    fn packetize_chunks() {
        let mut raw = &mut create_packet()[..];
        let reference_packet = IPv4Packet::parse(raw);

        let data = [ 0x11u8, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88 ];
        let mut cursor = io::Cursor::new(&data);

        let ipv4_header = reference_packet.ipv4_header();
        let transport_header = reference_packet.transport_header().unwrap();
        let mut packetizer = Packetizer::new(&ipv4_header, &transport_header);

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(2)).unwrap();
            assert_eq!(30, packet.ipv4_header().total_length());
            assert_eq!([0x11, 0x22], packet.payload().unwrap());
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(3)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!([0x33, 0x44, 0x55], packet.payload().unwrap());
        }

        {
            let packet = packetizer.packetize_read(&mut cursor, Some(1024)).unwrap();
            assert_eq!(31, packet.ipv4_header().total_length());
            assert_eq!([0x66, 0x77, 0x88], packet.payload().unwrap());
        }
    }
}
<|endoftext|>


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

//! Test utilities to help with unit testing of Oak SDK code.

use log::{debug, info};
use oak_abi::proto::oak::application::{
    node_configuration::ConfigType, ApplicationConfiguration, NodeConfiguration,
    WebAssemblyConfiguration,
};
use prost::Message;
use std::{collections::HashMap, path::PathBuf, process::Command};
use tonic::transport::Certificate;

// TODO(#544): re-enable unit tests of SDK functionality

/// Uses cargo to compile a Rust manifest to Wasm bytes.
pub fn compile_rust_wasm(cargo_path: &str, module_name: &str) -> std::io::Result<Vec<u8>> {
    // Use a separate target dir for Wasm build artifacts. The precise name is not relevant, but it
    // should end with `target` so that it gets automatically ignored by our `.gitignore`.
    let target_dir = PathBuf::from("oak_tests/target");

    Command::new("cargo")
        .args(&[
            "build",
            &format!(
                "--target-dir={}",
                target_dir.to_str().expect("invalid target dir")
            ),
            "--target=wasm32-unknown-unknown",
            &format!("--manifest-path={}", cargo_path),
        ])
        .env_remove("RUSTFLAGS")
        .spawn()?
        .wait()?;

    let mut module_path = target_dir;
    module_path.push("wasm32-unknown-unknown/debug");
    module_path.push(module_name);

    info!("compiled Wasm module path: {:?}", module_path);

    std::fs::read(module_path)
}

/// Default entrypoint name for the module under test.
const DEFAULT_ENTRYPOINT_NAME: &str = "oak_main";

const DEFAULT_MODULE_MANIFEST: &str = "Cargo.toml";
const MODULE_WASM_SUFFIX: &str = ".wasm";

/// Convenience helper to build and run a single-Node Application with the given module name, using
/// the default name "oak_main" for its entrypoint.
pub fn run_single_module_default(
    module_config_name: &str,
) -> Result<(oak_runtime::RuntimeProxy, oak_abi::Handle), oak::OakStatus> {
    run_single_module(module_config_name, DEFAULT_ENTRYPOINT_NAME)
}

/// Convenience helper to build and run a single-Node application with the given module name, using
/// the provided entrypoint name.
pub fn run_single_module(
    module_config_name: &str,
    entrypoint_name: &str,
) -> Result<(oak_runtime::RuntimeProxy, oak_abi::Handle), oak::OakStatus> {
    let (app_config, runtime_config, grpc_config) =
        runtime_configs(module_config_name, entrypoint_name);
    oak_runtime::configure_and_run(app_config, runtime_config, grpc_config)
}

/// Build the set of configurations needed to launch a test Runtime instance that
/// runs a single-Node application with the given module name and entrypoint.
pub fn runtime_configs(
    module_config_name: &str,
    entrypoint_name: &str,
) -> (
    ApplicationConfiguration,
    oak_runtime::RuntimeConfiguration,
    oak_runtime::GrpcConfiguration,
) {
    let wasm: HashMap<String, Vec<u8>> = [(
        module_config_name.to_owned(),
        compile_rust_wasm(
            DEFAULT_MODULE_MANIFEST,
            &(module_config_name.to_owned() + MODULE_WASM_SUFFIX),
        )
        .expect("failed to build wasm module"),
    )]
    .iter()
    .cloned()
    .collect();

    (
        ApplicationConfiguration {
            wasm_modules: wasm,
            initial_node_configuration: Some(NodeConfiguration {
                name: "test".to_string(),
                config_type: Some(ConfigType::WasmConfig(WebAssemblyConfiguration {
                    wasm_module_name: module_config_name.to_string(),
                    wasm_entrypoint_name: entrypoint_name.to_string(),
                })),
            }),
        },
        oak_runtime::RuntimeConfiguration::default(),
        oak_runtime::GrpcConfiguration {
            grpc_server_tls_identity: None,
            oidc_client_info: None,
            // Some of the tests require a gRPC client, so we populate the required certificate with
            // an invalid value here, even though it will still fail when instantiating the actual
            // gRPC client.
            grpc_client_root_tls_certificate: Some(Certificate::from_pem("invalid-cert")),
        },
    )
}

/// Inject a gRPC request into an Oak Application under test and return the corresponding
/// response (or error).
pub fn grpc_request<R, Q>(
    proxy: &oak_runtime::RuntimeProxy,
    handle: oak_abi::Handle,
    method_name: &str,
    req: &R,
) -> oak::grpc::Result<Q>
where
    R: prost::Message,
    Q: prost::Message + Default,
{
    // Put the request in a GrpcRequest wrapper and serialize into a message.
    let grpc_req =
        oak_abi::grpc::encap_request(req, method_name).expect("failed to build GrpcRequest");
    let mut req_msg = oak_runtime::NodeMessage {
        data: vec![],
        handles: vec![],
    };

    grpc_req
        .encode(&mut req_msg.data)
        .expect("failed to serialize GrpcRequest message");

    // Create a new channel to hold the request message.
    //
    // In most cases we do not care about labels, so we use the least privileged label for this
    // channel.
    let (req_write_handle, req_read_handle) = proxy
        .channel_create(&oak_abi::label::Label::public_untrusted())
        .expect("could not create channel");
    proxy
        .channel_write(req_write_handle, req_msg)
        .expect("could not write message");

    // Create a new channel for responses to arrive on and also attach that to the message.
    //
    // In most cases we do not care about labels, so we use the least privileged label for this
    // channel.
    let (rsp_write_handle, rsp_read_handle) = proxy
        .channel_create(&oak_abi::label::Label::public_untrusted())
        .expect("could not create channel");

    // Create a notification message and attach the method-invocation specific channels to it.
    let notify_msg = oak_runtime::NodeMessage {
        data: vec![],
        handles: vec![req_read_handle, rsp_write_handle],
    };

    // Send the notification message (with attached handles) into the Node under test.
    proxy
        .channel_write(handle, notify_msg)
        .expect("could not write message");
    proxy
        .channel_close(req_write_handle)
        .expect("failed to close channel");
    proxy
        .channel_close(req_read_handle)
        .expect("failed to close channel");
    proxy
        .channel_close(rsp_write_handle)
        .expect("failed to close channel");

    // Read the serialized, encapsulated response.
    loop {
        let rsp = match proxy.channel_read(rsp_read_handle) {
            Ok(Some(r)) => r,
            Ok(None) => {
                debug!("no pending gRPC response message; poll again soon");
                std::thread::sleep(std::time::Duration::from_millis(100));
                continue;
            }
            Err(e) => {
                panic!("failed to read from response channel: {:?}", e);
            }
        };
        proxy
            .channel_close(rsp_read_handle)
            .expect("failed to close channel");
        if rsp.data.is_empty() {
            info!("no pending message; poll again soon");
            std::thread::sleep(std::time::Duration::from_millis(100));
            continue;
        }
        let rsp = oak_abi::proto::oak::encap::GrpcResponse::decode(rsp.data.as_slice())
            .expect("failed to parse GrpcResponse message");
        if !rsp.last {
            panic!("Expected single final response");
        }

        let status = rsp.status.unwrap_or_default();
        if status.code != oak::grpc::Code::Ok as i32 {
            return Err(status);
        }
        let rsp =
            Q::decode(rsp.rsp_msg.as_slice()).expect("Failed to parse response protobuf message");
        return Ok(rsp);
    }
}
<commit_msg>sdk: turn on aux servers in tests<commit_after>//
// Copyright 2019 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! Test utilities to help with unit testing of Oak SDK code.

use log::{debug, info};
use oak_abi::proto::oak::application::{
    node_configuration::ConfigType, ApplicationConfiguration, NodeConfiguration,
    WebAssemblyConfiguration,
};
use prost::Message;
use std::{collections::HashMap, path::PathBuf, process::Command};
use tonic::transport::Certificate;

// TODO(#544): re-enable unit tests of SDK functionality

/// Uses cargo to compile a Rust manifest to Wasm bytes.
pub fn compile_rust_wasm(cargo_path: &str, module_name: &str) -> std::io::Result<Vec<u8>> {
    // Use a separate target dir for Wasm build artifacts. The precise name is not relevant, but it
    // should end with `target` so that it gets automatically ignored by our `.gitignore`.
    let target_dir = PathBuf::from("oak_tests/target");

    Command::new("cargo")
        .args(&[
            "build",
            &format!(
                "--target-dir={}",
                target_dir.to_str().expect("invalid target dir")
            ),
            "--target=wasm32-unknown-unknown",
            &format!("--manifest-path={}", cargo_path),
        ])
        .env_remove("RUSTFLAGS")
        .spawn()?
        .wait()?;

    let mut module_path = target_dir;
    module_path.push("wasm32-unknown-unknown/debug");
    module_path.push(module_name);

    info!("compiled Wasm module path: {:?}", module_path);

    std::fs::read(module_path)
}

/// Default entrypoint name for the module under test.
const DEFAULT_ENTRYPOINT_NAME: &str = "oak_main";

const DEFAULT_MODULE_MANIFEST: &str = "Cargo.toml";
const MODULE_WASM_SUFFIX: &str = ".wasm";

/// Convenience helper to build and run a single-Node Application with the given module name, using
/// the default name "oak_main" for its entrypoint.
pub fn run_single_module_default(
    module_config_name: &str,
) -> Result<(oak_runtime::RuntimeProxy, oak_abi::Handle), oak::OakStatus> {
    run_single_module(module_config_name, DEFAULT_ENTRYPOINT_NAME)
}

/// Convenience helper to build and run a single-Node application with the given module name, using
/// the provided entrypoint name.
pub fn run_single_module(
    module_config_name: &str,
    entrypoint_name: &str,
) -> Result<(oak_runtime::RuntimeProxy, oak_abi::Handle), oak::OakStatus> {
    let (app_config, runtime_config, grpc_config) =
        runtime_configs(module_config_name, entrypoint_name);
    oak_runtime::configure_and_run(app_config, runtime_config, grpc_config)
}

/// Build the set of configurations needed to launch a test Runtime instance that
/// runs a single-Node application with the given module name and entrypoint.
pub fn runtime_configs(
    module_config_name: &str,
    entrypoint_name: &str,
) -> (
    ApplicationConfiguration,
    oak_runtime::RuntimeConfiguration,
    oak_runtime::GrpcConfiguration,
) {
    let wasm: HashMap<String, Vec<u8>> = [(
        module_config_name.to_owned(),
        compile_rust_wasm(
            DEFAULT_MODULE_MANIFEST,
            &(module_config_name.to_owned() + MODULE_WASM_SUFFIX),
        )
        .expect("failed to build wasm module"),
    )]
    .iter()
    .cloned()
    .collect();

    (
        ApplicationConfiguration {
            wasm_modules: wasm,
            initial_node_configuration: Some(NodeConfiguration {
                name: "test".to_string(),
                config_type: Some(ConfigType::WasmConfig(WebAssemblyConfiguration {
                    wasm_module_name: module_config_name.to_string(),
                    wasm_entrypoint_name: entrypoint_name.to_string(),
                })),
            }),
        },
        oak_runtime::RuntimeConfiguration {
            metrics_port: Some(9090),
            introspect_port: Some(1909),
        },
        oak_runtime::GrpcConfiguration {
            grpc_server_tls_identity: None,
            oidc_client_info: None,
            // Some of the tests require a gRPC client, so we populate the required certificate with
            // an invalid value here, even though it will still fail when instantiating the actual
            // gRPC client.
            grpc_client_root_tls_certificate: Some(Certificate::from_pem("invalid-cert")),
        },
    )
}

/// Inject a gRPC request into an Oak Application under test and return the corresponding
/// response (or error).
pub fn grpc_request<R, Q>(
    proxy: &oak_runtime::RuntimeProxy,
    handle: oak_abi::Handle,
    method_name: &str,
    req: &R,
) -> oak::grpc::Result<Q>
where
    R: prost::Message,
    Q: prost::Message + Default,
{
    // Put the request in a GrpcRequest wrapper and serialize into a message.
    let grpc_req =
        oak_abi::grpc::encap_request(req, method_name).expect("failed to build GrpcRequest");
    let mut req_msg = oak_runtime::NodeMessage {
        data: vec![],
        handles: vec![],
    };

    grpc_req
        .encode(&mut req_msg.data)
        .expect("failed to serialize GrpcRequest message");

    // Create a new channel to hold the request message.
    //
    // In most cases we do not care about labels, so we use the least privileged label for this
    // channel.
    let (req_write_handle, req_read_handle) = proxy
        .channel_create(&oak_abi::label::Label::public_untrusted())
        .expect("could not create channel");
    proxy
        .channel_write(req_write_handle, req_msg)
        .expect("could not write message");

    // Create a new channel for responses to arrive on and also attach that to the message.
    //
    // In most cases we do not care about labels, so we use the least privileged label for this
    // channel.
    let (rsp_write_handle, rsp_read_handle) = proxy
        .channel_create(&oak_abi::label::Label::public_untrusted())
        .expect("could not create channel");

    // Create a notification message and attach the method-invocation specific channels to it.
    let notify_msg = oak_runtime::NodeMessage {
        data: vec![],
        handles: vec![req_read_handle, rsp_write_handle],
    };

    // Send the notification message (with attached handles) into the Node under test.
    proxy
        .channel_write(handle, notify_msg)
        .expect("could not write message");
    proxy
        .channel_close(req_write_handle)
        .expect("failed to close channel");
    proxy
        .channel_close(req_read_handle)
        .expect("failed to close channel");
    proxy
        .channel_close(rsp_write_handle)
        .expect("failed to close channel");

    // Read the serialized, encapsulated response.
    loop {
        let rsp = match proxy.channel_read(rsp_read_handle) {
            Ok(Some(r)) => r,
            Ok(None) => {
                debug!("no pending gRPC response message; poll again soon");
                std::thread::sleep(std::time::Duration::from_millis(100));
                continue;
            }
            Err(e) => {
                panic!("failed to read from response channel: {:?}", e);
            }
        };
        proxy
            .channel_close(rsp_read_handle)
            .expect("failed to close channel");
        if rsp.data.is_empty() {
            info!("no pending message; poll again soon");
            std::thread::sleep(std::time::Duration::from_millis(100));
            continue;
        }
        let rsp = oak_abi::proto::oak::encap::GrpcResponse::decode(rsp.data.as_slice())
            .expect("failed to parse GrpcResponse message");
        if !rsp.last {
            panic!("Expected single final response");
        }

        let status = rsp.status.unwrap_or_default();
        if status.code != oak::grpc::Code::Ok as i32 {
            return Err(status);
        }
        let rsp =
            Q::decode(rsp.rsp_msg.as_slice()).expect("Failed to parse response protobuf message");
        return Ok(rsp);
    }
}
<|endoftext|>