<commit_before>//! Checks for any invariants we expect to be true, and pages whoever is on call
//! if they are not.
//!
//! Usage:
//!     cargo run --bin monitor

#![deny(warnings)]

#[macro_use]
extern crate serde_derive;

mod on_call;

use cargo_registry::{db, util::CargoResult};
use diesel::prelude::*;

fn main() -> CargoResult<()> {
    let conn = db::connect_now()?;

    check_stalled_background_jobs(&conn)?;
    Ok(())
}

fn check_stalled_background_jobs(conn: &PgConnection) -> CargoResult<()> {
    use cargo_registry::schema::background_jobs::dsl::*;
    use diesel::dsl::*;

    const BACKGROUND_JOB_KEY: &str = "background_jobs";

    println!("Checking for stalled background jobs");

    let max_job_time = dotenv::var("MAX_JOB_TIME")
        .map(|s| s.parse::<i32>().unwrap())
        .unwrap_or(15);

    let stalled_job_count = background_jobs
        .filter(created_at.lt(now - max_job_time.minutes()))
        .count()
        .get_result::<i64>(conn)?;

    let event = if stalled_job_count > 0 {
        on_call::Event::Trigger {
            incident_key: Some(BACKGROUND_JOB_KEY.into()),
            description: format!(
                "{} jobs have been in the queue for more than {} minutes",
                stalled_job_count, max_job_time
            ),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: BACKGROUND_JOB_KEY.into(),
            description: Some("No stalled background jobs".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn log_and_trigger_event(event: on_call::Event) -> CargoResult<()> {
    match event {
        on_call::Event::Trigger {
            ref description, ..
        } => println!("Paging on-call: {}", description),
        on_call::Event::Resolve {
            description: Some(ref description),
            ..
        } => println!("{}", description),
        _ => {} // noop
    }
    event.send()
}
<commit_msg>Add monitoring for common spam patterns<commit_after>//! Checks for any invariants we expect to be true, and pages whoever is on call
//! if they are not.
//!
//! Usage:
//!     cargo run --bin monitor

#![deny(warnings)]

#[macro_use]
extern crate serde_derive;

mod on_call;

use cargo_registry::{db, schema::*, util::CargoResult};
use diesel::prelude::*;

fn main() -> CargoResult<()> {
    let conn = db::connect_now()?;

    check_stalled_background_jobs(&conn)?;
    check_spam_attack(&conn)?;
    Ok(())
}

fn check_stalled_background_jobs(conn: &PgConnection) -> CargoResult<()> {
    use cargo_registry::schema::background_jobs::dsl::*;
    use diesel::dsl::*;

    const EVENT_KEY: &str = "background_jobs";

    println!("Checking for stalled background jobs");

    let max_job_time = dotenv::var("MAX_JOB_TIME")
        .map(|s| s.parse::<i32>().unwrap())
        .unwrap_or(15);

    let stalled_job_count = background_jobs
        .filter(created_at.lt(now - max_job_time.minutes()))
        .count()
        .get_result::<i64>(conn)?;

    let event = if stalled_job_count > 0 {
        on_call::Event::Trigger {
            incident_key: Some(EVENT_KEY.into()),
            description: format!(
                "{} jobs have been in the queue for more than {} minutes",
                stalled_job_count, max_job_time
            ),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: EVENT_KEY.into(),
            description: Some("No stalled background jobs".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn check_spam_attack(conn: &PgConnection) -> CargoResult<()> {
    use cargo_registry::models::krate::canon_crate_name;
    use diesel::dsl::*;
    use diesel::sql_types::Bool;

    const EVENT_KEY: &str = "spam_attack";

    println!("Checking for crates indicating someone is spamming us");

    let bad_crate_names = dotenv::var("SPAM_CRATE_NAMES");
    let bad_crate_names = bad_crate_names
        .as_ref()
        .map(|s| s.split(",").collect())
        .unwrap_or(Vec::new());
    let bad_author_patterns = dotenv::var("SPAM_AUTHOR_PATTERNS");
    let bad_author_patterns = bad_author_patterns
        .as_ref()
        .map(|s| s.split(",").collect())
        .unwrap_or(Vec::new());

    let mut event_description = None;

    let bad_crate = crates::table
        .filter(canon_crate_name(crates::name).eq(any(bad_crate_names)))
        .select(crates::name)
        .first::<String>(conn)
        .optional()?;

    if let Some(bad_crate) = bad_crate {
        event_description = Some(
            format!("Crate named {} published", bad_crate)
        );
    }

    let mut query = version_authors::table
        .select(version_authors::name)
        .filter(false.into_sql::<Bool>()) // Never return anything if we have no patterns
        .into_boxed();
    for author_pattern in bad_author_patterns {
        query = query.or_filter(version_authors::name.like(author_pattern));
    }
    let bad_author = query.first::<String>(conn).optional()?;

    if let Some(bad_author) = bad_author {
        event_description = Some(
            format!("Crate with author {} published", bad_author)
        );
    }

    let event = if let Some(event_description) = event_description {
        on_call::Event::Trigger {
            incident_key: Some(EVENT_KEY.into()),
            description: format!(
                "{}, possible spam attack underway",
                event_description,
            ),
        }
    } else {
        on_call::Event::Resolve {
            incident_key: EVENT_KEY.into(),
            description: Some("No spam crates detected".into()),
        }
    };

    log_and_trigger_event(event)?;
    Ok(())
}

fn log_and_trigger_event(event: on_call::Event) -> CargoResult<()> {
    match event {
        on_call::Event::Trigger {
            ref description, ..
        } => println!("Paging on-call: {}", description),
        on_call::Event::Resolve {
            description: Some(ref description),
            ..
        } => println!("{}", description),
        _ => {} // noop
    }
    event.send()
}
<|endoftext|>


<commit_before>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use std::cmp;
use std::time::Duration;
use r2d2::Pool;
use r2d2::config::Config;
use r2d2_redis::{RedisConnectionManager, Error};
use redis::{self, Value, Connection, Commands, PipelineCommands};
use futures::future::Future;
use futures_cpupool::CpuPool;

use super::route::ROUTE_PREFIX;
use APP_CONF;

static KEY_BODY: &'static str = "b";
static KEY_FINGERPRINT: &'static str = "f";
static KEY_TAGS: &'static str = "t";
static KEY_TAGS_SEPARATOR: &'static str = ",";

lazy_static! {
    pub static ref EXECUTOR_POOL: CpuPool = CpuPool::new(APP_CONF.cache.executor_pool as usize);
}

pub struct CacheStoreBuilder;

pub struct CacheStore {
    pool: Pool<RedisConnectionManager>,
}

#[derive(Debug)]
pub enum CacheStoreError {
    Disconnected,
    Failed,
    Invalid,
    Corrupted,
    TooLarge,
}

#[derive(Debug)]
pub enum CachePurgeVariant {
    Bucket,
    Auth,
}

type CacheReadResultFuture = Box<Future<Item = Option<String>, Error = CacheStoreError>>;
type CacheWriteResult = Result<(String), (CacheStoreError, String)>;
type CacheWriteResultFuture = Box<Future<Item = CacheWriteResult, Error = ()>>;
type CachePurgeResult = Result<(), CacheStoreError>;

impl CacheStoreBuilder {
    pub fn new() -> CacheStore {
        info!(
            "binding to store backend at {}:{}",
            APP_CONF.redis.host,
            APP_CONF.redis.port
        );

        let addr_auth = match APP_CONF.redis.password {
            Some(ref password) => format!(":{}@", password),
            None => "".to_string(),
        };

        let tcp_addr_raw =
            format!(
            "redis://{}{}:{}/{}",
            &addr_auth,
            APP_CONF.redis.host,
            APP_CONF.redis.port,
            APP_CONF.redis.database,
        );

        debug!("will connect to redis at: {}", tcp_addr_raw);

        match RedisConnectionManager::new(tcp_addr_raw.as_ref()) {
            Ok(manager) => {
                let config = Config::<Connection, Error>::builder()
                    .initialization_fail_fast(false)
                    .test_on_check_out(false)
                    .pool_size(APP_CONF.redis.pool_size)
                    .max_lifetime(Some(
                        Duration::from_secs(APP_CONF.redis.max_lifetime_seconds),
                    ))
                    .idle_timeout(Some(
                        Duration::from_secs(APP_CONF.redis.idle_timeout_seconds),
                    ))
                    .connection_timeout(Duration::from_secs(
                        APP_CONF.redis.connection_timeout_seconds,
                    ))
                    .build();

                match Pool::new(config, manager) {
                    Ok(pool) => {
                        info!("bound to store backend");

                        CacheStore { pool: pool }
                    }
                    Err(_) => panic!("could not spawn redis pool"),
                }
            }
            Err(_) => panic!("could not create redis connection manager"),
        }
    }
}

impl CacheStore {
    pub fn get_meta(&self, shard: u8, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, (Value, Value)>(key, (KEY_FINGERPRINT, KEY_TAGS)) {
                        Ok(value) => {
                            match value {
                                (Value::Data(fingerprint_bytes), tags_bytes) => {
                                    // Parse tags and bump their last access time
                                    if let Value::Data(tags_bytes_data) = tags_bytes {
                                        if let Ok(tags_data) = String::from_utf8(
                                            tags_bytes_data) {
                                            if tags_data.is_empty() == false {
                                                let tags = tags_data.split(KEY_TAGS_SEPARATOR)
                                                    .map(|tag| {
                                                        format!(
                                                            "{}:{}:{}", ROUTE_PREFIX, shard, tag
                                                        )
                                                    })
                                                    .collect::<Vec<String>>();

                                                // Proceed a soft bump of last access time of \
                                                //   associated tag keys. This prevents a \
                                                //   frequently accessed cache namespace to \
                                                //   become 'orphan' (ie. one or more tag keys \
                                                //   are LRU-expired), and thus cache namespace \
                                                //   not to be properly removed on purge of an \
                                                //   associated tag.
                                                // Also, count bumped keys. It may happen that \
                                                //   some tag keys are incorrectly removed by \
                                                //   Redis LRU system, as it is probabilistic \
                                                //   and thus might sample some keys incorrectly.
                                                // The conditions explained above only happens on \
                                                //   Redis instances with used memory going over \
                                                //   the threshold of the max memory policy.
                                                let tags_count = tags.len();

                                                match redis::cmd("TOUCH").arg(tags)
                                                    .query::<usize>(&*client) {
                                                    Ok(bump_count) => {
                                                        // Partial bump count? Consider cache as \
                                                        //   non-existing
                                                        if bump_count < tags_count {
                                                            return Ok(None);
                                                        }
                                                    },
                                                    Err(err) => {
                                                        error!(
                                                            "error bumping access time of tags: {}",
                                                            err
                                                        );
                                                    }
                                                }
                                            }
                                        }
                                    }

                                    // Decode raw bytes to string
                                    if let Ok(fingerprint) = String::from_utf8(fingerprint_bytes) {
                                        Ok(Some(fingerprint))
                                    } else {
                                        Err(CacheStoreError::Corrupted)
                                    }
                                },
                                (Value::Nil, _) => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn get_body(&self, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, Value>(key, KEY_BODY) {
                        Ok(value) => {
                            match value {
                                Value::Data(body_bytes) => {
                                    // Decode raw bytes to string
                                    if let Ok(body) = String::from_utf8(body_bytes) {
                                        Ok(Some(body))
                                    } else {
                                        Err(CacheStoreError::Corrupted)
                                    }
                                },
                                Value::Nil => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn set(
        &self,
        key: String,
        key_mask: String,
        value: String,
        fingerprint: String,
        ttl: usize,
        key_tags: Vec<(String, String)>,
    ) -> CacheWriteResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            Ok(get_cache_store_client!(
                    pool,
                    (CacheStoreError::Disconnected, fingerprint),

                    client {
                        // Cap TTL to 'max_key_expiration'
                        let ttl_cap = cmp::min(ttl, APP_CONF.redis.max_key_expiration);

                        // Ensure value is not larger than 'max_key_size'
                        if value.len() > APP_CONF.redis.max_key_size {
                            Err((CacheStoreError::TooLarge, fingerprint))
                        } else {
                            let mut pipeline = redis::pipe();

                            // Append storage command
                            {
                                let key_tag_masks = key_tags.iter()
                                    .map(|key_tag| key_tag.1.as_ref())
                                    .collect::<Vec<&str>>();

                                pipeline.hset_multiple(
                                    &key, &[
                                        (KEY_FINGERPRINT, &fingerprint),
                                        (KEY_TAGS, &key_tag_masks.join(KEY_TAGS_SEPARATOR)),
                                        (KEY_BODY, &value)
                                    ]
                                ).ignore();
                            }

                            pipeline.expire(&key, ttl_cap).ignore();

                            for key_tag in key_tags {
                                pipeline.sadd(&key_tag.0, &key_mask).ignore();
                                pipeline.expire(&key_tag.0, APP_CONF.redis.max_key_expiration);
                            }

                            // Bucket (MULTI operation for main data + bucket marker)
                            match pipeline.query::<()>(&*client) {
                                Ok(_) => Ok(fingerprint),
                                Err(err) => {
                                    error!("got store error: {}", err);

                                    Err((CacheStoreError::Failed, fingerprint))
                                }
                            }
                        }
                    }
                ))
        }))
    }

    pub fn purge_tag(
        &self,
        variant: &CachePurgeVariant,
        shard: u8,
        key_tag: &str,
    ) -> CachePurgeResult {
        get_cache_store_client!(self.pool, CacheStoreError::Disconnected, client {
            // Invoke keyspace cleanup script for key tag
            let result = redis::Script::new(variant.get_script())
                .arg(ROUTE_PREFIX)
                .arg(shard)
                .arg(key_tag)
                .invoke::<()>(&*client);

            result
                .and(Ok(()))
                .or(Err(CacheStoreError::Failed))
        })
    }
}

impl CachePurgeVariant {
    fn get_script(&self) -> &'static str {
        match *self {
            CachePurgeVariant::Bucket |
            CachePurgeVariant::Auth => {
                r#"
                    local targets = {}

                    for _, tag in pairs(redis.call('SMEMBERS', ARGV[3])) do
                        table.insert(targets, ARGV[1] .. ":" .. ARGV[2] .. ":c:" .. tag)
                    end

                    table.insert(targets, ARGV[3])

                    redis.call('DEL', unpack(targets))
                "#
            }
        }
    }
}
<commit_msg>Log<commit_after>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use std::cmp;
use std::time::Duration;
use r2d2::Pool;
use r2d2::config::Config;
use r2d2_redis::{RedisConnectionManager, Error};
use redis::{self, Value, Connection, Commands, PipelineCommands};
use futures::future::Future;
use futures_cpupool::CpuPool;

use super::route::ROUTE_PREFIX;
use APP_CONF;

static KEY_BODY: &'static str = "b";
static KEY_FINGERPRINT: &'static str = "f";
static KEY_TAGS: &'static str = "t";
static KEY_TAGS_SEPARATOR: &'static str = ",";

lazy_static! {
    pub static ref EXECUTOR_POOL: CpuPool = CpuPool::new(APP_CONF.cache.executor_pool as usize);
}

pub struct CacheStoreBuilder;

pub struct CacheStore {
    pool: Pool<RedisConnectionManager>,
}

#[derive(Debug)]
pub enum CacheStoreError {
    Disconnected,
    Failed,
    Invalid,
    Corrupted,
    TooLarge,
}

#[derive(Debug)]
pub enum CachePurgeVariant {
    Bucket,
    Auth,
}

type CacheReadResultFuture = Box<Future<Item = Option<String>, Error = CacheStoreError>>;
type CacheWriteResult = Result<(String), (CacheStoreError, String)>;
type CacheWriteResultFuture = Box<Future<Item = CacheWriteResult, Error = ()>>;
type CachePurgeResult = Result<(), CacheStoreError>;

impl CacheStoreBuilder {
    pub fn new() -> CacheStore {
        info!(
            "binding to store backend at {}:{}",
            APP_CONF.redis.host,
            APP_CONF.redis.port
        );

        let addr_auth = match APP_CONF.redis.password {
            Some(ref password) => format!(":{}@", password),
            None => "".to_string(),
        };

        let tcp_addr_raw =
            format!(
            "redis://{}{}:{}/{}",
            &addr_auth,
            APP_CONF.redis.host,
            APP_CONF.redis.port,
            APP_CONF.redis.database,
        );

        debug!("will connect to redis at: {}", tcp_addr_raw);

        match RedisConnectionManager::new(tcp_addr_raw.as_ref()) {
            Ok(manager) => {
                let config = Config::<Connection, Error>::builder()
                    .initialization_fail_fast(false)
                    .test_on_check_out(false)
                    .pool_size(APP_CONF.redis.pool_size)
                    .max_lifetime(Some(
                        Duration::from_secs(APP_CONF.redis.max_lifetime_seconds),
                    ))
                    .idle_timeout(Some(
                        Duration::from_secs(APP_CONF.redis.idle_timeout_seconds),
                    ))
                    .connection_timeout(Duration::from_secs(
                        APP_CONF.redis.connection_timeout_seconds,
                    ))
                    .build();

                match Pool::new(config, manager) {
                    Ok(pool) => {
                        info!("bound to store backend");

                        CacheStore { pool: pool }
                    }
                    Err(_) => panic!("could not spawn redis pool"),
                }
            }
            Err(_) => panic!("could not create redis connection manager"),
        }
    }
}

impl CacheStore {
    pub fn get_meta(&self, shard: u8, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, (Value, Value)>(key, (KEY_FINGERPRINT, KEY_TAGS)) {
                        Ok(value) => {
                            match value {
                                (Value::Data(fingerprint_bytes), tags_bytes) => {
                                    // Parse tags and bump their last access time
                                    if let Value::Data(tags_bytes_data) = tags_bytes {
                                        if let Ok(tags_data) = String::from_utf8(
                                            tags_bytes_data) {
                                            if tags_data.is_empty() == false {
                                                let tags = tags_data.split(KEY_TAGS_SEPARATOR)
                                                    .map(|tag| {
                                                        format!(
                                                            "{}:{}:{}", ROUTE_PREFIX, shard, tag
                                                        )
                                                    })
                                                    .collect::<Vec<String>>();

                                                // Proceed a soft bump of last access time of \
                                                //   associated tag keys. This prevents a \
                                                //   frequently accessed cache namespace to \
                                                //   become 'orphan' (ie. one or more tag keys \
                                                //   are LRU-expired), and thus cache namespace \
                                                //   not to be properly removed on purge of an \
                                                //   associated tag.
                                                // Also, count bumped keys. It may happen that \
                                                //   some tag keys are incorrectly removed by \
                                                //   Redis LRU system, as it is probabilistic \
                                                //   and thus might sample some keys incorrectly.
                                                // The conditions explained above only happens on \
                                                //   Redis instances with used memory going over \
                                                //   the threshold of the max memory policy.
                                                let tags_count = tags.len();

                                                match redis::cmd("TOUCH").arg(tags)
                                                    .query::<usize>(&*client) {
                                                    Ok(bump_count) => {
                                                        // Partial bump count? Consider cache as \
                                                        //   non-existing
                                                        if bump_count < tags_count {
                                                            info!(
                                                                "got only partial tag count: {}/{}",
                                                                bump_count, tags_count
                                                            );

                                                            return Ok(None);
                                                        }
                                                    },
                                                    Err(err) => {
                                                        error!(
                                                            "error bumping access time of tags: {}",
                                                            err
                                                        );
                                                    }
                                                }
                                            }
                                        }
                                    }

                                    // Decode raw bytes to string
                                    if let Ok(fingerprint) = String::from_utf8(fingerprint_bytes) {
                                        Ok(Some(fingerprint))
                                    } else {
                                        Err(CacheStoreError::Corrupted)
                                    }
                                },
                                (Value::Nil, _) => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn get_body(&self, key: String) -> CacheReadResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            get_cache_store_client!(pool, CacheStoreError::Disconnected, client {
                    match (*client).hget::<_, _, Value>(key, KEY_BODY) {
                        Ok(value) => {
                            match value {
                                Value::Data(body_bytes) => {
                                    // Decode raw bytes to string
                                    if let Ok(body) = String::from_utf8(body_bytes) {
                                        Ok(Some(body))
                                    } else {
                                        Err(CacheStoreError::Corrupted)
                                    }
                                },
                                Value::Nil => Ok(None),
                                _ => Err(CacheStoreError::Invalid),
                            }
                        },
                        _ => Err(CacheStoreError::Failed),
                    }
                })
        }))
    }

    pub fn set(
        &self,
        key: String,
        key_mask: String,
        value: String,
        fingerprint: String,
        ttl: usize,
        key_tags: Vec<(String, String)>,
    ) -> CacheWriteResultFuture {
        let pool = self.pool.to_owned();

        Box::new(EXECUTOR_POOL.spawn_fn(move || {
            Ok(get_cache_store_client!(
                    pool,
                    (CacheStoreError::Disconnected, fingerprint),

                    client {
                        // Cap TTL to 'max_key_expiration'
                        let ttl_cap = cmp::min(ttl, APP_CONF.redis.max_key_expiration);

                        // Ensure value is not larger than 'max_key_size'
                        if value.len() > APP_CONF.redis.max_key_size {
                            Err((CacheStoreError::TooLarge, fingerprint))
                        } else {
                            let mut pipeline = redis::pipe();

                            // Append storage command
                            {
                                let key_tag_masks = key_tags.iter()
                                    .map(|key_tag| key_tag.1.as_ref())
                                    .collect::<Vec<&str>>();

                                pipeline.hset_multiple(
                                    &key, &[
                                        (KEY_FINGERPRINT, &fingerprint),
                                        (KEY_TAGS, &key_tag_masks.join(KEY_TAGS_SEPARATOR)),
                                        (KEY_BODY, &value)
                                    ]
                                ).ignore();
                            }

                            pipeline.expire(&key, ttl_cap).ignore();

                            for key_tag in key_tags {
                                pipeline.sadd(&key_tag.0, &key_mask).ignore();
                                pipeline.expire(&key_tag.0, APP_CONF.redis.max_key_expiration);
                            }

                            // Bucket (MULTI operation for main data + bucket marker)
                            match pipeline.query::<()>(&*client) {
                                Ok(_) => Ok(fingerprint),
                                Err(err) => {
                                    error!("got store error: {}", err);

                                    Err((CacheStoreError::Failed, fingerprint))
                                }
                            }
                        }
                    }
                ))
        }))
    }

    pub fn purge_tag(
        &self,
        variant: &CachePurgeVariant,
        shard: u8,
        key_tag: &str,
    ) -> CachePurgeResult {
        get_cache_store_client!(self.pool, CacheStoreError::Disconnected, client {
            // Invoke keyspace cleanup script for key tag
            let result = redis::Script::new(variant.get_script())
                .arg(ROUTE_PREFIX)
                .arg(shard)
                .arg(key_tag)
                .invoke::<()>(&*client);

            result
                .and(Ok(()))
                .or(Err(CacheStoreError::Failed))
        })
    }
}

impl CachePurgeVariant {
    fn get_script(&self) -> &'static str {
        match *self {
            CachePurgeVariant::Bucket |
            CachePurgeVariant::Auth => {
                r#"
                    local targets = {}

                    for _, tag in pairs(redis.call('SMEMBERS', ARGV[3])) do
                        table.insert(targets, ARGV[1] .. ":" .. ARGV[2] .. ":c:" .. tag)
                    end

                    table.insert(targets, ARGV[3])

                    redis.call('DEL', unpack(targets))
                "#
            }
        }
    }
}
<|endoftext|>


<commit_before>use super::internal_prelude::*;

pub static CREATE_CANVAS: CreateCanvasNodeDef = CreateCanvasNodeDef{};

#[derive(Debug,Clone)]
pub struct CreateCanvasNodeDef{}

impl CreateCanvasNodeDef{
    fn get(&self, n: &NodeParams) -> Result<(usize, usize, PixelFormat, s::Color)>{
        if let &NodeParams::Json(s::Node::CreateCanvas { format,
                                    w,
                                    h,
                                    ref color }) = n{
            let max_dimension = 2000000; // 2million

            if w < 1 || w > max_dimension {
                Err(nerror!(::ErrorKind::InvalidCoordinates, "canvas width ({}) must be greater than zero and less than {}.", w, max_dimension))
            } else if h < 1 || h > max_dimension {
                Err(nerror!(::ErrorKind::InvalidCoordinates, "canvas height ({}) must be greater than zero and less than {}.", w, max_dimension))
            } else if h * w > 100000000 {
                Err(nerror!(::ErrorKind::InvalidCoordinates, "canvas size ({}) cannot exceed 100 megapixels.", w))
            } else if format == ffi::PixelFormat::Gray8 {
                Err(nerror!(::ErrorKind::InvalidNodeParams, "canvas format cannot be grayscale; single-channel grayscale bitmaps are not yet supported in Imageflow"))
            }else if format != ffi::PixelFormat::Bgr24 && format != ffi::PixelFormat::Bgr32 && format != ffi::PixelFormat::Bgra32{
                Err(nerror!(::ErrorKind::InvalidNodeParams, "canvas format {:?} not recognized", format))
            } else {
                Ok((w,h,format,color.clone()))
            }
        }else {
            Err(nerror!(::ErrorKind::NodeParamsMismatch))
        }
    }
}
impl NodeDef for CreateCanvasNodeDef {
    fn fqn(&self) -> &'static str {
        "imazen.create_canvas"
    }

    fn edges_required(&self, p: &NodeParams) -> Result<(EdgesIn, EdgesOut)> {
        Ok((EdgesIn::NoInput, EdgesOut::Any))
    }

    fn validate_params(&self, n: &NodeParams) -> Result<()> {
        self.get(n).map(|_| ())
    }

    fn estimate(&self, ctx: &mut OpCtxMut, ix: NodeIndex) -> Result<FrameEstimate> {
        self.get(&ctx.weight(ix).params).map(|(w, h, format, _)| {
            FrameEstimate::Some(FrameInfo {
                w: w as i32,
                h: h as i32,
                fmt: format,
            })
        })
    }
    fn can_execute(&self) -> bool{
        true
    }
    fn execute(&self, ctx: &mut OpCtxMut, ix: NodeIndex) -> Result<NodeResult> {
        match self.get(&ctx.weight(ix).params) {
            Ok((w, h, format, color)) => {
                let flow_pointer = ctx.flow_c();
                let weight = &mut ctx.weight_mut(ix);

                unsafe {

                    let ptr =
                        ::ffi::flow_bitmap_bgra_create(flow_pointer, w as i32, h as i32, true, format);
                    if ptr.is_null() {
                        return Err(cerror!(ctx.c, "Failed to allocate {}x{}x{} bitmap ({} bytes). Reduce dimensions or increase RAM.", w, h, format.bytes(), w * h * format.bytes()))
                    }
                    let color_val = color.clone();
                    let color_srgb_argb = color_val.clone().to_u32_bgra().unwrap();
                    (*ptr).compositing_mode = ::ffi::BitmapCompositingMode::ReplaceSelf;
                    if color_val != s::Color::Transparent {
                        if !ffi::flow_bitmap_bgra_fill_rect(flow_pointer,
                                                            ptr,
                                                            0,
                                                            0,
                                                            w as u32,
                                                            h as u32,
                                                            color_srgb_argb) {
                            return Err(cerror!(ctx.c, "Failed to fill rectangle"))

                        }
                        (*ptr).compositing_mode = ::ffi::BitmapCompositingMode::BlendWithMatte;
                    }

                    (*ptr).matte_color = mem::transmute(color_srgb_argb);

                    weight.result = NodeResult::Frame(ptr);

                    Ok(NodeResult::Frame(ptr))
                }
            },
            Err(e) => Err(e)
        }
    }
}

<commit_msg>Prevent canvas creation with bgr24<commit_after>use super::internal_prelude::*;

pub static CREATE_CANVAS: CreateCanvasNodeDef = CreateCanvasNodeDef{};

#[derive(Debug,Clone)]
pub struct CreateCanvasNodeDef{}

impl CreateCanvasNodeDef{
    fn get(&self, n: &NodeParams) -> Result<(usize, usize, PixelFormat, s::Color)>{
        if let &NodeParams::Json(s::Node::CreateCanvas { format,
                                    w,
                                    h,
                                    ref color }) = n{
            let max_dimension = 2000000; // 2million

            if w < 1 || w > max_dimension {
                Err(nerror!(::ErrorKind::InvalidCoordinates, "canvas width ({}) must be greater than zero and less than {}.", w, max_dimension))
            } else if h < 1 || h > max_dimension {
                Err(nerror!(::ErrorKind::InvalidCoordinates, "canvas height ({}) must be greater than zero and less than {}.", w, max_dimension))
            } else if h * w > 100000000 {
                Err(nerror!(::ErrorKind::InvalidCoordinates, "canvas size ({}) cannot exceed 100 megapixels.", w))
            } else if format == ffi::PixelFormat::Gray8 {
                Err(nerror!(::ErrorKind::InvalidNodeParams, "canvas format cannot be grayscale; single-channel grayscale bitmaps are not yet supported in Imageflow"))
            }else if format == ffi::PixelFormat::Bgr24{
                Err(nerror!(::ErrorKind::InvalidNodeParams, "canvas format {:?} not permitted. Use Bgr32 instead", format))
            }else if format != ffi::PixelFormat::Bgr24 && format != ffi::PixelFormat::Bgr32 && format != ffi::PixelFormat::Bgra32 {
                Err(nerror!(::ErrorKind::InvalidNodeParams, "canvas format {:?} not recognized", format))
            } else {
                Ok((w,h,format,color.clone()))
            }
        }else {
            Err(nerror!(::ErrorKind::NodeParamsMismatch))
        }
    }
}
impl NodeDef for CreateCanvasNodeDef {
    fn fqn(&self) -> &'static str {
        "imazen.create_canvas"
    }

    fn edges_required(&self, p: &NodeParams) -> Result<(EdgesIn, EdgesOut)> {
        Ok((EdgesIn::NoInput, EdgesOut::Any))
    }

    fn validate_params(&self, n: &NodeParams) -> Result<()> {
        self.get(n).map(|_| ())
    }

    fn estimate(&self, ctx: &mut OpCtxMut, ix: NodeIndex) -> Result<FrameEstimate> {
        self.get(&ctx.weight(ix).params).map(|(w, h, format, _)| {
            FrameEstimate::Some(FrameInfo {
                w: w as i32,
                h: h as i32,
                fmt: format,
            })
        })
    }
    fn can_execute(&self) -> bool{
        true
    }
    fn execute(&self, ctx: &mut OpCtxMut, ix: NodeIndex) -> Result<NodeResult> {
        match self.get(&ctx.weight(ix).params) {
            Ok((w, h, format, color)) => {
                let flow_pointer = ctx.flow_c();
                let weight = &mut ctx.weight_mut(ix);

                unsafe {

                    let ptr =
                        ::ffi::flow_bitmap_bgra_create(flow_pointer, w as i32, h as i32, true, format);
                    if ptr.is_null() {
                        return Err(cerror!(ctx.c, "Failed to allocate {}x{}x{} bitmap ({} bytes). Reduce dimensions or increase RAM.", w, h, format.bytes(), w * h * format.bytes()))
                    }
                    let color_val = color.clone();
                    let color_srgb_argb = color_val.clone().to_u32_bgra().unwrap();
                    (*ptr).compositing_mode = ::ffi::BitmapCompositingMode::ReplaceSelf;
                    if color_val != s::Color::Transparent {
                        if !ffi::flow_bitmap_bgra_fill_rect(flow_pointer,
                                                            ptr,
                                                            0,
                                                            0,
                                                            w as u32,
                                                            h as u32,
                                                            color_srgb_argb) {
                            return Err(cerror!(ctx.c, "Failed to fill rectangle"))

                        }
                        (*ptr).compositing_mode = ::ffi::BitmapCompositingMode::BlendWithMatte;
                    }

                    (*ptr).matte_color = mem::transmute(color_srgb_argb);

                    weight.result = NodeResult::Frame(ptr);

                    Ok(NodeResult::Frame(ptr))
                }
            },
            Err(e) => Err(e)
        }
    }
}

<|endoftext|>


<commit_before>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * 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.
 */

use crate::config::EndPoint;
use crate::extensions::filter_registry::CreateFilterArgs;
use crate::extensions::{Error, Filter, FilterFactory};
use metrics::Metrics;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::oneshot::{channel, Sender};
use tokio::time::{self, Instant};

mod metrics;

/// RateLimitFilter applies rate limiting to packets flowing through the proxy
///
/// # Configuration
///
/// ```yaml
/// local:
///   port: 7000 # the port to receive traffic to locally
/// filters:
///   - name: quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit
///     config:
///       max_packets: 10
///       period: 500ms
/// client:
///   addresses:
///     - 127.0.0.1:7001
///   connection_id: 1x7ijy6
/// ```
///  `config.max_packets` is the maximum number of packets allowed
///  to be forwarded by the rate limiter in a given duration.
///  `config.period` (optional) is the duration during which config.max_packets applies.
///  If none is provided, it defaults to 1 second.
///

/// Config represents a RateLimitFilter's configuration.
#[derive(Serialize, Deserialize, Debug)]
struct Config {
    /// max_packets is the maximum number of packets allowed
    /// to be forwarded by the rate limiter in a given duration.
    max_packets: usize,
    /// period is the duration during which max_packets applies.
    /// If none is provided, it defaults to 1 second.
    #[serde(with = "humantime_serde")]
    period: Option<Duration>,
}

/// Creates instances of RateLimitFilter.
#[derive(Default)]
pub struct RateLimitFilterFactory;

/// A filter that implements rate limiting on packets based on
/// the token-bucket algorithm.
/// Packets that violate the rate limit are dropped.
/// It only applies rate limiting on packets that are destined for the
/// proxy's endpoints. All other packets flow through the filter untouched.
struct RateLimitFilter {
    /// available_tokens is how many tokens are left in the bucket any
    /// any given moment.
    available_tokens: Arc<AtomicUsize>,
    /// metrics reporter for this filter.
    metrics: Metrics,
    /// shutdown_tx signals the spawned token refill future to exit.
    shutdown_tx: Option<Sender<()>>,
}

impl FilterFactory for RateLimitFilterFactory {
    fn name(&self) -> String {
        "quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit".into()
    }

    fn create_filter(&self, args: CreateFilterArgs) -> Result<Box<dyn Filter>, Error> {
        let config: Config = serde_yaml::to_string(args.config)
            .and_then(|raw_config| serde_yaml::from_str(raw_config.as_str()))
            .map_err(|err| Error::DeserializeFailed(err.to_string()))?;

        match config.period {
            Some(period) if period.lt(&Duration::from_millis(100)) => Err(Error::FieldInvalid {
                field: "period".into(),
                reason: "value must be at least 100ms".into(),
            }),
            _ => Ok(Box::new(RateLimitFilter::new(
                config,
                Metrics::new(&args.metrics_registry)?,
            ))),
        }
    }
}

impl RateLimitFilter {
    /// new returns a new RateLimitFilter. It spawns a future in the background
    /// that periodically refills the rate limiter's tokens.
    fn new(config: Config, metrics: Metrics) -> Self {
        let (shutdown_tx, mut shutdown_rx) = channel();

        let tokens = Arc::new(AtomicUsize::new(config.max_packets));

        let max_tokens = config.max_packets;
        let period = config.period.unwrap_or(Duration::from_secs(1));
        let available_tokens = tokens.clone();
        let _ = tokio::spawn(async move {
            let mut interval = time::interval_at(Instant::now() + period, period);
            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        // Refill tokens.
                        let mut refilled = false;
                        while !refilled {
                            let remaining_tokens = available_tokens.load(Ordering::Relaxed);

                            refilled = available_tokens.compare_and_swap(
                                remaining_tokens,
                                max_tokens,
                                Ordering::Relaxed) == remaining_tokens;
                        }
                    },
                    _ = &mut shutdown_rx => {
                        return;
                    }
                }
            }
        });

        RateLimitFilter {
            available_tokens: tokens,
            metrics,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// acquire_token is called on behalf of every packet that is eligible
    /// for rate limiting. It returns whether there exists a token in the current
    /// period - determining whether or not the packet should be forwarded or dropped.
    fn acquire_token(&self) -> Option<()> {
        loop {
            let remaining_tokens = self.available_tokens.load(Ordering::Relaxed);

            if remaining_tokens == 0 {
                return None;
            }

            if self.available_tokens.compare_and_swap(
                remaining_tokens,
                remaining_tokens - 1,
                Ordering::Relaxed,
            ) == remaining_tokens
            {
                return Some(());
            }
        }
    }
}

impl Drop for RateLimitFilter {
    fn drop(&mut self) {
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            shutdown_tx.send(()).ok();
        }
    }
}

impl Filter for RateLimitFilter {
    fn on_downstream_receive(
        &self,
        endpoints: &[EndPoint],
        _from: SocketAddr,
        contents: Vec<u8>,
    ) -> Option<(Vec<EndPoint>, Vec<u8>)> {
        self.acquire_token()
            .map(|()| (endpoints.to_vec(), contents))
            .or_else(|| {
                self.metrics.packets_dropped_total.inc();
                None
            })
    }

    fn on_upstream_receive(
        &self,
        _endpoint: &EndPoint,
        _from: SocketAddr,
        _to: SocketAddr,
        contents: Vec<u8>,
    ) -> Option<Vec<u8>> {
        Some(contents)
    }
}

#[cfg(test)]
mod tests {
    use crate::config::EndPoint;
    use crate::extensions::filters::local_rate_limit::metrics::Metrics;
    use crate::extensions::filters::local_rate_limit::{Config, RateLimitFilter};
    use crate::extensions::Filter;
    use prometheus::Registry;
    use std::time::Duration;
    use tokio::time;

    fn rate_limiter(config: Config) -> RateLimitFilter {
        RateLimitFilter::new(config, Metrics::new(&Registry::default()).unwrap())
    }

    #[tokio::test]
    async fn initially_available_tokens() {
        // Test that we always start with the max number of tokens available.
        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_exhaustion_and_refill() {
        let r = rate_limiter(Config {
            max_packets: 2,
            period: Some(Duration::from_millis(100)),
        });

        // Exhaust tokens
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);

        // Wait for refill
        time::delay_for(Duration::from_millis(110)).await;

        // Exhaust tokens again.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_refill_maximum() {
        // Test that we never refill more than the max_tokens specified.

        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        // Use up some of the tokens.
        assert_eq!(r.acquire_token(), Some(()));

        // Wait for refill
        time::delay_for(Duration::from_millis(110)).await;

        // Refill should not go over max token limit.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn filter_with_no_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 0,
            period: Some(Duration::from_millis(100)),
        });

        // Check that other routes are not affected.
        assert_eq!(
            r.on_upstream_receive(
                &EndPoint::new("e".to_string(), "127.0.0.1:8081".parse().unwrap(), vec![]),
                "127.0.0.1:8080".parse().unwrap(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9]
            ),
            Some(vec![9])
        );

        // Check that we're rate limited.
        assert_eq!(
            r.on_downstream_receive(&vec![], "127.0.0.1:8080".parse().unwrap(), vec![9]),
            None
        );
    }

    #[tokio::test]
    async fn filter_with_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 1,
            period: Some(Duration::from_millis(100)),
        });

        assert_eq!(
            r.on_downstream_receive(&vec![], "127.0.0.1:8080".parse().unwrap(), vec![9]),
            Some((vec![], vec![9]))
        );
        // We should be out of tokens now.
        assert_eq!(None, r.acquire_token());

        // Check that other routes are not affected.
        assert_eq!(
            r.on_upstream_receive(
                &EndPoint::new("e".to_string(), "127.0.0.1:8081".parse().unwrap(), vec![]),
                "127.0.0.1:8080".parse().unwrap(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9]
            ),
            Some(vec![9])
        );
    }
}
<commit_msg>Documentation for LocalRateLimit Filter Metrics (#88)<commit_after>/*
 * Copyright 2020 Google LLC All Rights Reserved.
 *
 * 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.
 */

use crate::config::EndPoint;
use crate::extensions::filter_registry::CreateFilterArgs;
use crate::extensions::{Error, Filter, FilterFactory};
use metrics::Metrics;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::oneshot::{channel, Sender};
use tokio::time::{self, Instant};

mod metrics;

/// RateLimitFilter applies rate limiting to packets flowing through the proxy
///
/// # Configuration
///
/// ```yaml
/// local:
///   port: 7000 # the port to receive traffic to locally
/// filters:
///   - name: quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit
///     config:
///       max_packets: 10
///       period: 500ms
/// client:
///   addresses:
///     - 127.0.0.1:7001
///   connection_id: 1x7ijy6
/// ```
///  `config.max_packets` is the maximum number of packets allowed
///  to be forwarded by the rate limiter in a given duration.
///  `config.period` (optional) is the duration during which config.max_packets applies.
///  If none is provided, it defaults to 1 second.
///
/// # Metrics
///
/// `quilkin_filter_LocalRateLimit_packets_dropped`: Total number of packets dropped due to rate limiting
///

/// Config represents a RateLimitFilter's configuration.
#[derive(Serialize, Deserialize, Debug)]
struct Config {
    /// max_packets is the maximum number of packets allowed
    /// to be forwarded by the rate limiter in a given duration.
    max_packets: usize,
    /// period is the duration during which max_packets applies.
    /// If none is provided, it defaults to 1 second.
    #[serde(with = "humantime_serde")]
    period: Option<Duration>,
}

/// Creates instances of RateLimitFilter.
#[derive(Default)]
pub struct RateLimitFilterFactory;

/// A filter that implements rate limiting on packets based on
/// the token-bucket algorithm.
/// Packets that violate the rate limit are dropped.
/// It only applies rate limiting on packets that are destined for the
/// proxy's endpoints. All other packets flow through the filter untouched.
struct RateLimitFilter {
    /// available_tokens is how many tokens are left in the bucket any
    /// any given moment.
    available_tokens: Arc<AtomicUsize>,
    /// metrics reporter for this filter.
    metrics: Metrics,
    /// shutdown_tx signals the spawned token refill future to exit.
    shutdown_tx: Option<Sender<()>>,
}

impl FilterFactory for RateLimitFilterFactory {
    fn name(&self) -> String {
        "quilkin.extensions.filters.local_rate_limit.v1alpha1.LocalRateLimit".into()
    }

    fn create_filter(&self, args: CreateFilterArgs) -> Result<Box<dyn Filter>, Error> {
        let config: Config = serde_yaml::to_string(args.config)
            .and_then(|raw_config| serde_yaml::from_str(raw_config.as_str()))
            .map_err(|err| Error::DeserializeFailed(err.to_string()))?;

        match config.period {
            Some(period) if period.lt(&Duration::from_millis(100)) => Err(Error::FieldInvalid {
                field: "period".into(),
                reason: "value must be at least 100ms".into(),
            }),
            _ => Ok(Box::new(RateLimitFilter::new(
                config,
                Metrics::new(&args.metrics_registry)?,
            ))),
        }
    }
}

impl RateLimitFilter {
    /// new returns a new RateLimitFilter. It spawns a future in the background
    /// that periodically refills the rate limiter's tokens.
    fn new(config: Config, metrics: Metrics) -> Self {
        let (shutdown_tx, mut shutdown_rx) = channel();

        let tokens = Arc::new(AtomicUsize::new(config.max_packets));

        let max_tokens = config.max_packets;
        let period = config.period.unwrap_or(Duration::from_secs(1));
        let available_tokens = tokens.clone();
        let _ = tokio::spawn(async move {
            let mut interval = time::interval_at(Instant::now() + period, period);
            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        // Refill tokens.
                        let mut refilled = false;
                        while !refilled {
                            let remaining_tokens = available_tokens.load(Ordering::Relaxed);

                            refilled = available_tokens.compare_and_swap(
                                remaining_tokens,
                                max_tokens,
                                Ordering::Relaxed) == remaining_tokens;
                        }
                    },
                    _ = &mut shutdown_rx => {
                        return;
                    }
                }
            }
        });

        RateLimitFilter {
            available_tokens: tokens,
            metrics,
            shutdown_tx: Some(shutdown_tx),
        }
    }

    /// acquire_token is called on behalf of every packet that is eligible
    /// for rate limiting. It returns whether there exists a token in the current
    /// period - determining whether or not the packet should be forwarded or dropped.
    fn acquire_token(&self) -> Option<()> {
        loop {
            let remaining_tokens = self.available_tokens.load(Ordering::Relaxed);

            if remaining_tokens == 0 {
                return None;
            }

            if self.available_tokens.compare_and_swap(
                remaining_tokens,
                remaining_tokens - 1,
                Ordering::Relaxed,
            ) == remaining_tokens
            {
                return Some(());
            }
        }
    }
}

impl Drop for RateLimitFilter {
    fn drop(&mut self) {
        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            shutdown_tx.send(()).ok();
        }
    }
}

impl Filter for RateLimitFilter {
    fn on_downstream_receive(
        &self,
        endpoints: &[EndPoint],
        _from: SocketAddr,
        contents: Vec<u8>,
    ) -> Option<(Vec<EndPoint>, Vec<u8>)> {
        self.acquire_token()
            .map(|()| (endpoints.to_vec(), contents))
            .or_else(|| {
                self.metrics.packets_dropped_total.inc();
                None
            })
    }

    fn on_upstream_receive(
        &self,
        _endpoint: &EndPoint,
        _from: SocketAddr,
        _to: SocketAddr,
        contents: Vec<u8>,
    ) -> Option<Vec<u8>> {
        Some(contents)
    }
}

#[cfg(test)]
mod tests {
    use crate::config::EndPoint;
    use crate::extensions::filters::local_rate_limit::metrics::Metrics;
    use crate::extensions::filters::local_rate_limit::{Config, RateLimitFilter};
    use crate::extensions::Filter;
    use prometheus::Registry;
    use std::time::Duration;
    use tokio::time;

    fn rate_limiter(config: Config) -> RateLimitFilter {
        RateLimitFilter::new(config, Metrics::new(&Registry::default()).unwrap())
    }

    #[tokio::test]
    async fn initially_available_tokens() {
        // Test that we always start with the max number of tokens available.
        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_exhaustion_and_refill() {
        let r = rate_limiter(Config {
            max_packets: 2,
            period: Some(Duration::from_millis(100)),
        });

        // Exhaust tokens
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);

        // Wait for refill
        time::delay_for(Duration::from_millis(110)).await;

        // Exhaust tokens again.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn token_refill_maximum() {
        // Test that we never refill more than the max_tokens specified.

        let r = rate_limiter(Config {
            max_packets: 3,
            period: Some(Duration::from_millis(100)),
        });

        // Use up some of the tokens.
        assert_eq!(r.acquire_token(), Some(()));

        // Wait for refill
        time::delay_for(Duration::from_millis(110)).await;

        // Refill should not go over max token limit.
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), Some(()));
        assert_eq!(r.acquire_token(), None);
    }

    #[tokio::test]
    async fn filter_with_no_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 0,
            period: Some(Duration::from_millis(100)),
        });

        // Check that other routes are not affected.
        assert_eq!(
            r.on_upstream_receive(
                &EndPoint::new("e".to_string(), "127.0.0.1:8081".parse().unwrap(), vec![]),
                "127.0.0.1:8080".parse().unwrap(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9]
            ),
            Some(vec![9])
        );

        // Check that we're rate limited.
        assert_eq!(
            r.on_downstream_receive(&vec![], "127.0.0.1:8080".parse().unwrap(), vec![9]),
            None
        );
    }

    #[tokio::test]
    async fn filter_with_available_tokens() {
        let r = rate_limiter(Config {
            max_packets: 1,
            period: Some(Duration::from_millis(100)),
        });

        assert_eq!(
            r.on_downstream_receive(&vec![], "127.0.0.1:8080".parse().unwrap(), vec![9]),
            Some((vec![], vec![9]))
        );
        // We should be out of tokens now.
        assert_eq!(None, r.acquire_token());

        // Check that other routes are not affected.
        assert_eq!(
            r.on_upstream_receive(
                &EndPoint::new("e".to_string(), "127.0.0.1:8081".parse().unwrap(), vec![]),
                "127.0.0.1:8080".parse().unwrap(),
                "127.0.0.1:8080".parse().unwrap(),
                vec![9]
            ),
            Some(vec![9])
        );
    }
}
<|endoftext|>


<commit_before>//
// base32n.rs - new base32 module
//
// The Base32 Alphabet
//
// Value Encoding  Value Encoding  Value Encoding  Value Encoding
//     0 A             9 J            18 S            27 3
//     1 B            10 K            19 T            28 4
//     2 C            11 L            20 U            29 5
//     3 D            12 M            21 V            30 6
//     4 E            13 N            22 W            31 7
//     5 F            14 O            23 X
//     6 G            15 P            24 Y         (pad) =
//     7 H            16 Q            25 Z
//     8 I            17 R            26 2
//
// The "Extended Hex" Base 32 Alphabet
//
// Value Encoding  Value Encoding  Value Encoding  Value Encoding
//     0 0             9 9            18 I            27 R
//     1 1            10 A            19 J            28 S
//     2 2            11 B            20 K            29 T
//     3 3            12 C            21 L            30 U
//     4 4            13 D            22 M            31 V
//     5 5            14 E            23 N
//     6 6            15 F            24 O         (pad) =
//     7 7            16 G            25 P
//     8 8            17 H            26 Q
//

use std;

import vec::len;

const PAD: u8 = 61u8;

iface enc {
    fn encode(dst: [mutable u8], src: [u8]);
    fn encode_h(dst: [mutable u8], src: [u8]);
    fn encode_bytes(src: [u8]) -> [u8];
    fn encode_bytes_h(src: [u8]) -> [u8];
    fn decode(dst: [mutable u8], src: [u8]);
    fn decode_h(dst: [mutable u8], src: [u8]);
    fn decode_bytes(src: [u8]) -> [u8];
    fn decode_bytes_h(src: [u8]) -> [u8];
}

fn mk() -> enc {
    type _enc = {table: [u8], table_h: [u8],
                 decode_map: [u8], decode_map_h: [u8]};

    impl of enc for _enc {
        fn encode(dst: [mutable u8], src: [u8]) {
            b32encode(self.table, dst, src);
        }
        fn encode_h(dst: [mutable u8], src: [u8]) {
            b32encode(self.table_h, dst, src);
        }
        fn encode_bytes(src: [u8]) -> [u8] {
            []
        }
        fn encode_bytes_h(src: [u8]) -> [u8] {
            []
        }
        fn decode(dst: [mutable u8], src: [u8]) {
            b32decode(self.decode_map, dst, src);
        }
        fn decode_h(dst: [mutable u8], src: [u8]) {
            b32decode(self.decode_map_h, dst, src);
        }
        fn decode_bytes(src: [u8]) -> [u8] {
            []
        }
        fn decode_bytes_h(src: [u8]) -> [u8] {
            []
        }
    }

    {table: [],
     table_h: [],
     decode_map: [],
     decode_map_h: []} as enc
}

fn b32encode(table: [u8], dst: [mutable u8], src: [u8]) {
}

fn b32decode(decode_map: [u8], dst: [mutable u8], src: [u8]) {
}

#[cfg(test)]
mod tests {
}
<commit_msg>base32n: wip<commit_after>//
// base32n.rs - new base32 module
//
// The Base32 Alphabet
//
// Value Encoding  Value Encoding  Value Encoding  Value Encoding
//     0 A             9 J            18 S            27 3
//     1 B            10 K            19 T            28 4
//     2 C            11 L            20 U            29 5
//     3 D            12 M            21 V            30 6
//     4 E            13 N            22 W            31 7
//     5 F            14 O            23 X
//     6 G            15 P            24 Y         (pad) =
//     7 H            16 Q            25 Z
//     8 I            17 R            26 2
//
// The "Extended Hex" Base 32 Alphabet
//
// Value Encoding  Value Encoding  Value Encoding  Value Encoding
//     0 0             9 9            18 I            27 R
//     1 1            10 A            19 J            28 S
//     2 2            11 B            20 K            29 T
//     3 3            12 C            21 L            30 U
//     4 4            13 D            22 M            31 V
//     5 5            14 E            23 N
//     6 6            15 F            24 O         (pad) =
//     7 7            16 G            25 P
//     8 8            17 H            26 Q
//

use std;

import vec::len;

const PAD: u8 = 61u8;

iface enc {
    fn encode(dst: [mutable u8], src: [u8]);
    fn encode_h(dst: [mutable u8], src: [u8]);
    fn encode_bytes(src: [u8]) -> [u8];
    fn encode_bytes_h(src: [u8]) -> [u8];
    fn decode(dst: [mutable u8], src: [u8]);
    fn decode_h(dst: [mutable u8], src: [u8]);
    fn decode_bytes(src: [u8]) -> [u8];
    fn decode_bytes_h(src: [u8]) -> [u8];
}

fn mk() -> enc {
    type _enc = {table: [u8], table_h: [u8],
                 decode_map: [u8], decode_map_h: [u8]};

    impl of enc for _enc {
        fn encode(dst: [mutable u8], src: [u8]) {
            b32encode(self.table, dst, src);
        }
        fn encode_h(dst: [mutable u8], src: [u8]) {
            b32encode(self.table_h, dst, src);
        }
        fn encode_bytes(src: [u8]) -> [u8] {
            []
        }
        fn encode_bytes_h(src: [u8]) -> [u8] {
            []
        }
        fn decode(dst: [mutable u8], src: [u8]) {
            b32decode(self.decode_map, dst, src);
        }
        fn decode_h(dst: [mutable u8], src: [u8]) {
            b32decode(self.decode_map_h, dst, src);
        }
        fn decode_bytes(src: [u8]) -> [u8] {
            []
        }
        fn decode_bytes_h(src: [u8]) -> [u8] {
            []
        }
    }

    let mut i = 0u8;
    let table = vec::to_mut(vec::from_elem(32u, 0u8));
    u8::range(65u8, 91u8) { |j| table[i] = j; i += 1u8; }
    u8::range(50u8, 56u8) { |j| table[i] = j; i += 1u8; }

    i = 0u8;
    let table_h = vec::to_mut(vec::from_elem(32u, 0u8));
    u8::range(48u8, 58u8) { |j| table_h[i] = j; i += 1u8; }
    u8::range(65u8, 87u8) { |j| table_h[i] = j; i += 1u8; }

    let decode_map = vec::to_mut(vec::from_elem(256u, 0xff_u8));
    let decode_map_h = vec::to_mut(vec::from_elem(256u, 0xff_u8));

    i = 0u8;
    while i < 32u8 {
        decode_map[table[i]] = i;
        i += 1u8;
    }

    i = 0u8;
    while i < 32u8 {
        decode_map_h[table_h[i]] = i;
        i += 1u8;
    }

    {table: vec::from_mut(table),
     table_h: vec::from_mut(table_h),
     decode_map: vec::from_mut(decode_map),
     decode_map_h: vec::from_mut(decode_map_h)} as enc
}

fn b32encode(table: [u8], dst: [mutable u8], src: [u8]) {
}

fn b32decode(decode_map: [u8], dst: [mutable u8], src: [u8]) {
}

#[cfg(test)]
mod tests {
}
<|endoftext|>


<commit_before>use super::{
    LENGTH_BYTES,
    MAX_LIST_SIZE,
};

#[derive(Debug)]
pub enum EncodeError {
    ListTooLong,
}

pub trait Encodable {
    fn ssz_append(&self, s: &mut SszStream);
}

/// Provides a buffer for appending ssz-encodable values.
///
/// Use the `append()` fn to add a value to a list, then use
/// the `drain()` method to consume the struct and return the
/// ssz encoded bytes.
pub struct SszStream {
    buffer: Vec<u8>
}

impl SszStream {
    /// Create a new, empty stream for writing ssz values.
    pub fn new() -> Self {
        SszStream {
            buffer: Vec::new()
        }
    }

    /// Append some ssz encodable value to the stream.
    pub fn append<E>(&mut self, value: &E) -> &mut Self
        where E: Encodable
    {
        value.ssz_append(self);
        self
    }

    /// Append some ssz encoded bytes to the stream.
    ///
    /// The length of the supplied bytes will be concatenated
    /// to the stream before the supplied bytes.
    pub fn append_encoded_val(&mut self, vec: &[u8]) {
        self.buffer.extend_from_slice(
            &encode_length(vec.len(),
            LENGTH_BYTES));
        self.buffer.extend_from_slice(&vec);
    }

    /// Append some ssz encoded bytes to the stream without calculating length
    ///
    /// The raw bytes will be concatenated to the stream.
    pub fn append_encoded_raw(&mut self, vec: &[u8]) {
        self.buffer.extend_from_slice(&vec);
    }

    /// Append some vector (list) of encodable values to the stream.
    ///
    /// The length of the list will be concatenated to the stream, then
    /// each item in the vector will be encoded and concatenated.
    pub fn append_vec<E>(&mut self, vec: &[E])
        where E: Encodable
    {
        let mut list_stream = SszStream::new();
        for item in vec {
            item.ssz_append(&mut list_stream);
        }
        self.append_encoded_val(&list_stream.drain());
    }

    /// Consume the stream and return the underlying bytes.
    pub fn drain(self) -> Vec<u8> {
        self.buffer
    }
}

/// Encode some length into a ssz size prefix.
///
/// The ssz size prefix is 4 bytes, which is treated as a continuious
/// 32bit big-endian integer.
pub fn encode_length(len: usize, length_bytes: usize) -> Vec<u8> {
    assert!(length_bytes > 0);  // For sanity
    assert!((len as usize) < 2usize.pow(length_bytes as u32 * 8));
    let mut header: Vec<u8> = vec![0; length_bytes];
    for i in 0..length_bytes {
        let offset = (length_bytes - i - 1) * 8;
        header[i] = ((len >> offset) & 0xff) as u8;
    };
    header
}

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

    #[test]
    #[should_panic]
    fn test_encode_length_0_bytes_panic() {
        encode_length(0, 0);
    }

    #[test]
    fn test_encode_length_4_bytes() {
        assert_eq!(
            encode_length(0, LENGTH_BYTES),
            vec![0; 4]
        );
        assert_eq!(
            encode_length(1, LENGTH_BYTES),
            vec![0, 0, 0, 1]
        );
        assert_eq!(
            encode_length(255, LENGTH_BYTES),
            vec![0, 0, 0, 255]
        );
        assert_eq!(
            encode_length(256, LENGTH_BYTES),
            vec![0, 0, 1, 0]
        );
        assert_eq!(
            encode_length(4294967295, LENGTH_BYTES),  // 2^(3*8) - 1
            vec![255, 255, 255, 255]
        );
    }

    #[test]
    #[should_panic]
    fn test_encode_length_4_bytes_panic() {
        encode_length(4294967296, LENGTH_BYTES);  // 2^(3*8)
    }

    #[test]
    fn test_encode_list() {
        let test_vec: Vec<u16> = vec![256; 12];
        let mut stream = SszStream::new();
        stream.append_vec(&test_vec);
        let ssz = stream.drain();

        assert_eq!(ssz.len(), 4 + (12 * 2));
        assert_eq!(ssz[0..4], *vec![0, 0, 0, 24]);
        assert_eq!(ssz[4..6], *vec![1, 0]);
    }
}
<commit_msg>needless_range ignore and derive[default] tag on sszstream<commit_after>#![cfg_attr(not(feature = "clippy"), allow(needless_range_loop))]

use super::{
    LENGTH_BYTES,
    MAX_LIST_SIZE,
};

#[derive(Debug)]
pub enum EncodeError {
    ListTooLong,
}

pub trait Encodable {
    fn ssz_append(&self, s: &mut SszStream);
}

/// Provides a buffer for appending ssz-encodable values.
///
/// Use the `append()` fn to add a value to a list, then use
/// the `drain()` method to consume the struct and return the
/// ssz encoded bytes.
#[derive(Default)]
pub struct SszStream {
    buffer: Vec<u8>
}

impl SszStream {
    /// Create a new, empty stream for writing ssz values.
    pub fn new() -> Self {
        SszStream {
            buffer: Vec::new()
        }
    }

    /// Append some ssz encodable value to the stream.
    pub fn append<E>(&mut self, value: &E) -> &mut Self
        where E: Encodable
    {
        value.ssz_append(self);
        self
    }

    /// Append some ssz encoded bytes to the stream.
    ///
    /// The length of the supplied bytes will be concatenated
    /// to the stream before the supplied bytes.
    pub fn append_encoded_val(&mut self, vec: &[u8]) {
        self.buffer.extend_from_slice(
            &encode_length(vec.len(),
            LENGTH_BYTES));
        self.buffer.extend_from_slice(&vec);
    }

    /// Append some ssz encoded bytes to the stream without calculating length
    ///
    /// The raw bytes will be concatenated to the stream.
    pub fn append_encoded_raw(&mut self, vec: &[u8]) {
        self.buffer.extend_from_slice(&vec);
    }

    /// Append some vector (list) of encodable values to the stream.
    ///
    /// The length of the list will be concatenated to the stream, then
    /// each item in the vector will be encoded and concatenated.
    pub fn append_vec<E>(&mut self, vec: &[E])
        where E: Encodable
    {
        let mut list_stream = SszStream::new();
        for item in vec {
            item.ssz_append(&mut list_stream);
        }
        self.append_encoded_val(&list_stream.drain());
    }

    /// Consume the stream and return the underlying bytes.
    pub fn drain(self) -> Vec<u8> {
        self.buffer
    }
}

/// Encode some length into a ssz size prefix.
///
/// The ssz size prefix is 4 bytes, which is treated as a continuious
/// 32bit big-endian integer.
pub fn encode_length(len: usize, length_bytes: usize) -> Vec<u8> {
    assert!(length_bytes > 0);  // For sanity
    assert!((len as usize) < 2usize.pow(length_bytes as u32 * 8));
    let mut header: Vec<u8> = vec![0; length_bytes];
    for i in 0..length_bytes {
        let offset = (length_bytes - i - 1) * 8;
        header[i] = ((len >> offset) & 0xff) as u8;
    };
    header
}

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

    #[test]
    #[should_panic]
    fn test_encode_length_0_bytes_panic() {
        encode_length(0, 0);
    }

    #[test]
    fn test_encode_length_4_bytes() {
        assert_eq!(
            encode_length(0, LENGTH_BYTES),
            vec![0; 4]
        );
        assert_eq!(
            encode_length(1, LENGTH_BYTES),
            vec![0, 0, 0, 1]
        );
        assert_eq!(
            encode_length(255, LENGTH_BYTES),
            vec![0, 0, 0, 255]
        );
        assert_eq!(
            encode_length(256, LENGTH_BYTES),
            vec![0, 0, 1, 0]
        );
        assert_eq!(
            encode_length(4294967295, LENGTH_BYTES),  // 2^(3*8) - 1
            vec![255, 255, 255, 255]
        );
    }

    #[test]
    #[should_panic]
    fn test_encode_length_4_bytes_panic() {
        encode_length(4294967296, LENGTH_BYTES);  // 2^(3*8)
    }

    #[test]
    fn test_encode_list() {
        let test_vec: Vec<u16> = vec![256; 12];
        let mut stream = SszStream::new();
        stream.append_vec(&test_vec);
        let ssz = stream.drain();

        assert_eq!(ssz.len(), 4 + (12 * 2));
        assert_eq!(ssz[0..4], *vec![0, 0, 0, 24]);
        assert_eq!(ssz[4..6], *vec![1, 0]);
    }
}
<|endoftext|>


<commit_before>/// Type of a data packet.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum DataPacketOpcode {
    ConfirmedPacket,
    UnconfirmedPacket,
    ResponsePacket,
    TrunkingPacket,
}

impl DataPacketOpcode {
    /// Convert a symbolic type to its associated identifier.
    pub fn to_bits(&self) -> u8 {
        use self::DataPacketOpcode::*;

        match *self {
            ConfirmedPacket => 0b10110,
            UnconfirmedPacket => 0b10101,
            ResponsePacket => 0b00011,
            TrunkingPacket => 0b10111,
        }
    }

    /// Parse a packet type from the given 5 bits.
    pub fn from_bits(bits: u8) -> Option<DataPacketOpcode> {
        use self::DataPacketOpcode::*;

        assert!(bits >> 5 == 0);

        match bits {
            0b10110 => Some(ConfirmedPacket),
            0b10101 => Some(UnconfirmedPacket),
            0b00011 => Some(ResponsePacket),
            0b10111 => Some(TrunkingPacket),
            _ => None,
        }
    }
}

#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum ServiceAccessPoint {
    UnencryptedUserData,
    EncryptedUserData,
    CircuitData,
    CircuitDataControl,
    PacketData,
    ARP,
    SNDCPControl,
    ExtendedAddressing,
    RegistrationAuth,
    ChannelReassignment,
    SystemConfiguration,
    Loopback,
    Statistics,
    OutOfService,
    Paging,
    Configuration,
    UnencryptedKeyManagement,
    EncryptedKeyManagement,
    TrunkingControl,
    EncryptedTrunkingControl,
}

impl ServiceAccessPoint {
    pub fn from_bits(bits: u8) -> Option<ServiceAccessPoint> {
        use self::ServiceAccessPoint::*;

        assert!(bits >> 6 == 0);

        match bits {
            0x00 => Some(UnencryptedUserData),
            0x01 => Some(EncryptedUserData),
            0x02 => Some(CircuitData),
            0x03 => Some(CircuitDataControl),
            0x04 => Some(PacketData),
            0x05 => Some(ARP),
            0x06 => Some(SNDCPControl),
            0x1F => Some(ExtendedAddressing),
            0x20 => Some(RegistrationAuth),
            0x21 => Some(ChannelReassignment),
            0x22 => Some(SystemConfiguration),
            0x23 => Some(Loopback),
            0x24 => Some(Statistics),
            0x25 => Some(OutOfService),
            0x26 => Some(Paging),
            0x27 => Some(Configuration),
            0x28 => Some(UnencryptedKeyManagement),
            0x29 => Some(EncryptedKeyManagement),
            0x3D => Some(TrunkingControl),
            0x3F => Some(EncryptedTrunkingControl),
            _ => None,
        }
    }

    pub fn to_bits(&self) -> u8 {
        use self::ServiceAccessPoint::*;

        match *self {
            UnencryptedUserData => 0x00,
            EncryptedUserData => 0x01,
            CircuitData => 0x02,
            CircuitDataControl => 0x03,
            PacketData => 0x04,
            ARP => 0x05,
            SNDCPControl => 0x06,
            ExtendedAddressing => 0x1F,
            RegistrationAuth => 0x20,
            ChannelReassignment => 0x21,
            SystemConfiguration => 0x22,
            Loopback => 0x23,
            Statistics => 0x24,
            OutOfService => 0x25,
            Paging => 0x26,
            Configuration => 0x27,
            UnencryptedKeyManagement => 0x28,
            EncryptedKeyManagement => 0x29,
            TrunkingControl => 0x3D,
            EncryptedTrunkingControl => 0x3F,
        }
    }
}

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

    #[test]
    #[should_panic]
    fn test_sap_validate() {
        ServiceAccessPoint::from_bits(0b11111111);
    }
}
<commit_msg>Fix up docs for data::fields<commit_after>//! Data packet fields.

/// Data packet type present in every header.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum DataPacketOpcode {
    /// Confirmed packet that requires an acknowledgement response from the recipient.
    ConfirmedPacket,
    /// Unconfirmed packet that doesn't require an acknowledgement response from the
    /// recipient.
    UnconfirmedPacket,
    /// Response to the sender of a confirmed data packet.
    ResponsePacket,
    /// Multiblock trunking data packet.
    TrunkingPacket,
}

impl DataPacketOpcode {
    /// Convert packet type to its 5-bit representation.
    pub fn to_bits(&self) -> u8 {
        use self::DataPacketOpcode::*;

        match *self {
            ConfirmedPacket => 0b10110,
            UnconfirmedPacket => 0b10101,
            ResponsePacket => 0b00011,
            TrunkingPacket => 0b10111,
        }
    }

    /// Parse a packet type from the given 5 bits.
    pub fn from_bits(bits: u8) -> Option<DataPacketOpcode> {
        use self::DataPacketOpcode::*;

        assert!(bits >> 5 == 0);

        match bits {
            0b10110 => Some(ConfirmedPacket),
            0b10101 => Some(UnconfirmedPacket),
            0b00011 => Some(ResponsePacket),
            0b10111 => Some(TrunkingPacket),
            _ => None,
        }
    }
}

/// Destination service for data packet.
#[derive(Copy, Clone, Eq, PartialEq, Debug)]
pub enum ServiceAccessPoint {
    UnencryptedUserData,
    EncryptedUserData,
    CircuitData,
    CircuitDataControl,
    PacketData,
    ARP,
    SNDCPControl,
    ExtendedAddressing,
    RegistrationAuth,
    ChannelReassignment,
    SystemConfiguration,
    Loopback,
    Statistics,
    OutOfService,
    Paging,
    Configuration,
    UnencryptedKeyManagement,
    EncryptedKeyManagement,
    TrunkingControl,
    EncryptedTrunkingControl,
}

impl ServiceAccessPoint {
    /// Convert the given 6 bits to a SAP identifier.
    pub fn from_bits(bits: u8) -> Option<ServiceAccessPoint> {
        use self::ServiceAccessPoint::*;

        assert!(bits >> 6 == 0);

        match bits {
            0x00 => Some(UnencryptedUserData),
            0x01 => Some(EncryptedUserData),
            0x02 => Some(CircuitData),
            0x03 => Some(CircuitDataControl),
            0x04 => Some(PacketData),
            0x05 => Some(ARP),
            0x06 => Some(SNDCPControl),
            0x1F => Some(ExtendedAddressing),
            0x20 => Some(RegistrationAuth),
            0x21 => Some(ChannelReassignment),
            0x22 => Some(SystemConfiguration),
            0x23 => Some(Loopback),
            0x24 => Some(Statistics),
            0x25 => Some(OutOfService),
            0x26 => Some(Paging),
            0x27 => Some(Configuration),
            0x28 => Some(UnencryptedKeyManagement),
            0x29 => Some(EncryptedKeyManagement),
            0x3D => Some(TrunkingControl),
            0x3F => Some(EncryptedTrunkingControl),
            _ => None,
        }
    }

    /// Convert SAP identifier to its 6-bit representation.
    pub fn to_bits(&self) -> u8 {
        use self::ServiceAccessPoint::*;

        match *self {
            UnencryptedUserData => 0x00,
            EncryptedUserData => 0x01,
            CircuitData => 0x02,
            CircuitDataControl => 0x03,
            PacketData => 0x04,
            ARP => 0x05,
            SNDCPControl => 0x06,
            ExtendedAddressing => 0x1F,
            RegistrationAuth => 0x20,
            ChannelReassignment => 0x21,
            SystemConfiguration => 0x22,
            Loopback => 0x23,
            Statistics => 0x24,
            OutOfService => 0x25,
            Paging => 0x26,
            Configuration => 0x27,
            UnencryptedKeyManagement => 0x28,
            EncryptedKeyManagement => 0x29,
            TrunkingControl => 0x3D,
            EncryptedTrunkingControl => 0x3F,
        }
    }
}

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

    #[test]
    #[should_panic]
    fn test_sap_validate() {
        ServiceAccessPoint::from_bits(0b11111111);
    }
}
<|endoftext|>