use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use brotli::{CompressorWriter, Decompressor};
use std::fmt;
use std::io::{Read, Write};
const COMPRESSED_PREFIX: &str = "$br$";
const QUALITY: u32 = 9;
const WINDOW_BITS: u32 = 16;
const BUFFER_SIZE: usize = 4096;
const MAX_SCOPE_LEN: u64 = 64 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScopeDecodeError {
Base64DecodeFailed,
DecompressFailed,
TooLarge,
}
impl fmt::Display for ScopeDecodeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Base64DecodeFailed => write!(f, "Base64 decode of compressed scope failed"),
Self::DecompressFailed => write!(f, "Brotli decompression of scope failed"),
Self::TooLarge => write!(f, "Decompressed scope exceeds maximum length"),
}
}
}
impl std::error::Error for ScopeDecodeError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScopeEncodeError {
TooLarge,
}
impl fmt::Display for ScopeEncodeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::TooLarge => write!(f, "Scope exceeds maximum length"),
}
}
}
impl std::error::Error for ScopeEncodeError {}
fn brotli_compress(input: &str) -> Vec<u8> {
let mut writer = CompressorWriter::new(Vec::new(), BUFFER_SIZE, QUALITY, WINDOW_BITS);
writer
.write_all(input.as_bytes())
.expect("writing to a Vec cannot fail");
writer.into_inner()
}
fn brotli_decompress(input: &[u8]) -> Result<String, ScopeDecodeError> {
let mut output = String::new();
Decompressor::new(input, BUFFER_SIZE)
.take(MAX_SCOPE_LEN + 1)
.read_to_string(&mut output)
.map_err(|_| ScopeDecodeError::DecompressFailed)?;
if output.len() as u64 > MAX_SCOPE_LEN {
return Err(ScopeDecodeError::TooLarge);
}
Ok(output)
}
pub fn encode_scope(scope: &str) -> Result<String, ScopeEncodeError> {
if scope.len() as u64 > MAX_SCOPE_LEN {
return Err(ScopeEncodeError::TooLarge);
}
let tagged = format!(
"{COMPRESSED_PREFIX}{}",
URL_SAFE_NO_PAD.encode(brotli_compress(scope))
);
if tagged.len() < scope.len() || scope.starts_with(COMPRESSED_PREFIX) {
Ok(tagged)
} else {
Ok(scope.to_owned())
}
}
pub fn decode_scope(scope: &str) -> Result<String, ScopeDecodeError> {
let Some(encoded) = scope.strip_prefix(COMPRESSED_PREFIX) else {
return Ok(scope.to_owned());
};
let compressed = URL_SAFE_NO_PAD
.decode(encoded)
.map_err(|_| ScopeDecodeError::Base64DecodeFailed)?;
brotli_decompress(&compressed)
}
#[cfg(test)]
mod tests {
use super::*;
fn long_scope() -> String {
let mut scope = String::from("transition:generic transition:chat.bsky");
for collection in [
"social.colibri.message",
"social.colibri.community",
"social.colibri.reaction",
"social.colibri.member",
"social.colibri.channel.read",
] {
scope.push_str(&format!(" repo:{collection}?action=create&action=delete"));
}
scope
}
#[test]
fn long_scope_roundtrips_through_compression() {
let scope = long_scope();
let encoded = encode_scope(&scope).unwrap();
assert!(encoded.starts_with(COMPRESSED_PREFIX));
assert!(encoded.len() < scope.len());
assert_eq!(decode_scope(&encoded).unwrap(), scope);
}
#[test]
fn short_scope_stays_plaintext() {
let encoded = encode_scope("com.atproto.access").unwrap();
assert_eq!(encoded, "com.atproto.access");
assert_eq!(decode_scope(&encoded).unwrap(), "com.atproto.access");
}
#[test]
fn untagged_scope_passes_through() {
assert_eq!(
decode_scope("com.atproto.refresh").unwrap(),
"com.atproto.refresh"
);
assert_eq!(decode_scope("").unwrap(), "");
}
#[test]
fn malformed_compressed_scope_errors_instead_of_panicking() {
assert_eq!(
decode_scope("$br$not valid base64!"),
Err(ScopeDecodeError::Base64DecodeFailed)
);
assert_eq!(
decode_scope("$br$AAAAAAAAAAAAAAAA"),
Err(ScopeDecodeError::DecompressFailed)
);
}
#[test]
fn compression_bomb_is_rejected() {
let bomb = URL_SAFE_NO_PAD.encode(brotli_compress(&"a".repeat(MAX_SCOPE_LEN as usize * 2)));
assert_eq!(
decode_scope(&format!("{COMPRESSED_PREFIX}{bomb}")),
Err(ScopeDecodeError::TooLarge)
);
}
#[test]
fn plaintext_that_looks_compressed_roundtrips() {
let scope = "$br$repo:*";
let encoded = encode_scope(scope).unwrap();
assert!(encoded.starts_with(COMPRESSED_PREFIX));
assert_eq!(decode_scope(&encoded).unwrap(), scope);
}
#[test]
fn encode_rejects_oversized_scope() {
let oversized = "a".repeat(MAX_SCOPE_LEN as usize + 1);
assert_eq!(encode_scope(&oversized), Err(ScopeEncodeError::TooLarge));
assert!(encode_scope(&"a".repeat(MAX_SCOPE_LEN as usize)).is_ok());
}
}