main.rs

extern crate libc;

use std::io::{self, Read, Write};
use std::process::{Command, Stdio};
use std::{env, str};

#[macro_use]
extern crate pest_derive;
extern crate rustc_demangle;
extern crate stack_sizes;

use pest::Parser;
use rustc_demangle::demangle;

use std::fs::File;

use petgraph::{
    algo::is_cyclic_directed,
    dot::{Config, Dot},
    graph::NodeIndex,
    visit::DfsPostOrder,
    Graph, Incoming,
};
use std::collections::hash_map::HashMap;
use std::fmt;

#[derive(Parser)]
#[grammar = "ident.pest"]
struct IdentParser;

#[derive(Debug)]
struct Out {
    hash: Option<String>,
    crate_name: Option<String>,
    out_dir: Option<String>,
}

// Parse the output from rustc
// we skip the last line (Finished ...)
// and look for `--crate-name x`, `extra-filename=y` and `--out dir z`
//
// Notice, the parsing is a bit of a hack, implemnting look-ahead using flags.
fn parse_out(out_str: &str) -> Out {
    let mut out = Out {
        hash: None,
        crate_name: None,
        out_dir: None,
    };
    let output = out_str;
    let mut i = output.lines().into_iter();
    i.next_back(); // skip last line (Finished)
    if let Some(line) = i.next_back() {
        let mut crate_name = false;
        let mut out_dir = false;
        for part in line.split(' ') {
            if crate_name {
                out.crate_name = Some(part.to_string());
                crate_name = false;
            } else if out_dir {
                out.out_dir = Some(part.to_string());
                out_dir = false;
            } else if part.starts_with("--crate-name") {
                crate_name = true;
            } else if part.starts_with("--out-dir") {
                out_dir = true;
            } else if part.starts_with("extra-filename=") {
                out.hash = Some(part.split('=').nth(1).unwrap().to_string());
            }
        }
    }
    out
}

struct Node {
    name: String,
    local_stack: u32,
    call_stack: u32,
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{:#}\\n stack local = {}, call = {}",
            demangle(&self.name),
            self.local_stack,
            self.call_stack
        )
    }
}

fn new_node(
    hm: &mut HashMap<String, NodeIndex<u32>>,
    g: &mut Graph<Node, ()>,
    s: String,
    v: u32,
) -> NodeIndex<u32> {
    let n = Node {
        name: s.to_string(),
        local_stack: v,
        call_stack: v,
    };
    let i = g.add_node(n);
    hm.insert(s, i);
    i
}

fn add_edge<'a>(
    hm: &mut HashMap<String, NodeIndex<u32>>,
    g: &mut Graph<Node, ()>,
    from: &str,
    to: &str,
) {
    let i_from = hm.get(from).expect("from not found");
    if let Some(i_to) = hm.get(to) {
        g.add_edge(*i_from, *i_to, ());
    }
}

fn read_elf(path: &str, buffer: &mut Vec<u8>) {
    let mut elf = File::open(path).expect(format!("error opening '{}'", path).as_str());
    elf.read_to_end(buffer)
        .expect(format!("error opening '{}'", path).as_str());
}

fn parse_elf<'a>(
    hm: &mut HashMap<String, NodeIndex<u32>>,
    g: &mut Graph<Node, ()>,
    path: &str,
    vec: &'a mut Vec<u8>,
) {
    read_elf(path, vec);

    // read and analyse stack-sizes of elf
    let ss = stack_sizes::analyze(&vec).expect("error analysing elf");
    if let Some(v32) = ss.left() {
        // put into graph
        for s in v32.iter() {
            let d_name = s.name().to_string(); //  format!("{:#}", demangle(s.name()));
            let _ = new_node(hm, g, d_name, s.stack() as u32);
        }
    } else {
        panic!("64 bit elfs not supported");
    }
}

fn parse_callgraph(hm: &mut HashMap<String, NodeIndex<u32>>, g: &mut Graph<Node, ()>, input: &str) {
    let functions = IdentParser::parse(Rule::ident_list, &input)
        .unwrap_or_else(|e| panic!("parse error {}", e));

    for f in functions {
        match f.as_rule() {
            Rule::fun => {
                let mut p = f.into_inner();
                let mut i = p.next().unwrap().into_inner();
                let from = i.next().unwrap().into_inner();

                for c in p {
                    match c.as_rule() {
                        Rule::callext => {}
                        Rule::callfun => {
                            let mut p = c.into_inner();
                            p.next().unwrap();
                            let to = p.next().unwrap().into_inner().as_str().to_string();
                            add_edge(hm, g, from.as_str(), to.as_str());
                        }
                        _ => panic!("internal error"),
                    }
                }
            }
            Rule::null => (),
            Rule::EOI => (),
            _ => {
                panic!("internal error");
            }
        }
    }
}

fn main() {
    println!("start sub command");
    // first argument is the path to this binary; the second argument is always "call-stack" -- both can
    // be ignored
    let args = env::args().skip(2).collect::<Vec<_>>();

    // turn `cargo klee --example foo` into `cargo rustc --example foo -- (..)`
    let mut c = Command::new("cargo");
    c.arg("rustc")
        .args(&args)
        // verbose output for debugging purposes
        .arg("-v")
        .arg("--color=always")
        .arg("--")
        // ignore linking
        //.args(&["-C", "linker=true"])
        // link for embedded with retained stack section
        // .args(&["-C", "link-arg=-Tlink.x"])
        // .arg("codegen-units = 1") // for now assume that this is already in the user crate
        // force LTO
        .args(&["-C", "lto"])
        // also output the LLVM-IR (.ll file) for debugging purposes
        .arg("--emit=llvm-bc,llvm-ir")
        // emit stack sizes for the top level (not dependencies)
        .args(&["-Z", "emit-stack-sizes"])
        // force panic=abort in all crates
        //.env("RUSTFLAGS", "-C panic=abort")
        // collect stderr
        .stderr(Stdio::piped());

    // print full command for debugging purposes
    eprintln!("{:?}", c);

    // spawn process and drive it until completion while replicating its stderr in ours
    let mut p = c.spawn().unwrap();
    let mut pstderr = p.stderr.take().unwrap();
    let mut buf = [0; 1024];
    let mut output = vec![];

    let stderr = io::stderr();
    let mut stderr = stderr.lock();
    loop {
        if let Ok(n) = pstderr.read(&mut buf) {
            stderr.write(&buf[..n]).unwrap();
            output.extend_from_slice(&buf[..n]);
        }

        if let Some(status) = p.try_wait().unwrap() {
            assert!(status.success());
            break;
        }
    }

    let out = parse_out(str::from_utf8(&output).unwrap());
    println!("{:?}", out);

    use std::fmt::Write;
    let mut s = String::new();
    write!(
        &mut s,
        "{}/{}{}",
        out.out_dir.expect("--out-dir missing"),
        out.crate_name.expect("--crate-name missing"),
        out.hash.expect("extra-filename missing")
    )
    .expect("internal error");

    let mut ir = s.clone();
    write!(&mut ir, ".ll").expect("internal error");
    println!("s: {}", ir);

    let mut c = Command::new("opt");
    c.arg("-analyze")
        .arg("-print-callgraph")
        .arg(ir)
        .output()
        .expect("failed to execute process");
    let dump = c.output().expect("failed to execute process");

    println!("\nstatus: {:?}", dump.status);
    println!("\nstdout: {:?}", str::from_utf8(&dump.stdout).unwrap());
    println!("\nstderr: {:?}", str::from_utf8(&dump.stderr).unwrap());

    let mut g = Graph::<_, ()>::new();
    let mut hm = HashMap::new();

    let mut buffer = Vec::new();
    println!("elf = {}", &s);

    parse_elf(&mut hm, &mut g, &s, &mut buffer);

    parse_callgraph(&mut hm, &mut g, &str::from_utf8(&dump.stderr).unwrap());

    if is_cyclic_directed(&g) {
        println!("/* warning, cyclic graph */");
    }

    let mut v = Vec::new();
    for ext in g.externals(Incoming) {
        v.push(ext);
    }

    let mut dfs = DfsPostOrder::empty(&g);
    while let Some(ext) = v.pop() {
        dfs.move_to(ext);
        while let Some(node) = dfs.next(&g) {
            // use a detached neighbors walker
            let mut edges = g.neighbors_directed(node, Incoming).detach();
            while let Some(to) = edges.next_node(&g) {
                g[to].call_stack = g[to].call_stack.max(g[to].local_stack + g[node].call_stack);
            }
        }
    }

    println!("{:?}", Dot::with_config(&g, &[Config::EdgeNoLabel]));
}