aria2-rust-pro/crates/aria2-rust-pro-protocol/src/torrent/dht/codec.rs

use std::collections::BTreeMap;

use super::DhtBencodeValue;

/// Encodes one DHT bencode value into a byte buffer.
fn dht_encode_value(value: &DhtBencodeValue, out: &mut Vec<u8>) {
    match value {
        DhtBencodeValue::Int(number) => {
            out.push(b'i');
            out.extend_from_slice(number.to_string().as_bytes());
            out.push(b'e');
        }
        DhtBencodeValue::Bytes(bytes) => {
            out.extend_from_slice(bytes.len().to_string().as_bytes());
            out.push(b':');
            out.extend_from_slice(bytes);
        }
        DhtBencodeValue::List(values) => {
            out.push(b'l');
            for item in values {
                dht_encode_value(item, out);
            }
            out.push(b'e');
        }
        DhtBencodeValue::Dict(values) => {
            out.push(b'd');
            for (key, value) in values {
                out.extend_from_slice(key.len().to_string().as_bytes());
                out.push(b':');
                out.extend_from_slice(key);
                dht_encode_value(value, out);
            }
            out.push(b'e');
        }
    }
}

/// Encodes a DHT dictionary into canonical bencode bytes.
pub(super) fn dht_encode_dict(dict: &BTreeMap<Vec<u8>, DhtBencodeValue>) -> Vec<u8> {
    let mut out = Vec::new();
    dht_encode_value(&DhtBencodeValue::Dict(dict.clone()), &mut out);
    out
}

/// Parses one DHT bencode value and returns the decoded value plus next index.
pub(super) fn dht_parse_value(
    input: &[u8],
    index: usize,
) -> Result<(DhtBencodeValue, usize), String> {
    match input.get(index).copied() {
        Some(b'i') => dht_parse_int(input, index + 1),
        Some(b'l') => dht_parse_list(input, index + 1),
        Some(b'd') => dht_parse_dict(input, index + 1),
        Some(byte) if byte.is_ascii_digit() => dht_parse_bytes(input, index),
        Some(_) => Err("invalid dht bencode value".to_owned()),
        None => Err("unexpected end of dht bencode input".to_owned()),
    }
}

/// Parses one DHT bencode integer starting at `index`.
fn dht_parse_int(input: &[u8], index: usize) -> Result<(DhtBencodeValue, usize), String> {
    let mut cursor = index;
    while cursor < input.len() && input[cursor] != b'e' {
        cursor += 1;
    }
    if cursor >= input.len() {
        return Err("unterminated dht integer".to_owned());
    }
    let text = std::str::from_utf8(&input[index..cursor])
        .map_err(|_| "invalid dht integer bytes".to_owned())?;
    let value = text
        .parse::<i64>()
        .map_err(|_| "invalid dht integer value".to_owned())?;
    Ok((DhtBencodeValue::Int(value), cursor + 1))
}

/// Parses one DHT bencode list starting at `index`.
fn dht_parse_list(input: &[u8], index: usize) -> Result<(DhtBencodeValue, usize), String> {
    let mut values = Vec::new();
    let mut cursor = index;
    while cursor < input.len() {
        if input[cursor] == b'e' {
            return Ok((DhtBencodeValue::List(values), cursor + 1));
        }
        let (value, next) = dht_parse_value(input, cursor)?;
        values.push(value);
        cursor = next;
    }
    Err("unterminated dht list".to_owned())
}

/// Parses one DHT bencode dictionary starting at `index`.
fn dht_parse_dict(input: &[u8], index: usize) -> Result<(DhtBencodeValue, usize), String> {
    let mut map = BTreeMap::new();
    let mut cursor = index;
    while cursor < input.len() {
        if input[cursor] == b'e' {
            return Ok((DhtBencodeValue::Dict(map), cursor + 1));
        }
        let (key, key_end) = dht_parse_bytes_raw(input, cursor)?;
        let (value, value_end) = dht_parse_value(input, key_end)?;
        map.insert(key, value);
        cursor = value_end;
    }
    Err("unterminated dht dictionary".to_owned())
}

/// Parses one DHT bencode byte string starting at `index`.
fn dht_parse_bytes(input: &[u8], index: usize) -> Result<(DhtBencodeValue, usize), String> {
    let (bytes, next) = dht_parse_bytes_raw(input, index)?;
    Ok((DhtBencodeValue::Bytes(bytes), next))
}

/// Parses one raw DHT bencode byte string and returns its bytes plus next index.
fn dht_parse_bytes_raw(input: &[u8], index: usize) -> Result<(Vec<u8>, usize), String> {
    let mut cursor = index;
    while cursor < input.len() && input[cursor].is_ascii_digit() {
        cursor += 1;
    }
    if cursor == index || cursor >= input.len() || input[cursor] != b':' {
        return Err("invalid dht byte string".to_owned());
    }
    let length = std::str::from_utf8(&input[index..cursor])
        .map_err(|_| "invalid dht byte string length".to_owned())?
        .parse::<usize>()
        .map_err(|_| "invalid dht byte string length".to_owned())?;
    let start = cursor + 1;
    let end = start.saturating_add(length);
    if end > input.len() {
        return Err("truncated dht byte string".to_owned());
    }
    Ok((input[start..end].to_vec(), end))
}

/// Looks up a raw byte-string field inside a DHT dictionary.
pub(super) fn dht_dict_get_bytes(
    dict: &BTreeMap<Vec<u8>, DhtBencodeValue>,
    key: &[u8],
) -> Option<Vec<u8>> {
    match dict.get(key) {
        Some(DhtBencodeValue::Bytes(bytes)) => Some(bytes.clone()),
        _ => None,
    }
}

/// Looks up an integer field inside a DHT dictionary.
pub(super) fn dht_dict_get_int(
    dict: &BTreeMap<Vec<u8>, DhtBencodeValue>,
    key: &[u8],
) -> Option<i64> {
    match dict.get(key) {
        Some(DhtBencodeValue::Int(value)) => Some(*value),
        _ => None,
    }
}

/// Looks up a nested dictionary field inside a DHT dictionary.
pub(super) fn dht_dict_get_dict<'a>(
    dict: &'a BTreeMap<Vec<u8>, DhtBencodeValue>,
    key: &[u8],
) -> Option<&'a BTreeMap<Vec<u8>, DhtBencodeValue>> {
    match dict.get(key) {
        Some(DhtBencodeValue::Dict(value)) => Some(value),
        _ => None,
    }
}

/// Looks up a list field inside a DHT dictionary.
pub(super) fn dht_dict_get_list(
    dict: &BTreeMap<Vec<u8>, DhtBencodeValue>,
    key: &[u8],
) -> Option<Vec<DhtBencodeValue>> {
    match dict.get(key) {
        Some(DhtBencodeValue::List(values)) => Some(values.clone()),
        _ => None,
    }
}