candid 0.10.37

use crate::types::internal::{Field, Function, Label, Type, TypeInner};
use crate::types::{FuncMode, TypeEnv};
use anyhow::{anyhow, Context, Result};
use binrw::{BinRead, BinResult, Error as BError};
use std::convert::TryInto;

const MAX_TYPE_TABLE_LEN: u64 = 10_000; // Max type entries

/// Default bound on the byte length of the type-table header.
///
/// A header describes types, so its size follows the interface rather than the payload.
/// Every element it declares occupies at least one byte, so bounding its length bounds
/// each declared count at once. [`MAX_TYPE_TABLE_LEN`] bounds a different dimension:
/// the number of type-table entries, whatever their size.
pub(crate) const DEFAULT_MAX_HEADER_LEN: usize = 64 * 1024;

// Upper bound on a single allocation step while reading a length-prefixed byte
// blob. The buffer grows in steps of at most this size.
const READ_CHUNK: u64 = 8 * 1024;

/// Marks a header read that stopped because the reader ran out of bytes, rather than
/// one that found something malformed. Displays as the plain field name, so error text
/// is unchanged; the type is what makes the distinction recoverable.
#[derive(Debug)]
pub(crate) struct Truncated(pub(crate) &'static str);

impl std::fmt::Display for Truncated {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.0)
    }
}

fn leb_failure(e: leb128::read::Error, name: &'static str) -> Box<dyn binrw::error::CustomError> {
    match e {
        // Ran out of input: within a bounded prefix, that is the bound.
        leb128::read::Error::IoError(_) => Box::new(Truncated(name)),
        // A value too large for the target type is malformed, not truncated.
        leb128::read::Error::Overflow => Box::new(name),
    }
}

/// Whether a header parse stopped because it ran out of input, rather than because the
/// header was malformed. Inside a bounded prefix that is the signal that the header
/// needed more than the bound.
///
/// The failure offset cannot serve here: a length-prefixed field and a multi-byte
/// LEB128 are reported where they *start*, and binrw reports a failure inside an enum
/// variant at the start of the enum, all of which can be short of the boundary actually
/// crossed. Hence the marker, and the walk over the error tree.
pub(crate) fn is_truncation(e: &BError) -> bool {
    match e {
        // `Backtrace::error` is guaranteed not to be another backtrace.
        BError::Backtrace(bt) => is_truncation(&bt.error),
        BError::EnumErrors { variant_errors, .. } => {
            variant_errors.iter().any(|(_, inner)| is_truncation(inner))
        }
        BError::Io(io) => io.kind() == std::io::ErrorKind::UnexpectedEof,
        _ => e.custom_err::<Truncated>().is_some(),
    }
}

/// Read `len` bytes into a fresh `Vec`, growing the buffer in bounded steps.
///
/// A plain `#[br(count = len)]` on a `Vec<u8>` passes the wire-declared length
/// straight to `reserve_exact`, allocating `len` bytes up front before any are
/// read. A length prefix need not match the amount of data actually present, so
/// a value far larger than the input would reserve a correspondingly large
/// buffer instead of failing cleanly on the short read. Growing the buffer in
/// bounded chunks keeps the reservation proportional to the bytes available, so
/// an out-of-range or truncated length surfaces as an ordinary parse error.
fn read_len_prefixed<R: std::io::Read>(reader: &mut R, pos: u64, len: u64) -> BinResult<Vec<u8>> {
    use std::io::Read;
    let mut buf = Vec::new();
    let mut remaining = len;
    while remaining > 0 {
        let want = remaining.min(READ_CHUNK);
        // `read_to_end` on a bounded `Take` reserves at most `want`, so the
        // buffer never grows faster than the input is consumed.
        let read = reader.by_ref().take(want).read_to_end(&mut buf)? as u64;
        if read < want {
            // Fewer bytes were available than declared: short input.
            return Err(BError::Custom {
                pos,
                err: Box::new(Truncated("not enough bytes")),
            });
        }
        remaining -= read;
    }
    Ok(buf)
}

#[binrw::parser(reader)]
fn read_bytes(len: u64) -> BinResult<Vec<u8>> {
    let pos = reader.stream_position()?;
    read_len_prefixed(reader, pos, len)
}
#[binrw::parser(reader)]
fn read_text(len: u64) -> BinResult<String> {
    let pos = reader.stream_position()?;
    let bytes = read_len_prefixed(reader, pos, len)?;
    String::from_utf8(bytes).map_err(|_| BError::Custom {
        pos,
        err: Box::new("invalid utf8"),
    })
}

#[binrw::parser(reader)]
fn read_leb(name: &'static str) -> BinResult<u64> {
    let pos = reader.stream_position()?;
    leb128::read::unsigned(reader).map_err(|e| BError::Custom {
        pos,
        err: leb_failure(e, name),
    })
}
#[binrw::parser(reader)]
fn read_sleb(name: &'static str) -> BinResult<i64> {
    let pos = reader.stream_position()?;
    leb128::read::signed(reader).map_err(|e| BError::Custom {
        pos,
        err: leb_failure(e, name),
    })
}
#[binrw::parser(reader)]
fn read_leb_u32(name: &'static str, range_msg: &'static str) -> BinResult<u32> {
    let pos = reader.stream_position()?;
    let v = leb128::read::unsigned(reader).map_err(|e| BError::Custom {
        pos,
        err: leb_failure(e, name),
    })?;
    v.try_into().map_err(|_| BError::Custom {
        pos,
        err: Box::new(range_msg),
    })
}
#[binrw::parser(reader)]
fn read_leb_usize(name: &'static str, range_msg: &'static str) -> BinResult<usize> {
    let pos = reader.stream_position()?;
    let v = leb128::read::unsigned(reader).map_err(|e| BError::Custom {
        pos,
        err: leb_failure(e, name),
    })?;
    v.try_into().map_err(|_| BError::Custom {
        pos,
        err: Box::new(range_msg),
    })
}

#[derive(BinRead, Debug)]
#[br(magic = b"DIDL")]
#[br(import(max_type_len: Option<usize>))]
pub struct Header {
    #[br(args(max_type_len))]
    table: Table,
    #[br(parse_with = read_leb, args("len"))]
    len: u64,
    #[br(count = len)]
    args: Vec<IndexType>,
}

#[derive(BinRead, Debug)]
#[br(import(max_type_len: Option<usize>))]
struct Table {
    #[br(parse_with = read_leb, args("len"))]
    #[br(assert(len <= max_type_len.unwrap_or(MAX_TYPE_TABLE_LEN as usize) as u64, "type table size exceeded"))]
    len: u64,
    #[br(count = len)]
    table: Vec<ConsType>,
}
#[derive(BinRead, Debug)]
enum ConsType {
    #[br(magic = 0x6eu8)]
    Opt(Box<IndexType>),
    #[br(magic = 0x6du8)]
    Vec(Box<IndexType>),
    #[br(magic = 0x6cu8)]
    Record(Fields),
    #[br(magic = 0x6bu8)]
    Variant(Fields),
    #[br(magic = 0x6au8)]
    Func(FuncType),
    #[br(magic = 0x69u8)]
    Service(ServType),
    Future(FutureType),
}
#[derive(BinRead, Debug)]
struct IndexType {
    #[br(parse_with = read_sleb, args("index"), assert(index >= -17 || index == -24, "unknown opcode {}", index))]
    index: i64,
}
#[derive(BinRead, Debug)]
struct Fields {
    #[br(parse_with = read_leb_u32, args("len", "field length out of 32-bit range"))]
    len: u32,
    #[br(count = len)]
    inner: Vec<FieldType>,
}
#[derive(BinRead, Debug)]
struct FieldType {
    #[br(parse_with = read_leb_u32, args("id", "field id out of 32-bit range"))]
    id: u32,
    index: IndexType,
}
#[derive(BinRead, Debug)]
struct FuncType {
    #[br(parse_with = read_leb, args("arg_len"))]
    arg_len: u64,
    #[br(count = arg_len)]
    args: Vec<IndexType>,
    #[br(parse_with = read_leb, args("ret_len"))]
    ret_len: u64,
    #[br(count = ret_len)]
    rets: Vec<IndexType>,
    #[br(assert(ann_len <= 1u8, "function annotation length should be at most 1"))]
    ann_len: u8,
    #[br(count = ann_len)]
    ann: Vec<Mode>,
}
#[derive(BinRead, Debug)]
struct ServType {
    #[br(parse_with = read_leb, args("len"))]
    len: u64,
    #[br(count = len)]
    meths: Vec<Meths>,
}
#[derive(BinRead, Debug)]
struct FutureType {
    #[br(parse_with = read_sleb, args("opcode"), assert(opcode < -24, "{} is not a valid future type", opcode))]
    opcode: i64,
    #[br(parse_with = read_leb, args("len"))]
    len: u64,
    #[br(parse_with = read_bytes, args(len))]
    blob: Vec<u8>,
}
#[derive(BinRead, Debug)]
struct Meths {
    #[br(parse_with = read_leb, args("len"))]
    len: u64,
    #[br(parse_with = read_text, args(len))]
    name: String,
    ty: IndexType,
}
#[derive(BinRead, Debug)]
struct Mode {
    #[br(try_map = |x:u8| match x { 1u8 => Ok(FuncMode::Query), | 2u8 => Ok(FuncMode::Oneway), | 3u8 => Ok(FuncMode::CompositeQuery), | _ => Err("Unknown annotation") })]
    inner: FuncMode,
}

#[derive(BinRead)]
pub struct BoolValue(
    #[br(try_map = |x:u8| match x { 0u8 => Ok(false), | 1u8 => Ok(true), | _ => Err("Expect 00 or 01") } )]
    pub bool,
);
#[derive(BinRead)]
pub struct Len(
    #[br(parse_with = read_leb_usize, args("len", "length out of usize range"))] pub usize,
);
#[derive(BinRead)]
pub struct PrincipalBytes {
    #[br(assert(flag == 1u8, "Opaque reference not supported"))]
    pub flag: u8,
    #[br(parse_with = read_leb, args("len"), assert(len <= 29, "Principal is longer than 29 bytes"))]
    pub len: u64,
    #[br(count = len)]
    pub inner: Vec<u8>,
}

fn index_to_var(ind: i64) -> String {
    format!("table{ind}")
}
impl IndexType {
    fn to_type(&self, len: u64) -> Result<Type> {
        Ok(match self.index {
            v if v >= 0 => {
                if v >= len as i64 {
                    return Err(anyhow!("type index {} out of range", v));
                }
                TypeInner::Var(index_to_var(v))
            }
            -1 => TypeInner::Null,
            -2 => TypeInner::Bool,
            -3 => TypeInner::Nat,
            -4 => TypeInner::Int,
            -5 => TypeInner::Nat8,
            -6 => TypeInner::Nat16,
            -7 => TypeInner::Nat32,
            -8 => TypeInner::Nat64,
            -9 => TypeInner::Int8,
            -10 => TypeInner::Int16,
            -11 => TypeInner::Int32,
            -12 => TypeInner::Int64,
            -13 => TypeInner::Float32,
            -14 => TypeInner::Float64,
            -15 => TypeInner::Text,
            -16 => TypeInner::Reserved,
            -17 => TypeInner::Empty,
            -24 => TypeInner::Principal,
            _ => unreachable!(),
        }
        .into())
    }
}
impl ConsType {
    fn to_type(&self, len: u64) -> Result<Type> {
        Ok(match &self {
            ConsType::Opt(ref ind) => TypeInner::Opt(ind.to_type(len)?),
            ConsType::Vec(ref ind) => TypeInner::Vec(ind.to_type(len)?),
            ConsType::Record(fs) | ConsType::Variant(fs) => {
                let mut res = Vec::new();
                let mut prev = None;
                for f in &fs.inner {
                    if let Some(prev) = prev {
                        if prev >= f.id {
                            return Err(anyhow!("field id {} collision or not sorted", f.id));
                        }
                    }
                    prev = Some(f.id);
                    let field = Field {
                        id: Label::Id(f.id).into(),
                        ty: f.index.to_type(len)?,
                    };
                    res.push(field);
                }
                if matches!(&self, ConsType::Record(_)) {
                    TypeInner::Record(res)
                } else {
                    TypeInner::Variant(res)
                }
            }
            ConsType::Func(f) => {
                let mut args = Vec::new();
                let mut rets = Vec::new();
                for arg in &f.args {
                    args.push(arg.to_type(len)?);
                }
                for ret in &f.rets {
                    rets.push(ret.to_type(len)?);
                }
                TypeInner::Func(Function {
                    modes: f.ann.iter().map(|x| x.inner.clone()).collect(),
                    args,
                    rets,
                })
            }
            ConsType::Service(serv) => {
                let mut res = Vec::new();
                let mut prev = None;
                for m in &serv.meths {
                    if let Some(prev) = prev {
                        if prev >= &m.name {
                            return Err(anyhow!("method name {} duplicate or not sorted", m.name));
                        }
                    }
                    prev = Some(&m.name);
                    res.push((m.name.clone(), m.ty.to_type(len)?));
                }
                TypeInner::Service(res)
            }
            ConsType::Future(_) => TypeInner::Future,
        }
        .into())
    }
}
impl Table {
    fn to_env(&self, len: u64) -> Result<TypeEnv> {
        use std::collections::BTreeMap;
        let mut env = BTreeMap::new();
        for (i, t) in self.table.iter().enumerate() {
            let ty = t
                .to_type(len)
                .with_context(|| format!("Invalid table entry {i}: {t:?}"))?;
            env.insert(index_to_var(i as i64), ty);
        }
        // validate method has func type
        for t in env.values() {
            if let TypeInner::Service(ms) = t.as_ref() {
                for (name, ty) in ms {
                    if let TypeInner::Var(id) = ty.as_ref() {
                        if matches!(env.get(id).map(|t| t.as_ref()), Some(TypeInner::Func(_))) {
                            continue;
                        }
                    }
                    return Err(anyhow!("Method {name} has a non-function type {ty}"));
                }
            }
        }
        Ok(TypeEnv(env))
    }
}
impl Header {
    pub fn to_types(&self) -> Result<(TypeEnv, Vec<Type>)> {
        let len = self.table.len;
        let mut env = self.table.to_env(len)?;
        env.replace_empty()?;
        let mut args = Vec::new();
        for (i, t) in self.args.iter().enumerate() {
            args.push(
                t.to_type(len)
                    .with_context(|| format!("Invalid argument entry {i}: {t:?}"))?,
            );
        }
        Ok((env, args))
    }
}