Skip to content

Commit df2e75b

Browse files
wtdcodewtdcode
authored
Merge pull request #1764 from LG3696/fix-memory-leak
Fix Rust binding memory leak
2 parents 549f34f + aded30f commit df2e75b

File tree

2 files changed

+68
-55
lines changed

2 files changed

+68
-55
lines changed

‎bindings/rust/src/ffi.rs

+56-24
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
#![allow(non_camel_case_types)]
22
#![allow(dead_code)]
33

4-
use crate::Unicorn;
4+
use crate::{Unicorn, UnicornInner};
55

66
use super::unicorn_const::{uc_error, Arch, HookType, MemRegion, MemType, Mode, Query};
7-
use core::ffi::c_void;
7+
use alloc::rc::Weak;
8+
use core::{cell::UnsafeCell, ffi::c_void};
89
use libc::{c_char, c_int};
910

1011
pub type uc_handle = *mut c_void;
@@ -89,7 +90,7 @@ extern "C" {
8990

9091
pub struct UcHook<'a, D: 'a, F: 'a> {
9192
pub callback: F,
92-
pub uc: Unicorn<'a, D>,
93+
pub uc: Weak<UnsafeCell<UnicornInner<'a, D>>>,
9394
}
9495

9596
pub trait IsUcHook<'a> {}
@@ -106,8 +107,11 @@ where
106107
F: FnMut(&mut crate::Unicorn<D>, u64, usize) -> u64,
107108
{
108109
let user_data = unsafe { &mut *user_data };
109-
debug_assert_eq!(uc, user_data.uc.get_handle());
110-
(user_data.callback)(&mut user_data.uc, offset, size)
110+
let mut user_data_uc = Unicorn {
111+
inner: user_data.uc.upgrade().unwrap(),
112+
};
113+
debug_assert_eq!(uc, user_data_uc.get_handle());
114+
(user_data.callback)(&mut user_data_uc, offset, size)
111115
}
112116

113117
pub extern "C" fn mmio_write_callback_proxy<D, F>(
@@ -120,8 +124,11 @@ pub extern "C" fn mmio_write_callback_proxy<D, F>(
120124
F: FnMut(&mut crate::Unicorn<D>, u64, usize, u64),
121125
{
122126
let user_data = unsafe { &mut *user_data };
123-
debug_assert_eq!(uc, user_data.uc.get_handle());
124-
(user_data.callback)(&mut user_data.uc, offset, size, value);
127+
let mut user_data_uc = Unicorn {
128+
inner: user_data.uc.upgrade().unwrap(),
129+
};
130+
debug_assert_eq!(uc, user_data_uc.get_handle());
131+
(user_data.callback)(&mut user_data_uc, offset, size, value);
125132
}
126133

127134
pub extern "C" fn code_hook_proxy<D, F>(
@@ -133,8 +140,11 @@ pub extern "C" fn code_hook_proxy<D, F>(
133140
F: FnMut(&mut crate::Unicorn<D>, u64, u32),
134141
{
135142
let user_data = unsafe { &mut *user_data };
136-
debug_assert_eq!(uc, user_data.uc.get_handle());
137-
(user_data.callback)(&mut user_data.uc, address, size);
143+
let mut user_data_uc = Unicorn {
144+
inner: user_data.uc.upgrade().unwrap(),
145+
};
146+
debug_assert_eq!(uc, user_data_uc.get_handle());
147+
(user_data.callback)(&mut user_data_uc, address, size);
138148
}
139149

140150
pub extern "C" fn block_hook_proxy<D, F>(
@@ -146,8 +156,11 @@ pub extern "C" fn block_hook_proxy<D, F>(
146156
F: FnMut(&mut crate::Unicorn<D>, u64, u32),
147157
{
148158
let user_data = unsafe { &mut *user_data };
149-
debug_assert_eq!(uc, user_data.uc.get_handle());
150-
(user_data.callback)(&mut user_data.uc, address, size);
159+
let mut user_data_uc = Unicorn {
160+
inner: user_data.uc.upgrade().unwrap(),
161+
};
162+
debug_assert_eq!(uc, user_data_uc.get_handle());
163+
(user_data.callback)(&mut user_data_uc, address, size);
151164
}
152165

153166
pub extern "C" fn mem_hook_proxy<D, F>(
@@ -162,39 +175,52 @@ where
162175
F: FnMut(&mut crate::Unicorn<D>, MemType, u64, usize, i64) -> bool,
163176
{
164177
let user_data = unsafe { &mut *user_data };
165-
debug_assert_eq!(uc, user_data.uc.get_handle());
166-
(user_data.callback)(&mut user_data.uc, mem_type, address, size as usize, value)
178+
let mut user_data_uc = Unicorn {
179+
inner: user_data.uc.upgrade().unwrap(),
180+
};
181+
debug_assert_eq!(uc, user_data_uc.get_handle());
182+
(user_data.callback)(&mut user_data_uc, mem_type, address, size as usize, value)
167183
}
168184

169185
pub extern "C" fn intr_hook_proxy<D, F>(uc: uc_handle, value: u32, user_data: *mut UcHook<D, F>)
170186
where
171187
F: FnMut(&mut crate::Unicorn<D>, u32),
172188
{
173189
let user_data = unsafe { &mut *user_data };
174-
debug_assert_eq!(uc, user_data.uc.get_handle());
175-
(user_data.callback)(&mut user_data.uc, value);
190+
let mut user_data_uc = Unicorn {
191+
inner: user_data.uc.upgrade().unwrap(),
192+
};
193+
debug_assert_eq!(uc, user_data_uc.get_handle());
194+
(user_data.callback)(&mut user_data_uc, value);
176195
}
177196

178197
pub extern "C" fn insn_in_hook_proxy<D, F>(
179198
uc: uc_handle,
180199
port: u32,
181200
size: usize,
182201
user_data: *mut UcHook<D, F>,
183-
) where
202+
) -> u32
203+
where
184204
F: FnMut(&mut crate::Unicorn<D>, u32, usize) -> u32,
185205
{
186206
let user_data = unsafe { &mut *user_data };
187-
debug_assert_eq!(uc, user_data.uc.get_handle());
188-
(user_data.callback)(&mut user_data.uc, port, size);
207+
let mut user_data_uc = Unicorn {
208+
inner: user_data.uc.upgrade().unwrap(),
209+
};
210+
debug_assert_eq!(uc, user_data_uc.get_handle());
211+
(user_data.callback)(&mut user_data_uc, port, size)
189212
}
190213

191214
pub extern "C" fn insn_invalid_hook_proxy<D, F>(uc: uc_handle, user_data: *mut UcHook<D, F>) -> bool
192215
where
193216
F: FnMut(&mut crate::Unicorn<D>) -> bool,
194217
{
195218
let user_data = unsafe { &mut *user_data };
196-
debug_assert_eq!(uc, user_data.uc.get_handle());
197-
(user_data.callback)(&mut user_data.uc)
219+
let mut user_data_uc = Unicorn {
220+
inner: user_data.uc.upgrade().unwrap(),
221+
};
222+
debug_assert_eq!(uc, user_data_uc.get_handle());
223+
(user_data.callback)(&mut user_data_uc)
198224
}
199225

200226
pub extern "C" fn insn_out_hook_proxy<D, F>(
@@ -207,15 +233,21 @@ pub extern "C" fn insn_out_hook_proxy<D, F>(
207233
F: FnMut(&mut crate::Unicorn<D>, u32, usize, u32),
208234
{
209235
let user_data = unsafe { &mut *user_data };
210-
debug_assert_eq!(uc, user_data.uc.get_handle());
211-
(user_data.callback)(&mut user_data.uc, port, size, value);
236+
let mut user_data_uc = Unicorn {
237+
inner: user_data.uc.upgrade().unwrap(),
238+
};
239+
debug_assert_eq!(uc, user_data_uc.get_handle());
240+
(user_data.callback)(&mut user_data_uc, port, size, value);
212241
}
213242

214243
pub extern "C" fn insn_sys_hook_proxy<D, F>(uc: uc_handle, user_data: *mut UcHook<D, F>)
215244
where
216245
F: FnMut(&mut crate::Unicorn<D>),
217246
{
218247
let user_data = unsafe { &mut *user_data };
219-
debug_assert_eq!(uc, user_data.uc.get_handle());
220-
(user_data.callback)(&mut user_data.uc);
248+
let mut user_data_uc = Unicorn {
249+
inner: user_data.uc.upgrade().unwrap(),
250+
};
251+
debug_assert_eq!(uc, user_data_uc.get_handle());
252+
(user_data.callback)(&mut user_data_uc);
221253
}

‎bindings/rust/src/lib.rs

+12-31
Original file line numberDiff line numberDiff line change
@@ -369,17 +369,13 @@ impl<'a, D> Unicorn<'a, D> {
369369
let mut read_data = read_callback.map(|c| {
370370
Box::new(ffi::UcHook {
371371
callback: c,
372-
uc: Unicorn {
373-
inner: self.inner.clone(),
374-
},
372+
uc: Rc::downgrade(&self.inner),
375373
})
376374
});
377375
let mut write_data = write_callback.map(|c| {
378376
Box::new(ffi::UcHook {
379377
callback: c,
380-
uc: Unicorn {
381-
inner: self.inner.clone(),
382-
},
378+
uc: Rc::downgrade(&self.inner),
383379
})
384380
});
385381

@@ -586,7 +582,8 @@ impl<'a, D> Unicorn<'a, D> {
586582
return Err(uc_error::ARCH);
587583
}
588584

589-
let err: uc_error = unsafe { ffi::uc_reg_read(self.get_handle(), curr_reg_id, value.as_mut_ptr() as _) };
585+
let err: uc_error =
586+
unsafe { ffi::uc_reg_read(self.get_handle(), curr_reg_id, value.as_mut_ptr() as _) };
590587

591588
if err == uc_error::OK {
592589
boxed = value.into_boxed_slice();
@@ -622,9 +619,7 @@ impl<'a, D> Unicorn<'a, D> {
622619
let mut hook_ptr = core::ptr::null_mut();
623620
let mut user_data = Box::new(ffi::UcHook {
624621
callback,
625-
uc: Unicorn {
626-
inner: self.inner.clone(),
627-
},
622+
uc: Rc::downgrade(&self.inner),
628623
});
629624

630625
let err = unsafe {
@@ -654,9 +649,7 @@ impl<'a, D> Unicorn<'a, D> {
654649
let mut hook_ptr = core::ptr::null_mut();
655650
let mut user_data = Box::new(ffi::UcHook {
656651
callback,
657-
uc: Unicorn {
658-
inner: self.inner.clone(),
659-
},
652+
uc: Rc::downgrade(&self.inner),
660653
});
661654

662655
let err = unsafe {
@@ -697,9 +690,7 @@ impl<'a, D> Unicorn<'a, D> {
697690
let mut hook_ptr = core::ptr::null_mut();
698691
let mut user_data = Box::new(ffi::UcHook {
699692
callback,
700-
uc: Unicorn {
701-
inner: self.inner.clone(),
702-
},
693+
uc: Rc::downgrade(&self.inner),
703694
});
704695

705696
let err = unsafe {
@@ -730,9 +721,7 @@ impl<'a, D> Unicorn<'a, D> {
730721
let mut hook_ptr = core::ptr::null_mut();
731722
let mut user_data = Box::new(ffi::UcHook {
732723
callback,
733-
uc: Unicorn {
734-
inner: self.inner.clone(),
735-
},
724+
uc: Rc::downgrade(&self.inner),
736725
});
737726

738727
let err = unsafe {
@@ -763,9 +752,7 @@ impl<'a, D> Unicorn<'a, D> {
763752
let mut hook_ptr = core::ptr::null_mut();
764753
let mut user_data = Box::new(ffi::UcHook {
765754
callback,
766-
uc: Unicorn {
767-
inner: self.inner.clone(),
768-
},
755+
uc: Rc::downgrade(&self.inner),
769756
});
770757

771758
let err = unsafe {
@@ -796,9 +783,7 @@ impl<'a, D> Unicorn<'a, D> {
796783
let mut hook_ptr = core::ptr::null_mut();
797784
let mut user_data = Box::new(ffi::UcHook {
798785
callback,
799-
uc: Unicorn {
800-
inner: self.inner.clone(),
801-
},
786+
uc: Rc::downgrade(&self.inner),
802787
});
803788

804789
let err = unsafe {
@@ -830,9 +815,7 @@ impl<'a, D> Unicorn<'a, D> {
830815
let mut hook_ptr = core::ptr::null_mut();
831816
let mut user_data = Box::new(ffi::UcHook {
832817
callback,
833-
uc: Unicorn {
834-
inner: self.inner.clone(),
835-
},
818+
uc: Rc::downgrade(&self.inner),
836819
});
837820

838821
let err = unsafe {
@@ -870,9 +853,7 @@ impl<'a, D> Unicorn<'a, D> {
870853
let mut hook_ptr = core::ptr::null_mut();
871854
let mut user_data = Box::new(ffi::UcHook {
872855
callback,
873-
uc: Unicorn {
874-
inner: self.inner.clone(),
875-
},
856+
uc: Rc::downgrade(&self.inner),
876857
});
877858

878859
let err = unsafe {

0 commit comments

Comments
 (0)