aingle_id/

lib.rs

//! AIngle AI-ID base32 encoding utility.
//!
//! # Example
//!
//! ```
//! fn main() {
//!     let enc = aingle_id::AiidEncoding::with_kind("ais0").unwrap();
//!     let key = enc.encode(&[0; 32]).unwrap();
//!     assert_eq!("aiSciaaaa", &key[..9]);
//!     let buffer = enc.decode(&key).unwrap();
//!     assert_eq!([0; 32].to_vec(), buffer);
//! }
//! ```

mod error;
mod b32;
pub use error::{AiidError, AiidResult};

mod util;
use util::{b32_correct, cap_decode, cap_encode_bin, char_upper};

// AI_CODE_MAP: Maps third character of kind (e.g., 'k' in "aik0") to res byte
// The res byte encodes: bits 7-6 = 00 (for "ai" prefix), bits 5-1 = third char value, bit 0 = version
// Index = ASCII code of third char - 51
static AI_CODE_MAP: &[[u8; 2]] = &[
    [ 0x32, 0x33 ], // 51 '3': ai30/ai31 -> third char '3' = 25 = 11001 -> 0x32/0x33
    [ 0x34, 0x35 ], // 52 '4': ai40/ai41 -> third char '4' = 26 = 11010 -> 0x34/0x35
    [ 0x36, 0x37 ], // 53 '5': ai50/ai51 -> third char '5' = 27 = 11011 -> 0x36/0x37
    [ 0x38, 0x39 ], // 54 '6': ai60/ai61 -> third char '6' = 28 = 11100 -> 0x38/0x39
    [ 0x3a, 0x3b ], // 55 '7': ai70/ai71 -> third char '7' = 29 = 11101 -> 0x3a/0x3b
    [ 0x3c, 0x3d ], // 56 '8': ai80/ai81 -> third char '8' = 30 = 11110 -> 0x3c/0x3d
    [ 0x3e, 0x3f ], // 57 '9': ai90/ai91 -> third char '9' = 31 = 11111 -> 0x3e/0x3f

    // 58-61: reserved (: ; < =)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 62-65: reserved (> ? @ A)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 66-69: reserved (B C D E)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 70-73: reserved (F G H I)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 74-77: reserved (J K L M)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 78-81: reserved (N O P Q)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 82-85: reserved (R S T U)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 86-89: reserved (V W X Y)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 90-93: reserved (Z [ \ ])
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],
    // 94-96: reserved (^ _ `)
    [ 0xff, 0xff ], [ 0xff, 0xff ], [ 0xff, 0xff ],

    [ 0x00, 0x01 ], // 97 'a': aia0/aia1 -> third char 'A' = 0 = 00000 -> 0x00/0x01
    [ 0x02, 0x03 ], // 98 'b': aib0/aib1 -> third char 'B' = 1 = 00001 -> 0x02/0x03
    [ 0x04, 0x05 ], // 99 'c': aic0/aic1 -> third char 'C' = 2 = 00010 -> 0x04/0x05
    [ 0x06, 0x07 ], // 100 'd': aid0/aid1 -> third char 'D' = 3 = 00011 -> 0x06/0x07
    [ 0x08, 0x09 ], // 101 'e': aie0/aie1 -> third char 'E' = 4 = 00100 -> 0x08/0x09
    [ 0x0a, 0x0b ], // 102 'f': aif0/aif1 -> third char 'F' = 5 = 00101 -> 0x0a/0x0b
    [ 0x0c, 0x0d ], // 103 'g': aig0/aig1 -> third char 'G' = 6 = 00110 -> 0x0c/0x0d
    [ 0x0e, 0x0f ], // 104 'h': aih0/aih1 -> third char 'H' = 7 = 00111 -> 0x0e/0x0f
    [ 0x10, 0x11 ], // 105 'i': aii0/aii1 -> third char 'I' = 8 = 01000 -> 0x10/0x11
    [ 0x12, 0x13 ], // 106 'j': aij0/aij1 -> third char 'J' = 9 = 01001 -> 0x12/0x13
    [ 0x14, 0x15 ], // 107 'k': aik0/aik1 -> third char 'K' = 10 = 01010 -> 0x14/0x15

    [ 0xff, 0xff ], // 108 'l': reserved (no 'L' in base32 alphabet)

    [ 0x16, 0x17 ], // 109 'm': aim0/aim1 -> third char 'M' = 11 = 01011 -> 0x16/0x17
    [ 0x18, 0x19 ], // 110 'n': ain0/ain1 -> third char 'N' = 12 = 01100 -> 0x18/0x19
    [ 0x1a, 0x1b ], // 111 'o': aio0/aio1 -> third char 'O' = 13 = 01101 -> 0x1a/0x1b
    [ 0x1c, 0x1d ], // 112 'p': aip0/aip1 -> third char 'P' = 14 = 01110 -> 0x1c/0x1d
    [ 0x1e, 0x1f ], // 113 'q': aiq0/aiq1 -> third char 'Q' = 15 = 01111 -> 0x1e/0x1f
    [ 0x20, 0x21 ], // 114 'r': air0/air1 -> third char 'R' = 16 = 10000 -> 0x20/0x21
    [ 0x22, 0x23 ], // 115 's': ais0/ais1 -> third char 'S' = 17 = 10001 -> 0x22/0x23
    [ 0x24, 0x25 ], // 116 't': ait0/ait1 -> third char 'T' = 18 = 10010 -> 0x24/0x25
    [ 0x26, 0x27 ], // 117 'u': aiu0/aiu1 -> third char 'U' = 19 = 10011 -> 0x26/0x27
    [ 0x28, 0x29 ], // 118 'v': aiv0/aiv1 -> third char 'V' = 20 = 10100 -> 0x28/0x29
    [ 0x2a, 0x2b ], // 119 'w': aiw0/aiw1 -> third char 'W' = 21 = 10101 -> 0x2a/0x2b
    [ 0x2c, 0x2d ], // 120 'x': aix0/aix1 -> third char 'X' = 22 = 10110 -> 0x2c/0x2d
    [ 0x2e, 0x2f ], // 121 'y': aiy0/aiy1 -> third char 'Y' = 23 = 10111 -> 0x2e/0x2f
    [ 0x30, 0x31 ], // 122 'z': aiz0/aiz1 -> third char 'Z' = 24 = 11000 -> 0x30/0x31
];

/*
 * New prefix format for "ai" branding:
 * Prefix = [0x02, res, 0x24] produces "aiXc..." where X depends on res
 *
 * aiK v0 hex:     0x021424 -> "aiKc..."
 * aiK v1 hex:     0x021524 -> "aiKs..."
 * aiA v0 hex:     0x020024 -> "aiAc..."
 * aiA v1 hex:     0x020124 -> "aiAs..."
 * aiS v0 hex:     0x022224 -> "aiSc..."
 * aiS v1 hex:     0x022324 -> "aiSs..."
 */

/// represents an encoding configuration for aiid rendering and parsing
pub struct AiidEncodingConfig {
    /// byte count of actuall key data that will be encoded
    pub key_byte_count: usize,
    /// parity bytes that will be encoded directly into the base32 string (appended after key)
    pub base_parity_byte_count: usize,
    /// parity bytes that will be encoded in the alpha capitalization (appended after base parity)
    pub cap_parity_byte_count: usize,
    /// bytes to prefix before rendering to base32
    pub prefix: Vec<u8>,
    /// binary indication of the capitalization for prefix characters
    pub prefix_cap: Vec<u8>,
    /// how many characters are in a capitalization parity segment
    pub cap_segment_char_count: usize,
    /// how many characters long the fully rendered base32 string should be
    pub encoded_char_count: usize,
}

impl AiidEncodingConfig {
    /// create a new config given a kind token string
    ///
    /// # Example
    ///
    /// ```
    /// let aia0 = aingle_id::AiidEncodingConfig::new("aia0").unwrap();
    /// let aik0 = aingle_id::AiidEncodingConfig::new("aik0").unwrap();
    /// let ais0 = aingle_id::AiidEncodingConfig::new("ais0").unwrap();
    /// ```
    pub fn new(kind: &str) -> AiidResult<Self> {
        let kind_b = kind.as_bytes();
        // Check for "aiXY" format where X is a-z or 3-9, Y is 0 or 1
        // 'a' = 97, 'i' = 105
        if kind_b.len() != 4 || kind_b[0] != 97 || kind_b[1] != 105 ||
                (kind_b[3] != 48 && kind_b[3] != 49) ||
                kind_b[2] < 51 || kind_b[2] > 122 {
            return Err(format!("invalid kind: `{}`", kind).into());
        }

        let version = if kind_b[3] == 48 { 0 } else { 1 };
        let res = AI_CODE_MAP[(kind_b[2] - 51) as usize][version as usize];

        if res == 0xff {
            return Err(format!("invalid kind: `{}`", kind).into());
        }

        Ok(AiidEncodingConfig {
            key_byte_count: 32,
            base_parity_byte_count: 4,
            cap_parity_byte_count: 4,
            prefix: vec![0x02, res, 0x24],
            prefix_cap: b"001".to_vec(),  // "ai" lowercase, third char uppercase
            cap_segment_char_count: 15,
            encoded_char_count: 63,
        })
    }
}

/// an instance that can encode / decode a particular aiid encoding configuration
pub struct AiidEncoding {
    config: AiidEncodingConfig,
    rs_enc: reed_solomon::Encoder,
    rs_dec: reed_solomon::Decoder,
}

impl AiidEncoding {
    /// create a new AiidEncoding instance from given AiidEncodingConfig
    pub fn new(config: AiidEncodingConfig) -> AiidResult<Self> {
        // set up a reed-solomon encoder with proper parity count
        let rs_enc = reed_solomon::Encoder::new(
            config.base_parity_byte_count + config.cap_parity_byte_count,
        );

        // set up a reed-solomon decoder with proper parity count
        let rs_dec = reed_solomon::Decoder::new(
            config.base_parity_byte_count + config.cap_parity_byte_count,
        );

        Ok(Self {
            config,
            rs_enc,
            rs_dec,
        })
    }

    /// create a new config given a kind token string
    ///
    /// # Example
    ///
    /// ```
    /// let aia0 = aingle_id::AiidEncoding::with_kind("aia0").unwrap();
    /// let aik0 = aingle_id::AiidEncoding::with_kind("aik0").unwrap();
    /// let ais0 = aingle_id::AiidEncoding::with_kind("ais0").unwrap();
    /// ```
    pub fn with_kind(kind: &str) -> AiidResult<Self> {
        AiidEncoding::new(AiidEncodingConfig::new(kind)?)
    }

    /// encode a string to base32 with this instance's configuration
    pub fn encode(&self, data: &[u8]) -> AiidResult<String> {
        if data.len() != self.config.key_byte_count {
            return Err(AiidError(String::from(format!(
                "BadDataLen:{},Expected:{}",
                data.len(),
                self.config.key_byte_count
            ))));
        }

        // generate reed-solomon parity bytes
        let full_parity = self.rs_enc.encode(data);

        // extract the bytes that will be encoded as capitalization
        let cap_bytes = &full_parity[full_parity.len() - self.config.cap_parity_byte_count..];

        // base is the bytes that will be base32 encoded
        let mut base = self.config.prefix.clone();
        base.extend_from_slice(
            &full_parity[0..full_parity.len() - self.config.cap_parity_byte_count],
        );

        // do the base32 encoding
        let mut base32 = b32::encode(&base);

        if base32.len() != self.config.encoded_char_count {
            return Err(AiidError(String::from(format!(
                "InternalGeneratedBadLen:{},Expected:{}",
                base32.len(),
                self.config.encoded_char_count
            ))));
        }

        // capitalize the prefix with a fixed scheme
        cap_encode_bin(
            &mut base32[0..self.config.prefix_cap.len()],
            &self.config.prefix_cap,
            3,
        )?;

        // iterate over segments, applying parity capitalization
        for i in 0..cap_bytes.len() {
            let seg_start = self.config.prefix_cap.len() + (i * self.config.cap_segment_char_count);
            let seg = &mut base32[seg_start..seg_start + self.config.cap_segment_char_count];
            let bin = format!("{:08b}", cap_bytes[i]).into_bytes();
            cap_encode_bin(seg, &bin, 8)?;
        }
        
        // we only use ascii characters
        // use unchecked for performance / so we don't allocate again
        unsafe {
            // return the result as a String for ease of use
            Ok(String::from_utf8_unchecked(base32))
        }
    }

    /// decode the data from a base32 string with this instance's configuration.  Reed-Solomon can
    /// correct up to 1/2 its parity size worth of erasures (if no other errors are present).
    pub fn decode(&self, data: &str) -> AiidResult<Vec<u8>> {
        // get our parsed data with erasures
        let (data, erasures) = self.pre_decode(data)?;

        if erasures.len() > ( self.config.base_parity_byte_count + self.config.cap_parity_byte_count ) / 2 {
            // our reed-solomon library makes bad corrections once erasure count exceeds 1/2 the
            // parity count (it takes 2 parity symbols to find/correct one error, 1 parity symbol to
            // correct a known erasure)
            return Err(AiidError(String::from("TooManyErrors")));
        }

        // optimise for the case where there are no transcription errors
        // this makes correcting more expensive if there *are*,
        // but on average makes the system more efficient
        if self.pre_is_corrupt(&data, &erasures)? {
            // apply reed-solomon correction
            // will "throw" on too many errors
            Ok(
                self.rs_dec.correct(&data, Some(&erasures[..]))?[0..self.config.key_byte_count]
                    .to_vec(),
            )
        } else {
            Ok(data[0..self.config.key_byte_count].to_vec())
        }
    }

    /// a lighter-weight check to determine if a base32 string is corrupt
    pub fn is_corrupt(&self, data: &str) -> AiidResult<bool> {
        // get our parsed data with erasures
        let (data, erasures) = match self.pre_decode(data) {
            Ok(v) => v,
            Err(_) => return Ok(true),
        };

        match self.pre_is_corrupt(&data, &erasures) {
            Ok(v) => Ok(v),
            Err(_) => Ok(true),
        }
    }

    /// internal helper for is_corrupt checking
    fn pre_is_corrupt(&self, data: &[u8], erasures: &[u8]) -> AiidResult<bool> {
        // if we have any erasures, we can exit early
        if erasures.len() > 0 {
            return Ok(true);
        }

        // slightly more efficient reed-solomon corruption check
        Ok(self.rs_dec.is_corrupted(&data))
    }

    /// internal helper for preparing decoding
    fn pre_decode(&self, data: &str) -> AiidResult<(Vec<u8>, Vec<u8>)> {
        if data.len() != self.config.encoded_char_count {
            return Err(AiidError(String::from(format!(
                "BadIdLen:{},Expected:{}",
                data.len(),
                self.config.encoded_char_count
            ))));
        }

        let key_base_byte_size = self.config.key_byte_count + self.config.base_parity_byte_count;
        // All char_erasures are indexed from the 0th char of the full codeword w/ prefix, but
        // byte_erasures are indexed from the 0th byte of the key+parity (ie. without the prefix).
        // Any byte of key, or base/cap parity could be erased.
        let mut byte_erasures = vec![b'0'; key_base_byte_size + self.config.cap_parity_byte_count];
        let mut char_erasures = vec![b'0'; data.len()];

        // correct any transliteration errors into our base32 alphabet
        // marking any unrecognizable characters as char-level erasures
        let mut data = b32_correct(data.as_bytes(), &mut char_erasures);

        // Pull out the parity data that was encoded as capitalization.  If its erasure,
        // determine the 
        let mut cap_bytes: Vec<u8> = Vec::new();
        let mut all_zro = true;
        let mut all_one = true;
        for i in 0..self.config.cap_parity_byte_count {
            // For cap. parity, indexing starts after pre-defined Base-32 prefix
            let char_idx = self.config.prefix_cap.len() + (i * self.config.cap_segment_char_count);
            match cap_decode(
                char_idx,
                &data[char_idx..char_idx + self.config.cap_segment_char_count],
                &char_erasures
            )? {
                None => {
                    byte_erasures[key_base_byte_size + i] = b'1';
                    cap_bytes.push(0)
                }
                Some(parity) => {
                    if all_zro && parity != 0x00_u8 {
                        all_zro = false
                    }
                    if all_one && parity != 0xFF_u8 {
                        all_one = false
                    }
                    cap_bytes.push(parity)
                }
            }
        }

        // If either all caps or all lower case (or erasure), assume the casing was lost (eg. QR
        // code, or dns segment); mark all cap-derived parity as erasures.  This allows validation
        // of codeword if all remaining parity is intact and key is correct; since no parity
        // capacity remains, no correction will be attempted.  There is only a low probability that
        // any remaining errors will be detected, in this case.  However, we're no *worse* off than
        // if we had no R-S parity at all.
        if all_zro || all_one {
            for i in 0..self.config.cap_parity_byte_count {
                byte_erasures[key_base_byte_size + i] = b'1';
            }
        }

        // we have the cap data, uppercase everything
        for c in data.iter_mut() {
            char_upper(c);
        }

        // do the base32 decode
        let mut data = b32::decode(&data)?;

        if &data[0..self.config.prefix.len()] != self.config.prefix.as_slice() {
            return Err(AiidError(String::from("PrefixMismatch")));
        }

        // remove the prefix bytes
        data.drain(0..self.config.prefix.len());

        // append our cap parity bytes
        data.append(&mut cap_bytes);

        // Sort through the char-level erasures (5 bits), and associate them with byte-level data (8
        // bits) -- in the (now prefix-free) data buffer, so that we mark the proper erasures for
        // reed-solomon correction.  Some of these chars span multiple bytes... we need to mark both.
        for i in self.config.prefix_cap.len()..char_erasures.len() {
            let c = char_erasures[i];
            if c == b'1' {
                // 1st and last bit of 5-bit segment may index different bytes
                byte_erasures[( i * 5 + 0 ) / 8 - self.config.prefix.len()] = b'1';
                byte_erasures[( i * 5 + 4 ) / 8 - self.config.prefix.len()] = b'1';
            }
        }

        // translate erasures into the form expected by our reed-solomon lib
        let mut erasures: Vec<u8> = Vec::new();
        for i in 0..byte_erasures.len() {
            if byte_erasures[i] == b'1' {
                data[i] = 0;
                erasures.push(i as u8);
            }
        }

        Ok((data, erasures))
    }
}

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

    static TEST_HEX_1: &'static str =
        "0c71db50d35d760b0ea2002ff20147c7c3a8e8030d35ef28ed1adaec9e329aba";
    static TEST_ID_1: &'static str =
        "aiKciDds5OiogymxbnHKEabQ8iavqs8dwdVaGdJW76Vp4gx47tQDfGW4OWc9w5i";

    #[test]
    fn it_encodes_1() {
        let enc = AiidEncoding::with_kind("aik0").unwrap();

        let input = hex::decode(TEST_HEX_1.as_bytes()).unwrap();
        let id = enc.encode(&input).unwrap();
        assert_eq!(TEST_ID_1, id);
    }

    #[test]
    fn it_decodes_1() {
        let enc = AiidEncoding::with_kind("aik0").unwrap();

        let data = hex::encode(&enc.decode(TEST_ID_1).unwrap());
        assert_eq!(TEST_HEX_1, data);
    }

}