feat(tests): Add trace consistency tests

.github/workflows/ci.yml

@@ -26,6 +26,9 @@ jobs:
     steps:
     - uses: actions/checkout@v2
 
+    - name: System deps
+      run: sudo apt install -y nasm
+
     - name: Build
       run: cargo build
 

Cargo.toml

@@ -11,7 +11,7 @@ env_logger = "0.8"
 iced-x86 = "1.9.1"
 libc = "0.2"
 log = "0.4"
-nix = "0.19"
+nix = { git = "https://github.com/trailofbits/nix", branch = "ww/personality" }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 spawn-ptrace = "0.1.2"

README.md

@@ -38,3 +38,14 @@ Once you have the appropriate environment, just `cargo build`:
 $ cargo build
 $ ./target/debug/mttn -h
 ```
+
+### Testing
+
+*mttn*'s tests require some system depedencies to build test binaries with:
+`nasm`, (GNU) `ld`, and (GNU) `make`.
+
+Once you have those installed, running the tests should be as simple as:
+
+```bash
+$ cargo test
+```

src/main.rs

@@ -30,6 +30,12 @@ fn app<'a, 'b>() -> App<'a, 'b> {
                 .short("d")
                 .long("debug-on-fault"),
         )
+        .arg(
+            Arg::with_name("disable-aslr")
+                .help("Disable ASLR on the tracee")
+                .short("a")
+                .long("disable-aslr"),
+        )
         .arg(
             Arg::with_name("tracee-pid")
                 .help("Attach to the given PID for tracing")

src/trace.rs

@@ -3,6 +3,7 @@ use iced_x86::{
     Code, Decoder, DecoderOptions, Instruction, InstructionInfoFactory, InstructionInfoOptions,
     MemorySize, Mnemonic, OpAccess, Register,
 };
+use nix::sys::personality::{self, Persona};
 use nix::sys::ptrace;
 use nix::sys::signal;
 use nix::sys::uio;
@@ -12,16 +13,37 @@ use serde::Serialize;
 use spawn_ptrace::CommandPtraceSpawn;
 
 use std::convert::{TryFrom, TryInto};
+use std::io;
+use std::os::unix::process::CommandExt;
 use std::process::Command;
 
 const MAX_INSTR_LEN: usize = 15;
 
+pub trait CommandPersonality {
+    fn personality(&mut self, persona: Persona);
+}
+
+impl CommandPersonality for Command {
+    fn personality(&mut self, persona: Persona) {
+        unsafe {
+            self.pre_exec(move || match personality::set(persona) {
+                Ok(_) => Ok(()),
+                Err(nix::Error::Sys(e)) => Err(io::Error::from_raw_os_error(e as i32)),
+                Err(_) => Err(io::Error::new(
+                    io::ErrorKind::Other,
+                    "unknown personality error",
+                )),
+            })
+        };
+    }
+}
+
 /// Represents the width of a concrete memory operation.
 ///
 /// All `mttn` memory operations are 1, 2, 4, or 8 bytes.
 /// Larger operations are either modeled as multiple individual operations
 /// (if caused by a `REP` prefix), ignored (if configured), or cause a fatal error.
-#[derive(Clone, Copy, Debug, Serialize)]
+#[derive(Clone, Copy, Debug, PartialEq, Serialize)]
 pub enum MemoryMask {
     Byte = 1,
     Word = 2,
@@ -65,7 +87,7 @@ pub enum MemoryOp {
 
 /// Represents an entire traced memory operation, including its kind (`MemoryOp`),
 /// size (`MemoryMask`), concrete address, and actual read or written data.
-#[derive(Debug, Serialize)]
+#[derive(Debug, PartialEq, Serialize)]
 pub struct MemoryHint {
     address: u64,
     operation: MemoryOp,
@@ -76,7 +98,7 @@ pub struct MemoryHint {
 /// Represents an individual step in the trace, including the raw instruction bytes,
 /// the register file state before execution, and any memory operations that result
 /// from execution.
-#[derive(Debug, Serialize)]
+#[derive(Debug, PartialEq, Serialize)]
 pub struct Step {
     instr: Vec<u8>,
     regs: RegisterFile,
@@ -88,7 +110,7 @@ pub struct Step {
 /// Only the standard addressable registers, plus `RFLAGS`, are tracked.
 /// `mttn` assumes that all segment base addresses are `0` and therefore doesn't
 /// track them.
-#[derive(Clone, Copy, Debug, Default, Serialize)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize)]
 pub struct RegisterFile {
     rax: u64,
     rbx: u64,
@@ -507,7 +529,7 @@ impl<'a> Tracee<'a> {
         // fast-string-enable bit (0b) in the IA32_MISC_ENABLE MSR, but we just sleep
         // for a bit to give the CPU a chance to catch up.
         // TODO(ww): Longer term, we should just model REP'd instructions outright.
-        std::thread::sleep(std::time::Duration::from_millis(1));
+        std::thread::sleep(std::time::Duration::from_millis(2));
 
         for hint in hints.iter_mut() {
             if hint.operation != MemoryOp::Write {
@@ -538,6 +560,7 @@ impl Iterator for Tracee<'_> {
 pub struct Tracer {
     pub ignore_unsupported_memops: bool,
     pub debug_on_fault: bool,
+    pub disable_aslr: bool,
     pub bitness: u32,
     pub target: Target,
 }
@@ -560,6 +583,7 @@ impl From<clap::ArgMatches<'_>> for Tracer {
         Self {
             ignore_unsupported_memops: matches.is_present("ignore-unsupported-memops"),
             debug_on_fault: matches.is_present("debug-on-fault"),
+            disable_aslr: matches.is_present("disable-aslr"),
             bitness: matches.value_of("mode").unwrap().parse().unwrap(),
             target: target,
         }
@@ -570,7 +594,16 @@ impl Tracer {
     pub fn trace(&self) -> Result<Tracee> {
         let tracee_pid = match &self.target {
             Target::Program(name, args) => {
-                let child = Command::new(name).args(args).spawn_ptrace()?;
+                let child = {
+                    let mut cmd = Command::new(name);
+                    cmd.args(args);
+
+                    if self.disable_aslr {
+                        cmd.personality(Persona::ADDR_NO_RANDOMIZE);
+                    }
+
+                    cmd.spawn_ptrace()?
+                };
 
                 log::debug!("spawned {} for tracing as child {}", name, child.id());
 
@@ -594,7 +627,7 @@ impl Tracer {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use iced_x86::UsedMemory;
+    use std::path::PathBuf;
 
     fn dummy_regs() -> RegisterFile {
         RegisterFile {
@@ -604,6 +637,41 @@ mod tests {
         }
     }
 
+    fn build_test_program(program: &str) -> String {
+        let mut buf = {
+            let mut buf = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+            buf.push("test");
+
+            buf
+        };
+
+        let status = Command::new("make")
+            .arg("-C")
+            .arg(&buf)
+            .arg(program)
+            .status()
+            .expect(&format!("build failed: {}", program));
+
+        if !status.success() {
+            panic!("build failed: {}", program);
+        }
+
+        buf.push(program);
+        buf.into_os_string().into_string().unwrap()
+    }
+
+    fn test_program_tracer(program: &str) -> Tracer {
+        let target = Target::Program(program.into(), vec![]);
+
+        Tracer {
+            ignore_unsupported_memops: false,
+            debug_on_fault: false,
+            disable_aslr: true,
+            bitness: 32,
+            target: target,
+        }
+    }
+
     #[test]
     fn test_register_file_value() {
         let regs = dummy_regs();
@@ -632,4 +700,41 @@ mod tests {
         // Unaddressable and unsupported registers return an Err.
         assert!(regs.value(Register::ST0).is_err());
     }
+
+    macro_rules! trace_consistency_tests {
+        ($($name:ident,)*) => {
+            $(
+                #[test]
+                fn $name() {
+                    let program = build_test_program(concat!(stringify!($name), ".elf"));
+                    let tracer = test_program_tracer(&program);
+
+                    // TODO(ww): Don't collect these.
+                    let trace1 = tracer
+                        .trace()
+                        .expect("spawn failed")
+                        .collect::<Result<Vec<Step>>>()
+                        .expect("trace failed");
+
+                    let trace2 = tracer
+                        .trace()
+                        .expect("spawn failed")
+                        .collect::<Result<Vec<Step>>>()
+                        .expect("trace failed");
+
+                    assert_eq!(trace1.len(), trace2.len());
+                    for (step1, step2) in trace1.iter().zip(trace2.iter()) {
+                        assert_eq!(step1, step2);
+                    }
+                }
+            )*
+        }
+    }
+
+    trace_consistency_tests! {
+        memops,
+        stosb,
+        stosw,
+        stosd,
+    }
 }