@@ -181,9 +181,10 @@ fn assemble_jmp_to_address(address: usize, mut relative: isize) -> ([u8; JMP_MAX
181181/// When this structure is dropped (falls out of scope), the patch will be reverted and the function will return
182182/// to its original state.
183183pub struct PatchGuard {
184- ptr : * mut u8 ,
184+ pub ptr : * mut u8 ,
185185 len : usize ,
186186 data : [ u8 ; JMP_MAX_SIZE ] ,
187+ patch : [ u8 ; JMP_MAX_SIZE ] ,
187188}
188189
189190impl Drop for PatchGuard {
@@ -194,6 +195,20 @@ impl Drop for PatchGuard {
194195 }
195196}
196197
198+ impl PatchGuard {
199+ pub fn revert ( & self ) {
200+ unsafe {
201+ copy_to_protected_address ( self . ptr , & self . data [ ..self . len ] ) ;
202+ }
203+ }
204+
205+ pub fn restore ( & self ) {
206+ unsafe {
207+ copy_to_protected_address ( self . ptr , & self . patch [ ..self . len ] ) ;
208+ }
209+ }
210+ }
211+
197212#[ cfg( any( target_arch = "x86" , target_arch = "x86_64" ) ) ]
198213const UNSAFE_LEADING_BYTES : [ u8 ; 4 ] = [
199214 0xC3 , // ret near
@@ -238,6 +253,7 @@ macro_rules! define_patch {
238253 ptr: target,
239254 len,
240255 data: original,
256+ patch,
241257 }
242258 }
243259 ) ;
@@ -254,6 +270,26 @@ define_patch!(patch7(A, B, C, D, E, F, G,));
254270define_patch ! ( patch8( A , B , C , D , E , F , G , H , ) ) ;
255271define_patch ! ( patch9( A , B , C , D , E , F , G , H , I , ) ) ;
256272
273+ #[ macro_export]
274+ macro_rules! disable_patch {
275+ ( $guard: expr, async $( $body: tt) * ) => { {
276+ $guard. revert( ) ;
277+ let result = async { $( $body) * } . await ;
278+ $guard. restore( ) ;
279+ result
280+ } } ;
281+
282+ ( $guard: expr, $( $body: tt) * ) => { {
283+ $guard. revert( ) ;
284+ let result = ( || { $( $body) * } ) ( ) ;
285+ $guard. restore( ) ;
286+ result
287+ } } ;
288+ }
289+
290+ unsafe impl Send for PatchGuard { }
291+ unsafe impl Sync for PatchGuard { }
292+
257293#[ cfg( test) ]
258294#[ inline( never) ]
259295fn tiny ( ) { }
0 commit comments