tnj_pcc/encoding/

read.rs

//! Module containing deserialization for proofs.

use crate::encoding::iter::UntrustedSizeIterator;
use crate::Proof;
use arch::reg::Register;
use smol_str::SmolStr;
use std::collections::HashMap;
use std::io;
use std::io::Read;
use sym::{BinOpTy, BoolOpTy, CmpTy, Expr, Term};
use thiserror::Error;
use types::Type;

/// Read a proof from a reader.
pub fn read<R>(r: &mut R) -> Result<Proof, Error>
where
    R: ?Sized + Read,
{
    let mut proof = Proof::default();

    // first we read the strings
    let n_strings = leb128::read::unsigned(r)?;
    // we do not reserve space for the strings, as this would allow an attacker to easily crash our
    // program, making it run out of memory.
    let strings = UntrustedSizeIterator::new(n_strings)
        .map(|i| Ok((i, read_str(r)?)))
        .collect::<Result<_, Error>>()?;

    let assertions_len = leb128::read::unsigned(r)?;

    proof.assertions = UntrustedSizeIterator::new(assertions_len)
        .map(|_| {
            let addr = leb128::read::unsigned(r)?;
            let num_exprs = leb128::read::unsigned(r)?;
            let expressions = UntrustedSizeIterator::new(num_exprs)
                .map(|_| read_expr(r, &strings))
                .collect::<Result<Vec<_>, Error>>()?;
            Ok((addr, expressions))
        })
        .collect::<Result<HashMap<_, _>, Error>>()?;

    Ok(proof)
}

fn read_expr<R>(
    r: &mut R,
    strings: &HashMap<u64, SmolStr>,
) -> Result<Expr, crate::encoding::read::Error>
where
    R: ?Sized + Read,
{
    let ty = read_ty(r)?;
    let tag = leb128::read::unsigned(r)?;

    let expr = match tag {
        0 => {
            let id = leb128::read::unsigned(r)?
                .try_into()
                .map_err(|_| Error::Overflow("expected u32 unknown id"))?;
            Expr::Unknown(id, ty)
        }
        1 => {
            let val = leb128::read::unsigned(r)?;
            Expr::Const { val, ty }
        }
        2 => Expr::Load {
            addr: Box::new(read_expr(r, strings)?),
            ty,
        },
        3 => {
            let op = match leb128::read::unsigned(r)? {
                0 => BinOpTy::WrappingAdd,
                1 => BinOpTy::WrappingSub,
                n => return Err(Error::InvalidType(n, "BinOpTy")),
            };
            Expr::BinOp {
                lhs: Box::new(read_expr(r, strings)?),
                rhs: Box::new(read_expr(r, strings)?),
                op,
                ty,
            }
        }
        4 => Expr::Term(read_term(r, strings)?),
        5 => {
            let cmp_ty = match leb128::read::unsigned(r)? {
                0 => CmpTy::Equal,
                1 => CmpTy::NotEqual,
                2 => CmpTy::Less,
                3 => CmpTy::LessEqual,
                4 => CmpTy::Greater,
                5 => CmpTy::GreaterEqual,
                t => return Err(Error::InvalidType(t, "CmpTy")),
            };
            Expr::Cmp {
                lhs: Box::new(read_expr(r, strings)?),
                rhs: Box::new(read_expr(r, strings)?),
                ty: cmp_ty,
            }
        }
        6 => {
            let op = match leb128::read::unsigned(r)? {
                0 => BoolOpTy::And,
                1 => BoolOpTy::Or,
                2 => BoolOpTy::Imp,
                3 => BoolOpTy::BiImp,
                t => return Err(Error::InvalidType(t, "BoolOpTy")),
            };
            Expr::BoolOp {
                lhs: Box::new(read_expr(r, strings)?),
                rhs: Box::new(read_expr(r, strings)?),
                ty: op,
            }
        }
        7 => Expr::Not {
            inner: Box::new(read_expr(r, strings)?),
        },
        8 => {
            let string_id = leb128::read::unsigned(r)?;
            let name = strings
                .get(&string_id)
                .ok_or(Error::InvalidStringId(string_id))?
                .clone();
            Expr::Reg(Register { name, ty })
        }
        t => return Err(Error::InvalidTag(t)),
    };

    Ok(expr)
}

fn read_term<R>(
    r: &mut R,
    strings: &HashMap<u64, SmolStr>,
) -> Result<Term, crate::encoding::read::Error>
where
    R: ?Sized + Read,
{
    let string_id = leb128::read::unsigned(r)?;
    let name = strings
        .get(&string_id)
        .ok_or(Error::InvalidStringId(string_id))?
        .clone();
    let ty = read_ty(r)?;
    let len = leb128::read::unsigned(r)?;
    let args = UntrustedSizeIterator::new(len)
        .map(|_| read_expr(r, strings))
        .collect::<Result<Vec<_>, Error>>()?;

    Ok(Term { name, args, ty })
}

fn read_ty<R>(r: &mut R) -> Result<Type, Error>
where
    R: ?Sized + Read,
{
    let ty = leb128::read::unsigned(r)?;
    match ty {
        n if n <= 64 && n.is_power_of_two() => Ok(Type::int(n as u8)),
        n => Err(Error::InvalidType(n, "type")),
    }
}

fn read_str<R>(r: &mut R) -> Result<SmolStr, Error>
where
    R: ?Sized + Read,
{
    let mut len = leb128::read::unsigned(r)? as usize;

    // we read the string in chunks of max 64 bytes
    const CHUNK_SIZE: usize = 64;

    // Since we don't trust the length, we cannot pre-allocate using the full length.
    // We will pre-allocate using the chunk size instead
    let mut bytes = Vec::with_capacity(usize::min(len, CHUNK_SIZE));

    let mut buf = [0u8; CHUNK_SIZE];
    while len > 0 {
        let bytes_to_read = usize::min(len, CHUNK_SIZE);

        let buf_slice = &mut buf[0..bytes_to_read];
        r.read_exact(buf_slice)?;
        bytes.extend_from_slice(buf_slice);

        len -= bytes_to_read;
    }

    let s = std::str::from_utf8(&bytes)
        .map_err(|_| Error::Utf8Error(String::from_utf8_lossy(&bytes).to_string()))?;

    Ok(SmolStr::from(s))
}

/// Error type for reading proofs.
#[derive(Debug, Error)]
pub enum Error {
    /// An unexpected end of input was encountered.
    #[error("unexpected end of input")]
    UnexpectedEof,
    /// Wrapper around an leb128 error.
    #[error("leb128 error: {0}")]
    Leb128(#[from] leb128::read::Error),
    /// Wrapper around an io error.
    #[error("io error: {0}")]
    Io(#[from] io::Error),
    /// Invalid type
    #[error("invalid {1}: {0}")]
    InvalidType(u64, &'static str),
    /// Invalid expression tag
    #[error("invalid expression tag: {0}")]
    InvalidTag(u64),
    /// Overflow, value does not fit into type
    #[error("overflow: {0}")]
    Overflow(&'static str),
    /// Invalid Utf-8
    #[error("Invalid utf-8: {0}")]
    Utf8Error(String),
    /// Invalid string id
    #[error("Invalid string id {0}, does not exist in strings section")]
    InvalidStringId(u64),
}

#[cfg(test)]
mod tests {
    use crate::read::read;
    use arch::reg::Register;
    use sym::Expr;
    use types::I64;

    #[test]
    fn test_read() {
        let bytes = [
            // strings
            0x0, // no strings
            // assertions
            0x2, // two entries
            // first entry
            0x0, // at offset 0
            0x1, // one assertion
            // assertion
            0x40, // expression of type I64 (0x40)
            0x1,  // constant
            0x2a, // value 42
            // second entry
            0x1, // at offset 1
            0x1, // one assertion
            // assertion
            0x40, // expression of type I64 (0x40)
            0x1,  // constant
            0x2b, // value 43
        ];
        let mut cursor = std::io::Cursor::new(&bytes);
        let proof = read(&mut cursor).unwrap();
        assert_eq!(proof.assertions.len(), 2);
        assert_eq!(
            proof.assertions[&0],
            vec![Expr::Const {
                val: 42,
                ty: types::I64
            }]
        );
        assert_eq!(
            proof.assertions[&1],
            vec![Expr::Const {
                val: 43,
                ty: types::I64
            }]
        );
    }

    #[test]
    fn test_string() {
        let bytes = [
            // strings:
            0x1, // one string
            // string with id 0
            0xc, // string with length 12
            // h  e     l     l     o     ,
            0x68, 0x65, 0x6c, 0x6c, 0x6f, 0x2c, //
            // w  o     r     l     d     !
            0x20, 0x77, 0x6f, 0x72, 0x6c, 0x64, //
            // assertions
            0x1, // one entry
            0xc, // at offset 12
            // entry:
            0x1, // a single assertion
            // assertion:
            0x40, // expression of type I64 (0x40)
            0x8,  // expression of type reg
            0x0,  // register name: string with id 0
        ];
        let mut cursor = std::io::Cursor::new(&bytes);
        let proof = read(&mut cursor).unwrap();
        assert_eq!(proof.assertions.len(), 1);
        assert_eq!(
            proof.assertions[&12],
            vec![Expr::Reg(Register::new("hello, world", I64))]
        );
    }
}