path: root/uri-pack/src/lib.rs

/// Pack an URI ascii char
/// Panic if char not supported
fn pack_ascii(c: u8) -> u8 {
    [
        67, 68, 69, 70, 29, 71, 72, 73, 74, 75, 76, 77, 28, 26, 27, 57, 58, 59, 60, 61, 62, 63, 64,
        65, 66, 78, 79, 80, 81, 82, 83, 84, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,
        45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 85, 86, 87, 88, 30, 89, 0, 1, 2, 3, 4, 5,
        6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 90, 91, 92, 93,
    ][(c - b'!') as usize]
}

/// Unpack an URI ascii char
/// Panic if char not supported
fn unpack_ascii(c: u8) -> u8 {
    [
        b'a', b'b', b'c', b'd', b'e', b'f', b'g', b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o',
        b'p', b'q', b'r', b's', b't', b'u', b'v', b'w', b'x', b'y', b'z', b'.', b'/', b'-', b'%',
        b'_', b'A', b'B', b'C', b'D', b'E', b'F', b'G', b'H', b'I', b'J', b'K', b'L', b'M', b'N',
        b'O', b'P', b'Q', b'R', b'S', b'T', b'U', b'V', b'W', b'X', b'Y', b'Z', b'0', b'1', b'2',
        b'3', b'4', b'5', b'6', b'7', b'8', b'9', b'!', b'"', b'#', b'$', b'&', b'\'', b'(', b')',
        b'*', b'+', b',', b':', b';', b'<', b'=', b'>', b'?', b'@', b'[', b'\\', b']', b'^', b'`',
        b'{', b'|', b'}', b'~',
    ][c as usize]
}

/// Check if an ascii char is supported by the encoding
fn supported_ascii(c: &u8) -> bool {
    (b'!'..=b'~').contains(c)
}

/// Extended packing limit
pub const EXTENDED: u8 = 30;
/// EOF u5 encoding
pub const TERMINATOR: u8 = 31;
/// Ascii char that can be packed into 5 bits
pub const PACKED: [u8; 30] = [
    b'a', b'b', b'c', b'd', b'e', b'f', b'g', b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o', b'p',
    b'q', b'r', b's', b't', b'u', b'v', b'w', b'x', b'y', b'z', b'.', b'/', b'-', b'%',
];

#[derive(Debug, Clone, Copy, thiserror::Error)]
pub enum EncodeErr {
    #[error("{0} is not a valid uri char")]
    UnsupportedChar(u8),
}

#[derive(Debug, Clone, Copy, thiserror::Error)]
pub enum DecodeErr {
    #[error("An extended encoded char have been passed as an simple one")]
    ExpectedExtended,
    #[error("{0} is not an simple encoded char")]
    UnexpectedSimpleChar(u8),
    #[error("{0} is not an extended encoded char")]
    UnexpectedExtendedChar(u8),
    #[error("Missing bits")]
    UnexpectedEOF,
}

/// Pack an uri string into an optimized binary format
pub fn pack_uri(uri: &str) -> Result<Vec<u8>, EncodeErr> {
    let len = uri.as_bytes().len();
    let mut vec = Vec::with_capacity(len);

    if let Some(char) = uri.as_bytes().iter().find(|c| !supported_ascii(c)) {
        return Err(EncodeErr::UnsupportedChar(*char));
    }

    // Holds [buff_bits] pending bits beginning from the most significant bits
    let (mut buff, mut buff_bits) = (0u8, 0u8);

    // Write [nb_bits] less significant bits from [nb] to [buff]
    let mut write_bits = |nb: u8, mut nb_bits: u8| {
        while nb_bits > 0 {
            // Amount of bits we can write in buffer
            let writable = (8 - buff_bits).min(nb_bits);
            // Remove non writable bits
            let rmv_right = nb >> nb_bits - writable;
            let rmv_left = rmv_right << 8 - writable;
            // Align remaining bits with buff blank bits
            let align = rmv_left >> buff_bits;

            // Write bits in buffer
            buff = buff | align;
            buff_bits += writable;
            nb_bits -= writable;

            // Store buffer if full
            if buff_bits == 8 {
                vec.push(buff);
                buff = 0;
                buff_bits = 0;
            }
        }
    };

    for c in uri.bytes() {
        let nb = pack_ascii(c);
        if nb < EXTENDED {
            write_bits(nb, 5)
        } else {
            write_bits(EXTENDED, 5);
            write_bits(nb - EXTENDED, 6);
        }
    }
    write_bits(TERMINATOR, 5);

    // Push pending buffer if not empty
    if buff_bits > 0 {
        vec.push(buff);
    }

    return Ok(vec);
}

/// Unpack an uri string from its optimized binary format
pub fn unpack_uri(bytes: &[u8]) -> Result<String, DecodeErr> {
    let mut buf = String::with_capacity(bytes.len());
    let mut iter = bytes.iter();

    // Holds [buff_bits] pending bits beginning from the most significant bits
    let (mut buff, mut buff_bits) = (0u8, 0u8);

    // Write [nb_bits] less significant bits from [buff] to [nb]
    let mut read_nb = |mut nb_bits: u8| -> Result<u8, DecodeErr> {
        let mut nb = 0;
        while nb_bits > 0 {
            // Load buff if empty
            if buff_bits == 0 {
                buff = *iter.next().ok_or(DecodeErr::UnexpectedEOF)?;
                buff_bits = 8;
            }
            // Amount of bits we can read from buff
            let readable = buff_bits.min(nb_bits);
            // Remove non writable bits
            let rmv_left = buff << 8 - buff_bits;
            // Align remaining bits with nb blank bits
            let align = rmv_left >> (8 - readable);
            // Read bits from buff
            nb = (nb << readable) | align;
            buff_bits -= readable;
            nb_bits -= readable;
        }
        return Ok(nb);
    };

    loop {
        let encoded = match read_nb(5)? {
            TERMINATOR => break,
            EXTENDED => read_nb(6)? + EXTENDED,
            nb => nb,
        };
        buf.push(unpack_ascii(encoded) as char);
    }

    return Ok(buf);
}

#[cfg(test)]
#[macro_use(quickcheck)]
extern crate quickcheck_macros;

#[cfg(test)]
mod test {
    use std::str::FromStr;

    use serde_json::Value;

    use crate::{
        pack_ascii, pack_uri, supported_ascii, unpack_ascii, unpack_uri, EXTENDED, PACKED,
    };

    #[test]
    /// Check support every packable ascii character is packed
    fn packed() {
        for c in PACKED {
            assert!(pack_ascii(c) < EXTENDED);
        }
    }

    #[test]
    /// Check support every ascii graphic character and space
    fn supported() {
        for c in (0..=255u8).filter(supported_ascii) {
            assert_eq!(unpack_ascii(pack_ascii(c)), c);
        }
    }

    #[test]
    /// Check error on unsupported char
    fn unsupported() {
        for c in (0..=255u8).filter(|c| !supported_ascii(c)) {
            let string = String::from(c as char);
            assert!(pack_uri(&string).is_err());
        }
    }

    #[test]
    fn url_simple() {
        let mut majestic =
            csv::Reader::from_reader(include_str!("majestic_million.csv").as_bytes());
        for record in majestic.records() {
            let domain = &record.unwrap()[2];
            let encoded = pack_uri(domain).unwrap();
            let decoded = unpack_uri(&encoded).unwrap();
            assert_eq!(domain, decoded);
        }
    }

    #[test]
    fn url_complex() {
        let mut json = Value::from_str(include_str!("urltestdata.json"))
            .expect("JSON parse error in urltestdata.json");
        for entry in json.as_array_mut().unwrap() {
            if entry.is_string() {
                continue; // ignore comments
            }

            let href = entry.get("href").and_then(|it| it.as_str()).unwrap_or("");
            if href.chars().any(|c| !c.is_ascii_graphic() || c != ' ') {
                continue; // extended ascii
            }
            let encoded = pack_uri(&href).expect(&format!("Failed to encode {}", &href));
            let decoded =
                unpack_uri(&encoded).expect(&format!("Failed to decode encoded {}", &href));
            assert_eq!(href, decoded);
        }
    }

    #[quickcheck]
    fn fuzz(input: String) -> bool {
        if input.as_bytes().iter().all(supported_ascii) {
            let packed = pack_uri(&input).unwrap();
            let unpacked = unpack_uri(&packed).unwrap();
            input == unpacked
        } else {
            true
        }
    }
}