conduit/src/utils.rs

use argon2::{Config, Variant};
use cmp::Ordering;
use rand::prelude::*;
use ruma::serde::{try_from_json_map, CanonicalJsonError, CanonicalJsonObject};
use std::{
    cmp,
    convert::TryInto,
    str::FromStr,
    time::{SystemTime, UNIX_EPOCH},
};

#[tracing::instrument]
pub fn millis_since_unix_epoch() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("time is valid")
        .as_millis() as u64
}

pub fn increment(old: Option<&[u8]>) -> Option<Vec<u8>> {
    let number = match old.map(|bytes| bytes.try_into()) {
        Some(Ok(bytes)) => {
            let number = u64::from_be_bytes(bytes);
            number + 1
        }
        _ => 1, // Start at one. since 0 should return the first event in the db
    };

    Some(number.to_be_bytes().to_vec())
}

pub fn generate_keypair() -> Vec<u8> {
    let mut value = random_string(8).as_bytes().to_vec();
    value.push(0xff);
    value.extend_from_slice(
        &ruma::signatures::Ed25519KeyPair::generate()
            .expect("Ed25519KeyPair generation always works (?)"),
    );
    value
}

/// Parses the bytes into an u64.
#[tracing::instrument(skip(bytes))]
pub fn u64_from_bytes(bytes: &[u8]) -> Result<u64, std::array::TryFromSliceError> {
    let array: [u8; 8] = bytes.try_into()?;
    Ok(u64::from_be_bytes(array))
}

/// Parses the bytes into a string.
#[tracing::instrument(skip(bytes))]
pub fn string_from_bytes(bytes: &[u8]) -> Result<String, std::string::FromUtf8Error> {
    String::from_utf8(bytes.to_vec())
}

#[tracing::instrument(skip(length))]
pub fn random_string(length: usize) -> String {
    thread_rng()
        .sample_iter(&rand::distributions::Alphanumeric)
        .take(length)
        .map(char::from)
        .collect()
}

/// Calculate a new hash for the given password
#[tracing::instrument(skip(password))]
pub fn calculate_hash(password: &str) -> Result<String, argon2::Error> {
    let hashing_config = Config {
        variant: Variant::Argon2id,
        ..Default::default()
    };

    let salt = random_string(32);
    argon2::hash_encoded(password.as_bytes(), salt.as_bytes(), &hashing_config)
}

#[tracing::instrument(skip(iterators, check_order))]
pub fn common_elements(
    mut iterators: impl Iterator<Item = impl Iterator<Item = Vec<u8>>>,
    check_order: impl Fn(&[u8], &[u8]) -> Ordering,
) -> Option<impl Iterator<Item = Vec<u8>>> {
    let first_iterator = iterators.next()?;
    let mut other_iterators = iterators.map(|i| i.peekable()).collect::<Vec<_>>();

    Some(first_iterator.filter(move |target| {
        other_iterators.iter_mut().all(|it| {
            while let Some(element) = it.peek() {
                match check_order(element, target) {
                    Ordering::Greater => return false, // We went too far
                    Ordering::Equal => return true,    // Element is in both iters
                    Ordering::Less => {
                        // Keep searching
                        it.next();
                    }
                }
            }
            false
        })
    }))
}

/// Fallible conversion from any value that implements `Serialize` to a `CanonicalJsonObject`.
///
/// `value` must serialize to an `serde_json::Value::Object`.
#[tracing::instrument(skip(value))]
pub fn to_canonical_object<T: serde::Serialize>(
    value: T,
) -> Result<CanonicalJsonObject, CanonicalJsonError> {
    use serde::ser::Error;

    match serde_json::to_value(value).map_err(CanonicalJsonError::SerDe)? {
        serde_json::Value::Object(map) => try_from_json_map(map),
        _ => Err(CanonicalJsonError::SerDe(serde_json::Error::custom(
            "Value must be an object",
        ))),
    }
}

#[tracing::instrument(skip(deserializer))]
pub fn deserialize_from_str<
    'de,
    D: serde::de::Deserializer<'de>,
    T: FromStr<Err = E>,
    E: std::fmt::Display,
>(
    deserializer: D,
) -> std::result::Result<T, D::Error> {
    struct Visitor<T: FromStr<Err = E>, E>(std::marker::PhantomData<T>);
    impl<'de, T: FromStr<Err = Err>, Err: std::fmt::Display> serde::de::Visitor<'de>
        for Visitor<T, Err>
    {
        type Value = T;
        fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(formatter, "a parsable string")
        }
        fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
        where
            E: serde::de::Error,
        {
            v.parse().map_err(serde::de::Error::custom)
        }
    }
    deserializer.deserialize_str(Visitor(std::marker::PhantomData))
}

pub mod scale {
    use std::{collections::BTreeMap, mem::size_of, sync::Arc};

    use clru::WeightScale;
    use ruma::{EventId, RoomId, UserId};

    use crate::PduEvent;

    // This trait captures a plain size_of of a particular value,
    // but also grabs all "hidden owned memory" behind it dynamically
    pub trait Weighted {
        fn size_of() -> usize;

        // inner values that aren't taken into account with size_of
        //
        // total memory use = size_of() + inner()
        fn inner(&self) -> usize;
    }

    pub trait WeightedExt {
        fn weight(&self) -> usize;
    }

    impl<T> WeightedExt for T
    where
        T: Weighted + ?Sized,
    {
        fn weight(&self) -> usize {
            Self::size_of() + self.inner()
        }
    }

    impl<T> Weighted for Arc<T>
    where
        T: Weighted,
    {
        fn size_of() -> usize {
            size_of::<Arc<T>>()
        }

        fn inner(&self) -> usize {
            let i = self.as_ref();

            i.weight()
        }
    }

    impl<T> Weighted for Option<T>
    where
        T: Weighted,
    {
        fn size_of() -> usize {
            size_of::<Option<T>>()
        }

        fn inner(&self) -> usize {
            if let Some(i) = self {
                i.inner()
            } else {
                0
            }
        }
    }

    impl<T: ?Sized> Weighted for Box<T>
    where
        T: Weighted,
    {
        fn size_of() -> usize {
            size_of::<Box<T>>()
        }

        fn inner(&self) -> usize {
            (self as &T).weight()
        }
    }

    // impl Weighted for Box<ruma::ServerName> {
    //     fn size_of() -> usize {
    //         size_of::<Box<ruma::ServerName>>()
    //     }

    //     fn inner(&self) -> usize {
    //         (self as &ruma::ServerName).weight()
    //     }
    // }

    impl<T> Weighted for Vec<T>
    where
        T: Weighted,
    {
        fn size_of() -> usize {
            size_of::<Vec<T>>()
        }

        fn inner(&self) -> usize {
            (T::size_of() * self.len()) + self.iter().fold(0, |i, t| i + t.inner())
        }
    }

    impl<A, B> Weighted for BTreeMap<A, B>
    where
        A: Weighted,
        B: Weighted,
    {
        fn size_of() -> usize {
            size_of::<BTreeMap<A, B>>()
        }

        // FIXME: technically we have all values of key/value pairs,
        // but btree has a lot of internal nodeleafs andsoforth, so this isn't everything
        fn inner(&self) -> usize {
            self.iter().fold(0, |a, (k, v)| a + k.weight() + v.weight())
        }
    }

    impl Weighted for u64 {
        fn size_of() -> usize {
            size_of::<u64>()
        }

        fn inner(&self) -> usize {
            0
        }
    }

    impl Weighted for str {
        fn size_of() -> usize {
            0 // str is a transmutation of &[u8]
        }

        fn inner(&self) -> usize {
            self.as_bytes().len()
        }
    }

    impl Weighted for String {
        fn size_of() -> usize {
            size_of::<String>()
        }

        fn inner(&self) -> usize {
            self.as_str().inner()
        }
    }

    impl Weighted for ruma::ServerName {
        fn size_of() -> usize {
            0 // ServerName is effectively str
        }

        fn inner(&self) -> usize {
            self.as_str().inner()
        }
    }

    impl Weighted for EventId {
        fn size_of() -> usize {
            size_of::<EventId>()
        }

        // reconstructing the inner Box<str>
        fn inner(&self) -> usize {
            1 // stripped $ symbol
            + self.localpart().len() // everything up until :, or the end
            + if let Some(server) = self.server_name() {
                1 // stripped : symbol
                + server.inner() // everything from after : to the end
            } else {
                0
            }
        }
    }

    impl Weighted for RoomId {
        fn size_of() -> usize {
            size_of::<RoomId>()
        }

        // reconstructing the inner Box<str>
        fn inner(&self) -> usize {
            1 // stripped ! symbol
            + self.localpart().len() // everything up until :
            + 1 // stripped : symbol
            + self.server_name().inner() // everything from after : to the end
        }
    }

    impl Weighted for UserId {
        fn size_of() -> usize {
            size_of::<UserId>()
        }

        // reconstructing the inner Box<str>
        fn inner(&self) -> usize {
            1 // stripped @ symbol
            + self.localpart().len() // everything up until :
            + 1 // stripped : symbol
            + self.server_name().inner() // everything from after : to the end
        }
    }

    impl Weighted for ruma::events::pdu::EventHash {
        fn size_of() -> usize {
            size_of::<ruma::events::pdu::EventHash>()
        }

        // reconstructing the inner Box<str>
        fn inner(&self) -> usize {
            self.sha256.inner()
        }
    }

    impl<A, B: ?Sized> Weighted for ruma::KeyId<A, B> {
        fn size_of() -> usize {
            size_of::<ruma::KeyId<A, B>>()
        }

        fn inner(&self) -> usize {
            self.as_str().inner()
        }
    }

    impl Weighted for serde_json::Value {
        fn size_of() -> usize {
            size_of::<serde_json::Value>()
        }

        fn inner(&self) -> usize {
            use serde_json::Value;

            match self {
                Value::String(s) => s.inner(),
                Value::Array(a) => a.inner(),
                Value::Object(o) => o
                    .into_iter()
                    .fold(0, |a, (s, v)| a + s.weight() + v.weight()),
                _ => 0,
            }
        }
    }

    impl Weighted for PduEvent {
        fn size_of() -> usize {
            size_of::<PduEvent>()
        }

        fn inner(&self) -> usize {
            self.event_id.inner()
                + self.room_id.inner()
                + self.sender.inner()
                + self.content.inner()
                + self.prev_events.inner()
                + self.auth_events.inner()
                + self.redacts.inner()
                + self.unsigned.inner()
                + self.hashes.inner()
                + self.signatures.inner()
        }
    }

    pub(crate) struct ConduitScale;

    impl<A, B> WeightScale<A, B> for ConduitScale
    where
        A: Weighted,
        B: Weighted,
    {
        fn weight(&self, key: &A, value: &B) -> usize {
            key.weight() + value.weight()
        }
    }
}