ratproto-did / src / lib.rs

#![forbid(unsafe_op_in_unsafe_fn)]
//! This library implements a memory-efficient container for atproto DIDs.
//!
//! The [`Did`] type is guaranteed to take up exactly 16 bytes.
//! - `did:plc` is represented with just these 16 bytes
//! - `did:web` uses up to 255 bytes of heap
//! - other DID methods (not officially supported by atproto) allocate heap space as needed
//!
//! ---
//!
//! The most common DID in the Atmosphere (the atproto ecosystem) is `did:plc`.
//! Currently, its identifier is exactly 24 characters of base32. At 5 bits per characters,
//! that is 120 bits, or 15 bytes. This leaves just enough space for
//! a single-byte discriminator, making this implementation highly-optimized
//! for `did:plc` while also allowing it to represent other DID methods.
//! **The PLC identifier is expected to be exactly 24 characters long.**
//!
//! Atproto currently only supports host-level `did:web` DIDs, so the identifier is a web domain.
//! The maximum length of a web domain is 255 bytes, so a [`Did`] representing a `did:web` allocates
//! the necessary space on the heap. **Parsing a `did:web` with an unsupported format fails**.
//! _The only exception is `localhost`, which supports an optional port
//! (as in `localhost%3A12345`)._
//!
//! Other DID methods are currently unsupported by atproto, but not strictly disallowed.
//! DID methods that are not `plc` or `web` have their method and identifier strings
//! allocated on the heap.

mod plc_codec;
mod web_domain;

use std::{
    cmp::Ordering,
    fmt::{Debug, Display, Formatter},
    hash::{Hash, Hasher},
    mem::ManuallyDrop,
    str::FromStr,
};

use thiserror::Error;

use crate::web_domain::DidWebDomain;

/// A tightly-packed representation of DIDs for atproto.
///
/// Currently, `did:plc`, the most common kind of DID in the atmosphere,
/// has a base32 identifier with exactly 24 characters.
/// At 5 bits per characters, that is 120 bits, or 15 bytes.
/// This type is optimized for `did:plc`, and _is guaranteed_
/// to take up exactly 16 bytes.
///
/// did:web has its domain part heap-allocated as a string of up to 255 bytes.
/// Other DID methods (not currently supported by atproto)
/// allocate both the method and identifier as strings.
pub struct Did(DidInner);

impl Did {
    /// Which method does this [`Did`] represent?
    ///
    /// See [`Did::method()`] to get the string representation (especially for [`DidKind::Other`]).
    pub const fn kind(&self) -> DidKind {
        match &self.0 {
            DidInner::Plc(_) => DidKind::Plc,
            DidInner::Web { .. } => DidKind::Web,
            DidInner::Other(_) => DidKind::Other,
        }
    }

    /// The DID method string.
    pub const fn method(&self) -> &str {
        match &self.0 {
            DidInner::Plc(_) => "plc",
            DidInner::Web { .. } => "web",
            DidInner::Other(other) => other.method.as_str(),
        }
    }

    // TODO: did:plc: will now return the full 32 bytes - patch this to work with that
    // /// The identifier part of the DID.
    // ///
    // /// Allocates a String for `did:plc` (since it is stored as bytes rather than base32).
    // pub fn identifier<'a>(&'a self) -> Cow<'a, str> {
    //     match &self.0 {
    //         DidInner::Plc(bytes) => {
    //             let mut string = String::with_capacity(24);
    //             bytes_to_plc_ident_append(bytes, &mut string);
    //             Cow::Owned(string)
    //         }
    //         DidInner::Web { ptr, len } => {
    //             // SAFETY: ptr and len come from an unpacked DidWebDomain
    //             let web = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr, *len)) };
    //             // SAFETY: web is not dropped, so the string slice remains valid
    //             // The string slice is valid for as long as &self
    //             let ident: &'a str = unsafe { transmute::<_, &'a str>(web.as_str()) };
    //             ident.into()
    //         }
    //         DidInner::Other(other) => other.identifier.as_str().into(),
    //     }
    // }

    // /// Constructs a [`Did`] from a plain `did:plc` identifier (only the part after `did:plc:`).
    // pub fn from_plc_ident(_ident: &str) -> Option<Self> {
    //     todo!("Reimplement or remove")
    //     // plc_ident_to_bytes(ident).map(DidInner::Plc).map(Did)
    // }

    /// Constructs a [`Did`] directly from the identifier's bytes.
    pub fn from_plc_bytes(bytes: [u8; 15]) -> Self {
        Did(DidInner::Plc(bytes))
    }

    /// Constructs a [`Did`] from a web domain (the identifier after `did:web:`).
    pub fn from_web_domain(domain_str: &str) -> Option<Self> {
        // Try to parse the domain,
        // then unpack the raw data for a more packed repr
        // Storing the DidWebDomain directly in the enum makes it too large
        let web = DidWebDomain::from_str(domain_str).ok()?;
        let (ptr, len) = web.into_raw_parts();
        Some(Did(DidInner::Web { len, ptr }))
    }

    /// Constructs a [`Did`] from an arbitrary method and identifier.
    pub fn from_other_method_ident(method: &str, ident: &str) -> Option<Self> {
        // Method must be lowercase alphanumeric
        if !method.chars().all(|c: char| c.is_ascii_lowercase() || c.is_ascii_digit()) {
            return None;
        }

        // Identifier must be alphanumeric or one of: ._:%-
        if !ident.chars().all(|c| c.is_ascii_alphanumeric() || "._:%-".contains(c)) {
            return None;
        }

        // Identifier cannot end with % or :
        if ident.ends_with('%') || ident.ends_with(':') {
            return None;
        }

        // Identifier cannot be empty
        if ident.is_empty() {
            return None;
        }

        // Max identifier is (currently) 2048
        if ident.len() > 2048 {
            return None;
        }

        // Does not check percent sign encoding!

        Some(Did(DidInner::Other(Box::new(DidOther {
            method: method.to_owned(),
            identifier: ident.to_owned(),
        }))))
    }
}

#[derive(Debug, Error)]
pub enum DidParseError {
    #[error("missing `did:`")]
    MissingDid,
    #[error("missing identifier")]
    MissingIdent,
    #[error("invalid identifier")]
    InvalidIdent,
}

impl FromStr for Did {
    type Err = DidParseError;
    fn from_str(input: &str) -> Result<Self, Self::Err> {
        // Eagerly try to parse did:plc
        // This optimizes the code for the happy path
        if let Some(input_32) = input.as_bytes().as_array::<32>() {
            let plc_opt = plc_codec::decode_plc(input_32);
            if let Ok(did_inner) = plc_opt.try_into() {
                return Ok(Did(did_inner));
            }
        }

        let Some(method_ident) = input.strip_prefix("did:") else {
            return Err(DidParseError::MissingDid);
        };

        let Some((method, ident)) = method_ident.split_once(':') else {
            return Err(DidParseError::MissingIdent);
        };

        match method {
            "plc" => {
                // did:plc is parsed eagerly
                // If we got here, the DID was not 32 bytes long
                // We already checked for "did:" and the "plc" method,
                // so the identifier must be wrong
                Err(DidParseError::InvalidIdent)
            }
            "web" => Self::from_web_domain(ident).ok_or(DidParseError::InvalidIdent),
            _ => Self::from_other_method_ident(method, ident).ok_or(DidParseError::InvalidIdent),
        }
    }
}

/// Which method does this [`Did`] represent?
///
/// See [`Did::method()`] to get the string representation (especially for [`DidKind::Other`]).
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum DidKind {
    Plc,
    Web,
    Other,
}

#[repr(u8)]
enum DidInner {
    /// 1-byte discriminant + 15 identifier bytes
    Plc([u8; 15]),
    /// Contains an unpacked [`DidWebDomain`]
    Web { len: u8, ptr: *const u8 },
    /// A single-word pointer to two heap-allocated strings
    Other(Box<DidOther>),
}

#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
struct DidOther {
    method: String,
    identifier: String,
}

impl Drop for DidInner {
    fn drop(&mut self) {
        match self {
            DidInner::Web { len, ptr } => {
                // SAFETY: ptr and len come from an unpacked DidWebDomain
                // did:web needs to run its special drop code
                // This deallocates the memory at ptr!
                _ = unsafe { DidWebDomain::from_raw_parts(*ptr, *len) };
            }
            DidInner::Plc(_) | DidInner::Other(_) => {
                // plc and other don't require special drop logic
            }
        }
    }
}

impl PartialEq for Did {
    fn eq(&self, other: &Self) -> bool {
        match (&self.0, &other.0) {
            (DidInner::Plc(a), DidInner::Plc(b)) => a == b,
            (
                DidInner::Web { len: len_a, ptr: ptr_a },
                DidInner::Web { len: len_b, ptr: ptr_b },
            ) => {
                // SAFETY: ptr and len come from an unpacked DidWebDomain
                // ManuallyDrop prevents dropping the contents
                let a = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr_a, *len_a)) };
                let b = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr_b, *len_b)) };
                a.eq(&b)
            }
            (DidInner::Other(a), DidInner::Other(b)) => a.eq(b),
            _ => false,
        }
    }
}

impl Eq for Did {}

impl PartialOrd for Did {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Did {
    fn cmp(&self, other: &Self) -> Ordering {
        match (self, other) {
            // plc & plc -> compare bytes
            (Did(DidInner::Plc(a)), Did(DidInner::Plc(b))) => a.cmp(b),
            // web & web -> compare domains
            (
                Did(DidInner::Web { len: len_a, ptr: ptr_a }),
                Did(DidInner::Web { len: len_b, ptr: ptr_b }),
            ) => {
                // SAFETY: ptr and len come from an unpacked DidWebDomain
                // ManuallyDrop prevents dropping the contents
                let a = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr_a, *len_a)) };
                let b = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr_b, *len_b)) };
                a.cmp(&b)
            }
            // other & other -> compare methods and identifiers
            (Did(DidInner::Other(a)), Did(DidInner::Other(b))) => a.cmp(b),
            // different kinds -> just compare the method
            (a, b) => {
                debug_assert_ne!(a.kind(), b.kind());
                a.method().cmp(b.method())
            }
        }
    }
}

impl Hash for Did {
    fn hash<H: Hasher>(&self, state: &mut H) {
        match &self.0 {
            DidInner::Plc(bytes) => {
                state.write(b"plc");
                bytes.hash(state);
            }
            DidInner::Web { ptr, len } => {
                state.write(b"web");
                let web = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr, *len)) };
                web.hash(state);
            }
            DidInner::Other(other) => {
                other.method.hash(state);
            }
        }
    }
}

impl Clone for Did {
    fn clone(&self) -> Self {
        Did(match &self.0 {
            DidInner::Plc(data) => DidInner::Plc(*data),
            DidInner::Web { len, ptr } => {
                // SAFETY: ptr and len come from an unpacked DidWebDomain
                // ManuallyDrop prevents dropping the contents
                // The implementation of DidWebDomain::clone is used
                let web = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr, *len)) };
                let web = web.clone();
                // into_raw_parts doesn't drop the data
                let (ptr, len) = ManuallyDrop::into_inner(web).into_raw_parts();
                DidInner::Web { len, ptr }
            }
            DidInner::Other(other) => DidInner::Other(Box::clone(other)),
        })
    }
}

impl Debug for Did {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            DidInner::Plc(data) => f.debug_tuple("DidInner::Plc").field(data).finish(),
            DidInner::Web { len, ptr } => {
                // SAFETY: ptr and len come from an unpacked DidWebDomain
                // ManuallyDrop prevents dropping the contents
                let web = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr, *len)) };
                f.debug_tuple("DidInner::Web").field(&*web).finish()
            }
            DidInner::Other(other) => f.debug_tuple("DidInner::Other").field(other).finish(),
        }
    }
}

impl Display for Did {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match &self.0 {
            DidInner::Plc(data) => {
                let mut bytes = vec![0u8; 32];
                plc_codec::encode_plc(DidInner::Plc(*data), bytes.as_mut_array::<32>().unwrap());
                let string = String::from_utf8(bytes).expect("Encoded value should be UTF-8");
                write!(f, "{string}")
            }
            DidInner::Web { len, ptr } => {
                // SAFETY: ptr and len come from an unpacked DidWebDomain
                // ManuallyDrop prevents dropping the contents
                let web = unsafe { ManuallyDrop::new(DidWebDomain::from_raw_parts(*ptr, *len)) };
                write!(f, "did:web:{}", web.as_str())
            }
            DidInner::Other(other) => {
                write!(f, "did:{}:{}", other.method, other.identifier)
            }
        }
    }
}

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

    #[test]
    fn did_type_size() {
        assert_eq!(size_of::<Did>(), 16);
    }

    fn did_syntax_valid_lines() -> impl Iterator<Item = &'static str> {
        include_str!("./did_syntax_valid.txt")
            .lines()
            .filter(|l| !l.is_empty() && !l.starts_with('#'))
    }

    fn did_syntax_invalid_lines() -> impl Iterator<Item = &'static str> {
        include_str!("./did_syntax_invalid.txt")
            .lines()
            .filter(|l| !l.is_empty() && !l.starts_with('#'))
    }

    #[test_case::test_case("did:plc:c6te24qg5hx54qgegqylpqkx" => [0x17, 0xa6, 0x4d, 0x72, 0x06, 0xe9, 0xef, 0xde, 0x40, 0xc4, 0x34, 0x30, 0xb7, 0xc1, 0x57]
    )]
    #[test_case::test_case("did:plc:abcdefghijklmnopqrstuvwx" => [0x00, 0x44, 0x32, 0x14, 0xc7, 0x42, 0x54, 0xb6, 0x35, 0xcf, 0x84, 0x65, 0x3a, 0x56, 0xd7]
    )]
    #[test_case::test_case("did:plc:abcdefghabcdefghabcdefgh" => [0x00, 0x44, 0x32, 0x14, 0xc7, 0x00, 0x44, 0x32, 0x14, 0xc7, 0x00, 0x44, 0x32, 0x14, 0xc7]
    )]
    #[test_case::test_case("did:plc:234567234567234567234567" => [0xd6, 0xf9, 0xdf, 0x7f, 0x5b, 0xe7, 0x7d, 0xfd, 0x6f, 0x9d, 0xf7, 0xf5, 0xbe, 0x77, 0xdf]
    )]
    fn plc_str_to_bytes_ok(ident: &str) -> [u8; 15] {
        let DidInner::Plc(bytes) = Did::from_str(ident).unwrap().0 else { panic!() };
        bytes
    }

    #[test_case::test_case(&[0x17u8, 0xa6, 0x4d, 0x72, 0x06, 0xe9, 0xef, 0xde, 0x40, 0xc4, 0x34, 0x30, 0xb7, 0xc1, 0x57] => "did:plc:c6te24qg5hx54qgegqylpqkx"
    )]
    #[test_case::test_case(&[0x00u8, 0x44, 0x32, 0x14, 0xc7, 0x42, 0x54, 0xb6, 0x35, 0xcf, 0x84, 0x65, 0x3a, 0x56, 0xd7] => "did:plc:abcdefghijklmnopqrstuvwx"
    )]
    #[test_case::test_case(&[0x00u8, 0x44, 0x32, 0x14, 0xc7, 0x00, 0x44, 0x32, 0x14, 0xc7, 0x00, 0x44, 0x32, 0x14, 0xc7] => "did:plc:abcdefghabcdefghabcdefgh"
    )]
    #[test_case::test_case(&[0xd6u8, 0xf9, 0xdf, 0x7f, 0x5b, 0xe7, 0x7d, 0xfd, 0x6f, 0x9d, 0xf7, 0xf5, 0xbe, 0x77, 0xdf] => "did:plc:234567234567234567234567"
    )]
    fn bytes_to_plc_str_ok(bytes: &[u8; 15]) -> String {
        Did(DidInner::Plc(bytes.to_owned())).to_string()
    }

    #[test]
    fn did_syntax_valid() {
        for l in did_syntax_valid_lines() {
            Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
        }
    }

    #[test]
    fn did_syntax_invalid() {
        for l in did_syntax_invalid_lines() {
            Did::from_str(l).expect_err(&format!("Parsing {l} should have failed"));
        }
    }

    #[test]
    fn did_equals() {
        for l in did_syntax_valid_lines() {
            let a = Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
            let b = Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
            assert_eq!(a, b);
        }
    }

    #[test]
    fn did_clone_equals() {
        for l in did_syntax_valid_lines() {
            let a = Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
            let b = a.clone();
            assert_eq!(a, b, "cloned item should be equal");
        }
    }

    #[test]
    fn did_clone_drop_safe() {
        for l in did_syntax_valid_lines() {
            let a = Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
            let b = a.clone();
            drop(a);
            assert_eq!(b.to_string(), l, "cloned item roundtrip should work after drop");

            let a = Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
            let b = a.clone();
            drop(b);
            assert_eq!(a.to_string(), l, "original item roundtrip should work after drop");
        }
    }

    #[test]
    fn did_str_roundtrip() {
        for l in did_syntax_valid_lines() {
            let did = Did::from_str(l).unwrap_or_else(|_| panic!("Parsing {l} failed"));
            assert_eq!(did.to_string(), l);
        }
    }

    #[test]
    #[ignore]
    fn did_parts_roundtrip() {
        // TODO Fix once identifier() is fixed
        // for l in did_syntax_valid_lines() {
        //     let did = Did::from_str(l).expect(&format!("Parsing {l} failed"));
        //     assert_eq!(format!("did:{}:{}", did.method(), did.identifier()), l);
        // }
    }
}