<commit_before>use std::collections::{HashMap,BTreeSet};

use source::Src;
use var::Var;
use eval::Eval;

#[derive(Debug,PartialEq)]
pub struct SrcBlock {
    pub name: String,
    pub src: Vec<Src>,
    pub await_idx: usize,
    pub visited: bool,
}

#[derive(Debug,PartialEq)]
pub struct DefBlock {
    pub name: String,
    pub def: HashMap<String,Var>
}

#[derive(Debug,PartialEq)]
pub enum Block {
    Src(SrcBlock),
    Def(DefBlock),
}

/// Intermediate Representation
///
/// This is used during the parsing stage
#[derive(Debug,Clone,PartialEq)]
pub enum IR {
    String(String),
    /// A non-quoted string turns into a symbol/token
    Sym(String),
    /// Key-Value, pre-parsed
    Map(Vec<IR>),
}

impl From<IR> for String {
    fn from(t:IR) -> String {
        match t {
            IR::String(s) => s,
            IR::Sym(s) => s,
            IR::Map(mut v) => {
                let mut s = "{".to_owned();
                for n in v.drain(..) {
                    let n:String = n.into();
                    s.push_str(&n);
                }

                s.push('}');

                s
            }
        }
    }
}
impl From<Var> for IR {
    fn from(t:Var) -> IR {
        match t {
            Var::String(t) => { IR::String(t) }
            _ => { IR::Sym(t.to_string()) }
        }
    }
}

/// Map object for Selects
pub type Map = HashMap<String,Vec<Var>>;

pub struct Parser(Vec<Block>);

use std::ops::Deref;
impl Deref for Parser {
    type Target = Vec<Block>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

impl Parser {
    pub fn parse_blocks (src: &str) -> Parser {
        let mut v = vec!();
        let mut exp = String::new();
        let mut exps: Vec<IR> = vec!();
        let mut map_ir: Vec<IR> = vec!(); //contains pre-parsed map
        let mut block: Option<Block> = None;

        let mut in_string = false;
        let mut in_comment = false;
        let mut in_vec = false;
        let mut in_map = false;
        let mut was_if = false;

        let mut usyms = BTreeSet::new(); //unique set, remove dupes

        for c in src.chars() {
            if !in_comment && !in_string {
                if c == '[' { in_vec = true; continue }
                else if c == ']' { in_vec = false; }

            }

            if c == '#' && !in_string { in_comment = true; continue }
            else if  c == '\n' && in_comment && !in_string { in_comment = false; }

            if c == '\n' && (in_vec || in_map) && !in_comment && !in_string { continue }

            if (c == ']' ||
                c == '}' ||
                c == '#' ||
                c == '\n')
                && !in_string && !in_comment
            {
                for n in exp.split_whitespace() {
                    let sym = IR::Sym(n.trim().to_owned());
                    if in_map { map_ir.push(sym); }
                    else { exps.push(sym); }
                }

                if c == '}' && in_map && !in_comment{
                    in_map = false;
                    exps.push(IR::Map(map_ir));
                    map_ir = vec![];
                }

                exp = String::new();

                if exps.len() < 1 { continue }

                // determine block type
                if block.is_none() {
                    let name = exps.remove(0).into();
                    if name == "def" {
                        let b = DefBlock {
                            name: exps.pop().unwrap().into(),
                            def: HashMap::new(),
                        };

                        block = Some(Block::Def(b));
                    }
                    else {
                        let b = SrcBlock {
                            name: name,
                            src: vec!(),
                            await_idx: 0,
                            visited: false,
                        };

                        block = Some(Block::Src(b));
                    }

                }
                else { // build block type
                    let mut qsyms:Vec<(String,String)> = vec!();
                    let adjust_sym = |qsyms: &mut Vec<(String,String)>, s: &mut String| {
                        if s.chars().next().expect("ERROR: Empty Eager Symbol") == '!' {
                            let mut sym = "not_".to_owned();
                            sym.push_str(s[1..].trim());

                            let osym = s.trim().to_owned();

                            qsyms.push((sym.clone(),osym));
                            *s = sym;
                        }
                    };

                    // this builds symbol refs as a convenience
                    if exps.len() > 2 { //must be if/or/comp blocks
                        for n in exps.iter_mut() {
                            match n {
                                &mut IR::Sym(ref mut s) => {
                                    adjust_sym(&mut qsyms,s);
                                },
                                &mut IR::Map(ref mut v) => {
                                    for n in v.iter_mut() {
                                        match n {
                                            &mut IR::Sym(ref mut s) => {
                                                adjust_sym(&mut qsyms,s);
                                            },
                                            _ => {},
                                        }
                                    }
                                },
                                _ => {},
                            }
                        }
                    }

                    match block {
                        Some(Block::Def(ref mut b)) => {
                            let v = exps.pop().unwrap();
                            b.def.insert(exps.pop().unwrap().into(),
                                         Var::parse(v));
                        },
                        Some(Block::Src(ref mut b)) => {
                            let mut srcs = vec![];

                            for (qsym,sym) in qsyms.drain(..) {
                                if usyms.contains(&qsym) { continue }
                                usyms.insert(qsym.clone());

                                let src = Src::parse(vec![IR::Sym(qsym),
                                                          IR::Sym(sym)]);
                                srcs.push(src);
                            }

                            let src = Src::parse(exps);
                            srcs.push(src);

                            for src in srcs.drain(..) {
                                match &src {
                                    &Src::If(_,_,_) => { was_if = true; },
                                    &Src::Or(_,_) => {
                                        if !was_if {
                                            panic!("ERROR: IF must prepend OR");
                                        }
                                    },
                                    _ => { was_if = false; },
                                }

                                b.src.push(src);
                            }
                        },
                        _ => {}
                    }

                    exps = vec!();
                }
            }
            else if c == '"' && !in_comment {
                in_string = !in_string;
                if in_string { //starting a new quoted string? let's push this sym
                    for n in exp.split_whitespace() {
                        let sym = IR::Sym(n.trim().to_owned());
                        if in_map { map_ir.push(sym); }
                        else { exps.push(sym); }
                    }
                    exp = String::new();
                }
                else if !in_string { //finished the quoted string?
                    let sym = IR::String(exp);
                    if in_map { map_ir.push(sym); }
                    else { exps.push(sym); }
                    exp = String::new();
                }
            }
            else if c == ';' && !in_string && !in_comment {
                //fail otherwise, block should be built!
                v.push(block.expect("ERROR: Parse Block no built!"));
                usyms.clear(); //clear out on new block
                block = None;
            }
            else {
                if c == '{' && !in_comment && !in_string {
                    in_map = true;
                    // push previous symbol
                    let sym = IR::Sym(exp.trim().to_owned());
                    exps.push(sym);
                    exp.clear();
                }
                else if !in_comment {
                    exp.push(c);
                }
            }
        }

        Parser(v)
    }

    pub fn into_env (mut self) -> Env {
        let mut src = HashMap::new();
        let mut def = HashMap::new();

        for b in self.0.drain(..) {
            match b {
                Block::Def(db) => {
                    def.insert(db.name.clone(), db);
                },
                Block::Src(sb) => {
                    src.insert(sb.name.clone(), sb);
                },
            }

        }

        Env { def: def, src: src }
    }

    /// Parses a map from IR
    ///
    /// Parsed using commas for variable sized maps
    pub fn parse_map (map_ir: IR) -> Option<Map> {
        let mut map: Map = HashMap::new(); // optionally unbounded val-lengths

        match map_ir {
            IR::Map(mut exps) => {
                let mut key = "".to_owned();
                let mut vals = vec![];

                for n in exps.drain(..) {
                    if key.is_empty() { key = n.into(); continue }

                    match n {
                        IR::Sym(mut s) => {
                            if s.chars().last() == Some(',') {
                                let _ = s.pop();
                                vals.push(Var::parse(IR::Sym(s)));

                                map.insert(key,vals);
                                vals = vec![];
                                key = "".to_owned();

                                continue
                            }

                            vals.push(Var::parse(IR::Sym(s)));
                        },
                        _ => { vals.push(Var::parse(n)); },
                    }
                }

                if !key.is_empty() && vals.len() > 0 {
                    map.insert(key,vals);
                }
                else if !key.is_empty() {
                    panic!("ERROR: Unbalanced MAP at: {:?}",key);
                }

                Some(map)
            },
            _=> {None}
        }
    }
}

/// Def alias used for internal evaluation purposes
pub type Def = HashMap<String, DefBlock>;

impl Env {
    pub fn def_contains(def: &Def, path: Option<Vec<&str>>, lookup: &str) -> bool {
        if let Some(path) = path {
            if let Some(ref def) = def.get(path[0]) {
                return def.def.contains_key(lookup)
            }
        }

        false
    }
}

/// Environment containing all parsed definition and source blocks
pub struct Env {
    pub def: Def,
    pub src: HashMap<String, SrcBlock>
}

impl Eval for Def {
    fn get (&self, path: Option<Vec<&str>>, lookup: &str) -> Option<Var> {
        if let Some(path) = path {
            if let Some(ref def) = self.get(path[0]) {
                if let Some(v) = def.def.get(lookup) {
                    return Some(v.clone())
                }
            }
        }

        None
    }

    #[allow(unused_variables)]
    fn set (&mut self, path: Option<Vec<&str>>, lookup: &str, var: Var) {
        if let Some(path) = path {
            if let Some(ref mut def) = self.get_mut(path[0]) {
                if let Some(v) = def.def.get_mut(lookup) {
                    *v = var;
                }
            }
        }
    }

    #[allow(unused_variables)]
    fn call (&mut self, var: Var, fun: &str, vars: &Vec<Var>) -> Option<Var> {
        None
    }
}
<commit_msg>begin stream parser<commit_after>use std::collections::{HashMap,BTreeSet};
use std::io::prelude::*;

use source::Src;
use var::Var;
use eval::Eval;

#[derive(Debug,PartialEq)]
pub struct SrcBlock {
    pub name: String,
    pub src: Vec<Src>,
    pub await_idx: usize,
    pub visited: bool,
}

#[derive(Debug,PartialEq)]
pub struct DefBlock {
    pub name: String,
    pub def: HashMap<String,Var>
}

#[derive(Debug,PartialEq)]
pub enum Block {
    Src(SrcBlock),
    Def(DefBlock),
}

/// Intermediate Representation
///
/// This is used during the parsing stage
#[derive(Debug,Clone,PartialEq)]
pub enum IR {
    String(String),
    /// A non-quoted string turns into a symbol/token
    Sym(String),
    /// Key-Value, pre-parsed
    Map(Vec<IR>),
}

impl From<IR> for String {
    fn from(t:IR) -> String {
        match t {
            IR::String(s) => s,
            IR::Sym(s) => s,
            IR::Map(mut v) => {
                let mut s = "{".to_owned();
                for n in v.drain(..) {
                    let n:String = n.into();
                    s.push_str(&n);
                }

                s.push('}');

                s
            }
        }
    }
}
impl From<Var> for IR {
    fn from(t:Var) -> IR {
        match t {
            Var::String(t) => { IR::String(t) }
            _ => { IR::Sym(t.to_string()) }
        }
    }
}

/// Map object for Selects
pub type Map = HashMap<String,Vec<Var>>;

pub struct Parser(Vec<Block>);

use std::ops::Deref;
impl Deref for Parser {
    type Target = Vec<Block>;
    fn deref(&self) -> &Self::Target { &self.0 }
}

impl Parser {
    pub fn parse_blocks (src: &str) -> Parser {
        let mut v = vec!();
        let mut exp = String::new();
        let mut exps: Vec<IR> = vec!();
        let mut map_ir: Vec<IR> = vec!(); //contains pre-parsed map
        let mut block: Option<Block> = None;

        let mut in_string = false;
        let mut in_comment = false;
        let mut in_vec = false;
        let mut in_map = false;
        let mut was_if = false;

        let mut usyms = BTreeSet::new(); //unique set, remove dupes

        for c in src.chars() {
            if !in_comment && !in_string {
                if c == '[' { in_vec = true; continue }
                else if c == ']' { in_vec = false; }

            }

            if c == '#' && !in_string { in_comment = true; continue }
            else if  c == '\n' && in_comment && !in_string { in_comment = false; }

            if c == '\n' && (in_vec || in_map) && !in_comment && !in_string { continue }

            if (c == ']' ||
                c == '}' ||
                c == '#' ||
                c == '\n')
                && !in_string && !in_comment
            {
                for n in exp.split_whitespace() {
                    let sym = IR::Sym(n.trim().to_owned());
                    if in_map { map_ir.push(sym); }
                    else { exps.push(sym); }
                }

                if c == '}' && in_map && !in_comment{
                    in_map = false;
                    exps.push(IR::Map(map_ir));
                    map_ir = vec![];
                }

                exp = String::new();

                if exps.len() < 1 { continue }

                // determine block type
                if block.is_none() {
                    let name = exps.remove(0).into();
                    if name == "def" {
                        let b = DefBlock {
                            name: exps.pop().unwrap().into(),
                            def: HashMap::new(),
                        };

                        block = Some(Block::Def(b));
                    }
                    else {
                        let b = SrcBlock {
                            name: name,
                            src: vec!(),
                            await_idx: 0,
                            visited: false,
                        };

                        block = Some(Block::Src(b));
                    }

                }
                else { // build block type
                    let mut qsyms:Vec<(String,String)> = vec!();
                    let adjust_sym = |qsyms: &mut Vec<(String,String)>, s: &mut String| {
                        if s.chars().next().expect("ERROR: Empty Eager Symbol") == '!' {
                            let mut sym = "not_".to_owned();
                            sym.push_str(s[1..].trim());

                            let osym = s.trim().to_owned();

                            qsyms.push((sym.clone(),osym));
                            *s = sym;
                        }
                    };

                    // this builds symbol refs as a convenience
                    if exps.len() > 2 { //must be if/or/comp blocks
                        for n in exps.iter_mut() {
                            match n {
                                &mut IR::Sym(ref mut s) => {
                                    adjust_sym(&mut qsyms,s);
                                },
                                &mut IR::Map(ref mut v) => {
                                    for n in v.iter_mut() {
                                        match n {
                                            &mut IR::Sym(ref mut s) => {
                                                adjust_sym(&mut qsyms,s);
                                            },
                                            _ => {},
                                        }
                                    }
                                },
                                _ => {},
                            }
                        }
                    }

                    match block {
                        Some(Block::Def(ref mut b)) => {
                            let v = exps.pop().unwrap();
                            b.def.insert(exps.pop().unwrap().into(),
                                         Var::parse(v));
                        },
                        Some(Block::Src(ref mut b)) => {
                            let mut srcs = vec![];

                            for (qsym,sym) in qsyms.drain(..) {
                                if usyms.contains(&qsym) { continue }
                                usyms.insert(qsym.clone());

                                let src = Src::parse(vec![IR::Sym(qsym),
                                                          IR::Sym(sym)]);
                                srcs.push(src);
                            }

                            let src = Src::parse(exps);
                            srcs.push(src);

                            for src in srcs.drain(..) {
                                match &src {
                                    &Src::If(_,_,_) => { was_if = true; },
                                    &Src::Or(_,_) => {
                                        if !was_if {
                                            panic!("ERROR: IF must prepend OR");
                                        }
                                    },
                                    _ => { was_if = false; },
                                }

                                b.src.push(src);
                            }
                        },
                        _ => {}
                    }

                    exps = vec!();
                }
            }
            else if c == '"' && !in_comment {
                in_string = !in_string;
                if in_string { //starting a new quoted string? let's push this sym
                    for n in exp.split_whitespace() {
                        let sym = IR::Sym(n.trim().to_owned());
                        if in_map { map_ir.push(sym); }
                        else { exps.push(sym); }
                    }
                    exp = String::new();
                }
                else if !in_string { //finished the quoted string?
                    let sym = IR::String(exp);
                    if in_map { map_ir.push(sym); }
                    else { exps.push(sym); }
                    exp = String::new();
                }
            }
            else if c == ';' && !in_string && !in_comment {
                //fail otherwise, block should be built!
                v.push(block.expect("ERROR: Parse Block no built!"));
                usyms.clear(); //clear out on new block
                block = None;
            }
            else {
                if c == '{' && !in_comment && !in_string {
                    in_map = true;
                    // push previous symbol
                    let sym = IR::Sym(exp.trim().to_owned());
                    exps.push(sym);
                    exp.clear();
                }
                else if !in_comment {
                    exp.push(c);
                }
            }
        }

        Parser(v)
    }

    pub fn into_env (mut self) -> Env {
        let mut src = HashMap::new();
        let mut def = HashMap::new();

        for b in self.0.drain(..) {
            match b {
                Block::Def(db) => {
                    def.insert(db.name.clone(), db);
                },
                Block::Src(sb) => {
                    src.insert(sb.name.clone(), sb);
                },
            }

        }

        Env { def: def, src: src }
    }

    /// Parses a map from IR
    ///
    /// Parsed using commas for variable sized maps
    pub fn parse_map (map_ir: IR) -> Option<Map> {
        let mut map: Map = HashMap::new(); // optionally unbounded val-lengths

        match map_ir {
            IR::Map(mut exps) => {
                let mut key = "".to_owned();
                let mut vals = vec![];

                for n in exps.drain(..) {
                    if key.is_empty() { key = n.into(); continue }

                    match n {
                        IR::Sym(mut s) => {
                            if s.chars().last() == Some(',') {
                                let _ = s.pop();
                                vals.push(Var::parse(IR::Sym(s)));

                                map.insert(key,vals);
                                vals = vec![];
                                key = "".to_owned();

                                continue
                            }

                            vals.push(Var::parse(IR::Sym(s)));
                        },
                        _ => { vals.push(Var::parse(n)); },
                    }
                }

                if !key.is_empty() && vals.len() > 0 {
                    map.insert(key,vals);
                }
                else if !key.is_empty() {
                    panic!("ERROR: Unbalanced MAP at: {:?}",key);
                }

                Some(map)
            },
            _=> {None}
        }
    }
}

/// Def alias used for internal evaluation purposes
pub type Def = HashMap<String, DefBlock>;

impl Env {
    pub fn def_contains(def: &Def, path: Option<Vec<&str>>, lookup: &str) -> bool {
        if let Some(path) = path {
            if let Some(ref def) = def.get(path[0]) {
                return def.def.contains_key(lookup)
            }
        }

        false
    }
}

/// Environment containing all parsed definition and source blocks
pub struct Env {
    pub def: Def,
    pub src: HashMap<String, SrcBlock>
}

impl Eval for Def {
    fn get (&self, path: Option<Vec<&str>>, lookup: &str) -> Option<Var> {
        if let Some(path) = path {
            if let Some(ref def) = self.get(path[0]) {
                if let Some(v) = def.def.get(lookup) {
                    return Some(v.clone())
                }
            }
        }

        None
    }

    #[allow(unused_variables)]
    fn set (&mut self, path: Option<Vec<&str>>, lookup: &str, var: Var) {
        if let Some(path) = path {
            if let Some(ref mut def) = self.get_mut(path[0]) {
                if let Some(v) = def.def.get_mut(lookup) {
                    *v = var;
                }
            }
        }
    }

    #[allow(unused_variables)]
    fn call (&mut self, var: Var, fun: &str, vars: &Vec<Var>) -> Option<Var> {
        None
    }
}

pub struct StreamParser<'a, S:Read> {
    /// The non-parsed leftovers of a stream that is being buffered actively
    buf: Vec<&'a str>,
    stream: S,
}

impl<'a, S:Read> StreamParser<'a, S> {
    pub fn new (s: S) -> StreamParser<'a, S> {
        StreamParser {
            buf: vec![],
            stream: s,
        }
    }

}
<|endoftext|>


<commit_before>use ffi::h5i::hid_t;
use ffi::h5p::H5Pclose;

use object::{Handle, Object};

#[derive(Clone)]
pub struct PropertyList {
    handle: Handle,
}

impl Object for PropertyList {
    fn id(&self) -> hid_t {
        self.handle.id()
    }

    fn from_id(id: hid_t) -> PropertyList {
        PropertyList { handle: Handle::new(id) }
    }
}

impl Drop for PropertyList {
    fn drop(&mut self) {
        if self.refcount() == 1 {
            h5lock!(H5Pclose(self.id()));
        }
    }
}
<commit_msg>Remove Drop impl for PropertyList<commit_after>use ffi::h5i::hid_t;

use object::{Handle, Object};

#[derive(Clone)]
pub struct PropertyList {
    handle: Handle,
}

impl Object for PropertyList {
    fn id(&self) -> hid_t {
        self.handle.id()
    }

    fn from_id(id: hid_t) -> PropertyList {
        PropertyList { handle: Handle::new(id) }
    }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};
use log::{debug, trace};
use goblin::{Object};
use failure::{Error, Fail};
use std::ops::Range;

use super::super::super::arch::{RVA};
use super::super::super::loader::Permissions;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};
use std::collections::{HashSet};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
    #[fail(display = "Relocation target not found in image")]
    InvalidTargetAddress,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

// TODO: make this much faster
fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1, Permissions::R);
        }
    }
    return false;
}

fn is_zero(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(v) = ws.read_u32(rva) {
        return v == 0;
    }
    return false;
}

pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

fn sections_contain(sections: &[Range<RVA>], rva: RVA) -> bool {
    sections.iter()
        .find(|section| section.contains(&rva))
        .is_some()
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    /// lancelot::test::init_logging();
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    ///
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    ///
    /// assert!(ws.get_meta(RVA(0xC7F0)).unwrap().is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let x_sections: Vec<Range<RVA>> = ws.module.sections.iter()
            .filter(|section| section.perms.intersects(Permissions::X))
            .map(|section| Range {
                start: section.addr,
                end: section.end(),
            })
            .collect();

        let relocs: Vec<Reloc> = get_relocs(ws)?
            .into_iter()
            .filter(|r| match r.typ {
                RelocationType::ImageRelBasedHighLow => true,
                RelocationType::ImageRelBasedDir64 => true,
                _ => false,  // all other unsupported
            })
            .collect();
        debug!("found {} total relocs", relocs.len());

        for reloc in relocs.iter() {
            let target = ws.rva(ws.read_va(reloc.offset)?).ok_or(RelocAnalyzerError::InvalidTargetAddress)?;
            trace!("found reloc {} -> {} {}",
                   reloc.offset,
                   target,
                if sections_contain(&x_sections, target) {"(code)"} else {""}
            );
        }

        // scan for relocations to code.
        //
        // a relocation is a hardcoded offset that must be fixed up if the desired base address
        //  cannot be used.
        // for example, the function pointer passed to CreateThread will be a hardcoded address
        //  of the start of the function.
        //
        // relocated pointers may point to the .data section, e.g. strings or other constants.
        // we want to ignore these.
        // we are only interested in targets in executable sections.
        // we assume these are pointers to instructions/code.
        //
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.

        let mut unique_targets: HashSet<RVA> = HashSet::new();
        unique_targets.extend(relocs
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va)));

        debug!("reduced to {} unique reloc targets", unique_targets.len());

       let o: Vec<RVA> = unique_targets
            .iter()
            .filter(|&&rva| sections_contain(&x_sections, rva))
            .filter(|&&rva| !is_in_insn(ws, rva))
            .filter(|&&rva| !is_ptr(ws, rva))
            .filter(|&&rva| !is_zero(ws, rva))
            .map(|&rva| rva)
            // TODO: maybe ensure that the insn decodes.
            .collect();

        debug!("found {} relocs that point to instructions", o.len());

        o.iter().for_each(|&rva| {
            debug!("found pointer via relocations from executable section to {} (code)", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<commit_msg>reloc: warn when encountering unsupported reloc type<commit_after>use byteorder::{ByteOrder, LittleEndian};
use log::{debug, trace, warn};
use goblin::{Object};
use failure::{Error, Fail};
use std::ops::Range;

use super::super::super::arch::{RVA};
use super::super::super::loader::Permissions;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};
use std::collections::{HashSet};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
    #[fail(display = "Relocation target not found in image")]
    InvalidTargetAddress,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

// TODO: make this much faster
fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1, Permissions::R);
        }
    }
    return false;
}

fn is_zero(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(v) = ws.read_u32(rva) {
        return v == 0;
    }
    return false;
}

#[derive(Debug)]
pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

fn sections_contain(sections: &[Range<RVA>], rva: RVA) -> bool {
    sections.iter()
        .find(|section| section.contains(&rva))
        .is_some()
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    /// lancelot::test::init_logging();
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    ///
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    ///
    /// assert!(ws.get_meta(RVA(0xC7F0)).unwrap().is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let x_sections: Vec<Range<RVA>> = ws.module.sections.iter()
            .filter(|section| section.perms.intersects(Permissions::X))
            .map(|section| Range {
                start: section.addr,
                end: section.end(),
            })
            .collect();

        let relocs: Vec<Reloc> = get_relocs(ws)?
            .into_iter()
            .filter(|r| match &r.typ {
                RelocationType::ImageRelBasedHighLow => true,
                RelocationType::ImageRelBasedDir64 => true,
                reloc_type @ _ => {
                    // all other reloc types are currently unsupported (uncommon)
                    warn!("ignoring relocation with unsupported type: {:?}", reloc_type);
                    false
                },
            })
            .collect();
        debug!("found {} total relocs", relocs.len());

        for reloc in relocs.iter() {
            let target = ws.rva(ws.read_va(reloc.offset)?).ok_or(RelocAnalyzerError::InvalidTargetAddress)?;
            trace!("found reloc {} -> {} {}",
                   reloc.offset,
                   target,
                if sections_contain(&x_sections, target) {"(code)"} else {""}
            );
        }

        // scan for relocations to code.
        //
        // a relocation is a hardcoded offset that must be fixed up if the desired base address
        //  cannot be used.
        // for example, the function pointer passed to CreateThread will be a hardcoded address
        //  of the start of the function.
        //
        // relocated pointers may point to the .data section, e.g. strings or other constants.
        // we want to ignore these.
        // we are only interested in targets in executable sections.
        // we assume these are pointers to instructions/code.
        //
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.

        let mut unique_targets: HashSet<RVA> = HashSet::new();
        unique_targets.extend(relocs
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va)));

        debug!("reduced to {} unique reloc targets", unique_targets.len());

       let o: Vec<RVA> = unique_targets
            .iter()
            .filter(|&&rva| sections_contain(&x_sections, rva))
            .filter(|&&rva| !is_in_insn(ws, rva))
            .filter(|&&rva| !is_ptr(ws, rva))
            .filter(|&&rva| !is_zero(ws, rva))
            .map(|&rva| rva)
            // TODO: maybe ensure that the insn decodes.
            .collect();

        debug!("found {} relocs that point to instructions", o.len());

        o.iter().for_each(|&rva| {
            debug!("found pointer via relocations from executable section to {} (code)", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};
use log::{debug, trace};
use goblin::{Object};
use failure::{Error, Fail};
use std::ops::Range;

use super::super::super::arch::{RVA};
use super::super::super::loader::Permissions;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};
use std::collections::{HashSet};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

// TODO: make this much faster
fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

fn sections_contain(sections: &[Range<RVA>], rva: RVA) -> bool {
    sections.iter()
        .find(|section| section.contains(&rva))
        .is_some()
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    ///
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    ///
    /// assert!(ws.get_meta(RVA(0xC7F0)).unwrap().is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let x_sections: Vec<Range<RVA>> = ws.module.sections.iter()
            .filter(|section| section.perms.intersects(Permissions::X))
            .map(|section| Range {
                start: section.addr,
                end: section.end(),
            })
            .collect();

        let relocs = get_relocs(ws)?;
        debug!("found {} total relocs", relocs.len());

        for reloc in relocs.iter() {
            let target = ws.rva(ws.read_va(reloc.offset)?).expect("reloc target not in image");
            trace!("found reloc {} -> {} {}",
                   reloc.offset,
                   target,
                if sections_contain(&x_sections, target) {"(code)"} else {""}
            );
        }

        // scan for relocations to code.
        //
        // a relocation is a hardcoded offset that must be fixed up if the desired base address
        //  cannot be used.
        // for example, the function pointer passed to CreateThread will be a hardcoded address
        //  of the start of the function.
        //
        // relocated pointers may point to the .data section, e.g. strings or other constants.
        // we want to ignore these.
        // we are only interested in targets in executable sections.
        // we assume these are pointers to instructions/code.
        //
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.

        let mut unique_targets: HashSet<RVA> = HashSet::new();
        unique_targets.extend(relocs
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va)));

        debug!("reduced to {} unique reloc targets", unique_targets.len());

       let o: Vec<RVA> = unique_targets
            .iter()
            .filter(|&&rva| sections_contain(&x_sections, rva))
            .filter(|&&rva| !is_in_insn(ws, rva))
            .filter(|&&rva| !is_ptr(ws, rva))
            .map(|&rva| rva)
            // TODO: maybe ensure that the insn decodes.
            .collect();

        debug!("found {} relocs that point to instructions", o.len());

        o.iter().for_each(|&rva| {
            debug!("found pointer via relocations from executable section to {} (code)", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<commit_msg>analysis: reloc: explicitly support only some reloc types<commit_after>use byteorder::{ByteOrder, LittleEndian};
use log::{debug, trace};
use goblin::{Object};
use failure::{Error, Fail};
use std::ops::Range;

use super::super::super::arch::{RVA};
use super::super::super::loader::Permissions;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};
use std::collections::{HashSet};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

// TODO: make this much faster
fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

fn sections_contain(sections: &[Range<RVA>], rva: RVA) -> bool {
    sections.iter()
        .find(|section| section.contains(&rva))
        .is_some()
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    ///
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    ///
    /// assert!(ws.get_meta(RVA(0xC7F0)).unwrap().is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let x_sections: Vec<Range<RVA>> = ws.module.sections.iter()
            .filter(|section| section.perms.intersects(Permissions::X))
            .map(|section| Range {
                start: section.addr,
                end: section.end(),
            })
            .collect();

        let relocs: Vec<Reloc> = get_relocs(ws)?
            .into_iter()
            .filter(|r| match r.typ {
                RelocationType::ImageRelBasedHighLow => true,
                RelocationType::ImageRelBasedDir64 => true,
                _ => false,  // all other unsupported
            })
            .collect();
        debug!("found {} total relocs", relocs.len());

        for reloc in relocs.iter() {
            let target = ws.rva(ws.read_va(reloc.offset)?).expect("reloc target not in image");
            trace!("found reloc {} -> {} {}",
                   reloc.offset,
                   target,
                if sections_contain(&x_sections, target) {"(code)"} else {""}
            );
        }

        // scan for relocations to code.
        //
        // a relocation is a hardcoded offset that must be fixed up if the desired base address
        //  cannot be used.
        // for example, the function pointer passed to CreateThread will be a hardcoded address
        //  of the start of the function.
        //
        // relocated pointers may point to the .data section, e.g. strings or other constants.
        // we want to ignore these.
        // we are only interested in targets in executable sections.
        // we assume these are pointers to instructions/code.
        //
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.

        let mut unique_targets: HashSet<RVA> = HashSet::new();
        unique_targets.extend(relocs
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va)));

        debug!("reduced to {} unique reloc targets", unique_targets.len());

       let o: Vec<RVA> = unique_targets
            .iter()
            .filter(|&&rva| sections_contain(&x_sections, rva))
            .filter(|&&rva| !is_in_insn(ws, rva))
            .filter(|&&rva| !is_ptr(ws, rva))
            .map(|&rva| rva)
            // TODO: maybe ensure that the insn decodes.
            .collect();

        debug!("found {} relocs that point to instructions", o.len());

        o.iter().for_each(|&rva| {
            debug!("found pointer via relocations from executable section to {} (code)", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
#[macro_use]
extern crate approx;
#[macro_use]
extern crate quick_error;

mod wrapper;

macro_rules! scanlib_try {
    ($expr:expr, $scanlib:expr) => {{
        let result = $expr;
        if result != 0 {
            return Err(last_error(result, $scanlib));
        }
    }}
}

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        Scanlib(n: i32, message: String) {
            description("scanlib error")
            display("scanlib error code {}: {}", n, message)
        }
        FfiNul(err: std::ffi::NulError) {
            from()
            cause(err)
            description(err.description())
            display("ffi nul error: {}", err)
        }
    }
}

pub type Result<T> = std::result::Result<T, Error>;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct Inclination {
    pub time: f64,
    pub roll: f64,
    pub pitch: f64,
}

fn last_error(n: i32, scanlib: *const wrapper::Scanlib) -> Error {
    use std::ffi::CStr;
    let message = unsafe {
        CStr::from_ptr(wrapper::scanlib_last_error(scanlib))
            .to_string_lossy()
            .into_owned()
    };
    Error::Scanlib(n, message)
}

pub fn inclincations_from_path<P: AsRef<std::path::Path>>(
    path: P,
    sync_to_pps: bool,
) -> Result<Vec<Inclination>> {
    use std::ffi::CString;

    let path = CString::new(path.as_ref().to_string_lossy().as_ref())?;
    let mut inclinations = std::ptr::null_mut();
    let mut scanlib = std::ptr::null_mut();
    let slice = unsafe {
        wrapper::scanlib_new(&mut scanlib);
        scanlib_try!(
            wrapper::inclinations_from_path(scanlib, path.as_ptr(), sync_to_pps, &mut inclinations),
            scanlib
        );
        let mut p = std::ptr::null();
        let mut len = 0;
        wrapper::inclinations_pointer(inclinations, &mut p);
        wrapper::inclinations_len(inclinations, &mut len);
        std::slice::from_raw_parts(p, len as usize)
    };
    let inclinations_vec = slice.to_vec();
    unsafe {
        wrapper::inclinations_drop(inclinations);
        wrapper::scanlib_drop(scanlib);
    }
    Ok(inclinations_vec)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn inclincations_from_path_fixture() {
        let inclincations = inclincations_from_path("../scanifc/data/scan.rxp", false).unwrap();
        assert_eq!(36, inclincations.len());
        assert_relative_eq!(-8.442, inclincations[0].roll, epsilon = 1e-3);
        assert_relative_eq!(-0.981, inclincations[0].pitch, epsilon = 1e-3);
        assert_relative_eq!(67.7494, inclincations[35].time, epsilon = 1e-4);
        assert_relative_eq!(-8.451, inclincations[35].roll, epsilon = 1e-3);
        assert_relative_eq!(-1.004, inclincations[35].pitch, epsilon = 1e-3);
    }

    #[test]
    fn inclincations_from_path_dne() {
        assert!(inclincations_from_path("notafile", false).is_err());
    }
}
<commit_msg>Fix typo<commit_after>#[cfg(test)]
#[macro_use]
extern crate approx;
#[macro_use]
extern crate quick_error;

mod wrapper;

macro_rules! scanlib_try {
    ($expr:expr, $scanlib:expr) => {{
        let result = $expr;
        if result != 0 {
            return Err(last_error(result, $scanlib));
        }
    }}
}

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        Scanlib(n: i32, message: String) {
            description("scanlib error")
            display("scanlib error code {}: {}", n, message)
        }
        FfiNul(err: std::ffi::NulError) {
            from()
            cause(err)
            description(err.description())
            display("ffi nul error: {}", err)
        }
    }
}

pub type Result<T> = std::result::Result<T, Error>;

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct Inclination {
    pub time: f64,
    pub roll: f64,
    pub pitch: f64,
}

fn last_error(n: i32, scanlib: *const wrapper::Scanlib) -> Error {
    use std::ffi::CStr;
    let message = unsafe {
        CStr::from_ptr(wrapper::scanlib_last_error(scanlib))
            .to_string_lossy()
            .into_owned()
    };
    Error::Scanlib(n, message)
}

pub fn inclinations_from_path<P: AsRef<std::path::Path>>(
    path: P,
    sync_to_pps: bool,
) -> Result<Vec<Inclination>> {
    use std::ffi::CString;

    let path = CString::new(path.as_ref().to_string_lossy().as_ref())?;
    let mut inclinations = std::ptr::null_mut();
    let mut scanlib = std::ptr::null_mut();
    let slice = unsafe {
        wrapper::scanlib_new(&mut scanlib);
        scanlib_try!(
            wrapper::inclinations_from_path(scanlib, path.as_ptr(), sync_to_pps, &mut inclinations),
            scanlib
        );
        let mut p = std::ptr::null();
        let mut len = 0;
        wrapper::inclinations_pointer(inclinations, &mut p);
        wrapper::inclinations_len(inclinations, &mut len);
        std::slice::from_raw_parts(p, len as usize)
    };
    let inclinations_vec = slice.to_vec();
    unsafe {
        wrapper::inclinations_drop(inclinations);
        wrapper::scanlib_drop(scanlib);
    }
    Ok(inclinations_vec)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn inclinations_from_path_fixture() {
        let inclinations = inclinations_from_path("../scanifc/data/scan.rxp", false).unwrap();
        assert_eq!(36, inclinations.len());
        assert_relative_eq!(-8.442, inclinations[0].roll, epsilon = 1e-3);
        assert_relative_eq!(-0.981, inclinations[0].pitch, epsilon = 1e-3);
        assert_relative_eq!(67.7494, inclinations[35].time, epsilon = 1e-4);
        assert_relative_eq!(-8.451, inclinations[35].roll, epsilon = 1e-3);
        assert_relative_eq!(-1.004, inclinations[35].pitch, epsilon = 1e-3);
    }

    #[test]
    fn inclinations_from_path_dne() {
        assert!(inclinations_from_path("notafile", false).is_err());
    }
}
<|endoftext|>


<commit_before>use utils::bigint::M256;
use utils::opcode::Opcode;
use std::cmp::min;
use super::errors::PCError;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum Instruction {
    STOP, ADD, MUL, SUB, DIV, SDIV, MOD, SMOD, ADDMOD, MULMOD, EXP,
    SIGNEXTEND, LT, GT, SLT, SGT, EQ, ISZERO, AND, OR, XOR, NOT, BYTE,
    SHA3, ADDRESS, BALANCE, ORIGIN, CALLER, CALLVALUE, CALLDATALOAD,
    CALLDATASIZE, CALLDATACOPY, CODESIZE, CODECOPY, GASPRICE,
    EXTCODESIZE, EXTCODECOPY, BLOCKHASH, COINBASE, TIMESTAMP, NUMBER,
    DIFFICULTY, GASLIMIT, POP, MLOAD, MSTORE, MSTORE8, SLOAD, SSTORE,
    JUMP, JUMPI, PC, MSIZE, GAS, JUMPDEST, CREATE, CALL, CALLCODE,
    RETURN, DELEGATECALL, SUICIDE,

    PUSH(M256),
    DUP(usize),
    SWAP(usize),
    LOG(usize),
}

pub struct PC {
    position: usize,
    code: Vec<u8>,
    valids: Vec<bool>,
}

impl Default for PC {
    fn default() -> PC {
        PC {
            position: 0,
            code: Vec::new(),
            valids: Vec::new(),
        }
    }
}

impl PC {
    pub fn new(code: &[u8]) -> Self {
        let code: Vec<u8> = code.into();
        let mut valids: Vec<bool> = Vec::with_capacity(code.len());
        valids.resize(code.len(), false);

        let mut i = 0;
        while i < code.len() {
            let opcode: Opcode = code[i].into();
            match opcode {
                Opcode::JUMPDEST => {
                    valids[i] = true;
                    i = i + 1;
                },
                Opcode::PUSH(v) => {
                    i = i + v + 1;
                },
                _ => {
                    i = i + 1;
                }
            }
        }

        PC {
            position: 0,
            code: code,
            valids: valids,
        }
    }

    fn read_bytes(&self, from_position: usize, byte_count: usize) -> Result<M256, PCError> {
        if from_position > self.code.len() {
            return Err(PCError::Overflow);
        }
        let position = from_position;
        if position.checked_add(byte_count).is_none() {
            return Err(PCError::IndexNotSupported);
        }
        let max = min(position + byte_count, self.code.len());
        Ok(M256::from(&self.code[position..max]))
    }

    pub fn jump(&mut self, position: usize) -> Result<(), PCError> {
        if position >= self.code.len() {
            return Err(PCError::Overflow);
        }

        if !self.valids[position] {
            return Err(PCError::BadJumpDest);
        }

        self.position = position;
        Ok(())
    }

    pub fn position(&self) -> usize {
        self.position
    }

    pub fn check_valid(&self, position: usize) -> Result<(), PCError> {
        if self.is_valid(position) {
            Ok(())
        } else {
            Err(PCError::BadJumpDest)
        }
    }

    pub fn is_valid(&self, position: usize) -> bool {
        if position >= self.code.len() {
            return false;
        }

        if !self.valids[position] {
            return false;
        }

        return true;
    }

    pub fn is_end(&self) -> bool {
        self.position == self.code.len()
    }

    pub fn peek(&self) -> Result<Instruction, PCError> {
        let position = self.position;
        if position >= self.code.len() {
            return Err(PCError::Overflow);
        }
        let opcode: Opcode = self.code[position].into();
        Ok(match opcode {
            Opcode::STOP => Instruction::STOP,
            Opcode::ADD => Instruction::ADD,
            Opcode::MUL => Instruction::MUL,
            Opcode::SUB => Instruction::SUB,
            Opcode::DIV => Instruction::DIV,
            Opcode::SDIV => Instruction::SDIV,
            Opcode::MOD => Instruction::MOD,
            Opcode::SMOD => Instruction::SMOD,
            Opcode::ADDMOD => Instruction::ADDMOD,
            Opcode::MULMOD => Instruction::MULMOD,
            Opcode::EXP => Instruction::EXP,
            Opcode::SIGNEXTEND => Instruction::SIGNEXTEND,

            Opcode::LT => Instruction::LT,
            Opcode::GT => Instruction::GT,
            Opcode::SLT => Instruction::SLT,
            Opcode::SGT => Instruction::SGT,
            Opcode::EQ => Instruction::EQ,
            Opcode::ISZERO => Instruction::ISZERO,
            Opcode::AND => Instruction::AND,
            Opcode::OR => Instruction::OR,
            Opcode::XOR => Instruction::XOR,
            Opcode::NOT => Instruction::NOT,
            Opcode::BYTE => Instruction::BYTE,

            Opcode::SHA3 => Instruction::SHA3,

            Opcode::ADDRESS => Instruction::ADDRESS,
            Opcode::BALANCE => Instruction::BALANCE,
            Opcode::ORIGIN => Instruction::ORIGIN,
            Opcode::CALLER => Instruction::CALLER,
            Opcode::CALLVALUE => Instruction::CALLVALUE,
            Opcode::CALLDATALOAD => Instruction::CALLDATALOAD,
            Opcode::CALLDATASIZE => Instruction::CALLDATASIZE,
            Opcode::CALLDATACOPY => Instruction::CALLDATACOPY,
            Opcode::CODESIZE => Instruction::CODESIZE,
            Opcode::CODECOPY => Instruction::CODECOPY,
            Opcode::GASPRICE => Instruction::GASPRICE,
            Opcode::EXTCODESIZE => Instruction::EXTCODESIZE,
            Opcode::EXTCODECOPY => Instruction::EXTCODECOPY,

            Opcode::BLOCKHASH => Instruction::BLOCKHASH,
            Opcode::COINBASE => Instruction::COINBASE,
            Opcode::TIMESTAMP => Instruction::TIMESTAMP,
            Opcode::NUMBER => Instruction::NUMBER,
            Opcode::DIFFICULTY => Instruction::DIFFICULTY,
            Opcode::GASLIMIT => Instruction::GASLIMIT,

            Opcode::POP => Instruction::POP,
            Opcode::MLOAD => Instruction::MLOAD,
            Opcode::MSTORE => Instruction::MSTORE,
            Opcode::MSTORE8 => Instruction::MSTORE8,
            Opcode::SLOAD => Instruction::SLOAD,
            Opcode::SSTORE => Instruction::SSTORE,
            Opcode::JUMP => Instruction::JUMP,
            Opcode::JUMPI => Instruction::JUMPI,
            Opcode::PC => Instruction::PC,
            Opcode::MSIZE => Instruction::MSIZE,
            Opcode::GAS => Instruction::GAS,
            Opcode::JUMPDEST => Instruction::JUMPDEST,

            Opcode::PUSH(v) => {
                let param = self.read_bytes(position + 1, v)?;
                Instruction::PUSH(param)
            },

            Opcode::DUP(v) => Instruction::DUP(v),
            Opcode::SWAP(v) => Instruction::SWAP(v),
            Opcode::LOG(v) => Instruction::LOG(v),

            Opcode::CREATE => Instruction::CREATE,
            Opcode::CALL => Instruction::CALL,
            Opcode::CALLCODE => Instruction::CALLCODE,
            Opcode::RETURN => Instruction::RETURN,
            Opcode::DELEGATECALL => Instruction::DELEGATECALL,

            Opcode::INVALID => {
                return Err(PCError::InvalidOpcode);
            },
            Opcode::SUICIDE => Instruction::SUICIDE,
        })
    }

    pub fn read(&mut self) -> Result<Instruction, PCError> {
        let result = self.peek()?;
        let opcode: Opcode = self.code[self.position].into();
        match opcode {
            Opcode::PUSH(v) => {
                self.position = min(self.position + v + 1, self.code.len());
            },
            _ => {
                self.position = self.position + 1;
            },
        }
        Ok(result)
    }
}
<commit_msg>Problem: no docs for pc.rs Solution: add it<commit_after>//! EVM Program Counter.

use utils::bigint::M256;
use utils::opcode::Opcode;
use std::cmp::min;
use super::errors::PCError;

#[derive(Debug, Clone, Copy, Eq, PartialEq)]
/// Instructions for the program counter. This is the same as `Opcode`
/// except `PUSH`, which might take longer length.
pub enum Instruction {
    STOP, ADD, MUL, SUB, DIV, SDIV, MOD, SMOD, ADDMOD, MULMOD, EXP,
    SIGNEXTEND, LT, GT, SLT, SGT, EQ, ISZERO, AND, OR, XOR, NOT, BYTE,
    SHA3, ADDRESS, BALANCE, ORIGIN, CALLER, CALLVALUE, CALLDATALOAD,
    CALLDATASIZE, CALLDATACOPY, CODESIZE, CODECOPY, GASPRICE,
    EXTCODESIZE, EXTCODECOPY, BLOCKHASH, COINBASE, TIMESTAMP, NUMBER,
    DIFFICULTY, GASLIMIT, POP, MLOAD, MSTORE, MSTORE8, SLOAD, SSTORE,
    JUMP, JUMPI, PC, MSIZE, GAS, JUMPDEST, CREATE, CALL, CALLCODE,
    RETURN, DELEGATECALL, SUICIDE,

    PUSH(M256),
    DUP(usize),
    SWAP(usize),
    LOG(usize),
}

/// Represents a program counter in EVM.
pub struct PC {
    position: usize,
    code: Vec<u8>,
    valids: Vec<bool>,
}

impl Default for PC {
    fn default() -> PC {
        PC {
            position: 0,
            code: Vec::new(),
            valids: Vec::new(),
        }
    }
}

impl PC {
    /// Create a new program counter from the given code.
    pub fn new(code: &[u8]) -> Self {
        let code: Vec<u8> = code.into();
        let mut valids: Vec<bool> = Vec::with_capacity(code.len());
        valids.resize(code.len(), false);

        let mut i = 0;
        while i < code.len() {
            let opcode: Opcode = code[i].into();
            match opcode {
                Opcode::JUMPDEST => {
                    valids[i] = true;
                    i = i + 1;
                },
                Opcode::PUSH(v) => {
                    i = i + v + 1;
                },
                _ => {
                    i = i + 1;
                }
            }
        }

        PC {
            position: 0,
            code: code,
            valids: valids,
        }
    }

    fn read_bytes(&self, from_position: usize, byte_count: usize) -> Result<M256, PCError> {
        if from_position > self.code.len() {
            return Err(PCError::Overflow);
        }
        let position = from_position;
        if position.checked_add(byte_count).is_none() {
            return Err(PCError::IndexNotSupported);
        }
        let max = min(position + byte_count, self.code.len());
        Ok(M256::from(&self.code[position..max]))
    }

    /// Jump to a position in the code. The destination must be valid
    /// to jump.
    pub fn jump(&mut self, position: usize) -> Result<(), PCError> {
        if position >= self.code.len() {
            return Err(PCError::Overflow);
        }

        if !self.valids[position] {
            return Err(PCError::BadJumpDest);
        }

        self.position = position;
        Ok(())
    }

    /// Get the current program counter position.
    pub fn position(&self) -> usize {
        self.position
    }

    /// Check whether the position is a valid jump destination. If
    /// not, returns `PCError`.
    pub fn check_valid(&self, position: usize) -> Result<(), PCError> {
        if self.is_valid(position) {
            Ok(())
        } else {
            Err(PCError::BadJumpDest)
        }
    }

    /// Returns `true` if the position is a valid jump destination. If
    /// not, returns `false`.
    pub fn is_valid(&self, position: usize) -> bool {
        if position >= self.code.len() {
            return false;
        }

        if !self.valids[position] {
            return false;
        }

        return true;
    }

    /// Check whether the PC is ended. Next `read` on this PC would
    /// result in `PCError::PCOverflow`.
    pub fn is_end(&self) -> bool {
        self.position == self.code.len()
    }

    /// Peek the next instruction.
    pub fn peek(&self) -> Result<Instruction, PCError> {
        let position = self.position;
        if position >= self.code.len() {
            return Err(PCError::Overflow);
        }
        let opcode: Opcode = self.code[position].into();
        Ok(match opcode {
            Opcode::STOP => Instruction::STOP,
            Opcode::ADD => Instruction::ADD,
            Opcode::MUL => Instruction::MUL,
            Opcode::SUB => Instruction::SUB,
            Opcode::DIV => Instruction::DIV,
            Opcode::SDIV => Instruction::SDIV,
            Opcode::MOD => Instruction::MOD,
            Opcode::SMOD => Instruction::SMOD,
            Opcode::ADDMOD => Instruction::ADDMOD,
            Opcode::MULMOD => Instruction::MULMOD,
            Opcode::EXP => Instruction::EXP,
            Opcode::SIGNEXTEND => Instruction::SIGNEXTEND,

            Opcode::LT => Instruction::LT,
            Opcode::GT => Instruction::GT,
            Opcode::SLT => Instruction::SLT,
            Opcode::SGT => Instruction::SGT,
            Opcode::EQ => Instruction::EQ,
            Opcode::ISZERO => Instruction::ISZERO,
            Opcode::AND => Instruction::AND,
            Opcode::OR => Instruction::OR,
            Opcode::XOR => Instruction::XOR,
            Opcode::NOT => Instruction::NOT,
            Opcode::BYTE => Instruction::BYTE,

            Opcode::SHA3 => Instruction::SHA3,

            Opcode::ADDRESS => Instruction::ADDRESS,
            Opcode::BALANCE => Instruction::BALANCE,
            Opcode::ORIGIN => Instruction::ORIGIN,
            Opcode::CALLER => Instruction::CALLER,
            Opcode::CALLVALUE => Instruction::CALLVALUE,
            Opcode::CALLDATALOAD => Instruction::CALLDATALOAD,
            Opcode::CALLDATASIZE => Instruction::CALLDATASIZE,
            Opcode::CALLDATACOPY => Instruction::CALLDATACOPY,
            Opcode::CODESIZE => Instruction::CODESIZE,
            Opcode::CODECOPY => Instruction::CODECOPY,
            Opcode::GASPRICE => Instruction::GASPRICE,
            Opcode::EXTCODESIZE => Instruction::EXTCODESIZE,
            Opcode::EXTCODECOPY => Instruction::EXTCODECOPY,

            Opcode::BLOCKHASH => Instruction::BLOCKHASH,
            Opcode::COINBASE => Instruction::COINBASE,
            Opcode::TIMESTAMP => Instruction::TIMESTAMP,
            Opcode::NUMBER => Instruction::NUMBER,
            Opcode::DIFFICULTY => Instruction::DIFFICULTY,
            Opcode::GASLIMIT => Instruction::GASLIMIT,

            Opcode::POP => Instruction::POP,
            Opcode::MLOAD => Instruction::MLOAD,
            Opcode::MSTORE => Instruction::MSTORE,
            Opcode::MSTORE8 => Instruction::MSTORE8,
            Opcode::SLOAD => Instruction::SLOAD,
            Opcode::SSTORE => Instruction::SSTORE,
            Opcode::JUMP => Instruction::JUMP,
            Opcode::JUMPI => Instruction::JUMPI,
            Opcode::PC => Instruction::PC,
            Opcode::MSIZE => Instruction::MSIZE,
            Opcode::GAS => Instruction::GAS,
            Opcode::JUMPDEST => Instruction::JUMPDEST,

            Opcode::PUSH(v) => {
                let param = self.read_bytes(position + 1, v)?;
                Instruction::PUSH(param)
            },

            Opcode::DUP(v) => Instruction::DUP(v),
            Opcode::SWAP(v) => Instruction::SWAP(v),
            Opcode::LOG(v) => Instruction::LOG(v),

            Opcode::CREATE => Instruction::CREATE,
            Opcode::CALL => Instruction::CALL,
            Opcode::CALLCODE => Instruction::CALLCODE,
            Opcode::RETURN => Instruction::RETURN,
            Opcode::DELEGATECALL => Instruction::DELEGATECALL,

            Opcode::INVALID => {
                return Err(PCError::InvalidOpcode);
            },
            Opcode::SUICIDE => Instruction::SUICIDE,
        })
    }

    /// Read the next instruction and step the program counter.
    pub fn read(&mut self) -> Result<Instruction, PCError> {
        let result = self.peek()?;
        let opcode: Opcode = self.code[self.position].into();
        match opcode {
            Opcode::PUSH(v) => {
                self.position = min(self.position + v + 1, self.code.len());
            },
            _ => {
                self.position = self.position + 1;
            },
        }
        Ok(result)
    }
}
<|endoftext|>


<commit_before>use num::bigint::{BigUint, RandBigInt};
use num::One;

use rand::Rng;

use utils::modular::Inverse;
use utils::modular::Power;
use utils::primes::generate_prime;

pub struct SecretKeyExtra {
    p: BigUint,
    q: BigUint,
    dmp1: BigUint,
    dmq1: BigUint,
    qinv: BigUint
}

impl SecretKeyExtra {
    fn from_primes(p: BigUint, q: BigUint, d: &BigUint) -> Self {
        SecretKeyExtra {
            dmp1: d % (&p - BigUint::one()),
            dmq1: d % (&q - BigUint::one()),
            qinv: q.inverse(&p).unwrap(),
            p: p,
            q: q,
        }
    }
}

pub enum RSA {
    Public {
        /// Modulus
        n: BigUint,
        /// Public Exponent
        e: BigUint,
    },
    Private {
        /// Modulus
        n: BigUint,
        /// Private Exponent
        d: BigUint,
        extra: Option<SecretKeyExtra>,
    }
}

pub type KeyPair = (RSA, RSA);

impl RSA {
    pub fn keypair_from_primes<P, Q, E>(p: P, q: Q, e: E) -> KeyPair
        where P: Into<BigUint>, Q: Into<BigUint>, E: Into<BigUint> {
            let (p, q, e) = (p.into(), q.into(), e.into());

            let n = &p * &q;
            let phi_n = &n - (&p + &q - BigUint::one());

            let d = e.inverse(&phi_n).expect("e is irreversible in ring phi(pq) - error");

            let public = RSA::Public {
                n: n.clone(),
                e: e
            };
            let private = RSA::Private {
                n: n,
                extra: Some(SecretKeyExtra::from_primes(p, q, &d)),
                d: d,
            };

            (public, private)
        }

    pub fn generate_keypair<G, T>(mut rng: G, e: T, bits: usize) -> KeyPair
        where G: Rng + RandBigInt, T: Into<BigUint> {
            let e = e.into();

            let mut p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&p - BigUint::one()) % &e != BigUint::one() {
                p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }
            let mut q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&q - BigUint::one()) % &e != BigUint::one() {
                q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }

            Self::keypair_from_primes(p, q, e)
        }

    pub fn is_public(&self) -> bool {
        match *self {
            RSA::Public { .. } => true,
            _ => false
        }
    }

    pub fn is_private(&self) -> bool {
        match *self {
            RSA::Private { .. } => true,
            _ => false
        }
    }

    pub fn crypt(&self, msg: &BigUint) -> BigUint {
        match *self {
            RSA::Private { ref n, ref d, ref extra } => crypt(msg, n, d, extra.as_ref()),
            RSA::Public { ref n, ref e } => crypt(msg, n, e, None),
        }
    }
}

fn crypt(msg: &BigUint,
         modulus: &BigUint,
         exp: &BigUint,
         extra: Option<&SecretKeyExtra>) -> BigUint {
    if let Some(ref extra) = extra {
        chinese_reminders_power(msg, extra)
    } else {
        msg.pow_mod(exp, modulus)
    }
}

fn chinese_remainders_power(c: &BigUint, extra: &SecretKeyExtra) -> BigUint {
    let mut m1 = c.pow_mod(&extra.dmp1, &extra.p);
    let m2 = c.pow_mod(&extra.dmq1, &extra.q);

    while m1 < m2 {
        m1 = m1 + &extra.p;
    }

    let h = (&extra.qinv * (m1 - &m2)) % &extra.p;

    m2 + h * &extra.q
}

#[cfg(test)]
mod tests {
    use super::RSA;

    use num::bigint::{ToBigUint};

    fn keys() -> (RSA, RSA) {
        RSA::keypair_from_primes(
            61.to_biguint().unwrap(),
            53.to_biguint().unwrap(),
            17.to_biguint().unwrap())
    }

    #[test]
    fn test_encryption() {
        let (public, _) = keys();
        let c = public.crypt(&65.to_biguint().unwrap());

        assert_eq!(c, 2790.to_biguint().unwrap())
    }

    #[test]
    fn test_decryption() {
        let (_, private) = keys();
        let m = private.crypt(&2790.to_biguint().unwrap());

        assert_eq!(m, 65.to_biguint().unwrap())
    }
}
<commit_msg>fix(): more typos<commit_after>use num::bigint::{BigUint, RandBigInt};
use num::One;

use rand::Rng;

use utils::modular::Inverse;
use utils::modular::Power;
use utils::primes::generate_prime;

pub struct SecretKeyExtra {
    p: BigUint,
    q: BigUint,
    dmp1: BigUint,
    dmq1: BigUint,
    qinv: BigUint
}

impl SecretKeyExtra {
    fn from_primes(p: BigUint, q: BigUint, d: &BigUint) -> Self {
        SecretKeyExtra {
            dmp1: d % (&p - BigUint::one()),
            dmq1: d % (&q - BigUint::one()),
            qinv: q.inverse(&p).unwrap(),
            p: p,
            q: q,
        }
    }
}

pub enum RSA {
    Public {
        /// Modulus
        n: BigUint,
        /// Public Exponent
        e: BigUint,
    },
    Private {
        /// Modulus
        n: BigUint,
        /// Private Exponent
        d: BigUint,
        extra: Option<SecretKeyExtra>,
    }
}

pub type KeyPair = (RSA, RSA);

impl RSA {
    pub fn keypair_from_primes<P, Q, E>(p: P, q: Q, e: E) -> KeyPair
        where P: Into<BigUint>, Q: Into<BigUint>, E: Into<BigUint> {
            let (p, q, e) = (p.into(), q.into(), e.into());

            let n = &p * &q;
            let phi_n = &n - (&p + &q - BigUint::one());

            let d = e.inverse(&phi_n).expect("e is irreversible in ring phi(pq) - error");

            let public = RSA::Public {
                n: n.clone(),
                e: e
            };
            let private = RSA::Private {
                n: n,
                extra: Some(SecretKeyExtra::from_primes(p, q, &d)),
                d: d,
            };

            (public, private)
        }

    pub fn generate_keypair<G, T>(mut rng: G, e: T, bits: usize) -> KeyPair
        where G: Rng + RandBigInt, T: Into<BigUint> {
            let e = e.into();

            let mut p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&p - BigUint::one()) % &e != BigUint::one() {
                p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }
            let mut q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&q - BigUint::one()) % &e != BigUint::one() {
                q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }

            Self::keypair_from_primes(p, q, e)
        }

    pub fn is_public(&self) -> bool {
        match *self {
            RSA::Public { .. } => true,
            _ => false
        }
    }

    pub fn is_private(&self) -> bool {
        match *self {
            RSA::Private { .. } => true,
            _ => false
        }
    }

    pub fn crypt(&self, msg: &BigUint) -> BigUint {
        match *self {
            RSA::Private { ref n, ref d, ref extra } => crypt(msg, n, d, extra.as_ref()),
            RSA::Public { ref n, ref e } => crypt(msg, n, e, None),
        }
    }
}

fn crypt(msg: &BigUint,
         modulus: &BigUint,
         exp: &BigUint,
         extra: Option<&SecretKeyExtra>) -> BigUint {
    if let Some(ref extra) = extra {
        chinese_remainders_power(msg, extra)
    } else {
        msg.pow_mod(exp, modulus)
    }
}

fn chinese_remainders_power(c: &BigUint, extra: &SecretKeyExtra) -> BigUint {
    let mut m1 = c.pow_mod(&extra.dmp1, &extra.p);
    let m2 = c.pow_mod(&extra.dmq1, &extra.q);

    while m1 < m2 {
        m1 = m1 + &extra.p;
    }

    let h = (&extra.qinv * (m1 - &m2)) % &extra.p;

    m2 + h * &extra.q
}

#[cfg(test)]
mod tests {
    use super::RSA;

    use num::bigint::{ToBigUint};

    fn keys() -> (RSA, RSA) {
        RSA::keypair_from_primes(
            61.to_biguint().unwrap(),
            53.to_biguint().unwrap(),
            17.to_biguint().unwrap())
    }

    #[test]
    fn test_encryption() {
        let (public, _) = keys();
        let c = public.crypt(&65.to_biguint().unwrap());

        assert_eq!(c, 2790.to_biguint().unwrap())
    }

    #[test]
    fn test_decryption() {
        let (_, private) = keys();
        let m = private.crypt(&2790.to_biguint().unwrap());

        assert_eq!(m, 65.to_biguint().unwrap())
    }
}
<|endoftext|>