cryptid_rs/

codec.rs

use std::fmt;

use aes::Aes256;
use fpe::ff1::{BinaryNumeralString, FF1};
use hkdf::Hkdf;
use hmac::{Hmac, Mac};
use sha2::Sha256;
use uuid::Uuid;

use crate::Config;

type HmacSha256 = Hmac<Sha256>;

/// Error returned for encode/decode errors.
#[derive(Debug, PartialEq)]
pub enum Error {
    DecodingFailed,
    DecryptionFailed,
    EncryptionFailed,
    IncorrectMAC,
    InvalidDataLength,
    InvalidPrefix { received: String, expected: String },
    SentinelMismatch { received: u8, expected: u8 },
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::DecodingFailed => {
                write!(f, "Decoding string failed")
            }
            Error::DecryptionFailed => {
                write!(f, "FF1 decryption failed")
            }
            Error::EncryptionFailed => {
                write!(f, "FF1 encryption failed")
            }
            Error::IncorrectMAC => {
                write!(f, "Incorrect MAC")
            }
            Error::InvalidDataLength => {
                write!(f, "Invalid data length")
            }
            Error::SentinelMismatch { received, expected } => {
                write!(f, "Sentinel byte was {received}, expected {expected}")
            }
            Error::InvalidPrefix { received, expected } => {
                write!(f, "Prefix was {received}, expected {expected}")
            }
        }
    }
}

impl From<base62::DecodeError> for Error {
    fn from(_: base62::DecodeError) -> Error {
        Error::DecodingFailed
    }
}

impl std::error::Error for Error {}

// Maximum number of bytes we can base62 encode (an u128).
const MAX_BUFFER: usize = 16;

// The sentinel byte, in case we don't fill the full 16 bytes.
const SENTINEL: u8 = 1;

/// Core encoder/decoder.
pub struct Codec {
    ff1: FF1<Aes256>,
    hmac: HmacSha256,
    hmac_length: usize,
    prefix: String,
    zero_pad_length: usize,
}

impl Codec {
    /// Creates a new `Codec` instance with the given name and key.
    ///
    /// The `name` is used as a prefix in the encoded output and to derive a prefix-specifc
    /// key together with the master `key`.
    ///
    /// **Security note:** In order to be secure, you must provide a secure random `key`
    /// with sufficient entropy, and manage it appropriately.
    ///
    /// # Arguments
    ///
    /// * `name` - A string slice that holds the name of the codec.
    /// * `key` - A byte slice that holds the master key for encryption and MAC.
    ///
    /// # Returns
    ///
    /// A new instance of `Codec`.
    ///
    /// # Examples
    ///
    /// ```
    /// use cryptid_rs::{Config, Codec};
    ///
    /// let codec = Codec::new("example", &Config::new(b"your-secure-key"));
    /// ```
    pub fn new(name: &str, config: &Config) -> Codec {
        let hkdf = Hkdf::<Sha256>::new(None, &config.key);
        let mut ff1_key = [0u8; 32];
        let mut hmac_key = [0u8; 32];
        hkdf.expand(format!("{name}/ff1").as_bytes(), &mut ff1_key)
            .expect("Length 32 should be valid");
        hkdf.expand(format!("{name}/hmac").as_bytes(), &mut hmac_key)
            .expect("Length 32 should be valid");
        Codec {
            ff1: FF1::<Aes256>::new(&ff1_key, 2).expect("Radix 2 should be valid"),
            hmac: HmacSha256::new_from_slice(&hmac_key).expect("Key length 32 should be valid"),
            hmac_length: config.hmac_length as usize,
            prefix: format!("{name}_"),
            zero_pad_length: config.zero_pad_length as usize,
        }
    }

    /// Encodes a given numeric value into a secure string representation.
    ///
    /// This method applies format-preserving encryption to the number and
    /// then encodes it into a base62 string with a prefix. It also appends
    /// an HMAC for integrity verification.
    ///
    /// # Arguments
    ///
    /// * `num` - The 64-bit unsigned integer to be encoded.
    ///
    /// # Returns
    ///
    /// A `Result` which is `Ok` containing the encoded string if successful,
    /// or an `Error` if encoding fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use cryptid_rs::{Codec, Config};
    ///
    /// let codec = Codec::new("example", &Config::new(b"your-secure-key"));
    /// let encoded = codec.encode(12345);
    ///
    /// assert_eq!(encoded, "example_VgwPy6rwatl");
    /// ```
    pub fn encode(&self, num: u64) -> String {
        let encoded = base62::encode(self.encode_u128(num));
        format!("{}{}", self.prefix, encoded)
    }

    /// Encrypts `num` into a 128 bit value.  Note that high order bits may be zeroes,
    /// so that a short string representation can be made.
    fn encode_u128(&self, num: u64) -> u128 {
        let bytes = encrypt_number(
            &self.ff1,
            &self.hmac,
            self.hmac_length,
            self.zero_pad_length,
            num,
        );
        let mut num_array = [0u8; MAX_BUFFER];
        num_array[..bytes.len()].copy_from_slice(&bytes);
        if bytes.len() < num_array.len() {
            num_array[bytes.len()] = SENTINEL;
        }
        u128::from_le_bytes(num_array)
    }

    /// Encrypts `num` into an UUID.  This is useful in situations where using an UUID is
    /// forced by some interface.  Otherwise you should use the default type safe strings.
    pub fn encode_uuid(&self, num: u64) -> Uuid {
        // 8 bytes for hmac and 8 bytes for payload gets us a nice random 128 bit value.
        let vec = encrypt_number(&self.ff1, &self.hmac, 8, 8, num);
        let num = u128::from_le_bytes(vec.try_into().expect("Should have exactly 16 bytes"));
        Uuid::from_u128_le(num)
    }

    /// Decodes a previously encrypted UUID back into its original numeric value.
    ///
    /// # Arguments
    ///
    /// * `uuid` - The UUID to be decoded.
    ///
    /// # Returns
    ///
    /// A `Result` which is `Ok` containing the decoded 64-bit unsigned integer if successful,
    /// or an `Error` if decoding fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use cryptid_rs::{Codec, Config};
    /// use uuid::Uuid;
    ///
    /// let codec = Codec::new("test", &Config::new(b"Test key here"));
    /// let uuid = codec.encode_uuid(0);
    /// let decoded = codec.decode_uuid(uuid).unwrap();
    ///
    /// assert_eq!(decoded, 0);
    /// ```
    pub fn decode_uuid(&self, uuid: Uuid) -> Result<u64, Error> {
        let bytes = uuid.to_u128_le().to_le_bytes();
        decrypt_number(&self.ff1, &self.hmac, 8, 8, &bytes)
    }

    /// Decodes a previously encoded string back into its original numeric value.
    ///
    /// This method first verifies the integrity of the encoded data using HMAC,
    /// and then applies format-preserving decryption to retrieve the original number.
    /// It expects the encoded data to start with the correct prefix.
    ///
    /// # Arguments
    ///
    /// * `encoded` - A string slice representing the encoded data.
    ///
    /// # Returns
    ///
    /// A `Result` which is `Ok` containing the decoded 64-bit unsigned integer if successful,
    /// or an `Error` if decoding fails.
    ///
    /// # Examples
    ///
    /// ```
    /// use cryptid_rs::{Codec, Config};
    ///
    /// let codec = Codec::new("example", &Config::new(b"your-secure-key"));
    /// let decoded = codec.decode("example_VgwPy6rwatl").unwrap();
    ///
    /// assert_eq!(decoded, 12345);
    /// ```
    pub fn decode(&self, encoded: &str) -> Result<u64, Error> {
        // Ensure prefix matches (from last underscore).
        let received = match encoded.rfind('_') {
            None => "".to_string(),
            Some(i) => encoded[..i + 1].to_string(),
        };
        if received != self.prefix {
            let expected = self.prefix.clone();
            return Err(Error::InvalidPrefix { received, expected });
        }

        let tail = &encoded[self.prefix.len()..];
        let num = base62::decode(tail).map_err(Error::from)?;
        let num_array = num.to_le_bytes();

        let length;
        if self.hmac_length + self.zero_pad_length < MAX_BUFFER {
            length = last_nonzero(&num_array);
            if num_array[length] != SENTINEL {
                return Err(Error::SentinelMismatch {
                    received: num_array[length],
                    expected: SENTINEL,
                });
            }
        } else {
            length = MAX_BUFFER;
        }

        decrypt_number(
            &self.ff1,
            &self.hmac,
            self.hmac_length,
            self.zero_pad_length,
            &num_array[..length],
        )
    }
}

fn last_nonzero(bytes: &[u8]) -> usize {
    bytes.iter().rposition(|&b| b != 0).unwrap_or(0)
}

// Returns a memory representanion of `num` as a byte vector in little-endian byte
// order, leaving out trailing zero bytes beyond `min_length`.
fn num_to_le_vec(num: u64, min_length: usize) -> Vec<u8> {
    let bytes = num.to_le_bytes();
    let prefix_length = (last_nonzero(&bytes) + 1).max(min_length);
    bytes[..prefix_length].to_vec()
}

fn le_vec_to_num(bytes: &[u8]) -> u64 {
    let mut arr = [0; 8];
    arr[..bytes.len()].copy_from_slice(bytes);
    u64::from_le_bytes(arr)
}

fn encrypt_number(
    ff1: &FF1<Aes256>,
    hmac: &HmacSha256,
    hmac_length: usize,
    zero_pad_length: usize,
    num: u64,
) -> Vec<u8> {
    // Encrypt `num` using form-preserving encryption.
    let pt = num_to_le_vec(num, zero_pad_length);
    let encrypted_num = ff1
        .encrypt(&[], &BinaryNumeralString::from_bytes_le(&pt))
        .expect("Radix 2 should be valid")
        .to_bytes_le();

    // Compute a truncated MAC from the ciphertext.
    let mut hmac: HmacSha256 = hmac.clone();
    hmac.update(&encrypted_num);
    let truncated_mac = &hmac.finalize().into_bytes()[..hmac_length];

    // Return the combined bytes.
    let mut result = encrypted_num.to_vec();
    result.extend_from_slice(truncated_mac);

    result
}

fn decrypt_number(
    ff1: &FF1<Aes256>,
    hmac: &HmacSha256,
    hmac_length: usize,
    zero_pad_length: usize,
    encrypted_data: &[u8],
) -> Result<u64, Error> {
    if encrypted_data.len() < hmac_length + zero_pad_length {
        return Err(Error::InvalidDataLength);
    }
    let (encrypted_num, received_mac) = encrypted_data.split_at(encrypted_data.len() - hmac_length);

    // Verify MAC
    let mut hmac_clone: HmacSha256 = hmac.clone();
    hmac_clone.update(encrypted_num);
    let truncated_mac = &hmac_clone.finalize().into_bytes()[..hmac_length];
    if truncated_mac != received_mac {
        return Err(Error::IncorrectMAC);
    }

    // Decrypt the number
    let decrypted_num = ff1
        .decrypt(&[], &BinaryNumeralString::from_bytes_le(encrypted_num))
        .map_err(|_| Error::DecryptionFailed)?;

    // Convert decrypted bytes back to number
    let num: u64 = le_vec_to_num(&decrypted_num.to_bytes_le());
    Ok(num)
}

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

    #[test]
    fn test_defaults() {
        let codec = Codec::new("test", &Config::new(b"Test key here"));
        let test_cases = vec![
            (0, "test_g1HdsEGpXp5"),
            (1, "test_bTPc8uxHEwv"),
            (2, "test_dZ0iJdcLBgB"),
            (123, "test_hHLBCl4rZ3u"),
            (u64::MAX, "test_20cMzlnhTkILdJzWt"),
        ];

        for (input, expected) in test_cases {
            assert_eq!(codec.encode(input), expected);
            assert_eq!(codec.decode(expected).unwrap(), input);
        }
    }

    #[test]
    fn test_uuid() {
        let codec = Codec::new("test", &Config::new(b"Test key here"));
        let test_cases = [
            (0, "59142369-adeb-8ef9-a1be-28f61c05d4d6"),
            (1, "93196956-2d32-d8d2-54f7-9a86fc765f3a"),
            (2, "3c10f25c-005e-6f6f-87a9-781efe02d14d"),
            (123, "571fd9d5-e133-f7b0-b0df-f444e4dd1127"),
            (u64::MAX, "a3b06cf5-dd4d-3f09-4000-9d3519d4d6c2"),
        ];

        for &(input, expected) in &test_cases {
            let uuid = Uuid::parse_str(expected).unwrap();
            assert_eq!(codec.encode_uuid(input), uuid);

            // Test roundtrip
            let decoded = codec.decode_uuid(uuid).unwrap();
            assert_eq!(decoded, input, "Failed to decode UUID for input: {}", input);
        }
    }

    #[test]
    fn test_uuid_roundtrip() {
        let codec = Codec::new("test", &Config::new(b"Test key here"));
        let mut rng = rand::rng();

        for _ in 0..1_000 {
            let number: u64 = rng.random();
            let uuid = codec.encode_uuid(number);
            let decoded = codec.decode_uuid(uuid).expect("Decoding failed");

            assert_eq!(decoded, number, "Failed at number: {}", number);
        }
    }

    #[test]
    fn test_long() {
        let config = Config::new(b"Test key here")
            .hmac_length(8)
            .unwrap()
            .zero_pad_length(8)
            .unwrap();
        let codec = Codec::new("test", &config);
        assert_eq!(codec.encode(0), "test_6XNFaHOCeuIBNvRT4pIrVZ");
        assert_eq!(codec.encode(1), "test_1m9BJW23Jk5hSIlfPxoboZ");
        assert_eq!(codec.encode(2), "test_2MpvWPgnp5j1dIqFnJVOjU");
        assert_eq!(codec.encode(123), "test_1BirgT1ZJhfSsKFLgxA5gt");
        assert_eq!(codec.encode(u64::MAX), "test_5vegfyOLrrmwtgznQByI4J");
        assert_eq!(codec.decode("test_6XNFaHOCeuIBNvRT4pIrVZ").unwrap(), 0);
        assert_eq!(codec.decode("test_1m9BJW23Jk5hSIlfPxoboZ").unwrap(), 1);
        assert_eq!(codec.decode("test_2MpvWPgnp5j1dIqFnJVOjU").unwrap(), 2);
        assert_eq!(codec.decode("test_1BirgT1ZJhfSsKFLgxA5gt").unwrap(), 123);
        assert_eq!(
            codec.decode("test_5vegfyOLrrmwtgznQByI4J").unwrap(),
            u64::MAX
        );
    }

    #[test]
    fn test_short() {
        let config = Config::new(b"Test key here")
            .hmac_length(0)
            .unwrap()
            .zero_pad_length(3)
            .unwrap();
        let codec = Codec::new("test", &config);
        assert_eq!(codec.encode(0), "test_1zG8O");
        assert_eq!(codec.encode(1), "test_1R8PN");
        assert_eq!(codec.encode(2), "test_1nzgo");
        assert_eq!(codec.encode(123), "test_1YqNT");
        assert_eq!(codec.encode(u64::MAX), "test_Mlu72Yai97j");
        assert_eq!(codec.decode("test_1zG8O").unwrap(), 0);
        assert_eq!(codec.decode("test_1R8PN").unwrap(), 1);
        assert_eq!(codec.decode("test_1nzgo").unwrap(), 2);
        assert_eq!(codec.decode("test_1YqNT").unwrap(), 123);
        assert_eq!(codec.decode("test_Mlu72Yai97j").unwrap(), u64::MAX);

        // Without HMAC, pretty much anything decodes to some number.
        assert_eq!(codec.decode("test_1helloall").unwrap(), 20580488769766);
    }

    #[test]
    fn test_decode_errors() {
        let codec = Codec::new("test", &Config::new(b"Test key here"));

        assert_eq!(
            codec.decode("hHLBCl4rZ3u"),
            Err(Error::InvalidPrefix {
                received: "".to_string(),
                expected: "test_".to_string()
            })
        );

        assert_eq!(
            codec.decode("_hHLBCl4rZ3u"),
            Err(Error::InvalidPrefix {
                received: "_".to_string(),
                expected: "test_".to_string()
            })
        );

        assert_eq!(
            codec.decode("wrong_hHLBCl4rZ3u"),
            Err(Error::InvalidPrefix {
                received: "wrong_".to_string(),
                expected: "test_".to_string()
            })
        );

        assert_eq!(
            codec.decode("test_iHLBCl4rZ3u"),
            Err(Error::SentinelMismatch {
                received: 2,
                expected: SENTINEL,
            })
        );

        // Tampering with any part gives a MAC error.
        assert_eq!(codec.decode("test_hHLBCl4rZ3v"), Err(Error::IncorrectMAC));
        assert_eq!(codec.decode("test_hHMBCl4rZ3u"), Err(Error::IncorrectMAC));

        // Invalid characters aren't allowed.
        assert_eq!(codec.decode("test_hHLBCl+rZ3u"), Err(Error::DecodingFailed));

        // And just to validate the above, check that the correct string does decode.
        assert_eq!(codec.decode("test_hHLBCl4rZ3u"), Ok(123));
    }

    #[test]
    fn test_random_roundtrips() {
        let codec = Codec::new("test", &Config::new(b"Test key here"));
        let mut rng = rand::rng();

        for _ in 0..10_000 {
            let number: u64 = rng.random();
            let encoded = codec.encode(number);
            let decoded = codec.decode(&encoded).expect("Decoding failed");

            assert_eq!(decoded, number, "Failed at number: {}", number);
        }
    }

    #[test]
    fn test_thread_local_config_isolation() {
        use crate::{Config, Field, TypeMarker};
        use std::thread;

        #[derive(Clone, Copy, Debug)]
        pub struct TestIdMarker;
        impl TypeMarker for TestIdMarker {
            fn name() -> &'static str {
                "test"
            }
        }

        type TestId = Field<TestIdMarker>;

        // Set global config
        Config::set_global(Config::new(b"global-key-16bytes"));

        // Test that global config works
        let id = TestId::from(123);
        let global_encoded = id.to_string();

        // Test thread-local config in separate thread
        let handle = thread::spawn(|| {
            Config::set_thread_local(Config::new(b"thread-key-16bytes"));

            let id = TestId::from(123);
            id.to_string()
        });

        let thread_result = handle.join().unwrap();

        // Verify they're different (thread isolation works)
        assert_ne!(
            global_encoded, thread_result,
            "Thread-local config should produce different encoding than global config"
        );

        // Main thread should still use global config
        let main_again = TestId::from(123).to_string();
        assert_eq!(
            global_encoded, main_again,
            "Main thread should still use global config"
        );
    }
}