maarcadetweet/crates/at-mst/src/util.rs

use anyhow::{anyhow, Result};

use at_crypto::cid::sha256;

pub const DEFAULT_FANOUT: usize = 8;

pub fn max_layer_for_fanout(fanout: usize) -> usize {
    if fanout <= 1 {
        return 0;
    }
    (usize::ilog2(fanout) as usize).saturating_sub(1)
}

pub fn count_leading_zero_bits(hash: &[u8]) -> usize {
    let mut count = 0usize;
    for &byte in hash {
        if byte == 0 {
            count += 8;
        } else {
            count += byte.leading_zeros() as usize;
            break;
        }
    }
    count
}

/// Hash a record key to the 32-byte digest used as comparison input
/// throughout the MST. Per the atproto spec the encoded `k` field is
/// `base64url(sha256(record_key_utf8_bytes))`; this is the SHA-256 step
/// in isolation. Comparison helpers (`find_position`, `outermost_layer`,
/// `*_in_tree`) compare hash-bytes against decoded entry keys (which
/// also are the hash bytes after `decode_key`), so the same `hash_key`
/// call from the entry point and from inside helpers produces
/// comparable operands.
pub fn hash_key(raw_key: &str) -> [u8; 32] {
    sha256(raw_key.as_bytes())
}

/// Layer that an already-hashed key occupies in the tree of `fanout`.
/// Use after `hash_key` to avoid hashing twice.
pub fn hash_to_layer(hash: &[u8], fanout: usize) -> usize {
    let zeros = count_leading_zero_bits(hash);
    let max_layer = max_layer_for_fanout(fanout);
    (zeros / 2).min(max_layer)
}

pub fn key_to_layer(raw_key: &str, fanout: usize) -> usize {
    hash_to_layer(&hash_key(raw_key), fanout)
}

/// Encode a record key for storage in an MST entry.
///
/// Per the atproto MST spec
/// (<https://atproto.com/specs/data-model-repo#node-data>) the `k`
/// field is `base64url(sha256(record_key_utf8_bytes))`. Hashing
/// first ties the layer distribution to the cryptographic digest
/// of the key — under pre-image resistance, an attacker can't
/// craft keys that all land at the maximum layer by sorting their
/// bytes a certain way.
///
/// Decoding returns the raw 32-byte hash bytes; callers that need
/// the original key string have to keep it alongside. Cross-crate
/// callers passing `vec::Vec<u8>` vs `[u8; 32]` will need a trivial
/// .as_slice() conversion at the comparison site.
pub fn encode_key(raw_key: &str) -> String {
    use base64::Engine;
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(hash_key(raw_key))
}

/// Inverse of [`encode_key`]: round-trip the base64url string back
/// to the 32-byte SHA-256 hash. Rejects anything that doesn't decode
/// to exactly 32 bytes — i.e. catches the old `base64url(raw_key)`
/// encoding that predates this commit, which makes it easy to spot
/// incompatibilities during migration.
pub fn decode_key(encoded: &str) -> Result<Vec<u8>> {
    use base64::Engine;
    let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(encoded.as_bytes())
        .map_err(|e| anyhow!("invalid base64url key `{encoded}`: {e}"))?;
    if bytes.len() != 32 {
        return Err(anyhow!(
            "decoded key `{encoded}` is {} bytes; expected 32 (sha256 hash per the atproto MST spec)",
            bytes.len()
        ));
    }
    Ok(bytes)
}

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

    #[test]
    fn max_layer_for_fanout_8() {
        assert_eq!(max_layer_for_fanout(8), 2);
    }

    #[test]
    fn max_layer_for_fanout_16() {
        assert_eq!(max_layer_for_fanout(16), 3);
    }

    #[test]
    fn max_layer_for_fanout_1() {
        assert_eq!(max_layer_for_fanout(1), 0);
    }

    #[test]
    fn count_leading_zeros_all_zero() {
        let h = [0u8; 32];
        assert_eq!(count_leading_zero_bits(&h), 256);
    }

    #[test]
    fn count_leading_zeros_one_bit() {
        let mut h = [0u8; 32];
        h[0] = 0b0000_0001;
        assert_eq!(count_leading_zero_bits(&h), 7);
    }

    #[test]
    fn count_leading_zeros_one_nibble() {
        let mut h = [0u8; 32];
        h[0] = 0x0f;
        assert_eq!(count_leading_zero_bits(&h), 4);
    }

    #[test]
    fn count_leading_zeros_byte_boundary() {
        let mut h = [0u8; 32];
        h[2] = 0x80;
        assert_eq!(count_leading_zero_bits(&h), 16);
        let mut h = [0u8; 32];
        h[2] = 0x01;
        assert_eq!(count_leading_zero_bits(&h), 23);
    }

    #[test]
    fn key_to_layer_zero_layer() {
        let layer = key_to_layer("com.example.foo/abc", 8);
        assert!(layer <= 2);
    }
}