use crate::id::stable_id_bytes;
use super::estimator::ImageTokenEstimator;
use super::model::MediaRef;
#[derive(Debug, Clone, Copy)]
pub(crate) struct MediaAnalysis {
pub raw_hash: u64,
pub image_tokens: u64,
}
pub(crate) fn analyze(media: &MediaRef) -> MediaAnalysis {
let bytes = decode_base64(&media.bytes_b64).unwrap_or_default();
MediaAnalysis {
raw_hash: stable_id_bytes(&bytes),
image_tokens: ImageTokenEstimator::estimate(&media.mime, &bytes, &media.bytes_b64),
}
}
pub(crate) fn is_valid_base64(bytes_b64: &str) -> bool {
decode_base64(bytes_b64).is_ok()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct DecodeError;
pub(crate) fn decode_base64(input: &str) -> Result<Vec<u8>, DecodeError> {
let bytes = input.as_bytes();
if bytes.is_empty() {
return Ok(Vec::new());
}
if !bytes.len().is_multiple_of(4) {
return Err(DecodeError);
}
let quad_count = bytes.len() / 4;
let mut out = Vec::with_capacity(quad_count * 3);
for (index, quad) in bytes.chunks_exact(4).enumerate() {
let pad = decode_quad(quad, &mut out)?;
if pad > 0 && index + 1 != quad_count {
return Err(DecodeError);
}
}
Ok(out)
}
fn decode_quad(quad: &[u8], out: &mut Vec<u8>) -> Result<usize, DecodeError> {
let pad = quad.iter().rev().take_while(|&&b| b == b'=').count();
if pad > 2 {
return Err(DecodeError);
}
let mut sextets = [0_u8; 4];
for (index, &byte) in quad.iter().enumerate() {
sextets[index] = if byte == b'=' {
if index < 4 - pad {
return Err(DecodeError);
}
0
} else {
sextet(byte).ok_or(DecodeError)?
};
}
let word = (u32::from(sextets[0]) << 18)
| (u32::from(sextets[1]) << 12)
| (u32::from(sextets[2]) << 6)
| u32::from(sextets[3]);
out.extend_from_slice(&word.to_be_bytes()[1..4 - pad]);
Ok(pad)
}
fn sextet(byte: u8) -> Option<u8> {
match byte {
b'A'..=b'Z' => Some(byte - b'A'),
b'a'..=b'z' => Some(byte - b'a' + 26),
b'0'..=b'9' => Some(byte - b'0' + 52),
b'+' => Some(62),
b'/' => Some(63),
_ => None,
}
}
#[cfg(test)]
#[path = "media_tests.rs"]
mod tests;