<commit_before>use goblin::{container};
use mach;
use mach::header;
use mach::load_command;
use mach::exports::{Export};

use Opt;

use prettytable::cell::Cell;
use prettytable::row::Row;
use std::io::{self, Write};
use atty;
use termcolor::*;
use termcolor::Color::*;

use format::*;

pub struct Mach<'a>(pub mach::MachO<'a>, pub Opt);

impl<'a> Mach<'a> {
    pub fn print(&self) -> io::Result<()> {
        let mach = &self.0;
        let args = &self.1;

        let cc = if args.color || atty::is(atty::Stream::Stdout) { ColorChoice::Auto } else { ColorChoice::Never };
        let writer = BufferWriter::stdout(cc);
        let fmt = &mut writer.buffer();

        let header = &mach.header;
        let endianness = if header.is_little_endian() { "little-endian" } else { "big-endian" };
        let machine = header.cputype;
        let kind = |fmt: &mut Buffer, header: &header::Header| {
            let typ_cell = header.filetype;
            let kind_str = header::filetype_to_str(typ_cell);
            match typ_cell {
                header::MH_OBJECT =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Yellow)).set_fg(Some(Black)))?,
                header::MH_EXECUTE => fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Red)).set_fg(Some(Black)))?,
                header::MH_DYLIB =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Yellow)).set_fg(Some(Black)))?,
                header::MH_DYLINKER =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Yellow)).set_fg(Some(Black)))?,
                header::MH_DYLIB_STUB =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Blue)).set_fg(Some(Black)))?,
                header::MH_DSYM =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Green)).set_fg(Some(Black)))?,
                header::MH_CORE => fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(White)).set_fg(Some(Black)))?,
                _ => ()
            }
            write!(fmt, "{}", kind_str)?;
            fmt.reset()
        };
        fmt_hdr(fmt, "Mach-o ")?;
        kind(fmt, &mach.header)?;
        write!(fmt, " ")?;
        fmt_name_bold(fmt, mach::constants::cputype::get_arch_name_from_types(machine,header.cpusubtype).unwrap_or("None"))?;
        write!(fmt, "-{} @ ", endianness)?;
        fmt_addrx(fmt, mach.entry as u64)?;
        writeln!(fmt, ":")?;
        writeln!(fmt, "")?;

        let lcs = &mach.load_commands;
        fmt_header(fmt, "LoadCommands", lcs.len())?;
        for (i, lc) in lcs.into_iter().enumerate() {
            fmt_idx(fmt, i)?;
            write!(fmt, " ")?;
            let name = load_command::cmd_to_str(lc.command.cmd());
            let name = &format!("{:.27}", name);
            match lc.command {
                load_command::CommandVariant::Segment32        (_command) => fmt_name_color(fmt, name, Red)?,
                load_command::CommandVariant::Segment64        (_command) => fmt_name_color(fmt, name, Red)?,
                load_command::CommandVariant::Symtab           (_command) => fmt_name_color(fmt, name, Yellow)?,
                load_command::CommandVariant::Dysymtab         (_command) => fmt_name_color(fmt, name, Green)?,
                load_command::CommandVariant::LoadDylinker     (_command) => fmt_name_color(fmt, name, Yellow)?,
                load_command::CommandVariant::LoadDylib        (_command)
                    | load_command::CommandVariant::LoadUpwardDylib(_command)
                    | load_command::CommandVariant::ReexportDylib  (_command)
                    | load_command::CommandVariant::LazyLoadDylib  (_command) => fmt_name_color(fmt, name, Blue)?,
                load_command::CommandVariant::DyldInfo         (_command)
                    | load_command::CommandVariant::DyldInfoOnly   (_command) => fmt_name_color(fmt, name, Cyan)?,
                load_command::CommandVariant::Unixthread       (_command) => fmt_name_color(fmt, name, Red)?,
                load_command::CommandVariant::Main             (_command) => fmt_name_color(fmt, name, Red)?,
                _ => fmt_name_bold(fmt, name)?
            }
            writeln!(fmt, "")?;
        }

        writeln!(fmt, "")?;

        let segments = &mach.segments;
        fmt_header(fmt, "Segments", segments.len())?;
        for (ref _i, ref segment) in segments.into_iter().enumerate() {
            let name = segment.name().unwrap();
            let sections = &segment.sections().unwrap();
            let mut segment_table = new_table(row![b->"Segment", b->"# Sections"]);
            segment_table.add_row(Row::new(vec![
                str_cell(&name),
                Cell::new(&sections.len().to_string()),
            ]));
            flush(fmt, &writer, segment_table, args.color)?;

            let mut section_table = new_table(row![b->"", b->"Idx", b->"Name", b->"Addr", b->"Size", b->"Offset", b->"Align", b->"Reloff", b->"Nrelocs", b->"Flags"]);
            for (i, &(ref section, _)) in sections.into_iter().enumerate() {
                if let Ok(name) = section.name() {
                    section_table.add_row(Row::new(vec![
                        Cell::new(&format!("{:4}", "")), // filler
                        Cell::new(&i.to_string()),
                        Cell::new(name).style_spec("Fyb"),
                        addrx_cell(section.addr),
                        sz_cell(section.size),
                        offsetx_cell(section.offset as u64),
                        Cell::new(&format!("{}", section.align)),
                        offsetx_cell(section.reloff as u64),
                        Cell::new(&format!("{}", section.nreloc)),
                        Cell::new(&format!("{:#x}", section.flags)),
                    ]));
                } else {
                    section_table.add_row(Row::new(vec![
                        Cell::new(&i.to_string()),
                        Cell::new("BAD SECTION NAME"),
                    ]));
                }
            }
            flush(fmt, &writer, section_table, args.color)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        let mut relocations: Vec<_> = Vec::new();
        let mut nrelocs = 0;
        let r = mach.relocations().unwrap();
        for (_i, relocs, section) in r.into_iter() {
            let section_name = section.name().unwrap();
            let segment_name = section.segname().unwrap();
            let mut rs = Vec::new();
            for reloc in relocs {
                let reloc = reloc.unwrap();
                nrelocs += 1;
                rs.push(reloc);
            }
            if !rs.is_empty() { relocations.push((segment_name.to_owned(), section_name.to_owned(), rs)) };
        }

        fmt_header(fmt, "Relocations", nrelocs)?;
        for (n1, n2, relocs) in relocations {
            let mut section_table = new_table(row![b->"Segment", b->"Section", b->"Count"]);
            section_table.add_row(Row::new(vec![
                str_cell(&n1),
                str_cell(&n2),
                Cell::new(&relocs.len().to_string()),
            ]));
            flush(fmt, &writer, section_table, args.color)?;

            let mut reloc_table = new_table(row![b->"", b->"Type", b->"Offset", b->"SymbolNum", b->"Length", b->"PIC", b->"Extern"]);
            for reloc in relocs {
                reloc_table.add_row(Row::new(vec![
                    Cell::new(&format!("{:4}", "")),
                    cell(reloc.to_str(machine)),
                    addrx_cell(reloc.r_address as u64),
                    offsetx_cell(reloc.r_symbolnum() as u64),
                    cell(reloc.r_length()),
                    bool_cell(reloc.is_pic()),
                    bool_cell(reloc.is_extern()),
                ]));
            }
            flush(fmt, &writer, reloc_table, args.color)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        let sections = mach.segments.sections().flat_map(|x| x).map(|s| s.unwrap().0).collect::<Vec<_>>();
        let symbols = mach.symbols().collect::<Vec<_>>();
        fmt_header(fmt, "Symbols", symbols.len())?;
        let mut symbol_table = new_table(row![b->"Offset", b->"Name", b->"Section", b->"Global", b->"Undefined"]);
        for (i, symbol) in symbols.into_iter().enumerate() {
            match symbol {
                Ok((name, symbol)) => {
                    let section_cell = if symbol.get_type() == mach::symbols::N_SECT {
                        // we subtract 1 because when N_SECT it is an ordinal, and hence indexing starts from 1
                        cell(sections[symbol.n_sect - 1 as usize].name().unwrap()).style_spec("b")
                    } else {
                        cell("None").style_spec("i")
                    };
                    symbol_table.add_row(Row::new(vec![
                        addrx_cell(symbol.n_value as u64),
                        string_cell(&args, name),
                        section_cell,
                        bool_cell(symbol.is_global()),
                        bool_cell(symbol.is_undefined()),
                    ]));
                },
                Err(e) => {
                    writeln!(fmt, "  {}: {}", i, e)?;
                }
            }
        }
        flush(fmt, &writer, symbol_table, args.color)?;
        writeln!(fmt)?;

        let fmt_exports = |fmt: &mut Buffer, name: &str, syms: &[Export] | -> io::Result<()> {
            fmt_header(fmt, name, syms.len())?;
            for sym in syms {
                fmt_addr_right(fmt, sym.offset)?;
                write!(fmt, " ")?;
                fmt_string(fmt, args, &sym.name)?;
                write!(fmt, " (")?;
                fmt_sz(fmt, sym.size as u64)?;
                writeln!(fmt, ")")?;
            }
            writeln!(fmt, "")
        };

        let exports = match mach.exports () { Ok(exports) => exports, Err(_) => Vec::new() };
        fmt_exports(fmt, "Exports", &exports)?;

        let imports = match mach.imports () { Ok(imports) => imports, Err(_) => Vec::new() };
        fmt_header(fmt, "Imports", imports.len())?;
        for sym in imports {
            fmt_addr_right(fmt, sym.offset)?;
            write!(fmt, " ")?;
            fmt_string(fmt, args, &sym.name)?;
            write!(fmt, " (")?;
            fmt_sz(fmt, sym.size as u64)?;
            write!(fmt, ")")?;
            write!(fmt, " -> ")?;
            fmt_lib(fmt, sym.dylib)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        fmt_header(fmt, "Libraries", mach.libs.len() - 1)?;
        for lib in &mach.libs[1..] {
            fmt_lib_right(fmt, lib)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        write!(fmt, "Name: ")?;
        fmt_str_option(fmt, &mach.name)?;
        writeln!(fmt, "")?;
        write!(fmt, "is_64: ")?;
        fmt_bool(fmt, mach.header.container() == container::Container::Big)?;
        writeln!(fmt, "")?;
        write!(fmt, "is_lib: ")?;
        fmt_bool(fmt, mach.header.filetype == header::MH_DYLIB)?;
        writeln!(fmt, "")?;
        write!(fmt, "little_endian: ")?;
        fmt_bool(fmt, mach.header.is_little_endian())?;
        writeln!(fmt, "")?;
        write!(fmt, "entry: ")?;
        fmt_addr(fmt, mach.entry as u64)?;
        writeln!(fmt, "")?;

        writer.print(&fmt)?;
        Ok(())
    }
}
<commit_msg>mach: add relocation symbol and/or section reference printing<commit_after>use goblin::{container, error};
use mach;
use mach::header;
use mach::load_command;
use mach::exports::{Export};

use Opt;

use prettytable::cell::Cell;
use prettytable::row::Row;
use std::io::{self, Write};
use atty;
use termcolor::*;
use termcolor::Color::*;

use format::*;

pub struct Mach<'a>(pub mach::MachO<'a>, pub Opt);

impl<'a> Mach<'a> {
    pub fn print(&self) -> error::Result<()> {
        let mach = &self.0;
        let args = &self.1;

        let cc = if args.color || atty::is(atty::Stream::Stdout) { ColorChoice::Auto } else { ColorChoice::Never };
        let writer = BufferWriter::stdout(cc);
        let fmt = &mut writer.buffer();

        let header = &mach.header;
        let endianness = if header.is_little_endian() { "little-endian" } else { "big-endian" };
        let machine = header.cputype;
        let kind = |fmt: &mut Buffer, header: &header::Header| {
            let typ_cell = header.filetype;
            let kind_str = header::filetype_to_str(typ_cell);
            match typ_cell {
                header::MH_OBJECT =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Yellow)).set_fg(Some(Black)))?,
                header::MH_EXECUTE => fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Red)).set_fg(Some(Black)))?,
                header::MH_DYLIB =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Yellow)).set_fg(Some(Black)))?,
                header::MH_DYLINKER =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Yellow)).set_fg(Some(Black)))?,
                header::MH_DYLIB_STUB =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Blue)).set_fg(Some(Black)))?,
                header::MH_DSYM =>  fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(Green)).set_fg(Some(Black)))?,
                header::MH_CORE => fmt.set_color(::termcolor::ColorSpec::new().set_intense(true).set_bg(Some(White)).set_fg(Some(Black)))?,
                _ => ()
            }
            write!(fmt, "{}", kind_str)?;
            fmt.reset()
        };
        fmt_hdr(fmt, "Mach-o ")?;
        kind(fmt, &mach.header)?;
        write!(fmt, " ")?;
        fmt_name_bold(fmt, mach::constants::cputype::get_arch_name_from_types(machine,header.cpusubtype).unwrap_or("None"))?;
        write!(fmt, "-{} @ ", endianness)?;
        fmt_addrx(fmt, mach.entry as u64)?;
        writeln!(fmt, ":")?;
        writeln!(fmt, "")?;

        let lcs = &mach.load_commands;
        fmt_header(fmt, "LoadCommands", lcs.len())?;
        for (i, lc) in lcs.into_iter().enumerate() {
            fmt_idx(fmt, i)?;
            write!(fmt, " ")?;
            let name = load_command::cmd_to_str(lc.command.cmd());
            let name = &format!("{:.27}", name);
            match lc.command {
                load_command::CommandVariant::Segment32        (_command) => fmt_name_color(fmt, name, Red)?,
                load_command::CommandVariant::Segment64        (_command) => fmt_name_color(fmt, name, Red)?,
                load_command::CommandVariant::Symtab           (_command) => fmt_name_color(fmt, name, Yellow)?,
                load_command::CommandVariant::Dysymtab         (_command) => fmt_name_color(fmt, name, Green)?,
                load_command::CommandVariant::LoadDylinker     (_command) => fmt_name_color(fmt, name, Yellow)?,
                load_command::CommandVariant::LoadDylib        (_command)
                    | load_command::CommandVariant::LoadUpwardDylib(_command)
                    | load_command::CommandVariant::ReexportDylib  (_command)
                    | load_command::CommandVariant::LazyLoadDylib  (_command) => fmt_name_color(fmt, name, Blue)?,
                load_command::CommandVariant::DyldInfo         (_command)
                    | load_command::CommandVariant::DyldInfoOnly   (_command) => fmt_name_color(fmt, name, Cyan)?,
                load_command::CommandVariant::Unixthread       (_command) => fmt_name_color(fmt, name, Red)?,
                load_command::CommandVariant::Main             (_command) => fmt_name_color(fmt, name, Red)?,
                _ => fmt_name_bold(fmt, name)?
            }
            writeln!(fmt, "")?;
        }

        writeln!(fmt, "")?;

        let segments = &mach.segments;
        fmt_header(fmt, "Segments", segments.len())?;
        for (ref _i, ref segment) in segments.into_iter().enumerate() {
            let name = segment.name().unwrap();
            let sections = &segment.sections().unwrap();
            let mut segment_table = new_table(row![b->"Segment", b->"# Sections"]);
            segment_table.add_row(Row::new(vec![
                str_cell(&name),
                Cell::new(&sections.len().to_string()),
            ]));
            flush(fmt, &writer, segment_table, args.color)?;

            let mut section_table = new_table(row![b->"", b->"Idx", b->"Name", b->"Addr", b->"Size", b->"Offset", b->"Align", b->"Reloff", b->"Nrelocs", b->"Flags"]);
            for (i, &(ref section, _)) in sections.into_iter().enumerate() {
                if let Ok(name) = section.name() {
                    section_table.add_row(Row::new(vec![
                        Cell::new(&format!("{:4}", "")), // filler
                        Cell::new(&i.to_string()),
                        Cell::new(name).style_spec("Fyb"),
                        addrx_cell(section.addr),
                        sz_cell(section.size),
                        offsetx_cell(section.offset as u64),
                        Cell::new(&format!("{}", section.align)),
                        offsetx_cell(section.reloff as u64),
                        Cell::new(&format!("{}", section.nreloc)),
                        Cell::new(&format!("{:#x}", section.flags)),
                    ]));
                } else {
                    section_table.add_row(Row::new(vec![
                        Cell::new(&i.to_string()),
                        Cell::new("BAD SECTION NAME"),
                    ]));
                }
            }
            flush(fmt, &writer, section_table, args.color)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        let mut relocations: Vec<_> = Vec::new();
        let mut nrelocs = 0;
        let r = mach.relocations().unwrap();
        for (_i, relocs, section) in r.into_iter() {
            let section_name = section.name().unwrap();
            let segment_name = section.segname().unwrap();
            let mut rs = Vec::new();
            for reloc in relocs {
                let reloc = reloc.unwrap();
                nrelocs += 1;
                rs.push(reloc);
            }
            if !rs.is_empty() { relocations.push((segment_name.to_owned(), section_name.to_owned(), rs)) };
        }
        // need this to print relocation references
        let symbols = mach.symbols().collect::<Vec<_>>();
        let sections = mach.segments.sections().flat_map(|x| x).map(|s| s.unwrap().0).collect::<Vec<_>>();

        fmt_header(fmt, "Relocations", nrelocs)?;
        for (n1, n2, relocs) in relocations {
            let mut section_table = new_table(row![b->"Segment", b->"Section", b->"Count"]);
            section_table.add_row(Row::new(vec![
                str_cell(&n1),
                str_cell(&n2),
                Cell::new(&relocs.len().to_string()),
            ]));
            flush(fmt, &writer, section_table, args.color)?;
            let mut reloc_table = new_table(row![b->"", b->"Type", b->"Offset", b->"Length", b->"PIC", b->"Extern", b->"SymbolNum", b->"Symbol"]);
            for reloc in relocs {
                let idx = reloc.r_symbolnum();
                let name_cell = {
                    if reloc.is_extern() {
                        // FIXME: i cannot currently get this to compile without iterating and doing all this nonsense, otherwise move errors...
                        let mut maybe_name = None;
                        for (i, symbol) in symbols.iter().enumerate() {
                            match symbol {
                                &Ok((ref name, _)) => {
                                    let name: &str = name;
                                    if i == idx {
                                        maybe_name = Some(name);
                                    }
                                },
                                &Err(_) => ()
                            }
                        }
                        match maybe_name {
                            Some(name) => string_cell(&args, name),
                            None => {
                                cell("None").style_spec("b")
                            },
                        }
                    // not extern so the symbol num should reference a section
                    } else {
                        cell(sections[idx - 1 as usize].name().unwrap_or("None")).style_spec("bi")

                    }
                };
                reloc_table.add_row(Row::new(vec![
                    Cell::new(&format!("{:4}", "")),
                    cell(reloc.to_str(machine)),
                    addrx_cell(reloc.r_address as u64),
                    cell(reloc.r_length()),
                    bool_cell(reloc.is_pic()),
                    bool_cell(reloc.is_extern()),
                    offsetx_cell(idx as u64),
                    name_cell,
                ]));
            }

            flush(fmt, &writer, reloc_table, args.color)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        fmt_header(fmt, "Symbols", symbols.len())?;
        let mut symbol_table = new_table(row![b->"Offset", b->"Name", b->"Section", b->"Global", b->"Undefined"]);
        for (i, symbol) in symbols.into_iter().enumerate() {
            match symbol {
                Ok((name, symbol)) => {
                    let section_cell = if symbol.get_type() == mach::symbols::N_SECT {
                        // we subtract 1 because when N_SECT it is an ordinal, and hence indexing starts from 1
                        cell(sections[symbol.n_sect - 1 as usize].name().unwrap()).style_spec("b")
                    } else {
                        cell("None").style_spec("i")
                    };
                    symbol_table.add_row(Row::new(vec![
                        addrx_cell(symbol.n_value as u64),
                        string_cell(&args, name),
                        section_cell,
                        bool_cell(symbol.is_global()),
                        bool_cell(symbol.is_undefined()),
                    ]));
                },
                Err(e) => {
                    writeln!(fmt, "  {}: {}", i, e)?;
                }
            }
        }
        flush(fmt, &writer, symbol_table, args.color)?;
        writeln!(fmt)?;

        let fmt_exports = |fmt: &mut Buffer, name: &str, syms: &[Export] | -> io::Result<()> {
            fmt_header(fmt, name, syms.len())?;
            for sym in syms {
                fmt_addr_right(fmt, sym.offset)?;
                write!(fmt, " ")?;
                fmt_string(fmt, args, &sym.name)?;
                write!(fmt, " (")?;
                fmt_sz(fmt, sym.size as u64)?;
                writeln!(fmt, ")")?;
            }
            writeln!(fmt, "")
        };

        let exports = match mach.exports () { Ok(exports) => exports, Err(_) => Vec::new() };
        fmt_exports(fmt, "Exports", &exports)?;

        let imports = match mach.imports () { Ok(imports) => imports, Err(_) => Vec::new() };
        fmt_header(fmt, "Imports", imports.len())?;
        for sym in imports {
            fmt_addr_right(fmt, sym.offset)?;
            write!(fmt, " ")?;
            fmt_string(fmt, args, &sym.name)?;
            write!(fmt, " (")?;
            fmt_sz(fmt, sym.size as u64)?;
            write!(fmt, ")")?;
            write!(fmt, " -> ")?;
            fmt_lib(fmt, sym.dylib)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        fmt_header(fmt, "Libraries", mach.libs.len() - 1)?;
        for lib in &mach.libs[1..] {
            fmt_lib_right(fmt, lib)?;
            writeln!(fmt, "")?;
        }
        writeln!(fmt, "")?;

        write!(fmt, "Name: ")?;
        fmt_str_option(fmt, &mach.name)?;
        writeln!(fmt, "")?;
        write!(fmt, "is_64: ")?;
        fmt_bool(fmt, mach.header.container() == container::Container::Big)?;
        writeln!(fmt, "")?;
        write!(fmt, "is_lib: ")?;
        fmt_bool(fmt, mach.header.filetype == header::MH_DYLIB)?;
        writeln!(fmt, "")?;
        write!(fmt, "little_endian: ")?;
        fmt_bool(fmt, mach.header.is_little_endian())?;
        writeln!(fmt, "")?;
        write!(fmt, "entry: ")?;
        fmt_addr(fmt, mach.entry as u64)?;
        writeln!(fmt, "")?;

        writer.print(&fmt)?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Module that provides interaction with radare2.
//! Uses r2pipe.

use r2pipe::R2Pipe;

use rustc_serialize::json::{Json, ToJson, DecodeResult};
use rustc_serialize::json;
use std::collections::BTreeMap;

// All the struct required for json parsing,
use super::structs::*;

pub struct R2 {
	pipe:   R2Pipe,
	readin: String,
}

pub trait Search {
	fn get_values(&self, key: &[&str], deep: bool) -> Option<Json>;
}

impl Search for Json {
	fn get_values(&self, key: &[&str], deep: bool) -> Option<Json> {
		match *self {
			Json::Array(ref arr) => {
				let mut tmp = Vec::new();
				for x in arr.iter() {
					x.get_values(key, deep)
						.map_or((), |e| tmp.push(e));
				}
				if tmp.len() > 0 {
					return Some(tmp.to_json());
				}
			},

			Json::Object(ref obj) => {
				let mut result: BTreeMap<String, Json> = BTreeMap::new();
				for k in key {
					obj.get((*k)).map_or((), |val| {
						result.insert((*k).to_string(), val.clone());
					});
				}

				// If we have found something, return.
				if result.keys().len() > 0 {
					return Some(result.to_json());
				}

				// If we don't want a deep search, we just return None at this point.
				if !deep {
					return None;
				}

				// Manually iterate through the keys and check if the values have the key.
				for (_, val) in obj.iter() {
					let res = val.get_values(key, deep);
					if res.is_some() {
						return res;
					}
				}
			},

			_ => return None,
		}
		return None;
	}
}

// fn send and recv allow users to send their own commands, i.e. The ones that are not currently
// abstracted by the R2 API.
// Ideally, all commonly used commands must be supported for easier use.

impl R2 {
	pub fn new(path: &str) -> R2 {
		let pipe = open_pipe!(path).unwrap();
		R2 { pipe: pipe, readin: String::new() }
	}

	// Does some basic configurations (sane defaults).
	pub fn init(&mut self) {
		self.send("e asm.esil = true");
		self.send("e scr.color = false");
		self.send("aaa");
		self.flush();
	}

	pub fn close(&mut self) {
		self.send("q!");
	}

	pub fn send(&mut self, cmd: &str) {
		self.readin = self.pipe.cmd(cmd);
	}

	pub fn recv(&mut self) -> String {
		let res = self.readin.clone();
		self.flush();
		res
	}

	pub fn recv_json(&mut self) -> Json {
		let res = self.recv().replace("\n", "");
		Json::from_str(&*res).unwrap()
	}

	pub fn flush(&mut self) {
		self.readin = String::from("");
	}

	pub fn analyze(&mut self) {
		self.send("aa");
		self.flush();
	}

	pub fn get_function(&mut self, func: &str) -> DecodeResult<LFunctionInfo> {
		let cmd = format!("pdfj @ {}", func);
		self.send(&*cmd);
		let raw_json = self.recv();
		// Handle Error here.
		json::decode(&*raw_json)
	}

	// get 'n' (or 16) instructions at 'offset' (or current position if offset in `None`)
	pub fn get_insts(&mut self, n: Option<u64>, offset: Option<&str>) -> DecodeResult<Vec<LOpInfo>> {
		let n = n.unwrap_or(16);
		let offset: &str = offset.unwrap_or_default();
		let mut cmd = format!("pdj{}", n);
		if offset.len() > 0 {
			cmd = format!("{} @ {}", cmd, offset);
		}
		self.send(&*cmd);
		let raw_json = self.recv();
		json::decode(&*raw_json)
	}

	pub fn get_reg_info(&mut self) -> DecodeResult<LRegInfo> {
		self.send("drpj");
		let raw_json = self.recv();
		json::decode(&*raw_json)
	}

	pub fn get_flag_info(&mut self) -> DecodeResult<Vec<LFlagInfo>> {
		self.send("fj");
		let raw_json = self.recv();
		json::decode(&*raw_json)
	}

	pub fn get_fn_list(&mut self) -> Json {
		self.send("aflj");
		self.recv_json()
	}
}
<commit_msg>Allow to construct r2 from R2Pipe<commit_after>//! Module that provides interaction with radare2.
//! Uses r2pipe.

use r2pipe::R2Pipe;

use rustc_serialize::json::{Json, ToJson, DecodeResult};
use rustc_serialize::json;
use std::collections::BTreeMap;

// All the struct required for json parsing,
use super::structs::*;

pub struct R2 {
	pipe:   R2Pipe,
	readin: String,
}

pub trait Search {
	fn get_values(&self, key: &[&str], deep: bool) -> Option<Json>;
}

impl Search for Json {
	fn get_values(&self, key: &[&str], deep: bool) -> Option<Json> {
		match *self {
			Json::Array(ref arr) => {
				let mut tmp = Vec::new();
				for x in arr.iter() {
					x.get_values(key, deep)
						.map_or((), |e| tmp.push(e));
				}
				if tmp.len() > 0 {
					return Some(tmp.to_json());
				}
			},

			Json::Object(ref obj) => {
				let mut result: BTreeMap<String, Json> = BTreeMap::new();
				for k in key {
					obj.get((*k)).map_or((), |val| {
						result.insert((*k).to_string(), val.clone());
					});
				}

				// If we have found something, return.
				if result.keys().len() > 0 {
					return Some(result.to_json());
				}

				// If we don't want a deep search, we just return None at this point.
				if !deep {
					return None;
				}

				// Manually iterate through the keys and check if the values have the key.
				for (_, val) in obj.iter() {
					let res = val.get_values(key, deep);
					if res.is_some() {
						return res;
					}
				}
			},

			_ => return None,
		}
		return None;
	}
}

// fn send and recv allow users to send their own commands, i.e. The ones that are not currently
// abstracted by the R2 API.
// Ideally, all commonly used commands must be supported for easier use.

impl R2 {
	pub fn new(path: &str) -> R2 {
		let pipe = open_pipe!(path).unwrap();
		R2 { pipe: pipe, readin: String::new() }
	}

	pub fn from(r2p: R2Pipe) -> R2 {
		R2 { pipe: r2p, readin: String::new() }
	}

	// Does some basic configurations (sane defaults).
	pub fn init(&mut self) {
		self.send("e asm.esil = true");
		self.send("e scr.color = false");
		self.send("aaa");
		self.flush();
	}

	pub fn close(&mut self) {
		self.send("q!");
	}

	pub fn send(&mut self, cmd: &str) {
		self.readin = self.pipe.cmd(cmd);
	}

	pub fn recv(&mut self) -> String {
		let res = self.readin.clone();
		self.flush();
		res
	}

	pub fn recv_json(&mut self) -> Json {
		let res = self.recv().replace("\n", "");
		Json::from_str(&*res).unwrap()
	}

	pub fn flush(&mut self) {
		self.readin = String::from("");
	}

	pub fn analyze(&mut self) {
		self.send("aa");
		self.flush();
	}

	pub fn get_function(&mut self, func: &str) -> DecodeResult<LFunctionInfo> {
		let cmd = format!("pdfj @ {}", func);
		self.send(&*cmd);
		let raw_json = self.recv();
		// Handle Error here.
		json::decode(&*raw_json)
	}

	// get 'n' (or 16) instructions at 'offset' (or current position if offset in `None`)
	pub fn get_insts(&mut self, n: Option<u64>, offset: Option<&str>) -> DecodeResult<Vec<LOpInfo>> {
		let n = n.unwrap_or(16);
		let offset: &str = offset.unwrap_or_default();
		let mut cmd = format!("pdj{}", n);
		if offset.len() > 0 {
			cmd = format!("{} @ {}", cmd, offset);
		}
		self.send(&*cmd);
		let raw_json = self.recv();
		json::decode(&*raw_json)
	}

	pub fn get_reg_info(&mut self) -> DecodeResult<LRegInfo> {
		self.send("drpj");
		let raw_json = self.recv();
		json::decode(&*raw_json)
	}

	pub fn get_flag_info(&mut self) -> DecodeResult<Vec<LFlagInfo>> {
		self.send("fj");
		let raw_json = self.recv();
		json::decode(&*raw_json)
	}

	pub fn get_fn_list(&mut self) -> Json {
		self.send("aflj");
		self.recv_json()
	}
}
<|endoftext|>


<commit_before>use std::io;
use std::u64;

const ALPHABET: &'static[u8] =   b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

#[inline(always)]
pub fn encode_wr<W:io::Write>(wr: &mut W, id: u64) -> io::Result<()> {
    let mut n = id;
    let alphasize = ALPHABET.len() as u64;
    let mut buf: [u8; 11] = [0u8; 11];
    let mut i = 0;

    loop {
        let rem = n % alphasize;
        n = n / alphasize;
        buf[i] = ALPHABET[rem as usize];
        i += 1;
        if n == 0 { break; }
    }

    wr.write_all(&buf[..i])
}

pub fn encode(id: u64) -> String {
    let mut v = Vec::new();
    encode_wr(&mut v, id).unwrap();
    String::from_utf8(v).unwrap()
}

pub fn decode(bytes: &[u8]) -> Option<u64> {
    let mut inv: Vec<i8> = (0..256).map(|_| -1).collect();
    for (i, &b) in ALPHABET.iter().enumerate() {
        inv[b as usize] = i as i8;
    }

    let mut failure = false;
    let n = {
        bytes.into_iter().rev().filter_map(|&b| {
            match inv.get(b as usize) {
                Some(&-1) | None => {failure = true; None},
                Some(&val) => Some(val as u8)
            }
        }).fold(0u64, |acc, item| acc * ALPHABET.len() as u64 + item as u64)
    };

    if failure { None }
    else { Some(n) }
}

#[test]
fn test_encode() {
    assert_eq!("c4NP", encode(12343344));
    assert_eq!("fyha61AhGYl", encode(u64::MAX));
    assert_eq!("0", encode(0));
}

#[test]
fn test_decode() {
    assert_eq!(Some(12343344), decode(b"c4NP"));
    assert_eq!(Some(0), decode(b"0"));
    assert_eq!(Some(u64::MAX), decode(b"fyha61AhGYl"));
}

#[test]
fn test_cyclic() {
    use super::rand::random;
    assert_eq!(Some(0), decode(encode(0).as_bytes()));
    for _ in 0 .. 1_000 {
        let i = random::<u64>();
        let enc = encode(i);
        let dec = decode(enc.as_bytes());
        assert_eq!(Some(i), dec);
    }
}
<commit_msg>Add encode_into_vec() function<commit_after>use std::io;
#[cfg(test)]
use std::u64;

const ALPHABET: &'static[u8] =   b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";

#[inline(always)]
pub fn encode_wr<W:io::Write>(wr: &mut W, id: u64) -> io::Result<()> {
    let mut n = id;
    let alphasize = ALPHABET.len() as u64;
    let mut buf: [u8; 11] = [0u8; 11];
    let mut i = 0;

    loop {
        let rem = n % alphasize;
        n = n / alphasize;
        buf[i] = ALPHABET[rem as usize];
        i += 1;
        if n == 0 { break; }
    }

    wr.write_all(&buf[..i])
}

#[inline(always)]
pub fn encode_into_vec(vec: &mut Vec<u8>, id: u64) {
    let mut n = id;
    let alphasize = ALPHABET.len() as u64;

    loop {
        let rem = n % alphasize;
        n = n / alphasize;
        vec.push(ALPHABET[rem as usize]);
        if n == 0 { break; }
    }
}

pub fn encode(id: u64) -> String {
    let mut v = Vec::new();
    encode_into_vec(&mut v, id);
    String::from_utf8(v).unwrap()
}

pub fn decode(bytes: &[u8]) -> Option<u64> {
    let mut inv: Vec<i8> = (0..256).map(|_| -1).collect();
    for (i, &b) in ALPHABET.iter().enumerate() {
        inv[b as usize] = i as i8;
    }

    let mut failure = false;
    let n = {
        bytes.into_iter().rev().filter_map(|&b| {
            match inv.get(b as usize) {
                Some(&-1) | None => {failure = true; None},
                Some(&val) => Some(val as u8)
            }
        }).fold(0u64, |acc, item| acc * ALPHABET.len() as u64 + item as u64)
    };

    if failure { None }
    else { Some(n) }
}

#[test]
fn test_encode() {
    assert_eq!("c4NP", encode(12343344));
    assert_eq!("fyha61AhGYl", encode(u64::MAX));
    assert_eq!("0", encode(0));
}

#[test]
fn test_decode() {
    assert_eq!(Some(12343344), decode(b"c4NP"));
    assert_eq!(Some(0), decode(b"0"));
    assert_eq!(Some(u64::MAX), decode(b"fyha61AhGYl"));
}

#[test]
fn test_cyclic() {
    use super::rand::random;
    assert_eq!(Some(0), decode(encode(0).as_bytes()));
    for _ in 0 .. 1_000 {
        let i = random::<u64>();
        let enc = encode(i);
        let dec = decode(enc.as_bytes());
        assert_eq!(Some(i), dec);
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Helpers for reading directory structures from the local filesystem.

use std::path::PathBuf;
use std::fs::PathExt;

use std::fs::{read_dir};

use std::sync;
use std::sync::mpsc;

use threadpool;

pub trait PathHandler<D> {
  fn handle_path(&self, D, PathBuf) -> Option<D>;
}

pub fn iterate_recursively<P: 'static + Send + Clone, W: 'static + PathHandler<P> + Send + Clone>
  (root: (PathBuf, P), worker: &mut W)
{
  let threads = 10;
  let (push_ch, work_ch) = mpsc::sync_channel(threads);
  let pool = threadpool::ThreadPool::new(threads);

  // Insert the first task into the queue:
  push_ch.send(Some(root)).unwrap();
  let mut running_workers = 0 as i32;

  // Master thread:
  loop {
    match work_ch.recv() {
      Err(_) => unreachable!(),
      Ok(None) => {
        // A worker has completed a task.
        // We are done when no more workers are active (i.e. all tasks are done):
        running_workers -= 1;
        if running_workers == 0 {
          break
        }
      },
      Ok(Some((root, payload))) => {
        // Execute the task in a pool thread:
        running_workers += 1;
        let _worker = worker.clone();
        let _push_ch = push_ch.clone();
        pool.execute(move|| {
          let res = read_dir(&root);
          if res.is_ok() {
            for entry in res.unwrap() {
              if entry.is_ok() {
                let entry = entry.unwrap();
                let file = entry.path();
                let path = PathBuf::new(file.to_str().unwrap());
                let dir_opt = _worker.handle_path(payload.clone(), path);
                if dir_opt.is_some() {
                  _push_ch.send(Some((file.clone(), dir_opt.unwrap()))).unwrap();
                }
              }
            }
          }
          // Count this pool thread as idle:
          _push_ch.send(None).unwrap();
        });
      }
    }
  }
}

struct PrintPathHandler;

impl Clone for PrintPathHandler {
  fn clone(&self) -> PrintPathHandler { PrintPathHandler }
}

impl PathHandler<()> for PrintPathHandler {
  fn handle_path(&self, _: (), path: PathBuf) -> Option<()> {
    println!("{}", path.display());
    match path.metadata() {
      Ok(ref m) if m.is_dir() => Some(()),
      _ => None,
    }
  }
}
<commit_msg>Remove unused import std::sync.<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Helpers for reading directory structures from the local filesystem.

use std::path::PathBuf;
use std::fs::PathExt;

use std::fs::{read_dir};

use std::sync::mpsc;

use threadpool;

pub trait PathHandler<D> {
  fn handle_path(&self, D, PathBuf) -> Option<D>;
}

pub fn iterate_recursively<P: 'static + Send + Clone, W: 'static + PathHandler<P> + Send + Clone>
  (root: (PathBuf, P), worker: &mut W)
{
  let threads = 10;
  let (push_ch, work_ch) = mpsc::sync_channel(threads);
  let pool = threadpool::ThreadPool::new(threads);

  // Insert the first task into the queue:
  push_ch.send(Some(root)).unwrap();
  let mut running_workers = 0 as i32;

  // Master thread:
  loop {
    match work_ch.recv() {
      Err(_) => unreachable!(),
      Ok(None) => {
        // A worker has completed a task.
        // We are done when no more workers are active (i.e. all tasks are done):
        running_workers -= 1;
        if running_workers == 0 {
          break
        }
      },
      Ok(Some((root, payload))) => {
        // Execute the task in a pool thread:
        running_workers += 1;
        let _worker = worker.clone();
        let _push_ch = push_ch.clone();
        pool.execute(move|| {
          let res = read_dir(&root);
          if res.is_ok() {
            for entry in res.unwrap() {
              if entry.is_ok() {
                let entry = entry.unwrap();
                let file = entry.path();
                let path = PathBuf::new(file.to_str().unwrap());
                let dir_opt = _worker.handle_path(payload.clone(), path);
                if dir_opt.is_some() {
                  _push_ch.send(Some((file.clone(), dir_opt.unwrap()))).unwrap();
                }
              }
            }
          }
          // Count this pool thread as idle:
          _push_ch.send(None).unwrap();
        });
      }
    }
  }
}

struct PrintPathHandler;

impl Clone for PrintPathHandler {
  fn clone(&self) -> PrintPathHandler { PrintPathHandler }
}

impl PathHandler<()> for PrintPathHandler {
  fn handle_path(&self, _: (), path: PathBuf) -> Option<()> {
    println!("{}", path.display());
    match path.metadata() {
      Ok(ref m) if m.is_dir() => Some(()),
      _ => None,
    }
  }
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Helpers for reading directory structures from the local filesystem.

use std::sync;
use std::sync::atomic;
use std::os::{last_os_error};
use std::io::{TypeDirectory};
use std::io::fs::{lstat};

use std::c_str::CString;
use libc::funcs::posix88::dirent;
use libc::types::common::posix88::{DIR,dirent_t};
use libc::types::os::arch::c95::c_char;
use libc::{c_int};

pub struct DirIterator {
  fd: *mut DIR,
}

impl DirIterator {
  pub fn new(path: Path) -> Result<DirIterator, String> {
    let fd = path.with_c_str(|c_str| unsafe { dirent::opendir(c_str) });

    if fd as int > 0 { Ok(DirIterator{fd:fd}) }
    else { Err(last_os_error()) }
  }

  fn read(&mut self) -> String {
    extern {
      fn rust_dirent_t_size() -> c_int;
      fn rust_list_dir_val(ptr: *mut dirent_t) -> *const c_char;
    }

    let mut entry_ptr = 0 as *mut dirent_t;

    let size = unsafe { rust_dirent_t_size() };
    let mut buf = Vec::<u8>::with_capacity(size as uint);
    let buf_ptr = buf.as_mut_ptr() as *mut dirent_t;

    let retval = unsafe { dirent::readdir_r(self.fd, buf_ptr, &mut entry_ptr) };

    if retval == 0 && !entry_ptr.is_null() {
      let cstr = unsafe { CString::new(rust_list_dir_val(entry_ptr), false) };
      cstr.as_str().expect("Path not UTF8.").into_string()
    } else { "".to_string() }

  }

}

impl Drop for DirIterator {
  fn drop(&mut self) {
    unsafe {
      dirent::closedir(self.fd)
    };
  }
}

impl Iterator<String> for DirIterator {
  fn next(&mut self) -> Option<String> {
    let name = self.read();
    if name.len() == 0 { None }
    else { Some(name) }
  }
}

pub trait PathHandler<D> {
  fn handle_path(&mut self, D, Path) -> Option<D>;
}

pub fn iterate_recursively<P: Send + Clone, W: PathHandler<P> + Send + Clone>
  (root: (Path, P), worker: &mut W)
{
  let threads = 5;
  let (push_ch, work_ch) = sync_channel(threads);
  let mut pool = sync::TaskPool::new(threads, || proc(_){()});

  let seq_cst = atomic::SeqCst;  // Sequentially consistent access mode
  let running_workers = sync::Arc::new(atomic::AtomicInt::new(0));

  // Insert the first task into the queue:
  push_ch.send(Some(root));

  // Master thread:
  loop {
    match work_ch.recv() {
      None => {
        // A worker has reached completion and has decremented the counter.
        // We are done when no more workers are active (i.e. all tasks are done):
        if running_workers.load(seq_cst) == 0 {
          break
        }
      },
      Some((root, payload)) => {
        let t_worker = worker.clone();
        let t_push_ch = push_ch.clone();

        // Count the pool thread as active:
        let t_running_workers = running_workers.clone();
        t_running_workers.fetch_add(1, seq_cst);

        // Execute the task in a pool thread:
        pool.execute(proc(&()) {
          let mut root = root;
          let mut t_worker = t_worker;
          let res = DirIterator::new(root.clone());
          if res.is_ok() {
            let mut it = res.unwrap();
            for file in it {
              if file != ".".into_string() && file != "..".into_string() {
                let rel_path = Path::new(file);
                root.push(rel_path);
                let dir_opt = t_worker.handle_path(payload.clone(), root.clone());
                if dir_opt.is_some() {
                  t_push_ch.send(Some((root.clone(), dir_opt.unwrap())));
                }
                root.pop();
              }
            }
          }

          // Count this pool thread as idle:
          t_running_workers.fetch_sub(1, seq_cst);
          t_push_ch.send(None)
        });
      }
    }
  }
}

struct PrintPathHandler;

impl Clone for PrintPathHandler {
  fn clone(&self) -> PrintPathHandler { PrintPathHandler }
}

impl PathHandler<()> for PrintPathHandler {
  fn handle_path(&mut self, _: (), path: Path) -> Option<()> {
    let filename_opt = path.filename_str();
    println!("{}", path.display());
    match filename_opt {
      Some(".") | Some("..") | None => None,
      Some(_) =>
      match lstat(&path) {
        Ok(ref st) if st.kind == TypeDirectory => Some(()),
        _ => None,
      }
    }
  }
}
<commit_msg>Remove the atomic variable in iterate_recursively as the value can now reside in the master thread.<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Helpers for reading directory structures from the local filesystem.

use std::sync;
use std::os::{last_os_error};
use std::io::{TypeDirectory};
use std::io::fs::{lstat};

use std::c_str::CString;
use libc::funcs::posix88::dirent;
use libc::types::common::posix88::{DIR,dirent_t};
use libc::types::os::arch::c95::c_char;
use libc::{c_int};

pub struct DirIterator {
  fd: *mut DIR,
}

impl DirIterator {
  pub fn new(path: Path) -> Result<DirIterator, String> {
    let fd = path.with_c_str(|c_str| unsafe { dirent::opendir(c_str) });

    if fd as int > 0 { Ok(DirIterator{fd:fd}) }
    else { Err(last_os_error()) }
  }

  fn read(&mut self) -> String {
    extern {
      fn rust_dirent_t_size() -> c_int;
      fn rust_list_dir_val(ptr: *mut dirent_t) -> *const c_char;
    }

    let mut entry_ptr = 0 as *mut dirent_t;

    let size = unsafe { rust_dirent_t_size() };
    let mut buf = Vec::<u8>::with_capacity(size as uint);
    let buf_ptr = buf.as_mut_ptr() as *mut dirent_t;

    let retval = unsafe { dirent::readdir_r(self.fd, buf_ptr, &mut entry_ptr) };

    if retval == 0 && !entry_ptr.is_null() {
      let cstr = unsafe { CString::new(rust_list_dir_val(entry_ptr), false) };
      cstr.as_str().expect("Path not UTF8.").into_string()
    } else { "".to_string() }

  }

}

impl Drop for DirIterator {
  fn drop(&mut self) {
    unsafe {
      dirent::closedir(self.fd)
    };
  }
}

impl Iterator<String> for DirIterator {
  fn next(&mut self) -> Option<String> {
    let name = self.read();
    if name.len() == 0 { None }
    else { Some(name) }
  }
}

pub trait PathHandler<D> {
  fn handle_path(&mut self, D, Path) -> Option<D>;
}

pub fn iterate_recursively<P: Send + Clone, W: PathHandler<P> + Send + Clone>
  (root: (Path, P), worker: &mut W)
{
  let threads = 5;
  let (push_ch, work_ch) = sync_channel(threads);
  let mut pool = sync::TaskPool::new(threads, || proc(_){()});

  // Insert the first task into the queue:
  push_ch.send(Some(root));
  let mut running_workers = 0i;

  // Master thread:
  loop {
    match work_ch.recv() {
      None => {
        // A worker has reached completion and has decremented the counter.
        // We are done when no more workers are active (i.e. all tasks are done):
        running_workers -= 1;
        if running_workers == 0 {
          break
        }
      },
      Some((root, payload)) => {
        let t_worker = worker.clone();
        let t_push_ch = push_ch.clone();

        // Execute the task in a pool thread:
        running_workers += 1;
        pool.execute(proc(&()) {
          let mut root = root;
          let mut t_worker = t_worker;
          let res = DirIterator::new(root.clone());
          if res.is_ok() {
            let mut it = res.unwrap();
            for file in it {
              if file != ".".into_string() && file != "..".into_string() {
                let rel_path = Path::new(file);
                root.push(rel_path);
                let dir_opt = t_worker.handle_path(payload.clone(), root.clone());
                if dir_opt.is_some() {
                  t_push_ch.send(Some((root.clone(), dir_opt.unwrap())));
                }
                root.pop();
              }
            }
          }

          // Count this pool thread as idle:
          t_push_ch.send(None)
        });
      }
    }
  }
}

struct PrintPathHandler;

impl Clone for PrintPathHandler {
  fn clone(&self) -> PrintPathHandler { PrintPathHandler }
}

impl PathHandler<()> for PrintPathHandler {
  fn handle_path(&mut self, _: (), path: Path) -> Option<()> {
    let filename_opt = path.filename_str();
    println!("{}", path.display());
    match filename_opt {
      Some(".") | Some("..") | None => None,
      Some(_) =>
      match lstat(&path) {
        Ok(ref st) if st.kind == TypeDirectory => Some(()),
        _ => None,
      }
    }
  }
}
<|endoftext|>


<commit_before>use chan;
use chan::{Sender, Receiver};
use std;
use std::borrow::Cow;

use datatype::{AccessToken, Auth, ClientId, ClientSecret, Command, Config,
               Error, Event, UpdateState, UpdateRequestId};
use gateway::Interpret;
use http::{AuthClient, Client};
use oauth2::authenticate;
use ota_plus::OTA;

pub trait Interpreter<I, O> {
    fn interpret(&mut self, input: I, otx: &Sender<O>);

    fn run(&mut self, irx: Receiver<I>, otx: Sender<O>) {
        loop {
            self.interpret(irx.recv().expect("interpreter sender closed"), &otx);
        }
    }
}

pub struct EventInterpreter;

impl Interpreter<Event, Command> for EventInterpreter {
    fn interpret(&mut self, event: Event, ctx: &Sender<Command>) {
        info!("Event interpreter: {:?}", event);
        match event {
            Event::NotAuthenticated => {
                debug!("Trying to authenticate again...");
                ctx.send(Command::Authenticate(None));
            }

            _ => ()
        }
    }
}

pub struct CommandInterpreter;

impl Interpreter<Command, Interpret> for CommandInterpreter {
    fn interpret(&mut self, cmd: Command, itx: &Sender<Interpret>) {
        info!("Command interpreter: {:?}", cmd);
        itx.send(Interpret { command: cmd, response_tx: None });
    }
}

pub struct GlobalInterpreter<'t> {
    pub config:      Config,
    pub token:       Option<Cow<'t, AccessToken>>,
    pub http_client: Box<Client>,
    pub loopback_tx: Sender<Interpret>,
}

impl<'t> Interpreter<Interpret, Event> for GlobalInterpreter<'t> {
    fn interpret(&mut self, interpret: Interpret, etx: &Sender<Event>) {
        info!("Global interpreter started: {:?}", interpret.command);

        let (multi_tx, multi_rx) = chan::async::<Event>();
        let outcome = match (&self.token, self.config.auth.is_none()) {
            (&Some(_), _) | (_, true) => self.authenticated(&interpret.command, multi_tx),
            _                         => self.unauthenticated(&interpret.command, multi_tx)
        };

        let mut response_ev: Option<Event> = None;
        match outcome {
            Ok(_) => {
                for ev in multi_rx {
                    etx.send(ev.clone());
                    response_ev = Some(ev);
                }
                info!("Global interpreter success: {:?}", interpret.command);
            }

            Err(Error::AuthorizationError(_)) => {
                let ev = Event::NotAuthenticated;
                etx.send(ev.clone());
                response_ev = Some(ev);
                error!("Global interpreter authentication failed: {:?}", interpret.command);
            }

            Err(err) => {
                let ev = Event::Error(format!("{}", err));
                etx.send(ev.clone());
                response_ev = Some(ev);
                error!("Global interpreter failed: {:?}: {}", interpret.command, err);
            }
        }

        let ev = response_ev.expect("no response event to send back");
        interpret.response_tx.map(|tx| { tx.lock().unwrap().send(ev); });
    }
}

impl<'t> GlobalInterpreter<'t> {
    fn authenticated(&self, cmd: &Command, etx: Sender<Event>) -> Result<(), Error> {
        let mut ota = OTA::new(&self.config, self.http_client.as_ref());

        // always send at least one Event response
        match *cmd {
            Command::AcceptUpdates(ref ids) => {
                // TODO: let _ = self.upstream.lock().unwrap().send_start_download(ids[0].clone());
                for id in ids {
                    info!("Accepting ID: {}", id);
                    etx.send(Event::UpdateStateChanged(id.clone(), UpdateState::Downloading));
                    let report = try!(ota.install_package_update(&id, &etx));
                    try!(ota.send_install_report(&report));
                    info!("Install Report for {}: {:?}", id, report);
                    try!(ota.update_installed_packages())
                }
            }

            Command::Authenticate(_) => etx.send(Event::Ok),

            Command::GetPendingUpdates => {
                let mut updates = try!(ota.get_package_updates());
                if updates.len() > 0 {
                    updates.sort_by_key(|u| u.installPos);
                    info!("New package updates available: {:?}", updates);
                    let ids = updates.iter().map(|u| u.requestId.clone()).collect::<Vec<UpdateRequestId>>();
                    self.loopback_tx.send(Interpret { command: Command::AcceptUpdates(ids), response_tx: None });
                }
                etx.send(Event::Ok);
            }

            Command::ListInstalledPackages => {
                let pkgs = try!(self.config.ota.package_manager.installed_packages());
                info!("Installed packages: {:?}", pkgs);
                etx.send(Event::Ok);
            }

            Command::Shutdown => std::process::exit(0),

            Command::UpdateInstalledPackages => {
                try!(ota.update_installed_packages());
                etx.send(Event::Ok);
                info!("Posted installed packages to the server.")
            }

            Command::UpdateReport(_) => {
                // upstream = rvi::Upstream
                // TODO: let _ = self.upstream.lock().unwrap().send_update_report(report);
            }

            Command::ReportInstalledSoftware(_) => {
                // TODO: let _ = self.upstream.lock().unwrap().send_installed_software(installed);
            }
        }

        Ok(())
    }

    fn unauthenticated(&mut self, cmd: &Command, etx: Sender<Event>) -> Result<(), Error> {
        match *cmd {
            Command::Authenticate(_) => {
                let config = self.config.auth.clone().expect("trying to authenticate without auth config");
                self.set_client(Auth::Credentials(ClientId(config.client_id), ClientSecret(config.secret)));
                let server = config.server.join("/token").expect("couldn't build authentication url");
                let token  = try!(authenticate(server, self.http_client.as_ref()));
                self.set_client(Auth::Token(token.clone()));
                self.token = Some(token.into());
                etx.send(Event::Authenticated);
            }

            Command::AcceptUpdates(_)           |
            Command::GetPendingUpdates          |
            Command::ListInstalledPackages      |
            Command::UpdateReport(_)            |
            Command::ReportInstalledSoftware(_) |
            Command::UpdateInstalledPackages     => etx.send(Event::NotAuthenticated),

            Command::Shutdown => std::process::exit(0),
        }

        Ok(())
    }

    fn set_client(&mut self, auth: Auth) {
        if !self.http_client.is_testing() {
            self.http_client = Box::new(AuthClient::from(auth));
        }
    }
}

#[cfg(test)]
mod tests {
    use chan;
    use chan::{Sender, Receiver};
    use std::thread;

    use super::*;
    use datatype::{AccessToken, Command, Config, Event, UpdateState};
    use gateway::Interpret;
    use http::test_client::TestClient;
    use package_manager::PackageManager;
    use package_manager::tpm::assert_rx;

    fn new_interpreter(replies: Vec<String>, pkg_mgr: PackageManager) -> (Sender<Command>, Receiver<Event>) {
        let (etx, erx) = chan::sync::<Event>(0);
        let (ctx, crx) = chan::sync::<Command>(0);
        let (itx, _)   = chan::sync::<Interpret>(0);

        thread::spawn(move || {
            let mut gi = GlobalInterpreter {
                config:      Config::default(),
                token:       Some(AccessToken::default().into()),
                http_client: Box::new(TestClient::from(replies)),
                loopback_tx: itx,
            };
            gi.config.ota.package_manager = pkg_mgr;

            loop {
                match crx.recv() {
                    Some(cmd) => gi.interpret(Interpret { command: cmd, response_tx: None }, &etx),
                    None      => break
                }
            }
        });

        (ctx, erx)
    }

    #[test]
    fn already_authenticated() {
        let replies    = Vec::new();
        let pkg_mgr    = PackageManager::new_file(true);
        let (ctx, erx) = new_interpreter(replies, pkg_mgr);

        ctx.send(Command::Authenticate(None));
        assert_rx(erx, &[Event::Ok]);
    }

    #[test]
    fn accept_updates() {
        let replies    = vec!["[]".to_string(); 10];
        let pkg_mgr    = PackageManager::new_file(true);
        let (ctx, erx) = new_interpreter(replies, pkg_mgr);

        ctx.send(Command::AcceptUpdates(vec!["1".to_string(), "2".to_string()]));
        assert_rx(erx, &[
            Event::UpdateStateChanged("1".to_string(), UpdateState::Downloading),
            Event::UpdateStateChanged("1".to_string(), UpdateState::Installing),
            Event::UpdateStateChanged("1".to_string(), UpdateState::Installed),
            Event::UpdateStateChanged("2".to_string(), UpdateState::Downloading),
            Event::UpdateStateChanged("2".to_string(), UpdateState::Installing),
            Event::UpdateStateChanged("2".to_string(), UpdateState::Installed),
        ]);
    }

    #[test]
    fn failed_updates() {
        let replies    = vec!["[]".to_string(); 10];
        let pkg_mgr    = PackageManager::new_file(false);
        let (ctx, erx) = new_interpreter(replies, pkg_mgr);

        ctx.send(Command::AcceptUpdates(vec!["1".to_string()]));
        assert_rx(erx, &[Event::Error("IO error: No such file or directory (os error 2)".to_owned())]);
    }
}
<commit_msg>Send FoundInstalledPackages Event on ListInstalledPackages<commit_after>use chan;
use chan::{Sender, Receiver};
use std;
use std::borrow::Cow;

use datatype::{AccessToken, Auth, ClientId, ClientSecret, Command, Config,
               Error, Event, UpdateState, UpdateRequestId};
use gateway::Interpret;
use http::{AuthClient, Client};
use oauth2::authenticate;
use ota_plus::OTA;

pub trait Interpreter<I, O> {
    fn interpret(&mut self, input: I, otx: &Sender<O>);

    fn run(&mut self, irx: Receiver<I>, otx: Sender<O>) {
        loop {
            self.interpret(irx.recv().expect("interpreter sender closed"), &otx);
        }
    }
}

pub struct EventInterpreter;

impl Interpreter<Event, Command> for EventInterpreter {
    fn interpret(&mut self, event: Event, ctx: &Sender<Command>) {
        info!("Event interpreter: {:?}", event);
        match event {
            Event::NotAuthenticated => {
                debug!("Trying to authenticate again...");
                ctx.send(Command::Authenticate(None));
            }

            _ => ()
        }
    }
}

pub struct CommandInterpreter;

impl Interpreter<Command, Interpret> for CommandInterpreter {
    fn interpret(&mut self, cmd: Command, itx: &Sender<Interpret>) {
        info!("Command interpreter: {:?}", cmd);
        itx.send(Interpret { command: cmd, response_tx: None });
    }
}

pub struct GlobalInterpreter<'t> {
    pub config:      Config,
    pub token:       Option<Cow<'t, AccessToken>>,
    pub http_client: Box<Client>,
    pub loopback_tx: Sender<Interpret>,
}

impl<'t> Interpreter<Interpret, Event> for GlobalInterpreter<'t> {
    fn interpret(&mut self, interpret: Interpret, etx: &Sender<Event>) {
        info!("Global interpreter started: {:?}", interpret.command);

        let (multi_tx, multi_rx) = chan::async::<Event>();
        let outcome = match (&self.token, self.config.auth.is_none()) {
            (&Some(_), _) | (_, true) => self.authenticated(&interpret.command, multi_tx),
            _                         => self.unauthenticated(&interpret.command, multi_tx)
        };

        let mut response_ev: Option<Event> = None;
        match outcome {
            Ok(_) => {
                for ev in multi_rx {
                    etx.send(ev.clone());
                    response_ev = Some(ev);
                }
                info!("Global interpreter success: {:?}", interpret.command);
            }

            Err(Error::AuthorizationError(_)) => {
                let ev = Event::NotAuthenticated;
                etx.send(ev.clone());
                response_ev = Some(ev);
                error!("Global interpreter authentication failed: {:?}", interpret.command);
            }

            Err(err) => {
                let ev = Event::Error(format!("{}", err));
                etx.send(ev.clone());
                response_ev = Some(ev);
                error!("Global interpreter failed: {:?}: {}", interpret.command, err);
            }
        }

        let ev = response_ev.expect("no response event to send back");
        interpret.response_tx.map(|tx| { tx.lock().unwrap().send(ev); });
    }
}

impl<'t> GlobalInterpreter<'t> {
    fn authenticated(&self, cmd: &Command, etx: Sender<Event>) -> Result<(), Error> {
        let mut ota = OTA::new(&self.config, self.http_client.as_ref());

        // always send at least one Event response
        match *cmd {
            Command::AcceptUpdates(ref ids) => {
                // TODO: let _ = self.upstream.lock().unwrap().send_start_download(ids[0].clone());
                for id in ids {
                    info!("Accepting ID: {}", id);
                    etx.send(Event::UpdateStateChanged(id.clone(), UpdateState::Downloading));
                    let report = try!(ota.install_package_update(&id, &etx));
                    try!(ota.send_install_report(&report));
                    info!("Install Report for {}: {:?}", id, report);
                    try!(ota.update_installed_packages())
                }
            }

            Command::Authenticate(_) => etx.send(Event::Ok),

            Command::GetPendingUpdates => {
                let mut updates = try!(ota.get_package_updates());
                if updates.len() > 0 {
                    updates.sort_by_key(|u| u.installPos);
                    info!("New package updates available: {:?}", updates);
                    let ids = updates.iter().map(|u| u.requestId.clone()).collect::<Vec<UpdateRequestId>>();
                    self.loopback_tx.send(Interpret { command: Command::AcceptUpdates(ids), response_tx: None });
                }
                etx.send(Event::Ok);
            }

            Command::ListInstalledPackages => {
                let pkgs = try!(self.config.ota.package_manager.installed_packages());
                etx.send(Event::FoundInstalledPackages(pkgs));
            }

            Command::Shutdown => std::process::exit(0),

            Command::UpdateInstalledPackages => {
                try!(ota.update_installed_packages());
                etx.send(Event::Ok);
                info!("Posted installed packages to the server.")
            }

            Command::UpdateReport(_) => {
                // upstream = rvi::Upstream
                // TODO: let _ = self.upstream.lock().unwrap().send_update_report(report);
            }

            Command::ReportInstalledSoftware(_) => {
                // TODO: let _ = self.upstream.lock().unwrap().send_installed_software(installed);
            }
        }

        Ok(())
    }

    fn unauthenticated(&mut self, cmd: &Command, etx: Sender<Event>) -> Result<(), Error> {
        match *cmd {
            Command::Authenticate(_) => {
                let config = self.config.auth.clone().expect("trying to authenticate without auth config");
                self.set_client(Auth::Credentials(ClientId(config.client_id), ClientSecret(config.secret)));
                let server = config.server.join("/token").expect("couldn't build authentication url");
                let token  = try!(authenticate(server, self.http_client.as_ref()));
                self.set_client(Auth::Token(token.clone()));
                self.token = Some(token.into());
                etx.send(Event::Authenticated);
            }

            Command::AcceptUpdates(_)           |
            Command::GetPendingUpdates          |
            Command::ListInstalledPackages      |
            Command::UpdateReport(_)            |
            Command::ReportInstalledSoftware(_) |
            Command::UpdateInstalledPackages     => etx.send(Event::NotAuthenticated),

            Command::Shutdown => std::process::exit(0),
        }

        Ok(())
    }

    fn set_client(&mut self, auth: Auth) {
        if !self.http_client.is_testing() {
            self.http_client = Box::new(AuthClient::from(auth));
        }
    }
}

#[cfg(test)]
mod tests {
    use chan;
    use chan::{Sender, Receiver};
    use std::thread;

    use super::*;
    use datatype::{AccessToken, Command, Config, Event, UpdateState};
    use gateway::Interpret;
    use http::test_client::TestClient;
    use package_manager::PackageManager;
    use package_manager::tpm::assert_rx;

    fn new_interpreter(replies: Vec<String>, pkg_mgr: PackageManager) -> (Sender<Command>, Receiver<Event>) {
        let (etx, erx) = chan::sync::<Event>(0);
        let (ctx, crx) = chan::sync::<Command>(0);
        let (itx, _)   = chan::sync::<Interpret>(0);

        thread::spawn(move || {
            let mut gi = GlobalInterpreter {
                config:      Config::default(),
                token:       Some(AccessToken::default().into()),
                http_client: Box::new(TestClient::from(replies)),
                loopback_tx: itx,
            };
            gi.config.ota.package_manager = pkg_mgr;

            loop {
                match crx.recv() {
                    Some(cmd) => gi.interpret(Interpret { command: cmd, response_tx: None }, &etx),
                    None      => break
                }
            }
        });

        (ctx, erx)
    }

    #[test]
    fn already_authenticated() {
        let replies    = Vec::new();
        let pkg_mgr    = PackageManager::new_file(true);
        let (ctx, erx) = new_interpreter(replies, pkg_mgr);

        ctx.send(Command::Authenticate(None));
        assert_rx(erx, &[Event::Ok]);
    }

    #[test]
    fn accept_updates() {
        let replies    = vec!["[]".to_string(); 10];
        let pkg_mgr    = PackageManager::new_file(true);
        let (ctx, erx) = new_interpreter(replies, pkg_mgr);

        ctx.send(Command::AcceptUpdates(vec!["1".to_string(), "2".to_string()]));
        assert_rx(erx, &[
            Event::UpdateStateChanged("1".to_string(), UpdateState::Downloading),
            Event::UpdateStateChanged("1".to_string(), UpdateState::Installing),
            Event::UpdateStateChanged("1".to_string(), UpdateState::Installed),
            Event::UpdateStateChanged("2".to_string(), UpdateState::Downloading),
            Event::UpdateStateChanged("2".to_string(), UpdateState::Installing),
            Event::UpdateStateChanged("2".to_string(), UpdateState::Installed),
        ]);
    }

    #[test]
    fn failed_updates() {
        let replies    = vec!["[]".to_string(); 10];
        let pkg_mgr    = PackageManager::new_file(false);
        let (ctx, erx) = new_interpreter(replies, pkg_mgr);

        ctx.send(Command::AcceptUpdates(vec!["1".to_string()]));
        assert_rx(erx, &[Event::Error("IO error: No such file or directory (os error 2)".to_owned())]);
    }
}
<|endoftext|>


<commit_before>use header::*;
use constants::*;
use instructions::Instruction as Instr;
use num::FromPrimitive;
use std::fmt;

#[derive(Debug, Default)]
pub struct Runtime {
    interpreter: Interpreter,
    call_stack: Vec<StackFrame>,
}

impl Runtime {
    pub fn new(mut data: Vec<u8>) -> Runtime {
        let mut r = Runtime {
            interpreter: Interpreter::new(data),
            call_stack: Vec::new(),
        };
        r
    }
    pub fn run(&mut self, options: &::Options) {
        debug!("Running...");

        let debug = options.debug;

        loop {
            let res = self.dispatch_frame(debug);
            match res {
                None => {},
                Some(frm) => {self.call_stack.push(frm);},
            }
        }
    }

    fn dispatch_frame(&mut self, debug: bool) -> Option<StackFrame> {
        let ln = self.call_stack.len();
        let mut last_frame = &mut self.call_stack[ln-1];
        return last_frame.dispatch(&mut self.interpreter, debug);
    }
}

#[derive(Debug, Default)]
pub struct Interpreter {
    // File data
    header: RaptorHeader,
    const_table: ConstTable,

    // Rutime stuff
    op_stack: Vec<i32>,
    memory: Vec<i32>,
}

// Should this be here?

#[derive(Debug, Default, Clone)]
pub struct StackFrame {
    locals: Vec<i32>,
    // The index of the first op in the op_stack that should be kept
    return_addr: usize,
    bytecode: Vec<u8>,
    bc_counter: usize
}

impl Interpreter {
    pub fn new(mut data: Vec<u8>) -> Interpreter {
        debug!("Bytecode length: {} bytes", data.len());
        let header = read_header(&data);
        data.drain(..HEADER_SIZE);
        let const_table: ConstTable = read_const_table(data.as_slice());
        debug!("Constant table length: {} bytes", const_table.bc_counter);
        debug!("Bytecode length: {} bytes", data.len());
        data.drain(..const_table.bc_counter);
        let sf = StackFrame {
            bytecode: data,
            ..Default::default()
        };
        let mut i = Interpreter {
            header: header,
            const_table: const_table,
            op_stack: Vec::new(),
            memory: Vec::new(),
        };
        i.memory.resize(i.header.var_count as usize, 0);
        i
    }
}

impl StackFrame {

    fn get_next_4_bytes(&mut self) -> u32 {
        let val = (self.bytecode[self.bc_counter] as u32) << 24 |
        (self.bytecode[self.bc_counter + 1] as u32) << 16 |
        (self.bytecode[self.bc_counter + 2] as u32) << 8 |
        (self.bytecode[self.bc_counter + 3] as u32);
        self.bc_counter += 4;
        debug!("get_next_4_bytes: 0x{:04X}", val);
        val
    }

    fn dispatch(&mut self, inpr: &mut Interpreter, debug: bool) -> Option<StackFrame> {
        use std::ops::*;
        use std::cmp::*;

        // Main loop
        while self.bc_counter != self.bytecode.len() {
            // info!("PC: {}", bc_counter);

            // Use FromPrimitive trait to convert a value to its enum
            let instr = Instr::from_u8(self.bytecode[self.bc_counter]);
            self.bc_counter += 1;

            if instr.is_none() {
                warn!("Unimplemented instruction: {:04X}", self.bytecode[self.bc_counter]);
                continue;
            }

            let instr = instr.unwrap();   // We're sure it's Some here, so unpack it.

            if debug {
                debug!("{:?}", instr);
            }

            macro_rules! push {
                ( $x:expr ) => {
                    inpr.op_stack.push($x);
                };
            }
            macro_rules! pop {
                () => {
                    inpr.op_stack.pop();
                };
            }
            macro_rules! operation {
                ($op:ident) => ({
                    let final_length: usize = inpr.op_stack.len().saturating_sub(2);
                    let val = inpr.op_stack.drain(final_length..).fold(
                        0, |acc, x| acc.$op(x)
                    );
                    push!(val);
                })
            }
            macro_rules! reljump {
                ($op:ident) => ({
                    let top = pop!().unwrap();
                    if top.$op(&0) {
                        reljump!();
                    } else {
                        self.get_next_4_bytes();
                        if debug {debug!("Jump not taken"); }
                    }
                });
                () => ({
                    let offset = (self.get_next_4_bytes() - 1) as i32;
                    // Need this if because you can't have negative usizes
                    if offset > 0 {
                        if debug {debug!("RELJUMP: {}", offset);}
                        self.bc_counter += offset as usize;
                        if offset == 1 {debug!("RELJUMP 1 is redundant. This is a compiler bug")}
                    } else if offset < 0 {
                        if debug {debug!("RELJUMP: {}", offset);}
                        self.bc_counter -= (-offset) as usize;
                    } else {
                        warn!("Invalid reljump offset: 0");
                    }
                });
            }

            macro_rules! push_frame {
                ($id:expr) => ({
                    let func_const = &inpr.const_table.funcs[$id as usize];
                    let mut sf = StackFrame {
                        locals: Vec::new(),
                        ..Default::default()
                    };
                    for _ in 0..func_const.arg_count {
                        sf.locals.push(pop!().unwrap());
                    }
                    sf.return_addr = inpr.op_stack.len();
                    sf.locals.resize((func_const.arg_count + func_const.local_count) as usize, 0);
                    debug!("Pushed new frame: {}    {:?}", $id, sf);
                    debug!("Op stack: {:?}", inpr.op_stack);
                    return Some(sf);
                });
            }
            // TODO: More macros, less code

            match instr {
                Instr::NOP => {},
                Instr::HALT => {
                    println!("HALT issued, stopped execution.");
                    if debug {
                        debug!("Stack: {:?}", inpr.op_stack);
                        debug!("Memory: {:?}", inpr.memory);
                    }
                },
                Instr::ICONST => {
                    let b = self.get_next_4_bytes() as i32;
                    push!(b);
                },
                Instr::POP => { pop!(); },
                Instr::ADD =>       { operation!(add);    },
                Instr::SUB =>       { operation!(sub);    },
                Instr::MULTIPLY =>  { operation!(mul);    },
                Instr::DIVIDE =>    { operation!(div);    },
                Instr::MODULUS =>   { operation!(rem);    },
                Instr::RSHIFT =>    { operation!(shl);    },
                Instr::LSHIFT =>    { operation!(shr);    },
                Instr::AND =>       { operation!(bitand); },
                Instr::OR =>        { operation!(bitor);  },
                Instr::NOT =>       {
                    let val = pop!().unwrap();
                    push!(val.not());
                },
                Instr::COMP => {
                    let a = pop!(); let b = pop!();
                    push!(if a > b {1} else if a < b {-1} else {0});
                },
                Instr::COMP_LT => {
                    let a = pop!(); let b = pop!();
                    push!(if a < b {1} else {0});
                },
                Instr::COMP_EQ => {
                    let a = pop!(); let b = pop!();
                    push!(if a == b {1} else {0});
                },
                Instr::COMP_GT => {
                    let a = pop!(); let b = pop!();
                    push!(if a > b {1} else {0});
                },
                Instr::RELJUMP => {reljump!();},
                Instr::RELJUMP_GT => {reljump!(gt);},
                Instr::RELJUMP_LT => {reljump!(lt);},
                Instr::RELJUMP_EQ => {reljump!(eq);},
                Instr::STORE => {
                    let index = self.get_next_4_bytes() as usize;
                    self.locals[index] = pop!().unwrap();
                },
                Instr::LOAD => {
                    let index = self.get_next_4_bytes() as usize;
                    push!(self.locals[index]);
                },
                Instr::CALL => {
                    let id: u32 = self.get_next_4_bytes();
                    push_frame!(id);
                },
                Instr::RETURN => {
                    let val = pop!();
                    inpr.op_stack.resize(self.return_addr, 0)
                }
                Instr::PRINT => {
                    println!("PRINT: {}", pop!().unwrap());
                },
                Instr::DUMP_STACK => {
                    println!("{:?}", inpr.op_stack);
                },
                Instr::DUMP_GLOBALS => {
                    println!("{:?}", inpr.memory);},
            }
        }
        return None;    // Pop the current frame
    }
}

// TODO: Add tests
<commit_msg>Stack frames are in! buggy as shit, but theoretically working!<commit_after>use header::*;
use constants::*;
use instructions::Instruction as Instr;
use num::FromPrimitive;
use std::fmt;

#[derive(Debug, Default)]
pub struct Runtime {
    interpreter: Interpreter,
    call_stack: Vec<StackFrame>,
}

impl Runtime {
    pub fn new(mut data: Vec<u8>) -> Runtime {
        let mut r = Runtime {
            interpreter: Interpreter::new(data),
            call_stack: Vec::new(),
        };
        let prog_bc = r.interpreter.prog_bytecode.clone();
        r.call_stack.push(StackFrame {bytecode: prog_bc, ..Default::default()});
        r
    }
    pub fn run(&mut self, options: &::Options) {
        debug!("Running...");

        let debug = options.debug;

        while self.call_stack.len() > 0 {
            let res = self.dispatch_frame(debug);
            match res {
                None => {self.call_stack.pop();},
                Some(frm) => {self.call_stack.push(frm);},
            }
        }
    }

    fn dispatch_frame(&mut self, debug: bool) -> Option<StackFrame> {
        let ln = self.call_stack.len();
        let mut last_frame = &mut self.call_stack[ln-1];
        return last_frame.dispatch(&mut self.interpreter, debug);
    }
}

#[derive(Debug, Default)]
pub struct Interpreter {
    // File data
    header: RaptorHeader,
    const_table: ConstTable,

    // Rutime stuff
    op_stack: Vec<i32>,
    memory: Vec<i32>,
    prog_bytecode: Vec<u8>,
}

// Should this be here?

#[derive(Debug, Default, Clone)]
pub struct StackFrame {
    locals: Vec<i32>,
    // The index of the first op in the op_stack that should be kept
    return_addr: usize,
    bytecode: Vec<u8>,
    bc_counter: usize
}

impl Interpreter {
    pub fn new(mut data: Vec<u8>) -> Interpreter {
        debug!("Bytecode length: {} bytes", data.len());
        let header = read_header(&data);
        data.drain(..HEADER_SIZE);
        let const_table: ConstTable = read_const_table(data.as_slice());
        debug!("Constant table length: {} bytes", const_table.bc_counter);
        debug!("Bytecode length: {} bytes", data.len());
        data.drain(..const_table.bc_counter);
        let mut i = Interpreter {
            header: header,
            const_table: const_table,
            op_stack: Vec::new(),
            memory: Vec::new(),
            prog_bytecode: data,
        };
        i.memory.resize(i.header.var_count as usize, 0);
        i
    }
}

impl StackFrame {

    fn get_next_4_bytes(&mut self) -> u32 {
        let val = (self.bytecode[self.bc_counter] as u32) << 24 |
        (self.bytecode[self.bc_counter + 1] as u32) << 16 |
        (self.bytecode[self.bc_counter + 2] as u32) << 8 |
        (self.bytecode[self.bc_counter + 3] as u32);
        self.bc_counter += 4;
        debug!("get_next_4_bytes: 0x{:04X}", val);
        val
    }

    fn dispatch(&mut self, inpr: &mut Interpreter, debug: bool) -> Option<StackFrame> {
        use std::ops::*;
        use std::cmp::*;

        // Main loop
        while self.bc_counter != self.bytecode.len() {
            // info!("PC: {}", bc_counter);

            // Use FromPrimitive trait to convert a value to its enum
            let instr = Instr::from_u8(self.bytecode[self.bc_counter]);
            self.bc_counter += 1;

            if instr.is_none() {
                warn!("Unimplemented instruction: {:04X}", self.bytecode[self.bc_counter]);
                continue;
            }

            let instr = instr.unwrap();   // We're sure it's Some here, so unpack it.

            if debug {
                debug!("{:?}", instr);
            }

            macro_rules! push {
                ( $x:expr ) => {
                    inpr.op_stack.push($x);
                };
            }
            macro_rules! pop {
                () => {
                    inpr.op_stack.pop();
                };
            }
            macro_rules! operation {
                ($op:ident) => ({
                    let final_length: usize = inpr.op_stack.len().saturating_sub(2);
                    let val = inpr.op_stack.drain(final_length..).fold(
                        0, |acc, x| acc.$op(x)
                    );
                    push!(val);
                })
            }
            macro_rules! reljump {
                ($op:ident) => ({
                    let top = pop!().unwrap();
                    if top.$op(&0) {
                        reljump!();
                    } else {
                        self.get_next_4_bytes();
                        if debug {debug!("Jump not taken"); }
                    }
                });
                () => ({
                    let offset = (self.get_next_4_bytes() - 1) as i32;
                    // Need this if because you can't have negative usizes
                    if offset > 0 {
                        if debug {debug!("RELJUMP: {}", offset);}
                        self.bc_counter += offset as usize;
                        if offset == 1 {debug!("RELJUMP 1 is redundant. This is a compiler bug")}
                    } else if offset < 0 {
                        if debug {debug!("RELJUMP: {}", offset);}
                        self.bc_counter -= (-offset) as usize;
                    } else {
                        warn!("Invalid reljump offset: 0");
                    }
                });
            }

            macro_rules! push_frame {
                ($id:expr) => ({
                    let func_const = &inpr.const_table.funcs[$id as usize];
                    let mut sf = StackFrame {
                        locals: Vec::new(),
                        bytecode: func_const.body.clone(),
                        ..Default::default()
                    };
                    for _ in 0..func_const.arg_count {
                        sf.locals.push(pop!().unwrap());
                    }
                    sf.return_addr = inpr.op_stack.len();
                    sf.locals.resize((func_const.arg_count + func_const.local_count) as usize, 0);
                    debug!("Pushed new frame: {}    {:?}", $id, sf);
                    debug!("Op stack: {:?}", inpr.op_stack);
                    return Some(sf);
                });
            }
            // TODO: More macros, less code

            match instr {
                Instr::NOP => {},
                Instr::HALT => {
                    println!("HALT issued, stopped execution.");
                    if debug {
                        debug!("Stack: {:?}", inpr.op_stack);
                        debug!("Memory: {:?}", inpr.memory);
                    }
                },
                Instr::ICONST => {
                    let b = self.get_next_4_bytes() as i32;
                    push!(b);
                },
                Instr::POP => { pop!(); },
                Instr::ADD =>       { operation!(add);    },
                Instr::SUB =>       { operation!(sub);    },
                Instr::MULTIPLY =>  { operation!(mul);    },
                Instr::DIVIDE =>    { operation!(div);    },
                Instr::MODULUS =>   { operation!(rem);    },
                Instr::RSHIFT =>    { operation!(shl);    },
                Instr::LSHIFT =>    { operation!(shr);    },
                Instr::AND =>       { operation!(bitand); },
                Instr::OR =>        { operation!(bitor);  },
                Instr::NOT =>       {
                    let val = pop!().unwrap();
                    push!(val.not());
                },
                Instr::COMP => {
                    let a = pop!(); let b = pop!();
                    push!(if a > b {1} else if a < b {-1} else {0});
                },
                Instr::COMP_LT => {
                    let a = pop!(); let b = pop!();
                    push!(if a < b {1} else {0});
                },
                Instr::COMP_EQ => {
                    let a = pop!(); let b = pop!();
                    push!(if a == b {1} else {0});
                },
                Instr::COMP_GT => {
                    let a = pop!(); let b = pop!();
                    push!(if a > b {1} else {0});
                },
                Instr::RELJUMP => {reljump!();},
                Instr::RELJUMP_GT => {reljump!(gt);},
                Instr::RELJUMP_LT => {reljump!(lt);},
                Instr::RELJUMP_EQ => {reljump!(eq);},
                Instr::STORE => {
                    let index = self.get_next_4_bytes() as usize;
                    self.locals[index] = pop!().unwrap();
                },
                Instr::LOAD => {
                    let index = self.get_next_4_bytes() as usize;
                    push!(self.locals[index]);
                },
                Instr::CALL => {
                    let id: u32 = self.get_next_4_bytes();
                    return push_frame!(id);
                },
                Instr::RETURN => {
                    let val = pop!().unwrap();
                    inpr.op_stack.resize(self.return_addr + 1, 0);
                    push!(val);
                    return None;
                }
                Instr::PRINT => {
                    println!("PRINT: {}", pop!().unwrap());
                },
                Instr::DUMP_STACK => {
                    println!("{:?}", inpr.op_stack);
                },
                Instr::DUMP_GLOBALS => {
                    println!("{:?}", inpr.memory);},
            }
        }
        return None;    // Pop the current frame
    }
}

// TODO: Add tests
<|endoftext|>