<commit_before>use crate::agents::*;
use crate::config::LimitConfig;
use crate::crypto::SigningAlgorithm;
use crate::utils::{
    agent::*,
    redis::{locking, pubsub},
    SecureRandom,
};
use ::redis::{
    aio::MultiplexedConnection as RedisConn, pipe, AsyncCommands, Client as RedisClient,
    IntoConnectionInfo, RedisResult, Script,
};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// Internal message used to lock a key set.
struct LockKeys(SigningAlgorithm);
impl Message for LockKeys {
    type Reply = locking::LockGuard;
}

/// Internal message used to fetch a key set.
struct FetchKeys(SigningAlgorithm);
impl Message for FetchKeys {
    type Reply = RedisResult<KeySet>;
}

/// Internal message used to save a key set.
struct SaveKeys(KeySet);
impl Message for SaveKeys {
    type Reply = RedisResult<()>;
}

/// Internal message used to fetch keys and send an update to the key manager.
struct UpdateKeysLocked(SigningAlgorithm);
impl Message for UpdateKeysLocked {
    type Reply = ();
}

/// Store implementation using Redis.
pub struct RedisStore {
    /// A random unique ID for ourselves.
    id: Arc<Vec<u8>>,
    /// The connection.
    conn: RedisConn,
    /// Pubsub client.
    pubsub: pubsub::Subscriber,
    /// Locking client.
    locking: locking::LockClient,
    /// TTL of session keys
    expire_sessions: Duration,
    /// TTL of cache keys
    expire_cache: Duration,
    /// The agent used for fetching on cache miss.
    fetcher: Addr<FetchAgent>,
    /// Key manager if rotating keys are enabled.
    key_manager: Option<Addr<RotatingKeys>>,
    /// Script used to check a limit.
    limit_script: Arc<Script>,
    /// Configuration for per-email rate limiting.
    limit_per_email_config: LimitConfig,
}

impl RedisStore {
    pub async fn new(
        mut url: String,
        expire_sessions: Duration,
        expire_cache: Duration,
        limit_per_email_config: LimitConfig,
        fetcher: Addr<FetchAgent>,
        rng: SecureRandom,
    ) -> RedisResult<Self> {
        if url.starts_with("http://") {
            url = url.replace("http://", "redis://");
        } else if !url.starts_with("redis://") {
            url = format!("redis://{}", &url);
        }
        let id = Arc::new(rng.generate_async(16).await);
        let info = url.as_str().into_connection_info()?;
        let pubsub = pubsub::connect(&info).await?;
        let conn = RedisClient::open(info)?
            .get_multiplexed_tokio_connection()
            .await?;
        let locking = locking::LockClient::new(conn.clone(), pubsub.clone(), rng);

        log::warn!("Storing sessions and keys in Redis at {}", url);
        log::warn!("Please always double check this Redis and the connection to it are secure!");
        log::warn!("(This warning can't be fixed; it's a friendly reminder.)");

        let limit_script = Arc::new(Script::new(
            r"
            local count = redis.call('incr', KEYS[1])
            if count == 1 then
                redis.call('expire', KEYS[1], ARGV[1])
            end
            return count
            ",
        ));

        Ok(RedisStore {
            id,
            conn,
            pubsub,
            locking,
            expire_sessions,
            expire_cache,
            fetcher,
            key_manager: None,
            limit_script,
            limit_per_email_config,
        })
    }

    fn format_session_key(session_id: &str) -> String {
        format!("session:{}", session_id)
    }
}

impl Agent for RedisStore {
    fn started(&mut self, cx: Context<Self, AgentStarted>) {
        // Ping Redis at an interval.
        let mut conn = self.conn.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(20));
            loop {
                interval.tick().await;
                let _: String = ::redis::cmd("PING")
                    .query_async(&mut conn)
                    .await
                    .expect("Redis ping failed");
            }
        });
        cx.reply(());
    }
}

impl Handler<SaveSession> for RedisStore {
    fn handle(&mut self, message: SaveSession, cx: Context<Self, SaveSession>) {
        let mut conn = self.conn.clone();
        let ttl = self.expire_sessions;
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data = serde_json::to_string(&message.data)?;
            conn.set_ex(&key, data, ttl.as_secs() as usize).await?;
            Ok(())
        });
    }
}

impl Handler<GetSession> for RedisStore {
    fn handle(&mut self, message: GetSession, cx: Context<Self, GetSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data: String = conn.get(&key).await?;
            let data = serde_json::from_str(&data)?;
            Ok(data)
        });
    }
}

impl Handler<DeleteSession> for RedisStore {
    fn handle(&mut self, message: DeleteSession, cx: Context<Self, DeleteSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            conn.del(&key).await?;
            Ok(())
        });
    }
}

impl Handler<FetchUrlCached> for RedisStore {
    fn handle(&mut self, message: FetchUrlCached, cx: Context<Self, FetchUrlCached>) {
        let mut conn = self.conn.clone();
        let mut locking = self.locking.clone();
        let fetcher = self.fetcher.clone();
        let expire_cache = self.expire_cache;
        cx.reply_later(async move {
            let key = format!("cache:{}", message.url);
            let _lock = locking.lock(format!("lock:{}", key).as_bytes()).await;
            if let Some(data) = conn.get(key).await? {
                Ok(data)
            } else {
                let key = message.url.as_str().to_owned();
                let result = fetcher.send(FetchUrl { url: message.url }).await?;
                let ttl = std::cmp::max(expire_cache, result.max_age);
                conn.set_ex(key, result.data.clone(), ttl.as_secs() as usize)
                    .await?;
                Ok(result.data)
            }
        });
    }
}

impl Handler<IncrAndTestLimit> for RedisStore {
    fn handle(&mut self, message: IncrAndTestLimit, cx: Context<Self, IncrAndTestLimit>) {
        let mut conn = self.conn.clone();
        let script = self.limit_script.clone();
        let (key, config) = match message {
            IncrAndTestLimit::PerEmail { addr } => (
                format!("ratelimit:per-email:{}", addr),
                self.limit_per_email_config,
            ),
        };
        cx.reply_later(async move {
            let mut invocation = script.prepare_invoke();
            invocation.key(key).arg(config.duration.as_secs());
            let count: usize = invocation.invoke_async(&mut conn).await?;
            Ok(count <= config.max_count)
        });
    }
}

impl Handler<EnableRotatingKeys> for RedisStore {
    fn handle(&mut self, message: EnableRotatingKeys, cx: Context<Self, EnableRotatingKeys>) {
        let me = cx.addr().clone();
        let my_id = self.id.clone();
        let mut pubsub = self.pubsub.clone();
        self.key_manager = Some(message.key_manager.clone());
        cx.reply_later(async move {
            for signing_alg in &message.signing_algs {
                let signing_alg = *signing_alg;
                // Listen for key changes by other workers.
                let chan = format!("keys-updated:{}", signing_alg).into_bytes();
                let mut sub = pubsub.subscribe(chan).await;
                let me2 = me.clone();
                let my_id2 = my_id.clone();
                tokio::spawn(async move {
                    loop {
                        let from_id = sub.recv().await.expect("Redis keys subscription failed");
                        if from_id != *my_id2 {
                            me2.send(UpdateKeysLocked(signing_alg));
                        }
                    }
                });
                // Fetch current keys.
                me.send(UpdateKeysLocked(signing_alg)).await;
            }
        });
    }
}

impl Handler<RotateKeysLocked> for RedisStore {
    fn handle(&mut self, message: RotateKeysLocked, cx: Context<Self, RotateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let lock = me.send(LockKeys(message.0)).await;
            let key_set = me
                .send(FetchKeys(message.0))
                .await
                .expect("Failed to fetch keys from Redis");
            if let Some(key_set) = key_manager.send(RotateKeys(key_set)).await {
                me.send(SaveKeys(key_set.clone()))
                    .await
                    .expect("Failed to save keys to Redis");
                drop(lock);
                key_manager.send(UpdateKeys(key_set)).await;
            }
        });
    }
}

impl Handler<LockKeys> for RedisStore {
    fn handle(&mut self, message: LockKeys, cx: Context<Self, LockKeys>) {
        let mut locking = self.locking.clone();
        let lock_key = format!("lock:keys:{}", message.0);
        cx.reply_later(async move { locking.lock(lock_key.as_bytes()).await });
    }
}

impl Handler<FetchKeys> for RedisStore {
    fn handle(&mut self, message: FetchKeys, cx: Context<Self, FetchKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0;
        let current_key = format!("keys:{}:current", signing_alg);
        let next_key = format!("keys:{}:next", signing_alg);
        let previous_key = format!("keys:{}:previous", signing_alg);
        let mut pipe = pipe();
        pipe.atomic()
            .get(&current_key)
            .ttl(&current_key)
            .get(&next_key)
            .ttl(&next_key)
            .get(&previous_key);
        cx.reply_later(async move {
            let (current, current_ttl, next, next_ttl, previous): (
                Option<String>,
                i64,
                Option<String>,
                i64,
                Option<String>,
            ) = pipe.query_async(&mut conn).await?;
            let now = SystemTime::now();
            let parse_ttl = move |ttl: i64| -> SystemTime {
                if ttl <= 0 {
                    now
                } else {
                    now + Duration::from_secs(ttl as u64)
                }
            };
            Ok(KeySet {
                signing_alg,
                current: current.map(|value| Expiring {
                    value,
                    expires: parse_ttl(current_ttl),
                }),
                next: next.map(|value| Expiring {
                    value,
                    expires: parse_ttl(next_ttl),
                }),
                previous,
            })
        })
    }
}

impl Handler<SaveKeys> for RedisStore {
    fn handle(&mut self, message: SaveKeys, cx: Context<Self, SaveKeys>) {
        if let KeySet {
            signing_alg,
            current: Some(ref current),
            next: Some(ref next),
            ref previous,
            ..
        } = message.0
        {
            let mut conn = self.conn.clone();
            let current_key = format!("keys:{}:current", signing_alg);
            let next_key = format!("keys:{}:next", signing_alg);
            let previous_key = format!("keys:{}:previous", signing_alg);
            let mut pipe = pipe();
            fn to_unix(time: SystemTime) -> usize {
                time.duration_since(UNIX_EPOCH).unwrap().as_secs() as usize
            }
            pipe.atomic()
                .set(&current_key, &current.value)
                .expire_at(&current_key, to_unix(current.expires))
                .set(&next_key, &next.value)
                .expire_at(&next_key, to_unix(next.expires));
            if let Some(ref previous) = previous {
                pipe.set(&previous_key, previous);
            } else {
                pipe.del(&previous_key);
            }
            pipe.publish(format!("keys-updated:{}", signing_alg), &self.id[..]);
            cx.reply_later(async move { pipe.query_async(&mut conn).await });
        } else {
            panic!("Incomplete key pair in SaveKeys message");
        }
    }
}

impl Handler<UpdateKeysLocked> for RedisStore {
    fn handle(&mut self, message: UpdateKeysLocked, cx: Context<Self, UpdateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let key_set = {
                let _lock = me.send(LockKeys(message.0)).await;
                me.send(FetchKeys(message.0))
                    .await
                    .expect("Failed to fetch keys from Redis")
            };
            key_manager.send(UpdateKeys(key_set)).await;
        });
    }
}

impl StoreSender for Addr<RedisStore> {}
<commit_msg>Also use JSON to store keys in Redis, without TTL<commit_after>use crate::agents::*;
use crate::config::LimitConfig;
use crate::crypto::SigningAlgorithm;
use crate::utils::{
    agent::*,
    redis::{locking, pubsub},
    SecureRandom,
};
use ::redis::{
    aio::MultiplexedConnection as RedisConn, pipe, AsyncCommands, Client as RedisClient,
    IntoConnectionInfo, RedisResult, Script,
};
use std::sync::Arc;
use std::time::Duration;

/// Internal message used to lock a key set.
struct LockKeys(SigningAlgorithm);
impl Message for LockKeys {
    type Reply = locking::LockGuard;
}

/// Internal message used to fetch a key set.
struct FetchKeys(SigningAlgorithm);
impl Message for FetchKeys {
    type Reply = RedisResult<KeySet>;
}

/// Internal message used to save a key set.
struct SaveKeys(KeySet);
impl Message for SaveKeys {
    type Reply = RedisResult<()>;
}

/// Internal message used to fetch keys and send an update to the key manager.
struct UpdateKeysLocked(SigningAlgorithm);
impl Message for UpdateKeysLocked {
    type Reply = ();
}

/// Store implementation using Redis.
pub struct RedisStore {
    /// A random unique ID for ourselves.
    id: Arc<Vec<u8>>,
    /// The connection.
    conn: RedisConn,
    /// Pubsub client.
    pubsub: pubsub::Subscriber,
    /// Locking client.
    locking: locking::LockClient,
    /// TTL of session keys
    expire_sessions: Duration,
    /// TTL of cache keys
    expire_cache: Duration,
    /// The agent used for fetching on cache miss.
    fetcher: Addr<FetchAgent>,
    /// Key manager if rotating keys are enabled.
    key_manager: Option<Addr<RotatingKeys>>,
    /// Script used to check a limit.
    limit_script: Arc<Script>,
    /// Configuration for per-email rate limiting.
    limit_per_email_config: LimitConfig,
}

impl RedisStore {
    pub async fn new(
        mut url: String,
        expire_sessions: Duration,
        expire_cache: Duration,
        limit_per_email_config: LimitConfig,
        fetcher: Addr<FetchAgent>,
        rng: SecureRandom,
    ) -> RedisResult<Self> {
        if url.starts_with("http://") {
            url = url.replace("http://", "redis://");
        } else if !url.starts_with("redis://") {
            url = format!("redis://{}", &url);
        }
        let id = Arc::new(rng.generate_async(16).await);
        let info = url.as_str().into_connection_info()?;
        let pubsub = pubsub::connect(&info).await?;
        let conn = RedisClient::open(info)?
            .get_multiplexed_tokio_connection()
            .await?;
        let locking = locking::LockClient::new(conn.clone(), pubsub.clone(), rng);

        log::warn!("Storing sessions and keys in Redis at {}", url);
        log::warn!("Please always double check this Redis and the connection to it are secure!");
        log::warn!("(This warning can't be fixed; it's a friendly reminder.)");

        let limit_script = Arc::new(Script::new(
            r"
            local count = redis.call('incr', KEYS[1])
            if count == 1 then
                redis.call('expire', KEYS[1], ARGV[1])
            end
            return count
            ",
        ));

        Ok(RedisStore {
            id,
            conn,
            pubsub,
            locking,
            expire_sessions,
            expire_cache,
            fetcher,
            key_manager: None,
            limit_script,
            limit_per_email_config,
        })
    }

    fn format_session_key(session_id: &str) -> String {
        format!("session:{}", session_id)
    }
}

impl Agent for RedisStore {
    fn started(&mut self, cx: Context<Self, AgentStarted>) {
        // Ping Redis at an interval.
        let mut conn = self.conn.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(20));
            loop {
                interval.tick().await;
                let _: String = ::redis::cmd("PING")
                    .query_async(&mut conn)
                    .await
                    .expect("Redis ping failed");
            }
        });
        cx.reply(());
    }
}

impl Handler<SaveSession> for RedisStore {
    fn handle(&mut self, message: SaveSession, cx: Context<Self, SaveSession>) {
        let mut conn = self.conn.clone();
        let ttl = self.expire_sessions;
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data = serde_json::to_string(&message.data)?;
            conn.set_ex(&key, data, ttl.as_secs() as usize).await?;
            Ok(())
        });
    }
}

impl Handler<GetSession> for RedisStore {
    fn handle(&mut self, message: GetSession, cx: Context<Self, GetSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data: String = conn.get(&key).await?;
            let data = serde_json::from_str(&data)?;
            Ok(data)
        });
    }
}

impl Handler<DeleteSession> for RedisStore {
    fn handle(&mut self, message: DeleteSession, cx: Context<Self, DeleteSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            conn.del(&key).await?;
            Ok(())
        });
    }
}

impl Handler<FetchUrlCached> for RedisStore {
    fn handle(&mut self, message: FetchUrlCached, cx: Context<Self, FetchUrlCached>) {
        let mut conn = self.conn.clone();
        let mut locking = self.locking.clone();
        let fetcher = self.fetcher.clone();
        let expire_cache = self.expire_cache;
        cx.reply_later(async move {
            let key = format!("cache:{}", message.url);
            let _lock = locking.lock(format!("lock:{}", key).as_bytes()).await;
            if let Some(data) = conn.get(key).await? {
                Ok(data)
            } else {
                let key = message.url.as_str().to_owned();
                let result = fetcher.send(FetchUrl { url: message.url }).await?;
                let ttl = std::cmp::max(expire_cache, result.max_age);
                conn.set_ex(key, result.data.clone(), ttl.as_secs() as usize)
                    .await?;
                Ok(result.data)
            }
        });
    }
}

impl Handler<IncrAndTestLimit> for RedisStore {
    fn handle(&mut self, message: IncrAndTestLimit, cx: Context<Self, IncrAndTestLimit>) {
        let mut conn = self.conn.clone();
        let script = self.limit_script.clone();
        let (key, config) = match message {
            IncrAndTestLimit::PerEmail { addr } => (
                format!("ratelimit:per-email:{}", addr),
                self.limit_per_email_config,
            ),
        };
        cx.reply_later(async move {
            let mut invocation = script.prepare_invoke();
            invocation.key(key).arg(config.duration.as_secs());
            let count: usize = invocation.invoke_async(&mut conn).await?;
            Ok(count <= config.max_count)
        });
    }
}

impl Handler<EnableRotatingKeys> for RedisStore {
    fn handle(&mut self, message: EnableRotatingKeys, cx: Context<Self, EnableRotatingKeys>) {
        let me = cx.addr().clone();
        let my_id = self.id.clone();
        let mut pubsub = self.pubsub.clone();
        self.key_manager = Some(message.key_manager.clone());
        cx.reply_later(async move {
            for signing_alg in &message.signing_algs {
                let signing_alg = *signing_alg;
                // Listen for key changes by other workers.
                let chan = format!("keys-updated:{}", signing_alg).into_bytes();
                let mut sub = pubsub.subscribe(chan).await;
                let me2 = me.clone();
                let my_id2 = my_id.clone();
                tokio::spawn(async move {
                    loop {
                        let from_id = sub.recv().await.expect("Redis keys subscription failed");
                        if from_id != *my_id2 {
                            me2.send(UpdateKeysLocked(signing_alg));
                        }
                    }
                });
                // Fetch current keys.
                me.send(UpdateKeysLocked(signing_alg)).await;
            }
        });
    }
}

impl Handler<RotateKeysLocked> for RedisStore {
    fn handle(&mut self, message: RotateKeysLocked, cx: Context<Self, RotateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let lock = me.send(LockKeys(message.0)).await;
            let key_set = me
                .send(FetchKeys(message.0))
                .await
                .expect("Failed to fetch keys from Redis");
            if let Some(key_set) = key_manager.send(RotateKeys(key_set)).await {
                me.send(SaveKeys(key_set.clone()))
                    .await
                    .expect("Failed to save keys to Redis");
                drop(lock);
                key_manager.send(UpdateKeys(key_set)).await;
            }
        });
    }
}

impl Handler<LockKeys> for RedisStore {
    fn handle(&mut self, message: LockKeys, cx: Context<Self, LockKeys>) {
        let mut locking = self.locking.clone();
        let lock_key = format!("lock:keys:{}", message.0);
        cx.reply_later(async move { locking.lock(lock_key.as_bytes()).await });
    }
}

impl Handler<FetchKeys> for RedisStore {
    fn handle(&mut self, message: FetchKeys, cx: Context<Self, FetchKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0;
        let db_key = format!("keys:{}", signing_alg);
        cx.reply_later(async move {
            let key_set: Option<String> = conn.get(db_key).await?;
            let key_set = key_set
                .map(|data| serde_json::from_str(&data).expect("Invalid key set JSON in Redis"))
                .unwrap_or_else(|| KeySet::empty(signing_alg));
            Ok(key_set)
        })
    }
}

impl Handler<SaveKeys> for RedisStore {
    fn handle(&mut self, message: SaveKeys, cx: Context<Self, SaveKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0.signing_alg;
        let db_key = format!("keys:{}", signing_alg);
        let data = serde_json::to_string(&message.0).expect("Could not encode key set as JSON");
        let mut pipe = pipe();
        pipe.atomic()
            .set(db_key, data)
            .publish(format!("keys-updated:{}", signing_alg), &self.id[..]);
        cx.reply_later(async move { pipe.query_async(&mut conn).await });
    }
}

impl Handler<UpdateKeysLocked> for RedisStore {
    fn handle(&mut self, message: UpdateKeysLocked, cx: Context<Self, UpdateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let key_set = {
                let _lock = me.send(LockKeys(message.0)).await;
                me.send(FetchKeys(message.0))
                    .await
                    .expect("Failed to fetch keys from Redis")
            };
            key_manager.send(UpdateKeys(key_set)).await;
        });
    }
}

impl StoreSender for Addr<RedisStore> {}
<|endoftext|>


<commit_before>use std::result;

use env::Env;
use library::{self, Nullable};
use nameutil::crate_name;
use super::conversion_type::ConversionType;
use traits::*;

pub type Result = result::Result<String, String>;

impl AsStr for Result {
    #[inline]
    fn as_str(&self) -> &str {
        self.as_ref().unwrap_or_else(|s| s)
    }
}

pub fn rust_type(env: &Env, type_id: library::TypeId) -> Result {
    rust_type_full(env, type_id, Nullable(false), false)
}

fn rust_type_full(env: &Env, type_id: library::TypeId, nullable: Nullable, by_ref: bool) -> Result {
    use library::Type::*;
    use library::Fundamental::*;
    let mut skip_option = false;
    let type_ = env.library.type_(type_id);
    let mut rust_type = match *type_ {
        Fundamental(fund) => {
            let ok = |s: &str| Ok(s.into());
            let err = |s: &str| Err(s.into());
            match fund {
                None => err("()"),
                Boolean => ok("bool"),
                Int8 => ok("i8"),
                UInt8 => ok("u8"),
                Int16 => ok("i16"),
                UInt16 => ok("u16"),
                Int32 => ok("i32"),
                UInt32 => ok("u32"),
                Int64 => ok("i64"),
                UInt64 => ok("u64"),

                Int => ok("i32"),      //maybe dependent on target system
                UInt => ok("u32"),     //maybe dependent on target system

                Float => ok("f32"),
                Double => ok("f64"),

                Utf8 => if by_ref { ok("str") } else { ok("String") },
                Filename => if by_ref { ok("str") } else { ok("String") },

                Type => ok("glib::types::Type"),
                Unsupported => err("Unsupported"),
                _ => err(&format!("Fundamental: {:?}", fund)),
            }
        },

        Enumeration(ref enum_) => Ok(enum_.name.clone()),
        Bitfield(ref bitfield) => Ok(bitfield.name.clone()),
        Record(ref record) => Ok(record.name.clone()),
        Interface(ref interface) => Ok(interface.name.clone()),
        Class(ref klass) => Ok(klass.name.clone()),
        List(inner_tid) => {
            skip_option = true;
            rust_type(env, inner_tid)
                .map_any(|s| if by_ref { format!("[{}]", s) } else { format!("Vec<{}>", s) })
        }
        _ => Err(format!("Unknown rust type: {:?}", type_.get_name())),
        //TODO: check usage library::Type::get_name() when no _ in this
    };

    if type_id.ns_id != library::MAIN_NAMESPACE && type_id.ns_id != library::INTERNAL_NAMESPACE {
        let rust_type_with_prefix = rust_type.map(|s| format!("{}::{}",
            crate_name(&env.library.namespace(type_id.ns_id).name), s));
        if env.type_status(&type_id.full_name(&env.library)).ignored() {
            rust_type = Err(rust_type_with_prefix.as_str().into());
        } else {
            rust_type = rust_type_with_prefix;
        }
    }
    if by_ref {
        rust_type = rust_type.map_any(|s| format!("&{}", s));
    }
    if *nullable && !skip_option {
        match ConversionType::of(&env.library, type_id) {
            ConversionType::Pointer => rust_type = rust_type.map_any(|s| format!("Option<{}>", s)),
            _ => (),
        }
    }

    rust_type
}

pub fn used_rust_type(env: &Env, type_id: library::TypeId) -> Result {
    use library::Type::*;
    match *env.library.type_(type_id) {
        Fundamental(library::Fundamental::Type) |
            Bitfield(..) |
            Record(..) |
            Class(..) |
            Enumeration(..) |
            Interface(..) => rust_type(env, type_id),
        _ => Err("Don't need use".into()),
    }
}

pub fn parameter_rust_type(env: &Env, type_id:library::TypeId,
        direction: library::ParameterDirection, nullable: Nullable) -> Result {
    use library::Type::*;
    let type_ = env.library.type_(type_id);
    let by_ref = use_by_ref(type_, direction);
    let rust_type = rust_type_full(env, type_id, nullable, by_ref);
    match *type_ {
        Fundamental(fund) => {
            if fund == library::Fundamental::Utf8 || fund == library::Fundamental::Filename {
                match direction {
                    library::ParameterDirection::In |
                        library::ParameterDirection::Return => rust_type,
                    _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
                }
            } else {
                format_parameter(rust_type, direction)
            }
        },

        Enumeration(..) |
            Bitfield(..) => format_parameter(rust_type, direction),

        Record(..) => {
            match direction {
                _ if env.type_status(&type_id.full_name(&env.library)).ignored() => {
                    Err(format!("/*Ignored*/{}", rust_type.as_str()))
                }
                library::ParameterDirection::In |
                    library::ParameterDirection::Return => rust_type,
                _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
            }
        }

        Class(..) => {
            match direction {
                _ if env.type_status(&type_id.full_name(&env.library)).ignored() => {
                    Err(format!("/*Ignored*/{}", rust_type.as_str()))
                }
                library::ParameterDirection::In |
                    library::ParameterDirection::Return => rust_type,
                _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
            }
        }

        List(..) => {
            match direction {
                library::ParameterDirection::Return => rust_type,
                _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
            }
        }
        _ => Err(format!("Unknown rust type: {:?}", type_.get_name())),
        //TODO: check usage library::Type::get_name() when no _ in this
    }
}

#[inline]
fn format_parameter(rust_type: Result, direction: library::ParameterDirection) -> Result {
    if direction.is_out() {
        rust_type.map(|s| format!("&mut {}", s))
    } else {
        rust_type
    }
}

#[inline]
fn use_by_ref(type_: &library::Type, direction: library::ParameterDirection) -> bool {
    use library::Type::*;
    match *type_ {
        Fundamental(library::Fundamental::Utf8) |
            Fundamental(library::Fundamental::Filename) |
            Record(..) |
            Class(..) |
            List(..) => direction == library::ParameterDirection::In,
        _ => false,
    }
}
<commit_msg>Import types that are passed inside lists<commit_after>use std::result;

use env::Env;
use library::{self, Nullable};
use nameutil::crate_name;
use super::conversion_type::ConversionType;
use traits::*;

pub type Result = result::Result<String, String>;

impl AsStr for Result {
    #[inline]
    fn as_str(&self) -> &str {
        self.as_ref().unwrap_or_else(|s| s)
    }
}

pub fn rust_type(env: &Env, type_id: library::TypeId) -> Result {
    rust_type_full(env, type_id, Nullable(false), false)
}

fn rust_type_full(env: &Env, type_id: library::TypeId, nullable: Nullable, by_ref: bool) -> Result {
    use library::Type::*;
    use library::Fundamental::*;
    let mut skip_option = false;
    let type_ = env.library.type_(type_id);
    let mut rust_type = match *type_ {
        Fundamental(fund) => {
            let ok = |s: &str| Ok(s.into());
            let err = |s: &str| Err(s.into());
            match fund {
                None => err("()"),
                Boolean => ok("bool"),
                Int8 => ok("i8"),
                UInt8 => ok("u8"),
                Int16 => ok("i16"),
                UInt16 => ok("u16"),
                Int32 => ok("i32"),
                UInt32 => ok("u32"),
                Int64 => ok("i64"),
                UInt64 => ok("u64"),

                Int => ok("i32"),      //maybe dependent on target system
                UInt => ok("u32"),     //maybe dependent on target system

                Float => ok("f32"),
                Double => ok("f64"),

                Utf8 => if by_ref { ok("str") } else { ok("String") },
                Filename => if by_ref { ok("str") } else { ok("String") },

                Type => ok("glib::types::Type"),
                Unsupported => err("Unsupported"),
                _ => err(&format!("Fundamental: {:?}", fund)),
            }
        },

        Enumeration(ref enum_) => Ok(enum_.name.clone()),
        Bitfield(ref bitfield) => Ok(bitfield.name.clone()),
        Record(ref record) => Ok(record.name.clone()),
        Interface(ref interface) => Ok(interface.name.clone()),
        Class(ref klass) => Ok(klass.name.clone()),
        List(inner_tid) => {
            skip_option = true;
            rust_type(env, inner_tid)
                .map_any(|s| if by_ref { format!("[{}]", s) } else { format!("Vec<{}>", s) })
        }
        _ => Err(format!("Unknown rust type: {:?}", type_.get_name())),
        //TODO: check usage library::Type::get_name() when no _ in this
    };

    if type_id.ns_id != library::MAIN_NAMESPACE && type_id.ns_id != library::INTERNAL_NAMESPACE {
        let rust_type_with_prefix = rust_type.map(|s| format!("{}::{}",
            crate_name(&env.library.namespace(type_id.ns_id).name), s));
        if env.type_status(&type_id.full_name(&env.library)).ignored() {
            rust_type = Err(rust_type_with_prefix.as_str().into());
        } else {
            rust_type = rust_type_with_prefix;
        }
    }
    if by_ref {
        rust_type = rust_type.map_any(|s| format!("&{}", s));
    }
    if *nullable && !skip_option {
        match ConversionType::of(&env.library, type_id) {
            ConversionType::Pointer => rust_type = rust_type.map_any(|s| format!("Option<{}>", s)),
            _ => (),
        }
    }

    rust_type
}

pub fn used_rust_type(env: &Env, type_id: library::TypeId) -> Result {
    use library::Type::*;
    match *env.library.type_(type_id) {
        Fundamental(library::Fundamental::Type) |
            Bitfield(..) |
            Record(..) |
            Class(..) |
            Enumeration(..) |
            Interface(..) => rust_type(env, type_id),
        List(inner_tid) => used_rust_type(env, inner_tid),
        _ => Err("Don't need use".into()),
    }
}

pub fn parameter_rust_type(env: &Env, type_id:library::TypeId,
        direction: library::ParameterDirection, nullable: Nullable) -> Result {
    use library::Type::*;
    let type_ = env.library.type_(type_id);
    let by_ref = use_by_ref(type_, direction);
    let rust_type = rust_type_full(env, type_id, nullable, by_ref);
    match *type_ {
        Fundamental(fund) => {
            if fund == library::Fundamental::Utf8 || fund == library::Fundamental::Filename {
                match direction {
                    library::ParameterDirection::In |
                        library::ParameterDirection::Return => rust_type,
                    _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
                }
            } else {
                format_parameter(rust_type, direction)
            }
        },

        Enumeration(..) |
            Bitfield(..) => format_parameter(rust_type, direction),

        Record(..) => {
            match direction {
                _ if env.type_status(&type_id.full_name(&env.library)).ignored() => {
                    Err(format!("/*Ignored*/{}", rust_type.as_str()))
                }
                library::ParameterDirection::In |
                    library::ParameterDirection::Return => rust_type,
                _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
            }
        }

        Class(..) => {
            match direction {
                _ if env.type_status(&type_id.full_name(&env.library)).ignored() => {
                    Err(format!("/*Ignored*/{}", rust_type.as_str()))
                }
                library::ParameterDirection::In |
                    library::ParameterDirection::Return => rust_type,
                _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
            }
        }

        List(..) => {
            match direction {
                library::ParameterDirection::Return => rust_type,
                _ => Err(format!("/*Unimplemented*/{}", rust_type.as_str())),
            }
        }
        _ => Err(format!("Unknown rust type: {:?}", type_.get_name())),
        //TODO: check usage library::Type::get_name() when no _ in this
    }
}

#[inline]
fn format_parameter(rust_type: Result, direction: library::ParameterDirection) -> Result {
    if direction.is_out() {
        rust_type.map(|s| format!("&mut {}", s))
    } else {
        rust_type
    }
}

#[inline]
fn use_by_ref(type_: &library::Type, direction: library::ParameterDirection) -> bool {
    use library::Type::*;
    match *type_ {
        Fundamental(library::Fundamental::Utf8) |
            Fundamental(library::Fundamental::Filename) |
            Record(..) |
            Class(..) |
            List(..) => direction == library::ParameterDirection::In,
        _ => false,
    }
}
<|endoftext|>


<commit_before>// from: https://github.com/Maktm/FLIRTDB/blob/1f5763535e02d7cccf2f90a96a8ebaa36e9b2495/cmt/windows/libcmt_15_msvc_x86.pat#L355
//
//    3B0D........F27502F2C3F2E9...................................... 00 0000
// 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002
// ___security_cookie ^000D ___report_gsfailure    ^^^^        ^^^^^^^^^^^^^^
// bytes        ^^^^^^^^                 relocations
//
//                                                                     00 0000
// 0011                                                                     ^^
// number of bytes crc16'd
// ^^^^ crc16
// ^^^^ size of fn
//
//
// :0000 @__security_check_cookie@4
// ^^^^^ offset of match
// ^^^^^^^^^^^^^^^^^^^^^^^^^^ name
//
//
// :000B@ $failure$4
// ^^^^^^ local offset
// ^^^^^^^^^^ local name
//
//
// ^0002 ___security_cookie
// ^^^^^ reference offset
// ^^^^^^^^^^^^^^^^^^ reference name

use anyhow::Result;
use log::trace;
use nom::{
    branch::alt,
    bytes::complete::{tag, take_till1, take_while, take_while_m_n},
    combinator::{map, map_res, opt, peek},
    multi::{many0, many1},
    sequence::pair,
    IResult,
};
use thiserror::Error;

use super::{ByteSignature, FlirtSignature, Name, Offset, SigElement, Symbol, TailByte};

#[derive(Debug, Error)]
pub enum PatError {
    #[error("The pattern is not supported")]
    NotSupported,
    #[error("The .pat file is corrupt (or unsupported)")]
    CorruptPatFile,
}

fn whitespace(input: &str) -> IResult<&str, &str> {
    take_while(|c| c == ' ')(input)
}

fn is_not_newline(c: char) -> bool {
    c != '\r' && c != '\n'
}

fn is_hex_digit(c: char) -> bool {
    c.is_digit(16)
}

fn from_hex(input: &str) -> Result<u8, std::num::ParseIntError> {
    u8::from_str_radix(input, 16)
}

/// parse a single hex byte, like `AB`
fn hex(input: &str) -> IResult<&str, u8> {
    map_res(take_while_m_n(2, 2, is_hex_digit), from_hex)(input)
}

fn hex_byte(input: &str) -> IResult<&str, u8> {
    let (input, v) = hex(input)?;
    Ok((input, v))
}

fn hex_word(input: &str) -> IResult<&str, u16> {
    let (input, v1) = hex(input)?;
    let (input, v2) = hex(input)?;
    let v: u16 = ((v1 as u16) << 8) | (v2 as u16);
    Ok((input, v))
}

/// parse a single byte signature element, which is either a hex byte or a
/// wildcard.
fn sig_element(input: &str) -> IResult<&str, SigElement> {
    alt((map(hex, SigElement::Byte), map(tag(".."), |_| SigElement::Wildcard)))(input)
}

/// parse byte signature elements, hex or wildcard.
fn byte_signature(input: &str) -> IResult<&str, ByteSignature> {
    let (input, elems) = many1(sig_element)(input)?;
    Ok((input, ByteSignature(elems)))
}

/// parse a hex-encoded offset, like `0000`
/// max is 0x8000.
/// TODO: this can be negative.
fn hex_offset(input: &str) -> IResult<&str, u16> {
    hex_word(input)
}

/// parse a public offset, like `:0000`
fn public_offset(input: &str) -> IResult<&str, u16> {
    let (input, _) = tag(":")(input)?;
    let (input, offset) = hex_offset(input)?;
    let (input, _) = peek(tag(" "))(input)?;

    Ok((input, offset))
}

/// parse a local offset, like `:000B@`
fn local_offset(input: &str) -> IResult<&str, u16> {
    let (input, _) = tag(":")(input)?;
    let (input, offset) = hex_offset(input)?;
    let (input, _) = tag("@")(input)?;

    Ok((input, offset))
}

/// parse an external reference, like `^0002`
fn reference_offset(input: &str) -> IResult<&str, u16> {
    let (input, _) = tag("^")(input)?;
    let (input, offset) = hex_offset(input)?;

    Ok((input, offset))
}

fn offset(input: &str) -> IResult<&str, Offset> {
    alt((
        // this must go first, because it has trailing `@`,
        // otherwise, the same as public.
        map(local_offset, Offset::Local),
        map(public_offset, Offset::Public),
        map(reference_offset, Offset::Reference),
    ))(input)
}

fn symbol_name(input: &str) -> IResult<&str, &str> {
    take_till1(|c| c == ' ' || c == '\n')(input)
}

/// parse a (offset, symbol) pair.
///
/// note: this also consumes trailing spaces so that a sequence of these can be
/// easily parsed.
fn symbol(input: &str) -> IResult<&str, Symbol> {
    let (input, offset) = offset(input)?;
    let (input, _) = whitespace(input)?;
    let (input, name) = symbol_name(input)?;
    let (input, _) = opt(whitespace)(input)?;
    match offset {
        Offset::Public(v) => Ok((
            input,
            Symbol::Public(Name {
                name:   name.to_string(),
                offset: v as i64,
            }),
        )),
        Offset::Local(v) => Ok((
            input,
            Symbol::Local(Name {
                name:   name.to_string(),
                offset: v as i64,
            }),
        )),
        Offset::Reference(v) => Ok((
            input,
            Symbol::Reference(Name {
                name:   name.to_string(),
                offset: v as i64,
            }),
        )),
    }
}

fn symbols(input: &str) -> IResult<&str, Vec<Symbol>> {
    many1(symbol)(input)
}

// like: `(0012: 87)`
fn tail_byte(input: &str) -> IResult<&str, TailByte> {
    let (input, _) = tag("(")(input)?;
    let (input, offset) = hex_offset(input)?;
    let (input, _) = whitespace(input)?;
    let (input, value) = hex_byte(input)?;
    let (input, _) = tag(")")(input)?;

    Ok((
        input,
        TailByte {
            offset: offset as u64,
            value,
        },
    ))
}

fn tail_bytes(input: &str) -> IResult<&str, Vec<TailByte>> {
    match opt(many1(tail_byte))(input)? {
        (input, Some(tail_bytes)) => Ok((input, tail_bytes)),
        (input, None) => Ok((input, vec![])),
    }
}

fn pat_signature(input: &str) -> IResult<&str, FlirtSignature> {
    let (input, byte_sig) = byte_signature(input)?;
    let (input, _) = whitespace(input)?;
    trace!("sig: {:?}", byte_sig);

    let (input, size_of_bytes_crc16) = hex_byte(input)?;
    let (input, _) = whitespace(input)?;
    trace!("crc16 len: {:02x}", size_of_bytes_crc16);

    let (input, crc16) = hex_word(input)?;
    let (input, _) = whitespace(input)?;
    trace!("crc16: {:04x}", crc16);

    let (input, size_of_function) = hex_word(input)?;
    let (input, _) = whitespace(input)?;
    trace!("function size: {:04x}", size_of_function);

    let (input, names) = symbols(input)?;
    trace!("names: {:?}", names);

    // i'm not sure which of these comes first, footer pattern or tail bytes.
    // are they both actually valid in a .pat file!?!
    // pat.txt only describes the footer pattern (which it calls "tail bytes").
    // whereas the .sig file format uses (offset, value) tuples for tail bytes,
    // and dumpsig.exe uses a format like `(AAAA: BB)` for these.

    let (input, _) = opt(whitespace)(input)?;
    let (input, footer) = opt(byte_signature)(input)?;
    trace!("footer: {:?}", footer);

    let (input, _) = opt(whitespace)(input)?;
    let (input, tail_bytes) = tail_bytes(input)?;
    trace!("tail bytes: {:02x?}", tail_bytes);

    Ok((
        input,
        FlirtSignature {
            byte_sig,
            size_of_bytes_crc16,
            crc16,
            size_of_function: size_of_function as u64,
            names,
            footer,
            tail_bytes,
        },
    ))
}

fn comment(input: &str) -> IResult<&str, ()> {
    let (input, _) = tag("#")(input)?;
    let (input, _) = take_while(is_not_newline)(input)?;

    Ok((input, ()))
}

fn line(input: &str) -> IResult<&str, Option<FlirtSignature>> {
    if let Ok((input, _)) = comment(input) {
        return Ok((input, None));
    }

    match pat_signature(input) {
        Ok((input, sig)) => Ok((input, Some(sig))),
        Err(e) => Err(e),
    }
}

/// parse a .pat file into FLIRT signatures.
fn pat(input: &str) -> IResult<&str, Vec<FlirtSignature>> {
    // each signature is newline separated.
    // drop the newline after we've parsed it.
    let maybe_pat = map(pair(line, alt((tag("\r\n"), tag("\n")))), |p| p.0);
    let (input, sigs) = many0(maybe_pat)(input)?;

    // the file ends with `---`.
    let (input, _) = tag("---")(input)?;

    // TODO: assert there is nothing left in the file

    Ok((
        input,
        sigs.into_iter().filter(|s| s.is_some()).map(|s| s.unwrap()).collect(),
    ))
}

/// ```
/// use lancelot_flirt::pat;
/// let pat_buf = "3B0D........F27502F2C3F2E9...................................... 00 0000 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002 ___security_cookie ^000D ___report_gsfailure\n---";
/// assert_eq!(pat::parse(pat_buf).unwrap().len(), 1);
/// ```
pub fn parse(buf: &str) -> Result<Vec<FlirtSignature>> {
    if let Ok((_, sigs)) = pat(buf) {
        Ok(sigs)
    } else {
        Err(PatError::CorruptPatFile.into())
    }
}

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

    #[test]
    fn test_pat_parse() {
        parse("---").unwrap();

        parse("\
3B0D........F27502F2C3F2E9...................................... 00 0000 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002 ___security_cookie ^000D ___report_gsfailure
---
        ").unwrap();

        parse(
            "\
# foo
---
        ",
        )
        .unwrap();

        parse("\
# foo
3B0D........F27502F2C3F2E9...................................... 00 0000 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002 ___security_cookie ^000D ___report_gsfailure
---
        ").unwrap();
    }
}
<commit_msg>flirt: parse tail bytes from pat<commit_after>// from: https://github.com/Maktm/FLIRTDB/blob/1f5763535e02d7cccf2f90a96a8ebaa36e9b2495/cmt/windows/libcmt_15_msvc_x86.pat#L355
//
//    3B0D........F27502F2C3F2E9...................................... 00 0000
// 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002
// ___security_cookie ^000D ___report_gsfailure    ^^^^        ^^^^^^^^^^^^^^
// bytes        ^^^^^^^^                 relocations
//
//                                                                     00 0000
// 0011                                                                     ^^
// number of bytes crc16'd
// ^^^^ crc16
// ^^^^ size of fn
//
//
// :0000 @__security_check_cookie@4
// ^^^^^ offset of match
// ^^^^^^^^^^^^^^^^^^^^^^^^^^ name
//
//
// :000B@ $failure$4
// ^^^^^^ local offset
// ^^^^^^^^^^ local name
//
//
// ^0002 ___security_cookie
// ^^^^^ reference offset
// ^^^^^^^^^^^^^^^^^^ reference name

use anyhow::Result;
use log::trace;
use nom::{
    branch::alt,
    bytes::complete::{tag, take_till1, take_while, take_while_m_n},
    combinator::{map, map_res, opt, peek},
    multi::{many0, many1},
    sequence::pair,
    IResult,
};
use thiserror::Error;

use super::{ByteSignature, FlirtSignature, Name, Offset, SigElement, Symbol, TailByte};

#[derive(Debug, Error)]
pub enum PatError {
    #[error("The pattern is not supported")]
    NotSupported,
    #[error("The .pat file is corrupt (or unsupported)")]
    CorruptPatFile,
}

fn whitespace(input: &str) -> IResult<&str, &str> {
    take_while(|c| c == ' ')(input)
}

fn is_not_newline(c: char) -> bool {
    c != '\r' && c != '\n'
}

fn is_hex_digit(c: char) -> bool {
    c.is_digit(16)
}

fn from_hex(input: &str) -> Result<u8, std::num::ParseIntError> {
    u8::from_str_radix(input, 16)
}

/// parse a single hex byte, like `AB`
fn hex(input: &str) -> IResult<&str, u8> {
    map_res(take_while_m_n(2, 2, is_hex_digit), from_hex)(input)
}

fn hex_byte(input: &str) -> IResult<&str, u8> {
    let (input, v) = hex(input)?;
    Ok((input, v))
}

fn hex_word(input: &str) -> IResult<&str, u16> {
    let (input, v1) = hex(input)?;
    let (input, v2) = hex(input)?;
    let v: u16 = ((v1 as u16) << 8) | (v2 as u16);
    Ok((input, v))
}

/// parse a single byte signature element, which is either a hex byte or a
/// wildcard.
fn sig_element(input: &str) -> IResult<&str, SigElement> {
    alt((map(hex, SigElement::Byte), map(tag(".."), |_| SigElement::Wildcard)))(input)
}

/// parse byte signature elements, hex or wildcard.
fn byte_signature(input: &str) -> IResult<&str, ByteSignature> {
    let (input, elems) = many1(sig_element)(input)?;
    Ok((input, ByteSignature(elems)))
}

/// parse a hex-encoded offset, like `0000`
/// max is 0x8000.
/// TODO: this can be negative.
fn hex_offset(input: &str) -> IResult<&str, u16> {
    hex_word(input)
}

/// parse a public offset, like `:0000`
fn public_offset(input: &str) -> IResult<&str, u16> {
    let (input, _) = tag(":")(input)?;
    let (input, offset) = hex_offset(input)?;
    let (input, _) = peek(tag(" "))(input)?;

    Ok((input, offset))
}

/// parse a local offset, like `:000B@`
fn local_offset(input: &str) -> IResult<&str, u16> {
    let (input, _) = tag(":")(input)?;
    let (input, offset) = hex_offset(input)?;
    let (input, _) = tag("@")(input)?;

    Ok((input, offset))
}

/// parse an external reference, like `^0002`
fn reference_offset(input: &str) -> IResult<&str, u16> {
    let (input, _) = tag("^")(input)?;
    let (input, offset) = hex_offset(input)?;

    Ok((input, offset))
}

fn offset(input: &str) -> IResult<&str, Offset> {
    alt((
        // this must go first, because it has trailing `@`,
        // otherwise, the same as public.
        map(local_offset, Offset::Local),
        map(public_offset, Offset::Public),
        map(reference_offset, Offset::Reference),
    ))(input)
}

fn symbol_name(input: &str) -> IResult<&str, &str> {
    take_till1(|c| c == ' ' || c == '\n')(input)
}

/// parse a (offset, symbol) pair.
///
/// note: this also consumes trailing spaces so that a sequence of these can be
/// easily parsed.
fn symbol(input: &str) -> IResult<&str, Symbol> {
    let (input, offset) = offset(input)?;
    let (input, _) = whitespace(input)?;
    let (input, name) = symbol_name(input)?;
    let (input, _) = opt(whitespace)(input)?;
    match offset {
        Offset::Public(v) => Ok((
            input,
            Symbol::Public(Name {
                name:   name.to_string(),
                offset: v as i64,
            }),
        )),
        Offset::Local(v) => Ok((
            input,
            Symbol::Local(Name {
                name:   name.to_string(),
                offset: v as i64,
            }),
        )),
        Offset::Reference(v) => Ok((
            input,
            Symbol::Reference(Name {
                name:   name.to_string(),
                offset: v as i64,
            }),
        )),
    }
}

fn symbols(input: &str) -> IResult<&str, Vec<Symbol>> {
    many1(symbol)(input)
}

// like: `(0012: 87)`
fn tail_byte(input: &str) -> IResult<&str, TailByte> {
    let (input, _) = tag("(")(input)?;
    let (input, offset) = hex_offset(input)?;
    let (input, _) = tag(":")(input)?;
    let (input, _) = whitespace(input)?;
    let (input, value) = hex_byte(input)?;
    let (input, _) = tag(")")(input)?;

    Ok((
        input,
        TailByte {
            offset: offset as u64,
            value,
        },
    ))
}

fn tail_bytes(input: &str) -> IResult<&str, Vec<TailByte>> {
    match opt(many1(tail_byte))(input)? {
        (input, Some(tail_bytes)) => Ok((input, tail_bytes)),
        (input, None) => Ok((input, vec![])),
    }
}

fn pat_signature(input: &str) -> IResult<&str, FlirtSignature> {
    let (input, byte_sig) = byte_signature(input)?;
    let (input, _) = whitespace(input)?;
    trace!("sig: {:?}", byte_sig);

    let (input, size_of_bytes_crc16) = hex_byte(input)?;
    let (input, _) = whitespace(input)?;
    trace!("crc16 len: {:02x}", size_of_bytes_crc16);

    let (input, crc16) = hex_word(input)?;
    let (input, _) = whitespace(input)?;
    trace!("crc16: {:04x}", crc16);

    let (input, size_of_function) = hex_word(input)?;
    let (input, _) = whitespace(input)?;
    trace!("function size: {:04x}", size_of_function);

    let (input, names) = symbols(input)?;
    trace!("names: {:?}", names);

    // i'm not sure which of these comes first, footer pattern or tail bytes.
    // are they both actually valid in a .pat file!?!
    // pat.txt only describes the footer pattern (which it calls "tail bytes").
    // whereas the .sig file format uses (offset, value) tuples for tail bytes,
    // and dumpsig.exe uses a format like `(AAAA: BB)` for these.

    let (input, _) = opt(whitespace)(input)?;
    let (input, footer) = opt(byte_signature)(input)?;
    trace!("footer: {:?}", footer);

    let (input, _) = opt(whitespace)(input)?;
    let (input, tail_bytes) = tail_bytes(input)?;
    trace!("tail bytes: {:02x?}", tail_bytes);

    Ok((
        input,
        FlirtSignature {
            byte_sig,
            size_of_bytes_crc16,
            crc16,
            size_of_function: size_of_function as u64,
            names,
            footer,
            tail_bytes,
        },
    ))
}

fn comment(input: &str) -> IResult<&str, ()> {
    let (input, _) = tag("#")(input)?;
    let (input, _) = take_while(is_not_newline)(input)?;

    Ok((input, ()))
}

fn line(input: &str) -> IResult<&str, Option<FlirtSignature>> {
    if let Ok((input, _)) = comment(input) {
        return Ok((input, None));
    }

    match pat_signature(input) {
        Ok((input, sig)) => Ok((input, Some(sig))),
        Err(e) => Err(e),
    }
}

/// parse a .pat file into FLIRT signatures.
fn pat(input: &str) -> IResult<&str, Vec<FlirtSignature>> {
    // each signature is newline separated.
    // drop the newline after we've parsed it.
    let maybe_pat = map(pair(line, alt((tag("\r\n"), tag("\n")))), |p| p.0);
    let (input, sigs) = many0(maybe_pat)(input)?;

    // the file ends with `---`.
    let (input, _) = tag("---")(input)?;

    // TODO: assert there is nothing left in the file

    Ok((
        input,
        sigs.into_iter().filter(|s| s.is_some()).map(|s| s.unwrap()).collect(),
    ))
}

/// ```
/// use lancelot_flirt::pat;
/// let pat_buf = "3B0D........F27502F2C3F2E9...................................... 00 0000 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002 ___security_cookie ^000D ___report_gsfailure\n---";
/// assert_eq!(pat::parse(pat_buf).unwrap().len(), 1);
/// ```
pub fn parse(buf: &str) -> Result<Vec<FlirtSignature>> {
    if let Ok((_, sigs)) = pat(buf) {
        Ok(sigs)
    } else {
        Err(PatError::CorruptPatFile.into())
    }
}

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

    #[test]
    fn test_pat_parse() {
        parse("---").unwrap();

        parse("\
3B0D........F27502F2C3F2E9...................................... 00 0000 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002 ___security_cookie ^000D ___report_gsfailure
---
        ").unwrap();

        parse(
            "\
# foo
---
        ",
        )
        .unwrap();

        parse("\
# foo
3B0D........F27502F2C3F2E9...................................... 00 0000 0011 :0000 @__security_check_cookie@4 :000B@ $failure$4 ^0002 ___security_cookie ^000D ___report_gsfailure
---
        ").unwrap();
    }

    #[test]
    fn test_pat_tail_bytes() {
        parse("\
0330020028000000000000007f........16547f........1a5816547f 00 0000 0034 :0000 ???__E??_R3?$_Iosb@H@std@@8@@YMXXZ@?A0x7094697c@@$$FYMXXZ ^000d 040002DD (0003: 17)
0330020028000000000000007f........16547f........1a5816547f 00 0000 0034 :0000 ???__E??_R3?$_Iosb@H@std@@8@@YMXXZ@?A0xb8f7dd2a@@$$FYMXXZ ^000d 040002EB (0003: 17)
---
        ").unwrap();
    }
}
<|endoftext|>


<commit_before>use filetime::FileTime;
use serde::{Deserialize, Serialize};

use std::{
    collections::HashMap,
    fs, io,
    path::{Path, PathBuf},
    vec,
};

/// Helps identify changes in a set of files.
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(from = "DepGraphSerdeRepr", into = "DepGraphSerdeRepr")]
pub struct DepGraph {
    entries: HashMap<PathBuf, Option<DepGraphEntry>>,
}

#[derive(Debug, Clone, PartialEq)]
struct DepGraphEntry {
    mtime: FileTime,
    size: u64,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct DepGraphSerdeRepr {
    entries: Vec<DepGraphEntrySerdeRepr>,
}

impl From<DepGraphSerdeRepr> for DepGraph {
    fn from(from: DepGraphSerdeRepr) -> Self {
        DepGraph {
            entries: from
                .entries
                .iter()
                .map(|d| {
                    (
                        PathBuf::from(&d.path),
                        if d.has_graph_entry() {
                            Some(DepGraphEntry {
                                mtime: FileTime::from_unix_time(d.unix_seconds, d.nanos),
                                size: d.size,
                            })
                        } else {
                            None
                        },
                    )
                })
                .collect(),
        }
    }
}

impl From<DepGraph> for DepGraphSerdeRepr {
    fn from(dg: DepGraph) -> Self {
        DepGraphSerdeRepr {
            entries: dg
                .entries
                .iter()
                .map(|e| {
                    let path = e.0.to_str().expect("Only UTF names please").to_string();
                    if let Some(dep_entry) = e.1 {
                        DepGraphEntrySerdeRepr {
                            path,
                            unix_seconds: dep_entry.mtime.unix_seconds(),
                            nanos: dep_entry.mtime.nanoseconds(),
                            size: dep_entry.size,
                        }
                    } else {
                        DepGraphEntrySerdeRepr {
                            path,
                            unix_seconds: 0,
                            nanos: 0,
                            size: 0,
                        }
                    }
                })
                .collect(),
        }
    }
}

/// SystemTime lacks a platform independent representation we can
/// depend on so use FileTime's unix_seconds,nanos.
/// unix_seconds = nanos = size = 0 is used to represent lack of a DepGraphEntry.
#[derive(Serialize, Deserialize, Debug, Clone)]
struct DepGraphEntrySerdeRepr {
    path: String,
    unix_seconds: i64,
    nanos: u32,
    size: u64,
}

impl DepGraphEntrySerdeRepr {
    fn has_graph_entry(&self) -> bool {
        !(self.unix_seconds == 0 && self.nanos == 0 && self.size == 0)
    }
}

pub struct Change {
    path: PathBuf,
    updated_entry: DepGraphEntry,
}

pub enum InitialState {
    NotChanged,
    Changed,
}

impl DepGraphEntry {
    fn new(path: &Path) -> Result<DepGraphEntry, io::Error> {
        let metadata = path.metadata()?;
        Ok(DepGraphEntry {
            mtime: FileTime::from_system_time(metadata.modified()?),
            size: metadata.len(),
        })
    }
}

impl DepGraph {
    pub fn new() -> DepGraph {
        DepGraph {
            entries: HashMap::new(),
        }
    }
}

impl Default for DepGraph {
    fn default() -> Self {
        Self::new()
    }
}

impl DepGraph {
    pub fn is_tracked(&self, path: &PathBuf) -> bool {
        self.entries.contains_key(path)
    }

    pub fn track(&mut self, path: &PathBuf, default_state: InitialState) -> Result<(), io::Error> {
        if !self.entries.contains_key(path) {
            self.entries.insert(
                path.clone(),
                match default_state {
                    InitialState::NotChanged => Some(DepGraphEntry::new(path)?),
                    InitialState::Changed => None,
                },
            );

            // For a dir to register unchanged we need to add it's current contents
            if let InitialState::NotChanged = default_state {
                if path.is_dir() {
                    if let Some(new_paths) = self.new_paths(path) {
                        for new_path in new_paths {
                            self.entries
                                .insert(new_path.clone(), Some(DepGraphEntry::new(&new_path)?));
                        }
                    }
                }
            }
        }
        Ok(())
    }

    pub fn has_changed(&self, change: &Change) -> bool {
        return self
            .entries
            .get(&change.path)
            // If we are tracking with different state or we're not tracking at all then you've changed
            .map(|maybe_entry| {
                maybe_entry
                    .as_ref()
                    .map(|entry| change.updated_entry != *entry)
                    .unwrap_or(true)
            })
            .unwrap_or(true);
    }

    pub fn update(&mut self, changes: &Vec<Change>) {
        for change in changes {
            if self.has_changed(change) {
                self.entries
                    .insert(change.path.clone(), Some(change.updated_entry.clone()));
            }
        }
    }

    fn new_paths(&self, dir: &PathBuf) -> Option<Vec<PathBuf>> {
        assert!(dir.is_dir());

        // TODO: adding dir here makes Vec<&PathBuf> and then inside loop we can't push new entries. User
        let mut frontier = vec![dir.to_owned()];
        let mut results: Option<Vec<PathBuf>> = None;

        while let Some(dir) = frontier.pop() {
            let dir_entries = fs::read_dir(dir).expect("Unable to iterate directory");
            for dir_entry in dir_entries {
                let dir_entry = dir_entry.expect("Cannot read dir entry").path();

                // Queue subdir for processing. Bravely assume lack of cycles.
                if dir_entry.is_dir() {
                    frontier.push(dir_entry.to_owned());
                }

                // A path we've never seen before?!
                if !self.entries.contains_key(&dir_entry) {
                    let new_files = results.get_or_insert_with(Vec::new);
                    new_files.push(dir_entry);
                }
            }
        }
        results
    }

    /// Set of things that have either changed themselves or had
    /// something they depend on change.
    pub fn changed(&self) -> Result<Vec<Change>, io::Error> {
        let mut paths: Vec<&PathBuf> = self.entries.iter().map(|(k, _)| k).collect();

        let mut changes: Vec<Change> = Vec::new();

        // Add any new files in tracked dirs
        let new_paths: Vec<PathBuf> = paths
            .iter()
            .filter(|p| p.is_dir())
            .flat_map(|p| self.new_paths(p).unwrap_or_default())
            .collect();
        new_paths.iter().for_each(|p| paths.push(p));

        // Anything change 'round here?
        for path in paths {
            let updated_entry = DepGraphEntry::new(path)?;
            let has_changed = if self.entries.contains_key(path) {
                let maybe_current = self.entries.get(path).unwrap();
                maybe_current
                    .as_ref()
                    .map(|current_entry| updated_entry != *current_entry)
                    .unwrap_or(true)
            } else {
                true
            };

            if has_changed {
                changes.push(Change {
                    path: path.to_owned(),
                    updated_entry,
                });
            }
        }

        Ok(changes)
    }
}

#[cfg(test)]
mod tests {
    use std::{fs, ops::Add, path::PathBuf, time::Duration};

    use filetime::set_file_mtime;
    use tempfile::{tempdir, TempDir};

    use super::DepGraph;

    fn assert_no_changes(dg: &DepGraph) {
        assert_eq!(
            Vec::<&PathBuf>::new(),
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        )
    }

    #[test]
    fn detect_file_change() {
        let temp_dir = tempdir().unwrap();

        let file = temp_dir.path().join("a");
        fs::write(&file, "eh").unwrap();

        let mut dg = DepGraph::new();
        dg.track(&file, super::InitialState::NotChanged).unwrap();

        assert_no_changes(&dg);

        // Detect changed size
        fs::write(&file, "whoa").unwrap();
        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&file],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );

        // We're happy with that change
        dg.update(&changes);
        assert_eq!(
            Vec::<&PathBuf>::new(),
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        );

        // Detect changed mtime
        let new_mtime = file
            .metadata()
            .unwrap()
            .modified()
            .unwrap()
            .add(Duration::from_secs(1));
        set_file_mtime(&file, new_mtime.into()).unwrap();
        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&file],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );
    }

    #[test]
    fn detect_dir_change() {
        let temp_dir = tempdir().unwrap();

        let subdir = temp_dir.path().join("glif");
        fs::create_dir(&subdir).unwrap();

        let f1 = temp_dir.path().join("fileA");
        let f2 = subdir.join("fileB");
        fs::write(&f1, "eh").unwrap();
        fs::write(&f2, "eh").unwrap();

        // Track the parent dir from it's current state
        let mut dg = DepGraph::new();
        dg.track(
            &temp_dir.path().to_path_buf(),
            super::InitialState::NotChanged,
        )
        .unwrap();

        assert_eq!(
            Vec::<&PathBuf>::new(),
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        );
    }

    fn one_changed_file_one_not(temp_dir: &TempDir) -> (PathBuf, PathBuf, DepGraph) {
        let unchanged = temp_dir.path().join("a");
        let changed = temp_dir.path().join("b");

        fs::write(&unchanged, "eh").unwrap();
        fs::write(&changed, "eh").unwrap();

        let mut dg = DepGraph::new();
        dg.track(&unchanged, super::InitialState::NotChanged)
            .unwrap();
        dg.track(&changed, super::InitialState::Changed).unwrap();

        assert!(dg.is_tracked(&changed));
        assert!(dg.is_tracked(&unchanged));

        (changed, unchanged, dg)
    }

    #[test]
    fn track_with_initial_state() {
        let temp_dir = tempdir().unwrap();

        let (changed, _, mut dg) = one_changed_file_one_not(&temp_dir);

        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&changed],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );

        dg.update(&changes);
        let changes = dg.changed().unwrap();
        assert!(changes.is_empty());
    }

    #[test]
    fn read_write_bincode_entry() {}

    #[test]
    fn read_write_toml() {
        let temp_dir = tempdir().unwrap();

        let (changed, _, dg) = one_changed_file_one_not(&temp_dir);

        let toml = toml::ser::to_string_pretty(&dg).unwrap();
        let dg: DepGraph = toml::from_str(&toml).unwrap();

        assert_eq!(
            vec![&changed],
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        );
    }

    #[test]
    fn read_write_bincode() {
        let temp_dir = tempdir().unwrap();

        let (changed, _, dg) = one_changed_file_one_not(&temp_dir);

        let bc = bincode::serialize(&dg).unwrap();
        let dg: DepGraph = bincode::deserialize(&bc).unwrap();

        assert_eq!(
            vec![&changed],
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        );
    }
}
<commit_msg>Hopefully make test with dirs clearer, fix bug encountered along the way<commit_after>use filetime::FileTime;
use serde::{Deserialize, Serialize};

use std::{
    collections::{HashMap, HashSet},
    fs, io,
    path::{Path, PathBuf},
    vec,
};

/// Helps identify changes in a set of files.
#[derive(Serialize, Deserialize, Debug, Clone)]
#[serde(from = "DepGraphSerdeRepr", into = "DepGraphSerdeRepr")]
pub struct DepGraph {
    entries: HashMap<PathBuf, Option<DepGraphEntry>>,
}

#[derive(Debug, Clone, PartialEq)]
struct DepGraphEntry {
    mtime: FileTime,
    size: u64,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct DepGraphSerdeRepr {
    entries: Vec<DepGraphEntrySerdeRepr>,
}

impl From<DepGraphSerdeRepr> for DepGraph {
    fn from(from: DepGraphSerdeRepr) -> Self {
        DepGraph {
            entries: from
                .entries
                .iter()
                .map(|d| {
                    (
                        PathBuf::from(&d.path),
                        if d.has_graph_entry() {
                            Some(DepGraphEntry {
                                mtime: FileTime::from_unix_time(d.unix_seconds, d.nanos),
                                size: d.size,
                            })
                        } else {
                            None
                        },
                    )
                })
                .collect(),
        }
    }
}

impl From<DepGraph> for DepGraphSerdeRepr {
    fn from(dg: DepGraph) -> Self {
        DepGraphSerdeRepr {
            entries: dg
                .entries
                .iter()
                .map(|e| {
                    let path = e.0.to_str().expect("Only UTF names please").to_string();
                    if let Some(dep_entry) = e.1 {
                        DepGraphEntrySerdeRepr {
                            path,
                            unix_seconds: dep_entry.mtime.unix_seconds(),
                            nanos: dep_entry.mtime.nanoseconds(),
                            size: dep_entry.size,
                        }
                    } else {
                        DepGraphEntrySerdeRepr {
                            path,
                            unix_seconds: 0,
                            nanos: 0,
                            size: 0,
                        }
                    }
                })
                .collect(),
        }
    }
}

/// SystemTime lacks a platform independent representation we can
/// depend on so use FileTime's unix_seconds,nanos.
/// unix_seconds = nanos = size = 0 is used to represent lack of a DepGraphEntry.
#[derive(Serialize, Deserialize, Debug, Clone)]
struct DepGraphEntrySerdeRepr {
    path: String,
    unix_seconds: i64,
    nanos: u32,
    size: u64,
}

impl DepGraphEntrySerdeRepr {
    fn has_graph_entry(&self) -> bool {
        !(self.unix_seconds == 0 && self.nanos == 0 && self.size == 0)
    }
}

pub struct Change {
    path: PathBuf,
    updated_entry: DepGraphEntry,
}

pub enum InitialState {
    NotChanged,
    Changed,
}

impl DepGraphEntry {
    fn new(path: &Path) -> Result<DepGraphEntry, io::Error> {
        let metadata = path.metadata()?;
        Ok(DepGraphEntry {
            mtime: FileTime::from_system_time(metadata.modified()?),
            size: metadata.len(),
        })
    }
}

impl DepGraph {
    pub fn new() -> DepGraph {
        DepGraph {
            entries: HashMap::new(),
        }
    }
}

impl Default for DepGraph {
    fn default() -> Self {
        Self::new()
    }
}

impl DepGraph {
    pub fn is_tracked(&self, path: &PathBuf) -> bool {
        self.entries.contains_key(path)
    }

    pub fn track(&mut self, path: &PathBuf, default_state: InitialState) -> Result<(), io::Error> {
        if !self.entries.contains_key(path) {
            self.entries.insert(
                path.clone(),
                match default_state {
                    InitialState::NotChanged => Some(DepGraphEntry::new(path)?),
                    InitialState::Changed => None,
                },
            );

            // For a dir to register unchanged we need to add it's current contents
            if let InitialState::NotChanged = default_state {
                if path.is_dir() {
                    let mut dirs_visited = HashSet::new();
                    if let Some(new_paths) = self.new_paths(&mut dirs_visited, path) {
                        for new_path in new_paths {
                            self.entries
                                .insert(new_path.clone(), Some(DepGraphEntry::new(&new_path)?));
                        }
                    }
                }
            }
        }
        Ok(())
    }

    pub fn has_changed(&self, change: &Change) -> bool {
        return self
            .entries
            .get(&change.path)
            // If we are tracking with different state or we're not tracking at all then you've changed
            .map(|maybe_entry| {
                maybe_entry
                    .as_ref()
                    .map(|entry| change.updated_entry != *entry)
                    .unwrap_or(true)
            })
            .unwrap_or(true);
    }

    pub fn update(&mut self, changes: &Vec<Change>) {
        for change in changes {
            if self.has_changed(change) {
                self.entries
                    .insert(change.path.clone(), Some(change.updated_entry.clone()));
            }
        }
    }

    fn new_paths(
        &self,
        dirs_visited: &mut HashSet<PathBuf>,
        dir: &PathBuf,
    ) -> Option<Vec<PathBuf>> {
        assert!(dir.is_dir());
        if dirs_visited.contains(dir) {
            return None;
        }

        let mut frontier = vec![dir.to_owned()];
        let mut results: Option<Vec<PathBuf>> = None;

        while let Some(dir) = frontier.pop() {
            let dir_entries = fs::read_dir(&dir).expect("Unable to iterate directory");
            for dir_entry in dir_entries {
                let dir_entry = dir_entry.expect("Cannot read dir entry").path();

                // Queue subdir for processing. Bravely assume lack of cycles.
                if dir_entry.is_dir() && !dirs_visited.contains(&dir_entry) {
                    frontier.push(dir_entry.to_owned());
                }

                // A path we've never seen before?!
                if !self.entries.contains_key(&dir_entry) {
                    let new_files = results.get_or_insert_with(Vec::new);
                    new_files.push(dir_entry);
                }
            }
            dirs_visited.insert(dir.to_owned());
        }
        results
    }

    /// Set of things that have either changed themselves or had
    /// something they depend on change.
    pub fn changed(&self) -> Result<Vec<Change>, io::Error> {
        let mut paths: Vec<&PathBuf> = self.entries.iter().map(|(k, _)| k).collect();

        let mut changes: Vec<Change> = Vec::new();

        // Add any new files in tracked dirs
        let mut dirs_visited: HashSet<PathBuf> = HashSet::new();
        let new_paths: Vec<PathBuf> = paths
            .iter()
            .filter(|p| p.is_dir())
            .flat_map(|p| self.new_paths(&mut dirs_visited, p).unwrap_or_default())
            .collect();
        new_paths.iter().for_each(|p| paths.push(p));

        // Anything change 'round here?
        for path in paths {
            let updated_entry = DepGraphEntry::new(path)?;
            let has_changed = if self.entries.contains_key(path) {
                let maybe_current = self.entries.get(path).unwrap();
                maybe_current
                    .as_ref()
                    .map(|current_entry| updated_entry != *current_entry)
                    .unwrap_or(true)
            } else {
                true
            };

            if has_changed {
                changes.push(Change {
                    path: path.to_owned(),
                    updated_entry,
                });
            }
        }

        Ok(changes)
    }
}

#[cfg(test)]
mod tests {
    use std::{fs, ops::Add, path::PathBuf, time::Duration};

    use filetime::set_file_mtime;
    use tempfile::{tempdir, TempDir};

    use super::DepGraph;

    fn assert_no_changes(dg: &DepGraph) {
        assert_eq!(
            Vec::<&PathBuf>::new(),
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        )
    }

    #[test]
    fn detect_file_change() {
        let temp_dir = tempdir().unwrap();

        let file = temp_dir.path().join("a");
        fs::write(&file, "eh").unwrap();

        let mut dg = DepGraph::new();
        dg.track(&file, super::InitialState::NotChanged).unwrap();

        assert_no_changes(&dg);

        // Detect changed size
        fs::write(&file, "whoa").unwrap();
        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&file],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );

        // We're happy with that change
        dg.update(&changes);
        assert_no_changes(&dg);

        // Detect changed mtime
        let new_mtime = file
            .metadata()
            .unwrap()
            .modified()
            .unwrap()
            .add(Duration::from_secs(1));
        set_file_mtime(&file, new_mtime.into()).unwrap();
        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&file],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );
    }

    #[test]
    fn detect_dir_change() {
        let temp_dir = tempdir().unwrap();

        let subdir = temp_dir.path().join("glif");
        fs::create_dir(&subdir).unwrap();

        let f1 = temp_dir.path().join("fileA");
        let f2 = subdir.join("fileB");
        fs::write(&f1, "eh").unwrap();
        fs::write(&f2, "eh").unwrap();

        // Track the parent dir from it's current state
        let mut dg = DepGraph::new();
        dg.track(
            &temp_dir.path().to_path_buf(),
            super::InitialState::NotChanged,
        )
        .unwrap();

        // Notably, we should NOT report the files temp dir as changed
        // because we added the entire directory as unchanged and nothing
        // in it has changed since then
        assert_no_changes(&dg);

        // If we change or add files in the tracked dir that should count
        let f3 = subdir.join("fileC");
        fs::write(&f2, "eh+").unwrap();
        fs::write(&f3, "eh").unwrap();

        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&f2, &f3],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );
    }

    fn one_changed_file_one_not(temp_dir: &TempDir) -> (PathBuf, PathBuf, DepGraph) {
        let unchanged = temp_dir.path().join("a");
        let changed = temp_dir.path().join("b");

        fs::write(&unchanged, "eh").unwrap();
        fs::write(&changed, "eh").unwrap();

        let mut dg = DepGraph::new();
        dg.track(&unchanged, super::InitialState::NotChanged)
            .unwrap();
        dg.track(&changed, super::InitialState::Changed).unwrap();

        assert!(dg.is_tracked(&changed));
        assert!(dg.is_tracked(&unchanged));

        (changed, unchanged, dg)
    }

    #[test]
    fn track_with_initial_state() {
        let temp_dir = tempdir().unwrap();

        let (changed, _, mut dg) = one_changed_file_one_not(&temp_dir);

        let changes = dg.changed().unwrap();
        assert_eq!(
            vec![&changed],
            changes.iter().map(|c| &c.path).collect::<Vec<&PathBuf>>()
        );

        dg.update(&changes);
        let changes = dg.changed().unwrap();
        assert!(changes.is_empty());
    }

    #[test]
    fn read_write_bincode_entry() {}

    #[test]
    fn read_write_toml() {
        let temp_dir = tempdir().unwrap();

        let (changed, _, dg) = one_changed_file_one_not(&temp_dir);

        let toml = toml::ser::to_string_pretty(&dg).unwrap();
        let dg: DepGraph = toml::from_str(&toml).unwrap();

        assert_eq!(
            vec![&changed],
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        );
    }

    #[test]
    fn read_write_bincode() {
        let temp_dir = tempdir().unwrap();

        let (changed, _, dg) = one_changed_file_one_not(&temp_dir);

        let bc = bincode::serialize(&dg).unwrap();
        let dg: DepGraph = bincode::deserialize(&bc).unwrap();

        assert_eq!(
            vec![&changed],
            dg.changed()
                .unwrap()
                .iter()
                .map(|c| &c.path)
                .collect::<Vec<&PathBuf>>()
        );
    }
}
<|endoftext|>


<commit_before>//! Validate UUIDs.
//! A UUID is an unique number of 128bit.  UUIDs are used to  assign
//! unique identifiers to entities.
//! Check against 3 formats specified by the RFC: https://tools.ietf.org/html/rfc4122.
//! * simple:       9856BAC1F9CBD6b9d80C712AF75C832A7
//! * hyphenated:   551e7604-e35c-42b3-b825-416853441234
//! * urn: urn:uuid:551e7604-e35c-42b3-b825-416853441234
use std::result;

#[derive(Debug, PartialEq)]
pub enum UUIDError {
    InvalidChar(char),
    InvalidGroupCount(usize),
    InvalidLength(usize),
    InvalidGroupLength(u8),
}

type Result<T> = result::Result<T, UUIDError>;

// Length of each of the 5 hyphen delimited groups from the RFC.
const GROUP_LENS: [u8; 5] = [8, 4, 4, 4, 12];
const HYPHEN_FORMAT_GROUP_COUNT: usize = 4;
const HYPHEN_FORMAT_LEN: usize = 36;
const SIMPLE_FORMAT_LEN: usize = 32;
const URN_FORMAT_PREFIX_LEN: usize = 9;

/// Validate a `Uuid` from a string of hexadecimal digits against three of the
/// supported formats.
pub fn validate(mut input: &str) -> Result<()> {
    // Check against correct length of each format.
    let len = input.len();
    if input.starts_with("urn:uuid:") && len == (HYPHEN_FORMAT_LEN + URN_FORMAT_PREFIX_LEN) {
        input = &input[9..];
    } else if len != SIMPLE_FORMAT_LEN && len != HYPHEN_FORMAT_LEN {
        return Err(UUIDError::InvalidLength(len));
    }

    // Used to count digits from each group.
    let mut group_index = 0;
    let mut group_count = 0;

    // Iterate over index and value of each byte of the input string.
    for (index, value) in input.bytes().enumerate() {
        // If length is bigger than that of the simple format, check that the
        // groups are exactly 4 at this moment.
        if index as usize >= SIMPLE_FORMAT_LEN && group_index != HYPHEN_FORMAT_GROUP_COUNT {
            return Err(UUIDError::InvalidGroupCount(group_index + 1));
        }
        // Never let the UUID get more than 5 groups.
        if group_index > HYPHEN_FORMAT_GROUP_COUNT {
            return Err(UUIDError::InvalidGroupCount(group_index + 1));
        }
        match value {
            b'0'...b'9' => group_count += 1,
            b'a'...b'f' => group_count += 1,
            b'A'...b'F' => group_count += 1,
            // With each hyphen we need to make sure that the
            // number of digits from the current group meets the RFC.
            b'-' => {
                if GROUP_LENS[group_index] != group_count {
                    return Err(UUIDError::InvalidGroupLength(group_count));
                }
                group_index += 1;
                group_count = 0;
            }
            _ => {
                return Err(UUIDError::InvalidChar(
                    input[index..].chars().next().unwrap(),
                ))
            }
        }
    }
    if group_index > 0 && GROUP_LENS[group_index] != group_count {
        return Err(UUIDError::InvalidGroupLength(group_count));
    }
    Ok(())
}

#[test]
fn test_uuid_validate() {
    use super::*;

    // Testing invalid uuids.
    assert_eq!(validate(""), Err(UUIDError::InvalidLength(0)));
    assert_eq!(validate("a"), Err(UUIDError::InvalidLength(1)));
    assert_eq!(
        validate("551e7604-CDB3-5fbb-C6ADD-338BD51DB2F5"),
        Err(UUIDError::InvalidLength(37))
    );
    assert_eq!(
        validate("551e7604-DFC3A5abbAC7CAA431CA48AB2F3"),
        Err(UUIDError::InvalidGroupCount(2))
    );
    assert_eq!(
        validate("551e7604-DFC3-DFC3A5abbAC7CAA431CA48"),
        Err(UUIDError::InvalidGroupCount(3))
    );
    assert_eq!(
        validate("551e7604-DFC3-DFC3X5AbbAC7CAA431CA48"),
        Err(UUIDError::InvalidChar('X'))
    );

    assert_eq!(
        validate("551e7604-DAF-35ab-fC7CAA43-CA48AB2F5"),
        Err(UUIDError::InvalidGroupLength(3))
    );
    assert_eq!(
        validate("01020304-1112-2122-3132-41424344"),
        Err(UUIDError::InvalidGroupLength(8))
    );
    assert_eq!(
        validate("67e5604510c1436%9257bca00e6fe1da"),
        Err(UUIDError::InvalidChar('%'))
    );
    assert_eq!(
        validate("342342323323534535435434587353276861"),
        Err(UUIDError::InvalidGroupCount(1))
    );

    // Testing valid ones also.
    assert!(validate("00000000000000000000000000000000").is_ok());
    assert!(validate("68f56045-20c2-437a-8257-bb781e5ae0c9").is_ok());
    assert!(validate("F9168C5E-CEB2-4faa-B6BF-329BF39FA1E4").is_ok());
    assert!(validate("67e5504410b1426f9247bb680e5fe0c8").is_ok());
    assert!(validate("01020304-1112-2122-3132-414243444546").is_ok());
    assert!(validate("urn:uuid:67e55044-10b1-426f-9247-bb680e5fe0c8").is_ok());

    assert!(validate("00000000000000000000000000000000").is_ok());
    assert!(validate("00000000-0000-0000-0000-000000000000").is_ok());
}
<commit_msg>jailer: moved unit tests in tests module<commit_after>//! Validate UUIDs.
//! A UUID is an unique number of 128bit.  UUIDs are used to  assign
//! unique identifiers to entities.
//! Check against 3 formats specified by the RFC: https://tools.ietf.org/html/rfc4122.
//! * simple:       9856BAC1F9CBD6b9d80C712AF75C832A7
//! * hyphenated:   551e7604-e35c-42b3-b825-416853441234
//! * urn: urn:uuid:551e7604-e35c-42b3-b825-416853441234
use std::result;

#[derive(Debug, PartialEq)]
pub enum UUIDError {
    InvalidChar(char),
    InvalidGroupCount(usize),
    InvalidLength(usize),
    InvalidGroupLength(u8),
}

type Result<T> = result::Result<T, UUIDError>;

// Length of each of the 5 hyphen delimited groups from the RFC.
const GROUP_LENS: [u8; 5] = [8, 4, 4, 4, 12];
const HYPHEN_FORMAT_GROUP_COUNT: usize = 4;
const HYPHEN_FORMAT_LEN: usize = 36;
const SIMPLE_FORMAT_LEN: usize = 32;
const URN_FORMAT_PREFIX_LEN: usize = 9;

/// Validate a `Uuid` from a string of hexadecimal digits against three of the
/// supported formats.
pub fn validate(mut input: &str) -> Result<()> {
    // Check against correct length of each format.
    let len = input.len();
    if input.starts_with("urn:uuid:") && len == (HYPHEN_FORMAT_LEN + URN_FORMAT_PREFIX_LEN) {
        input = &input[9..];
    } else if len != SIMPLE_FORMAT_LEN && len != HYPHEN_FORMAT_LEN {
        return Err(UUIDError::InvalidLength(len));
    }

    // Used to count digits from each group.
    let mut group_index = 0;
    let mut group_count = 0;

    // Iterate over index and value of each byte of the input string.
    for (index, value) in input.bytes().enumerate() {
        // If length is bigger than that of the simple format, check that the
        // groups are exactly 4 at this moment.
        if index as usize >= SIMPLE_FORMAT_LEN && group_index != HYPHEN_FORMAT_GROUP_COUNT {
            return Err(UUIDError::InvalidGroupCount(group_index + 1));
        }
        // Never let the UUID get more than 5 groups.
        if group_index > HYPHEN_FORMAT_GROUP_COUNT {
            return Err(UUIDError::InvalidGroupCount(group_index + 1));
        }
        match value {
            b'0'...b'9' => group_count += 1,
            b'a'...b'f' => group_count += 1,
            b'A'...b'F' => group_count += 1,
            // With each hyphen we need to make sure that the
            // number of digits from the current group meets the RFC.
            b'-' => {
                if GROUP_LENS[group_index] != group_count {
                    return Err(UUIDError::InvalidGroupLength(group_count));
                }
                group_index += 1;
                group_count = 0;
            }
            _ => {
                return Err(UUIDError::InvalidChar(
                    input[index..].chars().next().unwrap(),
                ))
            }
        }
    }
    if group_index > 0 && GROUP_LENS[group_index] != group_count {
        return Err(UUIDError::InvalidGroupLength(group_count));
    }
    Ok(())
}

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

        // Testing invalid uuids.
        assert_eq!(validate(""), Err(UUIDError::InvalidLength(0)));
        assert_eq!(validate("a"), Err(UUIDError::InvalidLength(1)));
        assert_eq!(
            validate("551e7604-CDB3-5fbb-C6ADD-338BD51DB2F5"),
            Err(UUIDError::InvalidLength(37))
        );
        assert_eq!(
            validate("551e7604-DFC3A5abbAC7CAA431CA48AB2F3"),
            Err(UUIDError::InvalidGroupCount(2))
        );
        assert_eq!(
            validate("551e7604-DFC3-DFC3A5abbAC7CAA431CA48"),
            Err(UUIDError::InvalidGroupCount(3))
        );
        assert_eq!(
            validate("551e7604-DFC3-DFC3X5AbbAC7CAA431CA48"),
            Err(UUIDError::InvalidChar('X'))
        );

        assert_eq!(
            validate("551e7604-DAF-35ab-fC7CAA43-CA48AB2F5"),
            Err(UUIDError::InvalidGroupLength(3))
        );
        assert_eq!(
            validate("01020304-1112-2122-3132-41424344"),
            Err(UUIDError::InvalidGroupLength(8))
        );
        assert_eq!(
            validate("67e5604510c1436%9257bca00e6fe1da"),
            Err(UUIDError::InvalidChar('%'))
        );
        assert_eq!(
            validate("342342323323534535435434587353276861"),
            Err(UUIDError::InvalidGroupCount(1))
        );

        // Testing valid ones also.
        assert!(validate("00000000000000000000000000000000").is_ok());
        assert!(validate("68f56045-20c2-437a-8257-bb781e5ae0c9").is_ok());
        assert!(validate("F9168C5E-CEB2-4faa-B6BF-329BF39FA1E4").is_ok());
        assert!(validate("67e5504410b1426f9247bb680e5fe0c8").is_ok());
        assert!(validate("01020304-1112-2122-3132-414243444546").is_ok());
        assert!(validate("urn:uuid:67e55044-10b1-426f-9247-bb680e5fe0c8").is_ok());

        assert!(validate("00000000000000000000000000000000").is_ok());
        assert!(validate("00000000-0000-0000-0000-000000000000").is_ok());
    }
}
<|endoftext|>