<commit_before>use runtime::Error;
use values::Value;
use memory::Memory;

use std::rc::Rc;

pub static PRIMITIVES: [&'static str, ..18] = [
    "*",
    "+",
    "-",
    "/",
    "<",
    "<=",
    "=",
    ">",
    ">=",
    "car",
    "cdr",
    "cons",
    "length",
    "list",
    "list?",
    "not",
    "null?",
    "pair?",
];

pub fn apply(name: &String, args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    match name.as_slice() {
        "*"      => product(args, mem),
        "+"      => plus(args, mem),
        "-"      => minus(args, mem),
        "/"      => division(args, mem),
        "<"      => less_than(args, mem),
        "<="     => less_than_or_equal(args, mem),
        "="      => equals(args, mem),
        ">"      => greater_than(args, mem),
        ">="     => greater_than_or_equal(args, mem),
        "car"    => car(args, mem),
        "cdr"    => cdr(args, mem),
        "cons"   => cons(args, mem),
        "length" => length(args, mem),
        "list"   => list(args, mem),
        "list?"  => list_question_mark(args, mem),
        "not"    => not(args, mem),
        "null?"  => null(args, mem),
        "pair?"  => pair(args, mem),
        _        => Err(Error::PrimitiveFailed(name.to_string()))
    }
}

fn plus(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args_ = try!(list_of_integers(args));
    let sum = args_.iter().fold(0, |sum, n| sum + *n);

    Ok(mem.integer(sum))
}

fn minus(args_: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = try!(list_of_integers(args_));

    if args.len() == 0 {
        return Err(Error::BadArity(Some("-".to_string())))
    }

    let head = args.head().unwrap();
    let tail = args.tail();

    if tail.is_empty() {
        return Ok(mem.integer(- *head))
    }

    let tail_sum = tail.iter().fold(0, |sum, n| sum + *n);
    Ok(mem.integer(*head - tail_sum))
}

fn division(args_: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = try!(list_of_integers(args_));

    if args.len() == 0 {
        return Err(Error::BadArity(Some("/".to_string())))
    }

    let head = args.head().unwrap();
    let tail = args.tail();

    if tail.is_empty() {
        return Ok(mem.integer(1 / *head))
    }

    let div = tail.iter().fold(*head, |div, n| div / *n);
    Ok(mem.integer(div))
}

fn product(args_: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = try!(list_of_integers(args_));
    let product = args.iter().fold(1, |product, n| product * *n);

    Ok(mem.integer(product))
}

fn equals(args_: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args_.len() < 2 {
        return Ok(mem.b_true());
    }

    let args = try!(list_of_integers(args_));
    let head = args.head().unwrap();
    let outcome = args.iter().skip(1).all(|n| *n == *head);

    Ok(mem.boolean(outcome))
}

fn less_than(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a < b)
}

fn less_than_or_equal(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a <= b)
}

fn greater_than(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a > b)
}

fn greater_than_or_equal(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a >= b)
}

fn not(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("not".to_string())))
    }

    let outcome = match args.head().unwrap().deref() {
        &Value::Bool(false) => true,
        _                   => false,
    };

    Ok(mem.boolean(outcome))
}

fn list(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let list = args.iter().map(|a| a.deref().clone()).collect();

    Ok(mem.list(list))
}

fn length(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("length".to_string())));
    }

    match args[0].deref() {
        &Value::List(ref list) =>
            Ok(mem.integer(list.len() as i64)),
        value =>
            Err(Error::WrongArgumentType(value.clone())),
    }
}

fn pair(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("pair?".to_string())));
    }

    let outcome = if let &Value::List(ref list) = args[0].deref() {
        !list.is_empty()
    } else if let &Value::DottedList(ref _list, ref _el) = args[0].deref() {
        true
    } else {
        false
    };

    Ok(mem.boolean(outcome))
}

fn cons(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 2 {
        return Err(Error::BadArity(Some("cons".to_string())));
    }

    let mut list: Vec<Value> = Vec::new();

    if let &Value::List(ref l) = args[1].deref() {
        list.push(args[0].deref().clone());
        list.push_all(l.clone().as_slice());

        Ok(mem.list(list))
    } else {
        list.push(args[0].deref().clone());
        Ok(mem.dotted_list(list, args[1].deref().clone()))
    }
}

fn car(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("car".to_string())));
    }

    match args[0].deref() {
        &Value::List(ref l) if !l.is_empty() =>
            Ok(mem.store(l.head().unwrap().clone())),
        &Value::DottedList(ref l, ref _t) =>
            Ok(mem.store(l.head().unwrap().clone())),
        value =>
            Err(Error::WrongArgumentType(value.clone())),
    }
}

fn cdr(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("cdr".to_string())));
    }

    match args[0].deref() {
        &Value::List(ref l) if !l.is_empty() =>
            match l.tail() {
                tail if tail.is_empty() =>
                    Ok(mem.empty_list()),
                tail =>
                    Ok(mem.list(tail.to_vec())),
            },
        &Value::DottedList(ref _l, ref t) =>
            Ok(mem.store(*t.clone())),
        value =>
            Err(Error::WrongArgumentType(value.clone())),
    }
}

fn null(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("null?".to_string())));
    }

    let out = match args[0].deref() {
        &Value::List(ref l) if l.is_empty() => true,
        _ => false
    };

    Ok(mem.boolean(out))
}

fn list_question_mark(args: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("list?".to_string())));
    }

    let out = if let &Value::List(ref _l) = args[0].deref() {
        true
    } else {
        false
    };

    Ok(mem.boolean(out))
}

fn list_of_integers(list: Vec<Rc<Value>>) -> Result<Vec<i64>, Error> {
    let mut integers = Vec::with_capacity(list.len());

    for val in list.iter() {
        match val.deref() {
            &Value::Integer(n) =>
                integers.push(n),
            value =>
                return Err(Error::WrongArgumentType(value.clone())),
        }
    }

    Ok(integers)
}

fn ord(args_: Vec<Rc<Value>>, mem: &mut Memory, cmp: |i64, i64| -> bool) -> Result<Rc<Value>, Error> {
    if args_.len() < 2 {
        return Ok(mem.b_true())
    }

    let args = try!(list_of_integers(args_));
    let outcome = range(0, args.len() - 1).all(|i|
        cmp(args[i], args[i + 1u])
    );

    return Ok(mem.boolean(outcome))
}
<commit_msg>Passed around refs instead of Rc values in primitives.<commit_after>use runtime::Error;
use values::Value;
use memory::Memory;

use std::rc::Rc;

pub static PRIMITIVES: [&'static str, ..18] = [
    "*",
    "+",
    "-",
    "/",
    "<",
    "<=",
    "=",
    ">",
    ">=",
    "car",
    "cdr",
    "cons",
    "length",
    "list",
    "list?",
    "not",
    "null?",
    "pair?",
];

pub fn apply(name: &String, args_: Vec<Rc<Value>>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = args_.iter().map(|a| a.deref()).collect();

    match name.as_slice() {
        "*"      => product(args, mem),
        "+"      => plus(args, mem),
        "-"      => minus(args, mem),
        "/"      => division(args, mem),
        "<"      => less_than(args, mem),
        "<="     => less_than_or_equal(args, mem),
        "="      => equals(args, mem),
        ">"      => greater_than(args, mem),
        ">="     => greater_than_or_equal(args, mem),
        "car"    => car(args, mem),
        "cdr"    => cdr(args, mem),
        "cons"   => cons(args, mem),
        "length" => length(args, mem),
        "list"   => list(args, mem),
        "list?"  => list_question_mark(args, mem),
        "not"    => not(args, mem),
        "null?"  => null(args, mem),
        "pair?"  => pair(args, mem),
        _        => Err(Error::PrimitiveFailed(name.to_string()))
    }
}

fn plus(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args_ = try!(list_of_integers(args));
    let sum = args_.iter().fold(0, |sum, n| sum + *n);

    Ok(mem.integer(sum))
}

fn minus(args_: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = try!(list_of_integers(args_));

    if args.len() == 0 {
        return Err(Error::BadArity(Some("-".to_string())))
    }

    let head = args.head().unwrap();
    let tail = args.tail();

    if tail.is_empty() {
        return Ok(mem.integer(- *head))
    }

    let tail_sum = tail.iter().fold(0, |sum, n| sum + *n);
    Ok(mem.integer(*head - tail_sum))
}

fn division(args_: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = try!(list_of_integers(args_));

    if args.len() == 0 {
        return Err(Error::BadArity(Some("/".to_string())))
    }

    let head = args.head().unwrap();
    let tail = args.tail();

    if tail.is_empty() {
        return Ok(mem.integer(1 / *head))
    }

    let div = tail.iter().fold(*head, |div, n| div / *n);
    Ok(mem.integer(div))
}

fn product(args_: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let args = try!(list_of_integers(args_));
    let product = args.iter().fold(1, |product, n| product * *n);

    Ok(mem.integer(product))
}

fn equals(args_: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args_.len() < 2 {
        return Ok(mem.b_true());
    }

    let args = try!(list_of_integers(args_));
    let head = args.head().unwrap();
    let outcome = args.iter().skip(1).all(|n| *n == *head);

    Ok(mem.boolean(outcome))
}

fn less_than(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a < b)
}

fn less_than_or_equal(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a <= b)
}

fn greater_than(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a > b)
}

fn greater_than_or_equal(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    ord(args, mem, |a, b| a >= b)
}

fn not(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("not".to_string())))
    }

    let outcome = match *args.head().unwrap() {
        &Value::Bool(false) => true,
        _                   => false,
    };

    Ok(mem.boolean(outcome))
}

fn list(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    let list = args.iter().map(|a| a.deref().clone()).collect();

    Ok(mem.list(list))
}

fn length(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("length".to_string())));
    }

    match args[0] {
        &Value::List(ref list) =>
            Ok(mem.integer(list.len() as i64)),
        value =>
            Err(Error::WrongArgumentType(value.clone())),
    }
}

fn pair(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("pair?".to_string())));
    }

    let outcome = if let &Value::List(ref list) = args[0] {
        !list.is_empty()
    } else if let &Value::DottedList(ref _list, ref _el) = args[0] {
        true
    } else {
        false
    };

    Ok(mem.boolean(outcome))
}

fn cons(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 2 {
        return Err(Error::BadArity(Some("cons".to_string())));
    }

    let mut list: Vec<Value> = Vec::new();

    if let &Value::List(ref l) = args[1] {
        list.push(args[0].clone());
        list.push_all(l.clone().as_slice());

        Ok(mem.list(list))
    } else {
        list.push(args[0].clone());
        Ok(mem.dotted_list(list, args[1].clone()))
    }
}

fn car(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("car".to_string())));
    }

    match args[0] {
        &Value::List(ref l) if !l.is_empty() =>
            Ok(mem.store(l.head().unwrap().clone())),
        &Value::DottedList(ref l, ref _t) =>
            Ok(mem.store(l.head().unwrap().clone())),
        value =>
            Err(Error::WrongArgumentType(value.clone())),
    }
}

fn cdr(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("cdr".to_string())));
    }

    match args[0] {
        &Value::List(ref l) if !l.is_empty() =>
            match l.tail() {
                tail if tail.is_empty() =>
                    Ok(mem.empty_list()),
                tail =>
                    Ok(mem.list(tail.to_vec())),
            },
        &Value::DottedList(ref _l, ref t) =>
            Ok(mem.store(*t.clone())),
        value =>
            Err(Error::WrongArgumentType(value.clone())),
    }
}

fn null(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("null?".to_string())));
    }

    let out = match args[0] {
        &Value::List(ref l) if l.is_empty() => true,
        _ => false
    };

    Ok(mem.boolean(out))
}

fn list_question_mark(args: Vec<&Value>, mem: &mut Memory) -> Result<Rc<Value>, Error> {
    if args.len() != 1 {
        return Err(Error::BadArity(Some("list?".to_string())));
    }

    let out = if let &Value::List(ref _l) = args[0] {
        true
    } else {
        false
    };

    Ok(mem.boolean(out))
}

fn list_of_integers(list: Vec<&Value>) -> Result<Vec<i64>, Error> {
    let mut integers = Vec::with_capacity(list.len());

    for val in list.iter() {
        match *val {
            &Value::Integer(n) =>
                integers.push(n),
            value =>
                return Err(Error::WrongArgumentType(value.clone())),
        }
    }

    Ok(integers)
}

fn ord(args_: Vec<&Value>, mem: &mut Memory, cmp: |i64, i64| -> bool) -> Result<Rc<Value>, Error> {
    if args_.len() < 2 {
        return Ok(mem.b_true())
    }

    let args = try!(list_of_integers(args_));
    let outcome = range(0, args.len() - 1).all(|i|
        cmp(args[i], args[i + 1u])
    );

    return Ok(mem.boolean(outcome))
}
<|endoftext|>


<commit_before>use num::{ToPrimitive, FromPrimitive};
use byteorder::{ByteOrder, LittleEndian};
use failure::{Error, Fail};

use super::util;
use super::loader;
use super::arch::{Arch};
use super::loader::{LoadedModule, Loader};

#[derive(Debug, Fail)]
pub enum WorkspaceError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given address is not mapped")]
    InvalidAddress,
    #[fail(display = "Buffer overrun")]
    BufferOverrun,
}

pub struct WorkspaceBuilder<A: Arch> {
    filename: String,
    buf: Vec<u8>,

    loader: Option<Box<dyn Loader<A>>>,
}

impl<A: Arch + 'static> WorkspaceBuilder<A> {
    /// Override the default loader picker with the given loader.
    pub fn with_loader(self: WorkspaceBuilder<A>, loader: Box<dyn Loader<A>>) -> WorkspaceBuilder<A> {
        WorkspaceBuilder {
            loader: Some(loader),
            .. self
        }
    }

    /// Construct a workspace with the given builder configuration.
    ///
    /// This invokes the loaders, analyzers, and another other logic,
    ///  resulting in an initialized Workspace instance.
    ///
    /// Example (with default loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    ///
    /// Example (with specific loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn load(self: WorkspaceBuilder<A>) -> Result<Workspace<A>, Error> {
        // if the user provided a loader, use that.
        // otherwise, use the default detected loader.
        let (ldr, module) = match self.loader {
            Some(ldr) => {
                let module = ldr.load(&self.buf)?;
                (ldr, module)
            }
            None => loader::load::<A>(&self.buf)?,
        };

        Ok(Workspace::<A>{
            filename: self.filename,
            buf: self.buf,

            loader: ldr,
            module: module,
        })
    }
}

pub struct Workspace<A: Arch> {
    // name or source of the file
    pub filename: String,
    // raw bytes of the file
    pub buf: Vec<u8>,

    pub loader: Box<dyn Loader<A>>,
    pub module: LoadedModule<A>,

    // analysis {
    //   flowmeta,
    //   datameta,
    //   symbols,
    //   functions,
    // }
}

impl<A: Arch + 'static> Workspace<A> {
    /// Create a workspace and load the given bytes.
    ///
    /// See example on `WorkspaceBuilder::load()`
    pub fn from_bytes(filename: &str, buf: &[u8]) -> WorkspaceBuilder<A> {
        WorkspaceBuilder {
            filename: filename.to_string(),
            buf: buf.to_vec(),
            loader: None,
        }
    }

    pub fn from_file(filename: &str) -> Result<WorkspaceBuilder<A>, Error> {
        Ok(WorkspaceBuilder {
            filename: filename.to_string(),
            buf: util::read_file(filename)?,
            loader: None,
        })
    }

    /// Read bytes from the given RVA.
    ///
    /// Errors:
    ///
    ///   - InvalidAddress - if the address is not mapped.
    ///   - BufferOverrun - if the requested region runs beyond the matching section.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_bytes(0x0, 0x1).unwrap().to_vec(), b"\xEB");
    ///     assert_eq!(ws.read_bytes(0x1, 0x1).unwrap().to_vec(), b"\xFE");
    ///     assert_eq!(ws.read_bytes(0x0, 0x2).unwrap().to_vec(), b"\xEB\xFE");
    ///     assert_eq!(ws.read_bytes(0x0, 0x3).is_err(), true);
    ///     assert_eq!(ws.read_bytes(0x2, 0x1).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_bytes(&self, rva: A::RVA, length: usize) -> Result<&[u8], Error> {
        self.module.sections
            .iter()
            .filter(|section| section.contains(rva))
            .nth(0)
            .ok_or(WorkspaceError::InvalidAddress.into())
            .and_then(|section| -> Result<&[u8], Error> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .ok_or(WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + length > section.buf.len() {
                            Err(WorkspaceError::BufferOverrun.into())
                        } else {
                            Ok(&section.buf[offset..offset+length])
                        }
                    })
            })
    }

    /// Read a byte from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_u8(0x0).unwrap(), 0xEB);
    ///     assert_eq!(ws.read_u8(0x1).unwrap(), 0xFE);
    ///     assert_eq!(ws.read_u8(0x2).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_u8(&self, rva: A::RVA) -> Result<u8, Error> {
        self.read_bytes(rva, 1)
            .and_then(|buf| Ok(buf[0]))
    }

    /// Read a word from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_u16(0x0).unwrap(), 0xFEEB);
    ///     assert_eq!(ws.read_u16(0x1).is_err(), true);
    ///     assert_eq!(ws.read_u16(0x2).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_u16(&self, rva: A::RVA) -> Result<u16, Error> {
        self.read_bytes(rva, 2)
            .and_then(|buf| Ok(LittleEndian::read_u16(buf)))
    }

    /// Read a dword from the given RVA.
    pub fn read_u32(&self, rva: A::RVA) -> Result<u32, Error> {
        self.read_bytes(rva, 4)
            .and_then(|buf| Ok(LittleEndian::read_u32(buf)))
    }

    /// Read a qword from the given RVA.
    pub fn read_u64(&self, rva: A::RVA) -> Result<u64, Error> {
        self.read_bytes(rva, 8)
            .and_then(|buf| Ok(LittleEndian::read_u64(buf)))
    }

    /// Read an RVA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\x00\x11\x22\x33")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_rva(0x0).unwrap(), 0x33221100);
    ///     assert_eq!(ws.read_rva(0x1).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_rva(&self, rva: A::RVA) -> Result<A::RVA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::RVA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::RVA::from_u64(self.read_u64(rva)?).unwrap()),
            _  => panic!("unexpected architecture"),
        }
    }

    /// Read a VA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    pub fn read_va(&self, rva: A::RVA) -> Result<A::VA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::VA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::VA::from_u64(self.read_u64(rva)?).unwrap()),
            _  => panic!("unexpected architecture"),
        }
    }

    // API:
    //   get_insn
    //   get_xrefs_to
    //   get_xrefs_from
    //   get_functions

    // elsewhere:
    //   call graph
    //   control flow graph
}
<commit_msg>workspace: begin to sketch analysis runtime<commit_after>use num::{ToPrimitive, FromPrimitive};
use std::collections::VecDeque;
use failure::{Error, Fail};
use byteorder::{ByteOrder, LittleEndian};

use super::util;
use super::loader;
use super::arch::{Arch};
use super::loader::{LoadedModule, Loader};

#[derive(Debug, Fail)]
pub enum WorkspaceError {
    #[fail(display = "The given buffer is not supported (arch/plat/file format)")]
    NotSupported,
    #[fail(display = "The given address is not mapped")]
    InvalidAddress,
    #[fail(display = "Buffer overrun")]
    BufferOverrun,
}

pub enum AnalysisCommand<A: Arch> {
    MakeInsn(A::RVA),
}

pub struct WorkspaceBuilder<A: Arch> {
    filename: String,
    buf: Vec<u8>,

    loader: Option<Box<dyn Loader<A>>>,
}

impl<A: Arch + 'static> WorkspaceBuilder<A> {
    /// Override the default loader picker with the given loader.
    pub fn with_loader(self: WorkspaceBuilder<A>, loader: Box<dyn Loader<A>>) -> WorkspaceBuilder<A> {
        WorkspaceBuilder {
            loader: Some(loader),
            .. self
        }
    }

    /// Construct a workspace with the given builder configuration.
    ///
    /// This invokes the loaders, analyzers, and another other logic,
    ///  resulting in an initialized Workspace instance.
    ///
    /// Example (with default loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    ///
    /// Example (with specific loader):
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.loader.get_name(),       "Windows/32/Raw");
    ///     assert_eq!(ws.module.base_address,     0x0);
    ///     assert_eq!(ws.module.sections[0].name, "raw");
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn load(self: WorkspaceBuilder<A>) -> Result<Workspace<A>, Error> {
        // if the user provided a loader, use that.
        // otherwise, use the default detected loader.
        let (ldr, module) = match self.loader {
            Some(ldr) => {
                let module = ldr.load(&self.buf)?;
                (ldr, module)
            }
            None => loader::load::<A>(&self.buf)?,
        };

        Ok(Workspace::<A>{
            filename: self.filename,
            buf: self.buf,

            loader: ldr,
            module: module,

            analysis_queue: VecDeque::new(),
        })
    }
}

pub struct Workspace<A: Arch> {
    // name or source of the file
    pub filename: String,
    // raw bytes of the file
    pub buf: Vec<u8>,

    pub loader: Box<dyn Loader<A>>,
    pub module: LoadedModule<A>,

    // analysis {
    //   flowmeta,
    //   datameta,
    //   symbols,
    //   functions,
    // }

    analysis_queue: VecDeque<AnalysisCommand<A>>,
}

impl<A: Arch + 'static> Workspace<A> {
    /// Create a workspace and load the given bytes.
    ///
    /// See example on `WorkspaceBuilder::load()`
    pub fn from_bytes(filename: &str, buf: &[u8]) -> WorkspaceBuilder<A> {
        WorkspaceBuilder {
            filename: filename.to_string(),
            buf: buf.to_vec(),
            loader: None,
        }
    }

    pub fn from_file(filename: &str) -> Result<WorkspaceBuilder<A>, Error> {
        Ok(WorkspaceBuilder {
            filename: filename.to_string(),
            buf: util::read_file(filename)?,
            loader: None,
        })
    }

    /// Read bytes from the given RVA.
    ///
    /// Errors:
    ///
    ///   - InvalidAddress - if the address is not mapped.
    ///   - BufferOverrun - if the requested region runs beyond the matching section.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_bytes(0x0, 0x1).unwrap().to_vec(), b"\xEB");
    ///     assert_eq!(ws.read_bytes(0x1, 0x1).unwrap().to_vec(), b"\xFE");
    ///     assert_eq!(ws.read_bytes(0x0, 0x2).unwrap().to_vec(), b"\xEB\xFE");
    ///     assert_eq!(ws.read_bytes(0x0, 0x3).is_err(), true);
    ///     assert_eq!(ws.read_bytes(0x2, 0x1).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_bytes(&self, rva: A::RVA, length: usize) -> Result<&[u8], Error> {
        self.module.sections
            .iter()
            .filter(|section| section.contains(rva))
            .nth(0)
            .ok_or(WorkspaceError::InvalidAddress.into())
            .and_then(|section| -> Result<&[u8], Error> {
                // rva is guaranteed to be within this section,
                // so we can do an unchecked subtract here.
                let offset = rva - section.addr;
                A::RVA::to_usize(&offset)
                    .ok_or(WorkspaceError::InvalidAddress.into())
                    .and_then(|offset| {
                        if offset + length > section.buf.len() {
                            Err(WorkspaceError::BufferOverrun.into())
                        } else {
                            Ok(&section.buf[offset..offset+length])
                        }
                    })
            })
    }

    /// Read a byte from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_u8(0x0).unwrap(), 0xEB);
    ///     assert_eq!(ws.read_u8(0x1).unwrap(), 0xFE);
    ///     assert_eq!(ws.read_u8(0x2).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_u8(&self, rva: A::RVA) -> Result<u8, Error> {
        self.read_bytes(rva, 1)
            .and_then(|buf| Ok(buf[0]))
    }

    /// Read a word from the given RVA.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\xEB\xFE")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_u16(0x0).unwrap(), 0xFEEB);
    ///     assert_eq!(ws.read_u16(0x1).is_err(), true);
    ///     assert_eq!(ws.read_u16(0x2).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_u16(&self, rva: A::RVA) -> Result<u16, Error> {
        self.read_bytes(rva, 2)
            .and_then(|buf| Ok(LittleEndian::read_u16(buf)))
    }

    /// Read a dword from the given RVA.
    pub fn read_u32(&self, rva: A::RVA) -> Result<u32, Error> {
        self.read_bytes(rva, 4)
            .and_then(|buf| Ok(LittleEndian::read_u32(buf)))
    }

    /// Read a qword from the given RVA.
    pub fn read_u64(&self, rva: A::RVA) -> Result<u64, Error> {
        self.read_bytes(rva, 8)
            .and_then(|buf| Ok(LittleEndian::read_u64(buf)))
    }

    /// Read an RVA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    ///
    /// Errors: same as `read_bytes`.
    ///
    /// Example:
    ///
    /// ```
    /// use lancelot::arch::*;
    /// use lancelot::loader;
    /// use lancelot::workspace::Workspace;
    ///
    /// Workspace::<Arch32>::from_bytes("foo.bin", b"\x00\x11\x22\x33")
    ///   .with_loader(Box::new(loader::ShellcodeLoader::<Arch32>::new(loader::Platform::Windows)))
    ///   .load()
    ///   .map(|ws| {
    ///     assert_eq!(ws.read_rva(0x0).unwrap(), 0x33221100);
    ///     assert_eq!(ws.read_rva(0x1).is_err(), true);
    ///   })
    ///   .map_err(|e| panic!(e));
    /// ```
    pub fn read_rva(&self, rva: A::RVA) -> Result<A::RVA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::RVA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::RVA::from_u64(self.read_u64(rva)?).unwrap()),
            _  => panic!("unexpected architecture"),
        }
    }

    /// Read a VA from the given RVA.
    /// Note that the size of the read is dependent on the architecture.
    pub fn read_va(&self, rva: A::RVA) -> Result<A::VA, Error> {
        match A::get_bits() {
            // these conversions will never fail
            32 => Ok(A::VA::from_u32(self.read_u32(rva)?).unwrap()),
            64 => Ok(A::VA::from_u64(self.read_u64(rva)?).unwrap()),
            _  => panic!("unexpected architecture"),
        }
    }

    // API:
    //   get_insn
    //   get_xrefs_to
    //   get_xrefs_from
    //   get_functions

    // elsewhere:
    //   call graph
    //   control flow graph

    pub fn make_insn(&mut self, rva: A::RVA) -> Result<(), Error> {
        self.analysis_queue.push_back(AnalysisCommand::MakeInsn(rva));
        Ok(())
    }

    fn handle_make_insn(&mut self, rva: A::RVA) -> Result<(), Error> {
        // 1. ensure not exists
        // 2. compute len
        // 3. compute fallthrough
        // 4. compute flow ref
        Ok(())
    }

    pub fn analyze(&mut self) -> Result<(), Error> {
        while let Some(cmd) = self.analysis_queue.pop_front() {
            match cmd {
                AnalysisCommand::MakeInsn(rva) => self.handle_make_insn(rva)?,
            }
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>#![crate_name = "readlink"]
#![feature(collections, core, old_io, os, old_path, rustc_private)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Haitao Li <lihaitao@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

use getopts::{optflag, getopts, usage, OptGroup};
use std::old_io as io;
use std::old_io::fs;
use std::os;
use std::vec::Vec;

use CanonicalizeMode::{None, Normal, Existing, Missing};

#[path = "../common/util.rs"]
#[macro_use]
mod util;

const NAME: &'static str = "readlink";
const VERSION:  &'static str = "0.0.1";

#[derive(PartialEq)]
enum CanonicalizeMode {
    None,
    Normal,
    Existing,
    Missing,
}

fn resolve(original: &Path) -> io::IoResult<Path> {
    const MAX_LINKS_FOLLOWED: u32 = 255;
    let mut followed = 0;
    let mut result = original.clone();
    loop {
        if followed == MAX_LINKS_FOLLOWED {
            return Err(io::standard_error(io::InvalidInput));
        }

        match fs::lstat(&result) {
            Err(e) => return Err(e),
            Ok(ref stat) if stat.kind != io::FileType::Symlink => break,
            Ok(..) => {
                followed += 1;
                match fs::readlink(&result) {
                    Ok(path) => {
                        result.pop();
                        result.push(path);
                    },
                    Err(e) => {
                        return Err(e);
                    }
                }
            }
        }
    }
    return Ok(result);
}

fn canonicalize(original: &Path, can_mode: &CanonicalizeMode) -> io::IoResult<Path> {
    let original = os::make_absolute(original).unwrap();
    let result = original.root_path();
    let mut result = result.expect("make_absolute has no root_path");
    let mut parts = vec![];
    for part in original.components() {
        parts.push(part);
    }

    if parts.len() > 1 {
        for part in parts.init().iter() {
            result.push(part);

            if *can_mode == None {
                continue;
            }

            match resolve(&result) {
                Err(_) => match *can_mode {
                    Missing => continue,
                    _ => return Err(io::standard_error(io::InvalidInput)),
                },
                Ok(path) => {
                    result.pop();
                    result.push(path);
                }
            }
        }
    }

    result.push(parts.last().unwrap());
    match resolve(&result) {
        Err(e) => { if *can_mode == Existing { return Err(e); } },
        Ok(path) => {
            result.pop();
            result.push(path);
        }
    }
    return Ok(result);
}

pub fn uumain(args: Vec<String>) -> i32 {
    let program = args[0].as_slice();
    let opts = [
        optflag("f", "canonicalize",
                "canonicalize by following every symlink in every component of the \
                 given name recursively; all but the last component must exist"),
        optflag("e", "canonicalize-existing",
                "canonicalize by following every symlink in every component of the \
                 given name recursively, all components must exist"),
        optflag("m", "canonicalize-missing",
                "canonicalize by following every symlink in every component of the \
                 given name recursively, without requirements on components existence"),
        optflag("n", "no-newline", "do not output the trailing delimiter"),
        optflag("q", "quiet", "suppress most error messages"),
        optflag("s", "silent", "suppress most error messages"),
        optflag("v", "verbose", "report error message"),
        optflag("z", "zero", "separate output with NUL rather than newline"),
        optflag("", "help", "display this help and exit"),
        optflag("", "version", "output version information and exit"),
    ];

    let matches = match getopts(args.tail(), &opts) {
        Ok(m) => m,
        Err(f) => crash!(1, "Invalid options\n{}", f)
    };
    if matches.opt_present("help") {
        show_usage(program, &opts);
        return 0;
    }

    if matches.opt_present("version") {
        println!("{} v{}", NAME, VERSION);
        return 0;
    }

    let mut no_newline = matches.opt_present("no-newline");
    let use_zero = matches.opt_present("zero");
    let silent = matches.opt_present("silent") || matches.opt_present("quiet");

    let mut can_mode = None;
    if matches.opt_present("canonicalize") {
        can_mode = Normal;
    }

    if matches.opt_present("canonicalize-existing") {
        can_mode = Existing;
    }

    if matches.opt_present("canonicalize-missing") {
        can_mode = Missing;
    }

    let files = matches.free;
    if files.len() == 0 {
        crash!(1, "missing operand\nTry {} --help for more information", program);
    }

    if no_newline && files.len() > 1 {
        if !silent {
            eprintln!("{}: ignoring --no-newline with multiple arguments", program);
            no_newline = false;
        }
    }

    for f in files.iter() {
        let p = Path::new(f);
        if can_mode == None {
            match fs::readlink(&p) {
                Ok(path) => show(path.as_str().unwrap(), no_newline, use_zero),
                Err(_) => return 1
            }
        } else {
            match canonicalize(&p, &can_mode) {
                Ok(path) => show(path.as_str().unwrap(), no_newline, use_zero),
                Err(_) => return 1
            }
        }
    }
    return 0;
}

fn show(path: &str, no_newline: bool, use_zero: bool) {
    if use_zero {
        print!("{}\0", path);
    } else {
        if no_newline {
            io::print(path);
        } else {
            io::println(path);
        }
    }
}

fn show_usage(program: &str, opts: &[OptGroup]) {
    println!("Usage: {0} [OPTION]... [FILE]...", program);
    print!("Print value of a symbolic link or canonical file name");
    io::print(usage("", opts).as_slice());
}
<commit_msg>Update readlink.<commit_after>#![crate_name = "readlink"]
#![feature(file_type, rustc_private)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Haitao Li <lihaitao@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

extern crate getopts;

use getopts::{getopts, optflag, OptGroup, usage};
use std::env;
use std::fs::{metadata, read_link};
use std::io::{Error, ErrorKind, Result, Write};
use std::path::{Component, PathBuf};

use CanonicalizeMode::{None, Normal, Existing, Missing};

#[path = "../common/util.rs"]
#[macro_use]
mod util;

const NAME: &'static str = "readlink";
const VERSION: &'static str = "0.0.1";

#[derive(PartialEq)]
enum CanonicalizeMode {
    None,
    Normal,
    Existing,
    Missing,
}

fn resolve(original: &PathBuf) -> Result<PathBuf> {
    const MAX_LINKS_FOLLOWED: u32 = 255;
    let mut followed = 0;
    let mut result = original.clone();
    loop {
        if followed == MAX_LINKS_FOLLOWED {
            return Err(Error::new(ErrorKind::InvalidInput, "maximum links followed"));
        }

        match metadata(&result) {
            Err(e) => return Err(e),
            Ok(ref m) if !m.file_type().is_symlink() => break,
            Ok(..) => {
                followed += 1;
                match read_link(&result) {
                    Ok(path) => {
                        result.pop();
                        result.push(path);
                    },
                    Err(e) => {
                        return Err(e);
                    }
                }
            }
        }
    }
    Ok(result)
}

fn canonicalize(original: &PathBuf, can_mode: &CanonicalizeMode) -> Result<PathBuf> {
    // Create an absolute path
    let original = if original.as_path().is_absolute() {
        original.clone()
    } else {
        env::current_dir().unwrap().join(original)
    };

    let mut result = PathBuf::new();
    let mut parts = vec!();

    // Split path by directory separator; add prefix (Windows-only) and root
    // directory to final path buffer; add remaining parts to temporary
    // vector for canonicalization.
    for part in original.components() {
        match part {
            Component::Prefix(_) | Component::RootDir => {
                result.push(part.as_os_str());
            },
            Component::CurDir => {},
            Component::ParentDir => {
                parts.pop();
            },
            Component::Normal(_) => {
                parts.push(part.as_os_str());
            }
        }
    }

    // Resolve the symlinks where possible
    if parts.len() > 0 {
        for part in parts[..parts.len()-1].iter() {
            result.push(part);

            if *can_mode == None {
                continue;
            }

            match resolve(&result) {
                Err(e) => match *can_mode {
                    Missing => continue,
                    _ => return Err(e)
                },
                Ok(path) => {
                    result.pop();
                    result.push(path);
                }
            }
        }

        result.push(parts.last().unwrap());

        match resolve(&result) {
            Err(e) => { if *can_mode == Existing { return Err(e); } },
            Ok(path) => {
                result.pop();
                result.push(path);
            }
        }
    }
    Ok(result)
}

pub fn uumain(args: Vec<String>) -> i32 {
    let program = &args[0];
    let opts = [
        optflag("f", "canonicalize",
                "canonicalize by following every symlink in every component of the \
                 given name recursively; all but the last component must exist"),
        optflag("e", "canonicalize-existing",
                "canonicalize by following every symlink in every component of the \
                 given name recursively, all components must exist"),
        optflag("m", "canonicalize-missing",
                "canonicalize by following every symlink in every component of the \
                 given name recursively, without requirements on components existence"),
        optflag("n", "no-newline", "do not output the trailing delimiter"),
        optflag("q", "quiet", "suppress most error messages"),
        optflag("s", "silent", "suppress most error messages"),
        optflag("v", "verbose", "report error message"),
        optflag("z", "zero", "separate output with NUL rather than newline"),
        optflag("", "help", "display this help and exit"),
        optflag("", "version", "output version information and exit"),
    ];

    let matches = match getopts(&args[1..], &opts) {
        Ok(m) => m,
        Err(f) => crash!(1, "Invalid options\n{}", f)
    };
    if matches.opt_present("help") {
        show_usage(program, &opts);
        return 0;
    }

    if matches.opt_present("version") {
        println!("{} v{}", NAME, VERSION);
        return 0;
    }

    let mut no_newline = matches.opt_present("no-newline");
    let use_zero = matches.opt_present("zero");
    let silent = matches.opt_present("silent") || matches.opt_present("quiet");
    let verbose = matches.opt_present("verbose");

    let mut can_mode = None;
    if matches.opt_present("canonicalize") {
        can_mode = Normal;
    }

    if matches.opt_present("canonicalize-existing") {
        can_mode = Existing;
    }

    if matches.opt_present("canonicalize-missing") {
        can_mode = Missing;
    }

    let files = matches.free;
    if files.len() == 0 {
        crash!(1, "missing operand\nTry {} --help for more information", program);
    }

    if no_newline && files.len() > 1 {
        if !silent {
            eprintln!("{}: ignoring --no-newline with multiple arguments", program);
            no_newline = false;
        }
    }

    for f in files.iter() {
        let p = PathBuf::from(f);
        if can_mode == None {
            match read_link(&p) {
                Ok(path) => show(&path, no_newline, use_zero),
                Err(err) => {
                    if verbose {
                        eprintln!("{}: {}: errno {}", NAME, f, err.raw_os_error().unwrap());
                    }
                    return 1
                }
            }
        } else {
            match canonicalize(&p, &can_mode) {
                Ok(path) => show(&path, no_newline, use_zero),
                Err(err) => {
                    if verbose {
                        eprintln!("{}: {}: errno {:?}", NAME, f, err.raw_os_error().unwrap());
                    }
                    return 1
                }
            }
        }
    }

    0
}

fn show(path: &PathBuf, no_newline: bool, use_zero: bool) {
    let path = path.as_path().to_str().unwrap();
    if use_zero {
        print!("{}\0", path);
    } else if no_newline {
        print!("{}", path);
    } else {
        println!("{}", path);
    }
}

fn show_usage(program: &str, opts: &[OptGroup]) {
    println!("Usage: {0} [OPTION]... [FILE]...", program);
    print!("Print value of a symbolic link or canonical file name");
    print!("{}", usage("", opts));
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! Tasks that are executed by coordinators (short-lived compared to coordinators)
use crate::thread_pool::IO_POOL;
use crate::{
    core_mempool::{CoreMempool, TimelineState, TxnPointer},
    counters,
    logging::{LogEntry, LogEvent, LogSchema},
    network::{BroadcastError, MempoolSyncMsg},
    shared_mempool::types::{
        notify_subscribers, ScheduledBroadcast, SharedMempool, SharedMempoolNotification,
        SubmissionStatusBundle,
    },
    QuorumStoreRequest, QuorumStoreResponse, SubmissionStatus,
};
use anyhow::Result;
use aptos_config::network_id::PeerNetworkId;
use aptos_crypto::HashValue;
use aptos_infallible::{Mutex, RwLock};
use aptos_logger::prelude::*;
use aptos_metrics_core::HistogramTimer;
use aptos_types::{
    mempool_status::{MempoolStatus, MempoolStatusCode},
    on_chain_config::OnChainConfigPayload,
    transaction::SignedTransaction,
    vm_status::DiscardedVMStatus,
};
use consensus_types::common::TransactionSummary;
use futures::{channel::oneshot, stream::FuturesUnordered};
use network::application::interface::NetworkInterface;
use rayon::prelude::*;
use std::{
    cmp,
    collections::HashSet,
    sync::Arc,
    time::{Duration, Instant},
};
use storage_interface::state_view::LatestDbStateCheckpointView;
use tokio::runtime::Handle;
use vm_validator::vm_validator::{get_account_sequence_number, TransactionValidation};

// ============================== //
//  broadcast_coordinator tasks  //
// ============================== //

/// Attempts broadcast to `peer` and schedules the next broadcast.
pub(crate) async fn execute_broadcast<V>(
    peer: PeerNetworkId,
    backoff: bool,
    smp: &mut SharedMempool<V>,
    scheduled_broadcasts: &mut FuturesUnordered<ScheduledBroadcast>,
    executor: Handle,
) where
    V: TransactionValidation,
{
    let network_interface = &smp.network_interface.clone();
    // If there's no connection, don't bother to broadcast
    if network_interface.app_data().read(&peer).is_some() {
        if let Err(err) = network_interface
            .execute_broadcast(peer, backoff, smp)
            .await
        {
            match err {
                BroadcastError::NetworkError(peer, error) => error!(LogSchema::event_log(
                    LogEntry::BroadcastTransaction,
                    LogEvent::NetworkSendFail
                )
                .peer(&peer)
                .error(&error)),
                _ => {
                    trace!("{:?}", err)
                }
            }
        }
    } else {
        // Drop the scheduled broadcast, we're not connected anymore
        return;
    }
    let schedule_backoff = network_interface.is_backoff_mode(&peer);

    let interval_ms = if schedule_backoff {
        smp.config.shared_mempool_backoff_interval_ms
    } else {
        smp.config.shared_mempool_tick_interval_ms
    };

    scheduled_broadcasts.push(ScheduledBroadcast::new(
        Instant::now() + Duration::from_millis(interval_ms),
        peer,
        schedule_backoff,
        executor,
    ))
}

// =============================== //
// Tasks processing txn submission //
// =============================== //

/// Processes transactions directly submitted by client.
pub(crate) async fn process_client_transaction_submission<V>(
    smp: SharedMempool<V>,
    transaction: SignedTransaction,
    callback: oneshot::Sender<Result<SubmissionStatus>>,
    timer: HistogramTimer,
) where
    V: TransactionValidation,
{
    timer.stop_and_record();
    let _timer = counters::process_txn_submit_latency_timer_client();
    let statuses = process_incoming_transactions(&smp, vec![transaction], TimelineState::NotReady);
    log_txn_process_results(&statuses, None);

    if let Some(status) = statuses.get(0) {
        if callback.send(Ok(status.1.clone())).is_err() {
            error!(LogSchema::event_log(
                LogEntry::JsonRpc,
                LogEvent::CallbackFail
            ));
            counters::CLIENT_CALLBACK_FAIL.inc();
        }
    }
}

/// Processes get transaction by hash request by client.
pub(crate) async fn process_client_get_transaction<V>(
    smp: SharedMempool<V>,
    hash: HashValue,
    callback: oneshot::Sender<Option<SignedTransaction>>,
    timer: HistogramTimer,
) where
    V: TransactionValidation,
{
    timer.stop_and_record();
    let _timer = counters::process_get_txn_latency_timer_client();
    let txn = smp.mempool.lock().get_by_hash(hash);

    if callback.send(txn).is_err() {
        error!(LogSchema::event_log(
            LogEntry::GetTransaction,
            LogEvent::CallbackFail
        ));
        counters::CLIENT_CALLBACK_FAIL.inc();
    }
}

/// Processes transactions from other nodes.
pub(crate) async fn process_transaction_broadcast<V>(
    smp: SharedMempool<V>,
    transactions: Vec<SignedTransaction>,
    request_id: Vec<u8>,
    timeline_state: TimelineState,
    peer: PeerNetworkId,
    timer: HistogramTimer,
) where
    V: TransactionValidation,
{
    timer.stop_and_record();
    let _timer = counters::process_txn_submit_latency_timer(peer.network_id());
    let results = process_incoming_transactions(&smp, transactions, timeline_state);
    log_txn_process_results(&results, Some(peer));

    let ack_response = gen_ack_response(request_id, results, &peer);
    let network_sender = smp.network_interface.sender();
    if let Err(e) = network_sender.send_to(peer, ack_response) {
        counters::network_send_fail_inc(counters::ACK_TXNS);
        error!(
            LogSchema::event_log(LogEntry::BroadcastACK, LogEvent::NetworkSendFail)
                .peer(&peer)
                .error(&e.into())
        );
        return;
    }
    notify_subscribers(SharedMempoolNotification::ACK, &smp.subscribers);
}

/// If `MempoolIsFull` on any of the transactions, provide backpressure to the downstream peer.
fn gen_ack_response(
    request_id: Vec<u8>,
    results: Vec<SubmissionStatusBundle>,
    peer: &PeerNetworkId,
) -> MempoolSyncMsg {
    let mut backoff_and_retry = false;
    for (_, (mempool_status, _)) in results.into_iter() {
        if mempool_status.code == MempoolStatusCode::MempoolIsFull {
            backoff_and_retry = true;
            break;
        }
    }

    update_ack_counter(
        peer,
        counters::SENT_LABEL,
        backoff_and_retry,
        backoff_and_retry,
    );
    MempoolSyncMsg::BroadcastTransactionsResponse {
        request_id,
        retry: backoff_and_retry,
        backoff: backoff_and_retry,
    }
}

pub(crate) fn update_ack_counter(
    peer: &PeerNetworkId,
    direction_label: &str,
    retry: bool,
    backoff: bool,
) {
    if retry {
        counters::shared_mempool_ack_inc(
            peer.network_id(),
            direction_label,
            counters::RETRY_BROADCAST_LABEL,
        );
    }
    if backoff {
        counters::shared_mempool_ack_inc(
            peer.network_id(),
            direction_label,
            counters::BACKPRESSURE_BROADCAST_LABEL,
        );
    }
}

/// Submits a list of SignedTransaction to the local mempool
/// and returns a vector containing AdmissionControlStatus.
pub(crate) fn process_incoming_transactions<V>(
    smp: &SharedMempool<V>,
    transactions: Vec<SignedTransaction>,
    timeline_state: TimelineState,
) -> Vec<SubmissionStatusBundle>
where
    V: TransactionValidation,
{
    let mut statuses = vec![];

    let start_storage_read = Instant::now();
    let state_view = smp
        .db
        .latest_state_checkpoint_view()
        .expect("Failed to get latest state checkpoint view.");

    // Track latency: fetching seq number
    let seq_numbers = IO_POOL.install(|| {
        transactions
            .par_iter()
            .map(|t| {
                get_account_sequence_number(&state_view, t.sender()).map_err(|e| {
                    error!(LogSchema::new(LogEntry::DBError).error(&e));
                    counters::DB_ERROR.inc();
                    e
                })
            })
            .collect::<Vec<_>>()
    });
    // Track latency for storage read fetching sequence number
    let storage_read_latency = start_storage_read.elapsed();
    counters::PROCESS_TXN_BREAKDOWN_LATENCY
        .with_label_values(&[counters::FETCH_SEQ_NUM_LABEL])
        .observe(storage_read_latency.as_secs_f64() / transactions.len() as f64);

    let transactions: Vec<_> = transactions
        .into_iter()
        .enumerate()
        .filter_map(|(idx, t)| {
            if let Ok(crsn_or_seqno) = seq_numbers[idx] {
                if t.sequence_number() >= crsn_or_seqno.min_seq() {
                    return Some((t, crsn_or_seqno));
                } else {
                    statuses.push((
                        t,
                        (
                            MempoolStatus::new(MempoolStatusCode::VmError),
                            Some(DiscardedVMStatus::SEQUENCE_NUMBER_TOO_OLD),
                        ),
                    ));
                }
            } else {
                // Failed to get transaction
                statuses.push((
                    t,
                    (
                        MempoolStatus::new(MempoolStatusCode::VmError),
                        Some(DiscardedVMStatus::RESOURCE_DOES_NOT_EXIST),
                    ),
                ));
            }
            None
        })
        .collect();

    // Track latency: VM validation
    let vm_validation_timer = counters::PROCESS_TXN_BREAKDOWN_LATENCY
        .with_label_values(&[counters::VM_VALIDATION_LABEL])
        .start_timer();
    let validation_results = transactions
        .iter()
        .map(|t| smp.validator.read().validate_transaction(t.0.clone()))
        .collect::<Vec<_>>();
    vm_validation_timer.stop_and_record();
    {
        let mut mempool = smp.mempool.lock();
        for (idx, (transaction, crsn_or_seqno)) in transactions.into_iter().enumerate() {
            if let Ok(validation_result) = &validation_results[idx] {
                match validation_result.status() {
                    None => {
                        let ranking_score = validation_result.score();
                        let mempool_status = mempool.add_txn(
                            transaction.clone(),
                            ranking_score,
                            crsn_or_seqno,
                            timeline_state,
                        );
                        statuses.push((transaction, (mempool_status, None)));
                    }
                    Some(validation_status) => {
                        statuses.push((
                            transaction.clone(),
                            (
                                MempoolStatus::new(MempoolStatusCode::VmError),
                                Some(validation_status),
                            ),
                        ));
                    }
                }
            }
        }
    }
    notify_subscribers(SharedMempoolNotification::NewTransactions, &smp.subscribers);
    statuses
}

fn log_txn_process_results(results: &[SubmissionStatusBundle], sender: Option<PeerNetworkId>) {
    let network = match sender {
        Some(peer) => peer.network_id().to_string(),
        None => counters::CLIENT_LABEL.to_string(),
    };
    for (txn, (mempool_status, maybe_vm_status)) in results.iter() {
        if let Some(vm_status) = maybe_vm_status {
            trace!(
                SecurityEvent::InvalidTransactionMempool,
                failed_transaction = txn,
                vm_status = vm_status,
                sender = sender,
            );
            counters::shared_mempool_transactions_processed_inc(
                counters::VM_VALIDATION_LABEL,
                &network,
            );
            continue;
        }
        match mempool_status.code {
            MempoolStatusCode::Accepted => counters::shared_mempool_transactions_processed_inc(
                counters::SUCCESS_LABEL,
                &network,
            ),
            _ => counters::shared_mempool_transactions_processed_inc(
                &mempool_status.code.to_string(),
                &network,
            ),
        }
    }
}

// ================================= //
// intra-node communication handlers //
// ================================= //

/// Only applies to Validators. Either provides transactions to consensus [`GetBlockRequest`] or
/// handles rejecting transactions [`RejectNotification`]
pub(crate) fn process_quorum_store_request<V: TransactionValidation>(
    smp: &SharedMempool<V>,
    req: QuorumStoreRequest,
) {
    // Start latency timer
    let start_time = Instant::now();
    debug!(LogSchema::event_log(LogEntry::QuorumStore, LogEvent::Received).quorum_store_msg(&req));

    let (resp, callback, counter_label) = match req {
        QuorumStoreRequest::GetBatchRequest(max_batch_size, transactions, callback) => {
            let exclude_transactions: HashSet<TxnPointer> = transactions
                .iter()
                .map(|txn| (txn.sender, txn.sequence_number))
                .collect();
            let txns;
            {
                let mut mempool = smp.mempool.lock();
                // gc before pulling block as extra protection against txns that may expire in consensus
                // Note: this gc operation relies on the fact that consensus uses the system time to determine block timestamp
                let curr_time = aptos_infallible::duration_since_epoch();
                mempool.gc_by_expiration_time(curr_time);
                let batch_size = cmp::max(max_batch_size, 1);
                txns = mempool.get_batch(batch_size, exclude_transactions);
            }
            counters::mempool_service_transactions(counters::GET_BLOCK_LABEL, txns.len());

            (
                QuorumStoreResponse::GetBatchResponse(txns),
                callback,
                counters::GET_BLOCK_LABEL,
            )
        }
        QuorumStoreRequest::RejectNotification(transactions, callback) => {
            counters::mempool_service_transactions(
                counters::COMMIT_CONSENSUS_LABEL,
                transactions.len(),
            );
            process_committed_transactions(&smp.mempool, transactions, 0, true);
            (
                QuorumStoreResponse::CommitResponse(),
                callback,
                counters::COMMIT_CONSENSUS_LABEL,
            )
        }
    };
    // Send back to callback
    let result = if callback.send(Ok(resp)).is_err() {
        error!(LogSchema::event_log(
            LogEntry::QuorumStore,
            LogEvent::CallbackFail
        ));
        counters::REQUEST_FAIL_LABEL
    } else {
        counters::REQUEST_SUCCESS_LABEL
    };
    let latency = start_time.elapsed();
    counters::mempool_service_latency(counter_label, result, latency);
}

/// Remove transactions that are committed (or rejected) so that we can stop broadcasting them.
pub(crate) fn process_committed_transactions(
    mempool: &Mutex<CoreMempool>,
    transactions: Vec<TransactionSummary>,
    block_timestamp_usecs: u64,
    is_rejected: bool,
) {
    let mut pool = mempool.lock();

    for transaction in transactions {
        pool.remove_transaction(
            &transaction.sender,
            transaction.sequence_number,
            is_rejected,
        );
    }

    if block_timestamp_usecs > 0 {
        pool.gc_by_expiration_time(Duration::from_micros(block_timestamp_usecs));
    }
}

/// Processes on-chain reconfiguration notifications.  Restarts validator with the new info.
pub(crate) async fn process_config_update<V>(
    config_update: OnChainConfigPayload,
    validator: Arc<RwLock<V>>,
) where
    V: TransactionValidation,
{
    info!(
        LogSchema::event_log(LogEntry::ReconfigUpdate, LogEvent::Process)
            .reconfig_update(config_update.clone())
    );

    if let Err(e) = validator.write().restart(config_update) {
        counters::VM_RECONFIG_UPDATE_FAIL_COUNT.inc();
        error!(LogSchema::event_log(LogEntry::ReconfigUpdate, LogEvent::VMUpdateFail).error(&e));
    }
}
<commit_msg>Check broadcast_within_validator_network() for transactions directly submitted to validators (#1957)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

//! Tasks that are executed by coordinators (short-lived compared to coordinators)
use crate::thread_pool::IO_POOL;
use crate::{
    core_mempool::{CoreMempool, TimelineState, TxnPointer},
    counters,
    logging::{LogEntry, LogEvent, LogSchema},
    network::{BroadcastError, MempoolSyncMsg},
    shared_mempool::types::{
        notify_subscribers, ScheduledBroadcast, SharedMempool, SharedMempoolNotification,
        SubmissionStatusBundle,
    },
    QuorumStoreRequest, QuorumStoreResponse, SubmissionStatus,
};
use anyhow::Result;
use aptos_config::network_id::PeerNetworkId;
use aptos_crypto::HashValue;
use aptos_infallible::{Mutex, RwLock};
use aptos_logger::prelude::*;
use aptos_metrics_core::HistogramTimer;
use aptos_types::{
    mempool_status::{MempoolStatus, MempoolStatusCode},
    on_chain_config::OnChainConfigPayload,
    transaction::SignedTransaction,
    vm_status::DiscardedVMStatus,
};
use consensus_types::common::TransactionSummary;
use futures::{channel::oneshot, stream::FuturesUnordered};
use network::application::interface::NetworkInterface;
use rayon::prelude::*;
use std::{
    cmp,
    collections::HashSet,
    sync::Arc,
    time::{Duration, Instant},
};
use storage_interface::state_view::LatestDbStateCheckpointView;
use tokio::runtime::Handle;
use vm_validator::vm_validator::{get_account_sequence_number, TransactionValidation};

// ============================== //
//  broadcast_coordinator tasks  //
// ============================== //

/// Attempts broadcast to `peer` and schedules the next broadcast.
pub(crate) async fn execute_broadcast<V>(
    peer: PeerNetworkId,
    backoff: bool,
    smp: &mut SharedMempool<V>,
    scheduled_broadcasts: &mut FuturesUnordered<ScheduledBroadcast>,
    executor: Handle,
) where
    V: TransactionValidation,
{
    let network_interface = &smp.network_interface.clone();
    // If there's no connection, don't bother to broadcast
    if network_interface.app_data().read(&peer).is_some() {
        if let Err(err) = network_interface
            .execute_broadcast(peer, backoff, smp)
            .await
        {
            match err {
                BroadcastError::NetworkError(peer, error) => error!(LogSchema::event_log(
                    LogEntry::BroadcastTransaction,
                    LogEvent::NetworkSendFail
                )
                .peer(&peer)
                .error(&error)),
                _ => {
                    trace!("{:?}", err)
                }
            }
        }
    } else {
        // Drop the scheduled broadcast, we're not connected anymore
        return;
    }
    let schedule_backoff = network_interface.is_backoff_mode(&peer);

    let interval_ms = if schedule_backoff {
        smp.config.shared_mempool_backoff_interval_ms
    } else {
        smp.config.shared_mempool_tick_interval_ms
    };

    scheduled_broadcasts.push(ScheduledBroadcast::new(
        Instant::now() + Duration::from_millis(interval_ms),
        peer,
        schedule_backoff,
        executor,
    ))
}

// =============================== //
// Tasks processing txn submission //
// =============================== //

/// Processes transactions directly submitted by client.
pub(crate) async fn process_client_transaction_submission<V>(
    smp: SharedMempool<V>,
    transaction: SignedTransaction,
    callback: oneshot::Sender<Result<SubmissionStatus>>,
    timer: HistogramTimer,
) where
    V: TransactionValidation,
{
    timer.stop_and_record();
    let _timer = counters::process_txn_submit_latency_timer_client();
    let ineligible_for_broadcast =
        !smp.broadcast_within_validator_network() && smp.network_interface.is_validator();
    let timeline_state = if ineligible_for_broadcast {
        TimelineState::NonQualified
    } else {
        TimelineState::NotReady
    };
    let statuses = process_incoming_transactions(&smp, vec![transaction], timeline_state);
    log_txn_process_results(&statuses, None);

    if let Some(status) = statuses.get(0) {
        if callback.send(Ok(status.1.clone())).is_err() {
            error!(LogSchema::event_log(
                LogEntry::JsonRpc,
                LogEvent::CallbackFail
            ));
            counters::CLIENT_CALLBACK_FAIL.inc();
        }
    }
}

/// Processes get transaction by hash request by client.
pub(crate) async fn process_client_get_transaction<V>(
    smp: SharedMempool<V>,
    hash: HashValue,
    callback: oneshot::Sender<Option<SignedTransaction>>,
    timer: HistogramTimer,
) where
    V: TransactionValidation,
{
    timer.stop_and_record();
    let _timer = counters::process_get_txn_latency_timer_client();
    let txn = smp.mempool.lock().get_by_hash(hash);

    if callback.send(txn).is_err() {
        error!(LogSchema::event_log(
            LogEntry::GetTransaction,
            LogEvent::CallbackFail
        ));
        counters::CLIENT_CALLBACK_FAIL.inc();
    }
}

/// Processes transactions from other nodes.
pub(crate) async fn process_transaction_broadcast<V>(
    smp: SharedMempool<V>,
    transactions: Vec<SignedTransaction>,
    request_id: Vec<u8>,
    timeline_state: TimelineState,
    peer: PeerNetworkId,
    timer: HistogramTimer,
) where
    V: TransactionValidation,
{
    timer.stop_and_record();
    let _timer = counters::process_txn_submit_latency_timer(peer.network_id());
    let results = process_incoming_transactions(&smp, transactions, timeline_state);
    log_txn_process_results(&results, Some(peer));

    let ack_response = gen_ack_response(request_id, results, &peer);
    let network_sender = smp.network_interface.sender();
    if let Err(e) = network_sender.send_to(peer, ack_response) {
        counters::network_send_fail_inc(counters::ACK_TXNS);
        error!(
            LogSchema::event_log(LogEntry::BroadcastACK, LogEvent::NetworkSendFail)
                .peer(&peer)
                .error(&e.into())
        );
        return;
    }
    notify_subscribers(SharedMempoolNotification::ACK, &smp.subscribers);
}

/// If `MempoolIsFull` on any of the transactions, provide backpressure to the downstream peer.
fn gen_ack_response(
    request_id: Vec<u8>,
    results: Vec<SubmissionStatusBundle>,
    peer: &PeerNetworkId,
) -> MempoolSyncMsg {
    let mut backoff_and_retry = false;
    for (_, (mempool_status, _)) in results.into_iter() {
        if mempool_status.code == MempoolStatusCode::MempoolIsFull {
            backoff_and_retry = true;
            break;
        }
    }

    update_ack_counter(
        peer,
        counters::SENT_LABEL,
        backoff_and_retry,
        backoff_and_retry,
    );
    MempoolSyncMsg::BroadcastTransactionsResponse {
        request_id,
        retry: backoff_and_retry,
        backoff: backoff_and_retry,
    }
}

pub(crate) fn update_ack_counter(
    peer: &PeerNetworkId,
    direction_label: &str,
    retry: bool,
    backoff: bool,
) {
    if retry {
        counters::shared_mempool_ack_inc(
            peer.network_id(),
            direction_label,
            counters::RETRY_BROADCAST_LABEL,
        );
    }
    if backoff {
        counters::shared_mempool_ack_inc(
            peer.network_id(),
            direction_label,
            counters::BACKPRESSURE_BROADCAST_LABEL,
        );
    }
}

/// Submits a list of SignedTransaction to the local mempool
/// and returns a vector containing AdmissionControlStatus.
pub(crate) fn process_incoming_transactions<V>(
    smp: &SharedMempool<V>,
    transactions: Vec<SignedTransaction>,
    timeline_state: TimelineState,
) -> Vec<SubmissionStatusBundle>
where
    V: TransactionValidation,
{
    let mut statuses = vec![];

    let start_storage_read = Instant::now();
    let state_view = smp
        .db
        .latest_state_checkpoint_view()
        .expect("Failed to get latest state checkpoint view.");

    // Track latency: fetching seq number
    let seq_numbers = IO_POOL.install(|| {
        transactions
            .par_iter()
            .map(|t| {
                get_account_sequence_number(&state_view, t.sender()).map_err(|e| {
                    error!(LogSchema::new(LogEntry::DBError).error(&e));
                    counters::DB_ERROR.inc();
                    e
                })
            })
            .collect::<Vec<_>>()
    });
    // Track latency for storage read fetching sequence number
    let storage_read_latency = start_storage_read.elapsed();
    counters::PROCESS_TXN_BREAKDOWN_LATENCY
        .with_label_values(&[counters::FETCH_SEQ_NUM_LABEL])
        .observe(storage_read_latency.as_secs_f64() / transactions.len() as f64);

    let transactions: Vec<_> = transactions
        .into_iter()
        .enumerate()
        .filter_map(|(idx, t)| {
            if let Ok(crsn_or_seqno) = seq_numbers[idx] {
                if t.sequence_number() >= crsn_or_seqno.min_seq() {
                    return Some((t, crsn_or_seqno));
                } else {
                    statuses.push((
                        t,
                        (
                            MempoolStatus::new(MempoolStatusCode::VmError),
                            Some(DiscardedVMStatus::SEQUENCE_NUMBER_TOO_OLD),
                        ),
                    ));
                }
            } else {
                // Failed to get transaction
                statuses.push((
                    t,
                    (
                        MempoolStatus::new(MempoolStatusCode::VmError),
                        Some(DiscardedVMStatus::RESOURCE_DOES_NOT_EXIST),
                    ),
                ));
            }
            None
        })
        .collect();

    // Track latency: VM validation
    let vm_validation_timer = counters::PROCESS_TXN_BREAKDOWN_LATENCY
        .with_label_values(&[counters::VM_VALIDATION_LABEL])
        .start_timer();
    let validation_results = transactions
        .iter()
        .map(|t| smp.validator.read().validate_transaction(t.0.clone()))
        .collect::<Vec<_>>();
    vm_validation_timer.stop_and_record();
    {
        let mut mempool = smp.mempool.lock();
        for (idx, (transaction, crsn_or_seqno)) in transactions.into_iter().enumerate() {
            if let Ok(validation_result) = &validation_results[idx] {
                match validation_result.status() {
                    None => {
                        let ranking_score = validation_result.score();
                        let mempool_status = mempool.add_txn(
                            transaction.clone(),
                            ranking_score,
                            crsn_or_seqno,
                            timeline_state,
                        );
                        statuses.push((transaction, (mempool_status, None)));
                    }
                    Some(validation_status) => {
                        statuses.push((
                            transaction.clone(),
                            (
                                MempoolStatus::new(MempoolStatusCode::VmError),
                                Some(validation_status),
                            ),
                        ));
                    }
                }
            }
        }
    }
    notify_subscribers(SharedMempoolNotification::NewTransactions, &smp.subscribers);
    statuses
}

fn log_txn_process_results(results: &[SubmissionStatusBundle], sender: Option<PeerNetworkId>) {
    let network = match sender {
        Some(peer) => peer.network_id().to_string(),
        None => counters::CLIENT_LABEL.to_string(),
    };
    for (txn, (mempool_status, maybe_vm_status)) in results.iter() {
        if let Some(vm_status) = maybe_vm_status {
            trace!(
                SecurityEvent::InvalidTransactionMempool,
                failed_transaction = txn,
                vm_status = vm_status,
                sender = sender,
            );
            counters::shared_mempool_transactions_processed_inc(
                counters::VM_VALIDATION_LABEL,
                &network,
            );
            continue;
        }
        match mempool_status.code {
            MempoolStatusCode::Accepted => counters::shared_mempool_transactions_processed_inc(
                counters::SUCCESS_LABEL,
                &network,
            ),
            _ => counters::shared_mempool_transactions_processed_inc(
                &mempool_status.code.to_string(),
                &network,
            ),
        }
    }
}

// ================================= //
// intra-node communication handlers //
// ================================= //

/// Only applies to Validators. Either provides transactions to consensus [`GetBlockRequest`] or
/// handles rejecting transactions [`RejectNotification`]
pub(crate) fn process_quorum_store_request<V: TransactionValidation>(
    smp: &SharedMempool<V>,
    req: QuorumStoreRequest,
) {
    // Start latency timer
    let start_time = Instant::now();
    debug!(LogSchema::event_log(LogEntry::QuorumStore, LogEvent::Received).quorum_store_msg(&req));

    let (resp, callback, counter_label) = match req {
        QuorumStoreRequest::GetBatchRequest(max_batch_size, transactions, callback) => {
            let exclude_transactions: HashSet<TxnPointer> = transactions
                .iter()
                .map(|txn| (txn.sender, txn.sequence_number))
                .collect();
            let txns;
            {
                let mut mempool = smp.mempool.lock();
                // gc before pulling block as extra protection against txns that may expire in consensus
                // Note: this gc operation relies on the fact that consensus uses the system time to determine block timestamp
                let curr_time = aptos_infallible::duration_since_epoch();
                mempool.gc_by_expiration_time(curr_time);
                let batch_size = cmp::max(max_batch_size, 1);
                txns = mempool.get_batch(batch_size, exclude_transactions);
            }
            counters::mempool_service_transactions(counters::GET_BLOCK_LABEL, txns.len());

            (
                QuorumStoreResponse::GetBatchResponse(txns),
                callback,
                counters::GET_BLOCK_LABEL,
            )
        }
        QuorumStoreRequest::RejectNotification(transactions, callback) => {
            counters::mempool_service_transactions(
                counters::COMMIT_CONSENSUS_LABEL,
                transactions.len(),
            );
            process_committed_transactions(&smp.mempool, transactions, 0, true);
            (
                QuorumStoreResponse::CommitResponse(),
                callback,
                counters::COMMIT_CONSENSUS_LABEL,
            )
        }
    };
    // Send back to callback
    let result = if callback.send(Ok(resp)).is_err() {
        error!(LogSchema::event_log(
            LogEntry::QuorumStore,
            LogEvent::CallbackFail
        ));
        counters::REQUEST_FAIL_LABEL
    } else {
        counters::REQUEST_SUCCESS_LABEL
    };
    let latency = start_time.elapsed();
    counters::mempool_service_latency(counter_label, result, latency);
}

/// Remove transactions that are committed (or rejected) so that we can stop broadcasting them.
pub(crate) fn process_committed_transactions(
    mempool: &Mutex<CoreMempool>,
    transactions: Vec<TransactionSummary>,
    block_timestamp_usecs: u64,
    is_rejected: bool,
) {
    let mut pool = mempool.lock();

    for transaction in transactions {
        pool.remove_transaction(
            &transaction.sender,
            transaction.sequence_number,
            is_rejected,
        );
    }

    if block_timestamp_usecs > 0 {
        pool.gc_by_expiration_time(Duration::from_micros(block_timestamp_usecs));
    }
}

/// Processes on-chain reconfiguration notifications.  Restarts validator with the new info.
pub(crate) async fn process_config_update<V>(
    config_update: OnChainConfigPayload,
    validator: Arc<RwLock<V>>,
) where
    V: TransactionValidation,
{
    info!(
        LogSchema::event_log(LogEntry::ReconfigUpdate, LogEvent::Process)
            .reconfig_update(config_update.clone())
    );

    if let Err(e) = validator.write().restart(config_update) {
        counters::VM_RECONFIG_UPDATE_FAIL_COUNT.inc();
        error!(LogSchema::event_log(LogEntry::ReconfigUpdate, LogEvent::VMUpdateFail).error(&e));
    }
}
<|endoftext|>


<commit_before>mod parsers;
mod parsing_types;
pub mod preprocessing;
mod utils;

const LINESEP: &str = "\n";
const LATEX_TO_INDENTEX_FACTOR: f64 = 1.5;
const DO_NOT_EDIT_NOTICE: &str =
    "\
     % ============================================================== %\n\
     %                                                                %\n\
     % THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY. %\n\
     %                                                                %\n\
     % ============================================================== %\n";

pub struct TranspileOptions {
    pub flatten_output: bool,
    pub prepend_do_not_edit_notice: bool,
}

pub fn transpile(mut lines: Vec<String>, options: &TranspileOptions) -> String {
    use crate::parsers::process_line;
    use crate::parsing_types::Environment;
    use crate::parsing_types::Hashline::{OpenEnv, PlainLine};

    // The number of environments is not known beforehand
    let mut env_stack: Vec<Environment> = Vec::new();

    // Input size is the sum of all line lengths plus the number of lines (for lineseps)
    let input_size = lines.iter().fold(0, |sum, l| sum + l.len()) + lines.len();
    // We do not know how much larger the transpiled LaTeX file will be, but we can guess...
    let indentex_size = (LATEX_TO_INDENTEX_FACTOR * (input_size as f64)).round() as usize;
    let mut transpiled = String::with_capacity(indentex_size);

    let adjusted_indents = crate::utils::scan_indents(lines.as_ref());

    if options.prepend_do_not_edit_notice {
        transpiled.push_str(DO_NOT_EDIT_NOTICE);
    }

    for (line_num, line) in lines.drain(..).enumerate() {
        let list_like_active = match env_stack.last() {
            None => false, // No environment is active at all
            Some(d) => d.is_list_like(),
        };

        let tl = match process_line(line, list_like_active) {
            PlainLine(l) => l,
            OpenEnv(e) => {
                let tag_begin = e.latex_begin();
                env_stack.push(e);
                tag_begin
            }
        };
        if options.flatten_output {
            transpiled.push_str(tl.trim_start());
        } else {
            transpiled.push_str(&tl);
        }
        transpiled.push_str(LINESEP);

        // Check if we are in an environment and close as many as needed
        while match env_stack.last() {
            None => false,
            Some(d) => d.indent_depth() >= adjusted_indents[line_num + 1],
        } {
            // `unwrap()` is safe here since we have already checked if the stack is empty
            let tag_end = env_stack.pop().unwrap().latex_end();
            if options.flatten_output {
                transpiled.push_str(tag_end.trim_start());
            } else {
                transpiled.push_str(&tag_end);
            }
            transpiled.push_str(LINESEP);
        }
    }

    transpiled
}
<commit_msg>Add unit tests for the transpilation fn<commit_after>mod parsers;
mod parsing_types;
pub mod preprocessing;
mod utils;

const LINESEP: &str = "\n";
const LATEX_TO_INDENTEX_FACTOR: f64 = 1.5;
const DO_NOT_EDIT_NOTICE: &str =
    "\
     % ============================================================== %\n\
     %                                                                %\n\
     % THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY. %\n\
     %                                                                %\n\
     % ============================================================== %\n";

pub struct TranspileOptions {
    pub flatten_output: bool,
    pub prepend_do_not_edit_notice: bool,
}

pub fn transpile(mut lines: Vec<String>, options: &TranspileOptions) -> String {
    use crate::parsers::process_line;
    use crate::parsing_types::Environment;
    use crate::parsing_types::Hashline::{OpenEnv, PlainLine};

    // The number of environments is not known beforehand
    let mut env_stack: Vec<Environment> = Vec::new();

    // Input size is the sum of all line lengths plus the number of lines (for lineseps)
    let input_size = lines.iter().fold(0, |sum, l| sum + l.len()) + lines.len();
    // We do not know how much larger the transpiled LaTeX file will be, but we can guess...
    let indentex_size = (LATEX_TO_INDENTEX_FACTOR * (input_size as f64)).round() as usize;
    let mut transpiled = String::with_capacity(indentex_size);

    let adjusted_indents = crate::utils::scan_indents(lines.as_ref());

    if options.prepend_do_not_edit_notice {
        transpiled.push_str(DO_NOT_EDIT_NOTICE);
    }

    for (line_num, line) in lines.drain(..).enumerate() {
        let list_like_active = match env_stack.last() {
            None => false, // No environment is active at all
            Some(d) => d.is_list_like(),
        };

        let tl = match process_line(line, list_like_active) {
            PlainLine(l) => l,
            OpenEnv(e) => {
                let tag_begin = e.latex_begin();
                env_stack.push(e);
                tag_begin
            }
        };
        if options.flatten_output {
            transpiled.push_str(tl.trim_start());
        } else {
            transpiled.push_str(&tl);
        }
        transpiled.push_str(LINESEP);

        // Check if we are in an environment and close as many as needed
        while match env_stack.last() {
            None => false,
            Some(d) => d.indent_depth() >= adjusted_indents[line_num + 1],
        } {
            // `unwrap()` is safe here since we have already checked if the stack is empty
            let tag_end = env_stack.pop().unwrap().latex_end();
            if options.flatten_output {
                transpiled.push_str(tag_end.trim_start());
            } else {
                transpiled.push_str(&tag_end);
            }
            transpiled.push_str(LINESEP);
        }
    }

    transpiled
}

// LCOV_EXCL_START
#[cfg(test)]
mod tests {
    #[cfg(test)]
    mod transpilation_spec {
        use super::super::{transpile, TranspileOptions};

        #[test]
        fn vanilla_latex_is_not_modified() {
            let input = (vec![
                r"\newcommand{\foo}[1]{foo #1}",
                r"  \begin{equation}1+1\end{equation}",
                r" $ 1 + \frac{1}{2}$",
            ])
            .into_iter()
            .map(|s| s.to_string())
            .collect();

            let expected = "\
                            \\newcommand{\\foo}[1]{foo #1}\
                            \n  \\begin{equation}1+1\\end{equation}\
                            \n $ 1 + \\frac{1}{2}$\n";
            let to = TranspileOptions {
                flatten_output: false,
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn single_line_commands() {
            let input = (vec!["# foo: bar % qux", "  # foo[opts]: bar"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\foo{bar} % qux\n  \\foo[opts]{bar}\n";
            let to = TranspileOptions {
                flatten_output: false,
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly() {
            let input = (vec!["# foo:", "  bar"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n  bar\n\\end{foo}\n";
            let to = TranspileOptions {
                flatten_output: false,
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly_if_empty_body() {
            let input = (vec!["# foo:", "  # bar:"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n  \\begin{bar}\n  \\end{bar}\n\\end{foo}\n";
            let to = TranspileOptions {
                flatten_output: false,
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn open_and_close_envs_correctly_with_flattened_output() {
            let input = (vec!["# foo:", "  # bar:"])
                .into_iter()
                .map(|s| s.to_string())
                .collect();

            let expected = "\\begin{foo}\n\\begin{bar}\n\\end{bar}\n\\end{foo}\n";
            let to = TranspileOptions {
                flatten_output: true,
                prepend_do_not_edit_notice: false,
            };
            assert_eq!(&transpile(input, &to), expected);
        }

        #[test]
        fn do_not_edit_notice() {
            let to = TranspileOptions {
                flatten_output: false,
                prepend_do_not_edit_notice: true,
            };
            assert_eq!(
                &transpile(vec!["".to_string()], &to),
                &format!("{}\n", super::super::DO_NOT_EDIT_NOTICE)
            );
        }
    }
}
// LCOV_EXCL_STOP
<|endoftext|>