libinject/cmp/
cmp.rs

1//! **Fine-grained coverage instrumentation for CMP and TEST instructions:** Instructions are instrumented _iff_ they are in the target module, which is passed via the CLI.
2//!
3
4use crate::drcore::{self, log};
5use crate::utils::{self, Boolean};
6use crate::{cli, cmp, ffi};
7use std::os::raw::{c_uint, c_void};
8use std::ptr;
9
10/// **Instrumentation function called on each instruction in the basic block**, optionally transforming
11/// the block, just before it is added to the code cache. `instrlist` is the list of instructions in
12/// the basic block, and `instr` is the instruction focused on.
13///
14/// Only code blocks in the target module are instrumented.
15#[unsafe(no_mangle)]
16pub unsafe extern "C" fn cmp_coverage(
17    drcontext: *mut c_void,
18    tag: *mut c_void,
19    instrlist: *mut ffi::instrlist_t,
20    instr: *mut ffi::instr_t,
21    _for_trace: ffi::bool_,
22    _translating: ffi::bool_,
23    _user_data: *mut c_void,
24) -> ffi::dr_emit_flags_t {
25    unsafe {
26        let cli_args = cli::get_args().expect("args were parsed in libinject_init");
27
28        let bb_start_addr = ffi::dr_fragment_app_pc(tag);
29
30        // Only instrument CMP instructions that are in the target module.
31        if !drcore::app_pc_is_in_target_module(bb_start_addr, cli_args) {
32            return ffi::dr_emit_flags_t_DR_EMIT_DEFAULT;
33        }
34
35        match ffi::instr_get_opcode(instr) {
36            ffi::OP_cmp | ffi::OP_test => instrument_cmp(drcontext, instrlist, instr),
37            _ => ffi::dr_emit_flags_t_DR_EMIT_DEFAULT,
38        }
39    }
40}
41
42/// **The instrumentation applied to a CMP or TEST instruction:** The instruction operands may be
43/// registers, immediates, or memory references (or any combination). The correct instrumentation
44/// must be added (conditional on the operand types) to load the operand values into two scratch registers. These scratch registers are then passed as arguments to a clean call that does a byte-by-byte comparison and writes coverage information to the AFL bitmap.
45#[unsafe(no_mangle)]
46pub unsafe extern "C" fn instrument_cmp(
47    drcontext: *mut c_void,
48    instrlist: *mut ffi::instrlist_t,
49    instr: *mut ffi::instr_t,
50) -> ffi::dr_emit_flags_t {
51    unsafe {
52        if ffi::instr_num_srcs(instr) < 2 {
53            return ffi::dr_emit_flags_t_DR_EMIT_DEFAULT;
54        }
55
56        // Reserve 2 scratch registers, one for each operand.
57        let mut scratch_regs: [ffi::reg_id_t; 2] = [0; 2];
58        let mut reserved = 0_usize;
59        for reg in scratch_regs.iter_mut() {
60            let status = ffi::drreg_reserve_register(
61                drcontext,
62                instrlist,
63                instr,
64                ptr::null_mut(),
65                reg as *mut ffi::reg_id_t,
66            );
67            if status != ffi::drreg_status_t_DRREG_SUCCESS {
68                for &r in &scratch_regs[..reserved] {
69                    ffi::drreg_unreserve_register(drcontext, instrlist, instr, r);
70                }
71                log(&format!(
72                    "[cmp] drreg_reserve_register failed with status {}",
73                    status as i32
74                ));
75                return ffi::dr_emit_flags_t_DR_EMIT_DEFAULT;
76            }
77            reserved += 1;
78        }
79
80        // Insert instrumentation that parses the CMP operands (e.g. as register operands,
81        // immediates, or memory references) and inserts instrumentation to load them into the two
82        // scratch registers.
83        let (left_operand, right_operand) = match scratch_regs
84            .iter()
85            .enumerate()
86            .map(|(i, scratch_reg)| {
87                parse_operand(
88                    ffi::instr_get_src(instr, i as c_uint),
89                    *scratch_reg,
90                    drcontext,
91                    instrlist,
92                    instr,
93                )
94            })
95            .collect::<Result<Vec<Operand>, _>>()
96        {
97            Ok(operands) => (operands[0], operands[1]),
98            Err(reason) => {
99                log(&format!("Failed parsing left operand of cmp: {reason}"));
100                return ffi::dr_emit_flags_t_DR_EMIT_DEFAULT;
101            }
102        };
103
104        // Insert a clean call to collect the coverage information. Pass the IDs of the two scratch
105        // registers which will have the values of the two operands loaded into by the time the
106        // clean call fires. The scratch registers are passed as register operand arguments to the
107        // clean call.
108        ffi::dr_insert_clean_call(
109            drcontext,
110            instrlist,
111            instr,
112            cmp_clean_call as *mut c_void,
113            false as ffi::bool_,
114            5,
115            utils::opnd_create_uintptr(ffi::instr_get_app_pc(instr) as usize),
116            ffi::opnd_create_reg(left_operand.reg_id),
117            utils::opnd_create_uintptr(left_operand.size as usize),
118            ffi::opnd_create_reg(right_operand.reg_id),
119            utils::opnd_create_uintptr(right_operand.size as usize),
120        );
121
122        scratch_regs.iter().for_each(|reg_id| {
123            ffi::drreg_unreserve_register(drcontext, instrlist, instr, *reg_id);
124        });
125
126        return ffi::dr_emit_flags_t_DR_EMIT_DEFAULT;
127    }
128}
129
130/// A function inserted as instrumentation that takes two values (via registers) as arguments,
131/// along with the size (in bytes) of those values, and the current application program counter.
132///
133/// The input values are compared byte-by-byte. For every byte that matches, a corresponding byte
134/// on the CMP partition of the AFL bitmap is incremented (saturating at 8, rather than wrapping).
135/// The routine exits as soon as the first non-matching bytes occur.
136#[unsafe(no_mangle)]
137unsafe extern "C" fn cmp_clean_call(
138    app_pc: *const u8,
139    left_reg: ffi::reg_t,
140    left_size: usize,
141    right_reg: ffi::reg_t,
142    right_size: usize,
143) {
144    unsafe {
145        let afl_cmp_area = match cmp::utils::get_afl_cmp_area() {
146            Ok(cmp_area_ptr) => cmp_area_ptr,
147            Err(reason) => {
148                log(&reason);
149                return;
150            }
151        };
152
153        let base_idx = cmp::utils::base_idx_of_app_pc(app_pc);
154
155        // CMP must compare operands of the same size, but it is possible that smaller operand is
156        // compared by first promoting to the size of the other.
157        let cmp_len = usize::max(left_size, right_size);
158
159        // Operands will be <= full-register width (4 bytes on 32-bit, 8 bytes on 64-bit).
160        const MAX_OPND_SIZE: usize = if cfg!(target_pointer_width = "32") {
161            4
162        } else {
163            8
164        };
165        let buf_len = cmp_len.min(MAX_OPND_SIZE);
166
167        let mut left_buf = [0u8; MAX_OPND_SIZE];
168        let mut right_buf = [0u8; MAX_OPND_SIZE];
169        Operand::read_into(left_reg, buf_len, &mut left_buf[..buf_len]);
170        Operand::read_into(right_reg, buf_len, &mut right_buf[..buf_len]);
171
172        cmp::utils::write_coverage(
173            afl_cmp_area,
174            base_idx,
175            &left_buf[..buf_len],
176            &right_buf[..buf_len],
177            buf_len,
178        )
179    }
180}
181
182/// The id of a register holding the value, and the size of the operand value in bytes.
183#[derive(Debug, Copy, Clone)]
184pub struct Operand {
185    pub reg_id: ffi::reg_id_t,
186    pub size: u32,
187}
188
189impl Operand {
190    /// Fill a buffer with the operand’s value at runtime
191    fn read_into(reg_value: ffi::reg_t, buf_len: usize, buf: &mut [u8]) {
192        let bytes = reg_value.to_le_bytes();
193        buf[..buf_len].copy_from_slice(&bytes[..buf_len]);
194    }
195}
196
197/// Parse the operand type and insert instrumentation to load the value into the provided
198/// `dst_reg`.
199///
200/// The inserted instrumentation must be sensitive to size of the src operand, normalising it to
201/// the full width scratch register passed in as `dst_reg`.
202///
203/// In all cases the ID of the full-width scratch register is returned - even if the inserted
204/// instrumentation moves the value into the corresponding half-width sub-register. Returning with
205/// the full-width register (that is e.g. passed as an arg to a clean call) completes the
206/// normalisation, ensuring that downstream instrumentation can safely read the value as
207/// a usize.
208unsafe fn parse_operand(
209    operand: ffi::opnd_t,
210    dst_reg: ffi::reg_id_t,
211    drcontext: *mut c_void,
212    instrlist: *mut ffi::instrlist_t,
213    instr: *mut ffi::instr_t,
214) -> Result<Operand, String> {
215    unsafe {
216        let size: c_uint = ffi::opnd_size_in_bytes(ffi::opnd_get_size(operand));
217
218        // Valid normalising move op code is dependent on the size of the src operand.
219        let norm_move_op = match size {
220            1 | 2 => ffi::OP_movzx,
221            _ => ffi::OP_mov_ld,
222        };
223
224        if ffi::opnd_is_reg(operand).as_bool() {
225            let resized_dst_reg = get_resized_reg(dst_reg, size);
226
227            let move_ld = ffi::instr_create_1dst_1src(
228                drcontext,
229                norm_move_op,
230                ffi::opnd_create_reg(resized_dst_reg),
231                operand,
232            );
233            ffi::instr_set_meta(move_ld);
234            ffi::instrlist_meta_preinsert(instrlist, instr, move_ld);
235            Ok(Operand {
236                reg_id: dst_reg,
237                size,
238            })
239        } else if ffi::opnd_is_immed(operand).as_bool() {
240            // Noramlise the immediate to a full-width immediate.
241            const OPSZ_PTR: ffi::opnd_size_t = if cfg!(target_pointer_width = "64") {
242                ffi::OPSZ_8 as u8
243            } else {
244                ffi::OPSZ_4 as u8
245            };
246
247            let norm_opnd = ffi::opnd_create_immed_int(
248                ffi::opnd_get_immed_int(operand),
249                OPSZ_PTR as ffi::opnd_size_t,
250            );
251            let move_imm = ffi::instr_create_1dst_1src(
252                drcontext,
253                ffi::OP_mov_imm,
254                ffi::opnd_create_reg(dst_reg),
255                norm_opnd,
256            );
257            ffi::instr_set_meta(move_imm);
258            ffi::instrlist_meta_preinsert(instrlist, instr, move_imm);
259            Ok(Operand {
260                reg_id: dst_reg,
261                size,
262            })
263        } else if ffi::opnd_is_memory_reference(operand).as_bool() {
264            let resized_dst_reg = get_resized_reg(dst_reg, size);
265            let mov_ld = ffi::instr_create_1dst_1src(
266                drcontext,
267                norm_move_op,
268                ffi::opnd_create_reg(resized_dst_reg),
269                operand,
270            );
271            ffi::instr_set_meta(mov_ld);
272            ffi::instrlist_meta_preinsert(instrlist, instr, mov_ld);
273            Ok(Operand {
274                reg_id: dst_reg,
275                size,
276            })
277        } else {
278            Err("Unhandled cmp operand type".to_string())
279        }
280    }
281}
282
283unsafe fn get_resized_reg(reg: ffi::reg_id_t, size: u32) -> ffi::reg_id_t {
284    unsafe {
285        #[cfg(target_pointer_width = "64")]
286        if size == 4 {
287            ffi::reg_64_to_32(reg)
288        } else {
289            reg
290        }
291
292        #[cfg(target_pointer_width = "32")]
293        reg
294    }
295}
libinject/cmp/cmp.rs

libinject/cmp/
cmp.rs
