diff options
| -rw-r--r-- | src/display.rs | 34 | ||||
| -rw-r--r-- | src/lib.rs | 79 | ||||
| -rw-r--r-- | test/test.rs | 1 |
3 files changed, 112 insertions, 2 deletions
diff --git a/src/display.rs b/src/display.rs index 50375f2..76db507 100644 --- a/src/display.rs +++ b/src/display.rs @@ -10,7 +10,20 @@ use std::hint::unreachable_unchecked; use yaxpeax_arch::{Arch, Colorize, ColorSettings, Decodable, LengthedInstruction, ShowContextual, YaxColors}; use yaxpeax_arch::display::*; -use ::{RegSpec, RegisterBank, Opcode, Operand, Instruction}; +use ::{RegSpec, RegisterBank, Opcode, Operand, Instruction, Segment}; + +impl fmt::Display for Segment { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + Segment::CS => write!(f, "cs"), + Segment::DS => write!(f, "ds"), + Segment::ES => write!(f, "es"), + Segment::FS => write!(f, "fs"), + Segment::GS => write!(f, "gs"), + Segment::SS => write!(f, "ss"), + } + } +} impl fmt::Display for RegSpec { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -537,6 +550,7 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi write!(out, "lock ")?; } self.opcode.colorize(colors, out)?; + match context.and_then(|xs| xs[0].as_ref()) { Some(s) => { write!(out, " {}", s)?; }, None => { @@ -546,6 +560,9 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi }, ref x @ _ => { write!(out, " ")?; + if let Some(prefix) = self.segment_override_for_op(0) { + write!(out, "{}:", prefix); + } x.colorize(colors, out)?; } } @@ -562,10 +579,16 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi }, x @ &Operand::Register(_) => { write!(out, ", ")?; + if let Some(prefix) = self.segment_override_for_op(1) { + write!(out, "{}:", prefix); + } x.colorize(colors, out) } x @ _ => { write!(out, ", byte ")?; + if let Some(prefix) = self.segment_override_for_op(1) { + write!(out, "{}:", prefix); + } x.colorize(colors, out) } } @@ -582,10 +605,16 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi }, x @ &Operand::Register(_) => { write!(out, ", ")?; + if let Some(prefix) = self.segment_override_for_op(1) { + write!(out, "{}:", prefix); + } x.colorize(colors, out) } x @ _ => { write!(out, ", word ")?; + if let Some(prefix) = self.segment_override_for_op(1) { + write!(out, "{}:", prefix); + } x.colorize(colors, out) } } @@ -602,6 +631,9 @@ impl <T: std::fmt::Write> ShowContextual<u64, [Option<String>], T> for Instructi }, x @ _ => { write!(out, ", ")?; + if let Some(prefix) = self.segment_override_for_op(1) { + write!(out, "{}:", prefix); + } x.colorize(colors, out) } } @@ -172,6 +172,38 @@ pub enum Operand { Nothing } +impl Operand { + pub fn is_memory(&self) -> bool { + match self { + Operand::DisplacementU32(_) | + Operand::DisplacementU64(_) | + Operand::RegDeref(_) | + Operand::RegDisp(_, _) | + Operand::RegScale(_, _) | + Operand::RegIndexBase(_, _) | + Operand::RegIndexBaseDisp(_, _, _) | + Operand::RegScaleDisp(_, _, _) | + Operand::RegIndexBaseScale(_, _, _) | + Operand::RegIndexBaseScaleDisp(_, _, _, _) => { + true + }, + Operand::ImmediateI8(_) | + Operand::ImmediateU8(_) | + Operand::ImmediateI16(_) | + Operand::ImmediateU16(_) | + Operand::ImmediateU32(_) | + Operand::ImmediateI32(_) | + Operand::ImmediateU64(_) | + Operand::ImmediateI64(_) | + Operand::Register(_) | + Operand::Many(_) | + Operand::Nothing => { + false + } + } + } +} + #[cfg(feature="use-serde")] #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] pub enum RegisterBank { @@ -189,7 +221,7 @@ pub enum RegisterBank { ST, MM, // ST, MM regs (x87, mmx) } -#[derive(Debug)] +#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)] pub enum Segment { CS, DS, ES, FS, GS, SS } @@ -421,6 +453,51 @@ impl Instruction { _ => false } } + pub fn segment_override_for_op(&self, op: u8) -> Option<Segment> { + match self.opcode { + Opcode::STOS => { + if op == 0 { + Some(Segment::ES) + } else { + None + } + } + Opcode::LODS => { + if op == 1 { + Some(self.prefixes.segment) + } else { + None + } + } + Opcode::MOVS => { + if op == 0 { + Some(Segment::ES) + } else if op == 1 { + Some(self.prefixes.segment) + } else { + None + } + } + Opcode::CMPS => { + if op == 0 { + Some(self.prefixes.segment) + } else if op == 1 { + Some(Segment::ES) + } else { + None + } + }, + _ => { + // most operands are pretty simple: + if self.operands[op as usize].is_memory() && + self.prefixes.segment != Segment::DS { + Some(self.prefixes.segment) + } else { + None + } + } + } + } } #[derive(Debug)] diff --git a/test/test.rs b/test/test.rs index 1e73a2e..f8dadbb 100644 --- a/test/test.rs +++ b/test/test.rs @@ -83,6 +83,7 @@ fn test_mov() { test_display(&[0x0f, 0xb6, 0x06], "movzx eax, byte [rsi]"); test_display(&[0x0f, 0xb7, 0x06], "movzx eax, word [rsi]"); test_display(&[0x89, 0x55, 0x94], "mov [rbp - 0x6c], edx"); + test_display(&[0x65, 0x4c, 0x89, 0x04, 0x25, 0xa8, 0x01, 0x00, 0x00], "mov gs:[0x1a8], r8"); } #[test] |