ib_hook/inject/dll/
app.rs

1/*!
2Inject DLL into target processes with an opinioned RPC schema.
3
4## API
5- [`DllApp`]: DLL RPC schema.
6- [`DllInjection::inject()`]: Inject a DLL into a process and optionally appliy it (call `APPLY()`).
7- [`DllInjection::eject()`]: Eject a DLL (automatically unapply first if applied before).
8- [`DllInjection::leak()`]: Prevent automatic ejection on drop.
9- [`DllInjection::drop()`]: Automatically unapply and eject if not already ejected (or [`DllInjection::leak`]ed).
10- [`DllInjectionVec`]: Manages multiple injections with batch eject support.
11  - [`DllInjectionVecWithInput`]: With owned `dll_path` and `input` for `apply()`.
12
13## Example: Single process injection
14
15```no_run
16use ib_hook::inject::dll::app::{DllApp, DllInjection, OwnedProcess};
17
18// Define your DLL app trait implementation
19struct MyDll;
20impl DllApp for MyDll {
21    const APPLY: &str = "apply_hook";
22    type Input = String;
23    type Output = ();
24}
25
26// Inject into a single process
27let process = OwnedProcess::find_first_by_name("Notepad.exe").unwrap();
28let mut injection = DllInjection::<MyDll>::inject(process)
29    .dll_path(std::path::Path::new("hook.dll"))
30    .apply(&"input".into())
31    .call()
32    .unwrap();
33
34// Eject manually or let drop handle it
35injection.eject().unwrap();
36```
37
38## Example: Multiple processes by name
39
40```no_run
41use ib_hook::inject::dll::app::{DllApp, DllInjectionVec};
42
43struct MyDll;
44impl DllApp for MyDll {
45    const APPLY: &str = "apply_hook";
46    type Input = String;
47    type Output = ();
48}
49
50// Inject into all processes named Notepad.exe
51let mut injections = DllInjectionVec::<MyDll>::new();
52injections.inject_with_process_name("Notepad.exe")
53    .dll_path(std::path::Path::new("hook.dll"))
54    .apply(&"input".into())
55    .on_error(|pid, err| ())
56    .call()
57    .unwrap();
58
59// Eject all manually or let drop handle it
60injections.eject().on_error(|pid, err| ()).call();
61```
62
63## Example: Custom process iterator
64
65```no_run
66use ib_hook::inject::dll::app::{DllApp, DllInjectionVec, OwnedProcess};
67
68struct MyDll;
69impl DllApp for MyDll {
70    const APPLY: &str = "apply_hook";
71    type Input = ();
72    type Output = ();
73}
74
75let processes = vec![
76    OwnedProcess::find_first_by_name("Notepad.exe").unwrap(),
77    OwnedProcess::find_first_by_name("Time.exe").unwrap(),
78];
79
80let mut injections = DllInjectionVec::<MyDll>::new();
81injections.inject(processes.into_iter())
82    .dll_path(std::path::Path::new("hook.dll"))
83    .on_error(|pid, err| ())
84    .call()
85    .unwrap();
86```
87
88## Disclaimer
89This is currently implemented as a wrapper of [`dll_syringe`],
90for object ownership (avoiding self-references), RAII (drop guard) and `Send`.
91
92Ref: https://github.com/Chaoses-Ib/ib-shell/blob/7dc099ea07a9c0a0e2db6aea10a74b2b53c9373e/ib-shell-item/src/hook/inject.rs
93*/
94use std::{
95    mem::transmute,
96    path::{Path, PathBuf},
97};
98
99use bon::bon;
100use derive_more::{Deref, DerefMut};
101use dll_syringe::{
102    Syringe,
103    process::{BorrowedProcessModule, Process},
104    rpc::RemotePayloadProcedure,
105};
106use thiserror::Error;
107
108use crate::{log::*, process::Pid};
109
110#[doc(hidden)]
111pub use dll_syringe::payload_utils::__payload_procedure_helper;
112pub use dll_syringe::process::OwnedProcess;
113
114#[derive(Error, Debug)]
115pub enum InjectError {
116    #[error("dll not found: {0}")]
117    DllNotFound(PathBuf),
118    #[error("cannot find any {0} process")]
119    ProcessNotFound(String),
120    #[error("inject failed: {0}")]
121    InjectFailed(#[from] dll_syringe::error::InjectError),
122    #[error("get apply failed: {0}")]
123    GetApplyFailed(#[from] dll_syringe::error::LoadProcedureError),
124    #[error("apply not found")]
125    ApplyNotFound,
126    #[error("apply: {0}")]
127    ApplyFailed(#[from] dll_syringe::rpc::PayloadRpcError),
128    #[error("eject failed: {0}")]
129    EjectFailed(#[from] dll_syringe::error::EjectError),
130}
131
132/// DLL RPC schema.
133///
134/// Call `APPLY(Some(Input))` on [`inject`](DllInjection::inject),
135/// and `APPLY(None)` on [`eject`](DllInjection::eject)
136/// (and on drop if not [`leak`](DllInjection::leak)ed).
137/**
138<div class="warning">
139
140`APPLY(None)` should clean up all the references to the DLL's code,
141including hooks and threads. Otherwise, when the DLL is unloaded
142the process will crash due to memory access violation.
143</div>
144*/
145pub trait DllApp {
146    /// The name of the exported function for RPC.
147    const APPLY: &str;
148
149    type Input: serde::Serialize + 'static;
150    type Output: serde::de::DeserializeOwned + 'static;
151}
152
153/**
154Usage:
155```ignore
156ib_hook::inject::dll::app::export_apply!(apply_hook, "apply_hook");
157```
158
159## Implementation
160For example:
161```ignore
162#[derive(Debug, Clone, Copy)]
163#[repr(C)]
164pub(crate) struct ArgAndResultBufInfo {
165    pub data: u64,
166    pub len: u64,
167    pub is_error: bool,
168}
169
170const _: () = {
171    #[unsafe(export_name = "ib_shell_apply")]
172    pub unsafe extern "system" fn _ib_hook_inject_dll_app_apply(
173        __args_and_params: *mut ::core::ffi::c_void,
174    ) {
175        let buf_info_ptr = __args_and_params;
176        let buf_info_ptr = buf_info_ptr.cast::<ArgAndResultBufInfo>();
177        let buf_info = unsafe { &mut *buf_info_ptr };
178        // buf_info.len = buf_info.len - 1;
179        let buf =
180            unsafe { slice::from_raw_parts_mut(buf_info.data as *mut u8, buf_info.len as usize) };
181        // dbg!(&buf_info);
182
183        let config = bincode::config::standard();
184
185        // eprintln!("decode_from_slice {:02X?}", buf);
186        let args: Result<((Option<Input>,), usize), bincode::error::DecodeError> =
187            bincode::serde::decode_from_slice::<(Option<Input>,), _>(buf, config);
188        // eprintln!("{:X}\n{:02X?}\nargs: {:?}", buf_info.data, buf, &args);
189
190        ib_hook::inject::dll::app::__payload_procedure_helper(__args_and_params, |__args| {
191            let (input,) = __args;
192
193            // panic!(
194            //     "{:X}\n{:02X?}\n{:?}\n{:?}",
195            //     buf_info.data, buf, &input, &args
196            // );
197            apply_hook(input)
198        });
199    }
200};
201```
202
203TODO: https://github.com/rust-lang/rust/issues/52393
204TODO: https://github.com/rust-lang/rust/issues/143547 or proc macro
205*/
206#[macro_export]
207macro_rules! export_apply {
208    ($apply:ident, $export_name:literal) => {
209        const _: () = {
210            #[unsafe(export_name = $export_name)]
211            pub unsafe extern "system" fn _ib_hook_inject_dll_app_apply(
212                __args_and_params: *mut ::core::ffi::c_void,
213            ) {
214                $crate::inject::dll::app::__payload_procedure_helper(__args_and_params, |__args| {
215                    let (input,) = __args;
216                    $apply(input)
217                });
218            }
219        };
220    };
221}
222pub use export_apply;
223
224/// Represents an injected DLL with its syringe, payload, and remote apply function.
225pub struct DllInjection<D: DllApp> {
226    syringe: Syringe,
227    /// The injected DLL module (borrowed from the syringe).
228    payload: BorrowedProcessModule<'static>,
229    /// Remote procedure to call apply on the injected DLL.
230    remote_apply: RemotePayloadProcedure<fn(Option<&'static D::Input>) -> D::Output>,
231    /// PID of the target process.
232    pid: Pid,
233    /// Whether APPLY was successfully called.
234    applied: bool,
235    /// Whether the injection has been ejected (prevents cleanup on drop).
236    ejected: bool,
237}
238
239/**
240[`Syringe`] contains [`RemoteBoxAllocator`] which is [`Rc`] inner and thus `!Send`.
241But [`Syringe`] itself is `!Clone`, it's actually `Send`.
242*/
243unsafe impl<D: DllApp> Send for DllInjection<D> {}
244
245/*
246Only [`apply()`] works with `&self`.
247Unfortunately, it uses a `!Sync` allocator.
248*/
249// unsafe impl<D: DllApp> Sync for DllInjection<D> {}
250
251#[bon]
252impl<D: DllApp> DllInjection<D> {
253    /// Inject the DLL into the given process and optionally appliy it (call `APPLY()`).
254    #[builder]
255    pub fn inject(
256        #[builder(start_fn)] process: OwnedProcess,
257        dll_path: &Path,
258        apply: Option<&D::Input>,
259    ) -> Result<Self, InjectError> {
260        let pid = Pid(process.pid().unwrap().get());
261        let syringe = Syringe::for_process(process);
262
263        info!(%pid, ?dll_path, "Injecting");
264        let payload = syringe.find_or_inject(dll_path)?;
265        let eject = || {
266            if let Err(e) = syringe.eject(payload) {
267                warn!(?e, "eject");
268            }
269        };
270
271        // Eject if ApplyNotFound
272        let remote_apply = unsafe { syringe.get_payload_procedure(payload, D::APPLY) }
273            .map_err(InjectError::from)
274            .inspect_err(|_| eject())?
275            .ok_or(InjectError::ApplyNotFound)
276            .inspect_err(|_| eject())?;
277
278        // Transmute payload to 'static since syringe (owner of process) is returned
279        let payload = unsafe { transmute(payload) };
280
281        let mut injection = Self {
282            payload,
283            syringe,
284            remote_apply,
285            pid,
286            applied: false,
287            ejected: false,
288        };
289
290        if let Some(input) = apply {
291            // Drop & eject on error
292            injection.apply(input)?;
293            injection.applied = true;
294        }
295
296        info!(%pid, "Successfully injected");
297
298        Ok(injection)
299    }
300
301    pub fn pid(&self) -> Pid {
302        self.pid
303    }
304
305    /// Call [`DllApp::APPLY`] with the given input.
306    pub fn maybe_apply(
307        &self,
308        input: Option<&D::Input>,
309    ) -> Result<D::Output, dll_syringe::rpc::PayloadRpcError> {
310        if let Some(input) = input {
311            self.apply(input)
312        } else {
313            self.unapply()
314        }
315    }
316
317    /// Call [`DllApp::APPLY`] with the given input.
318    /**
319    ## Implementation
320    ```ignore
321    let args = (input,);
322    let config = bincode::config::standard();
323
324    let mut size_writer = bincode::enc::write::SizeWriter::default();
325    bincode::serde::encode_into_writer(&args, &mut size_writer, config)?;
326    let arg_bytes = size_writer.bytes_written;
327    let mut local_arg_buf = Vec::with_capacity(arg_bytes);
328    bincode::serde::encode_into_std_write(&args, &mut local_arg_buf, config)?;
329    unsafe { local_arg_buf.set_len(arg_bytes) };
330    // println!("{:02X?}", local_arg_buf);
331
332    /*
333    let buf = &local_arg_buf[..arg_bytes].to_owned();
334    println!("decode_from_slice {:02X?}", buf);
335    let args: Result<((Option<Input>,), usize), bincode::error::DecodeError> =
336        bincode::serde::decode_from_slice::<(Option<Input>,), _>(buf, config);
337    println!("{:?}", args);
338    */
339    ...
340    ```
341    */
342    pub fn apply(&self, input: &D::Input) -> Result<D::Output, dll_syringe::rpc::PayloadRpcError> {
343        // call() doen't really need 'static.
344        // &Option<D::Input> can be used instead and was used before,
345        // but it made life pathetic.
346        let input: &'static D::Input = unsafe { transmute(input) };
347        self.remote_apply.call(&Some(input))
348    }
349
350    pub fn unapply(&self) -> Result<D::Output, dll_syringe::rpc::PayloadRpcError> {
351        self.remote_apply.call(&None)
352    }
353
354    /// Eject the DLL from the target process.
355    ///
356    /// This will first unapply if it was applied before.
357    pub fn eject(&mut self) -> Result<(), InjectError> {
358        if self.ejected {
359            return Ok(());
360        }
361        if self.applied {
362            self.unapply()?;
363        }
364        // Panics if the given module was not loaded in the target process.
365        self.syringe.eject(self.payload)?;
366        self.ejected = true;
367        Ok(())
368    }
369
370    /// Leak the injection, preventing both manual and automatic `eject()`.
371    pub fn leak(&mut self) {
372        self.ejected = true;
373    }
374}
375
376impl<D: DllApp> Drop for DllInjection<D> {
377    fn drop(&mut self) {
378        if let Err(e) = self.eject() {
379            error!(pid = self.pid.0, ?e, "Failed to eject on drop");
380        }
381    }
382}
383
384/// A collection of injected processes that can be ejected together.
385pub struct DllInjectionVec<D: DllApp> {
386    injections: Vec<DllInjection<D>>,
387    ejected: bool,
388}
389
390/// #[derive(Debug)] will require D: Default
391impl<D: DllApp> Default for DllInjectionVec<D> {
392    fn default() -> Self {
393        Self {
394            injections: Default::default(),
395            ejected: Default::default(),
396        }
397    }
398}
399
400#[bon]
401impl<D: DllApp> DllInjectionVec<D> {
402    pub fn new() -> Self {
403        Self::default()
404    }
405
406    pub fn injections(&self) -> &[DllInjection<D>] {
407        &self.injections
408    }
409
410    pub fn injections_mut(&mut self) -> &mut [DllInjection<D>] {
411        &mut self.injections
412    }
413
414    /// Inject the DLL into the given processes.
415    ///
416    /// Before [`DllInjectionVec::eject()`], the DLL file will be locked and can't be deleted.
417    ///
418    /// # Returns
419    /// - `Ok(DllInjectionVec)`: Successfully injected processes
420    /// - `Err(InjectError)`: Error during injection
421    #[builder]
422    pub fn inject(
423        &mut self,
424        /// Processes to inject into.
425        #[builder(start_fn)]
426        processes: impl Iterator<Item = OwnedProcess>,
427        /// Path to the DLL
428        dll_path: &Path,
429        /// Optionally apply with the given input after injection.
430        apply: Option<&D::Input>,
431        /// Optional callback for errors during injection (called in the middle of the loop).
432        ///
433        /// Errors will always be logged.
434        mut on_error: Option<impl FnMut(Pid, InjectError)>,
435    ) -> Result<&mut Self, InjectError> {
436        if !dll_path.exists() {
437            return Err(InjectError::DllNotFound(dll_path.to_path_buf()));
438        }
439
440        // Store injected processes for later eject
441        for target_process in processes {
442            let pid = Pid(target_process.pid().unwrap().get());
443            match DllInjection::inject(target_process)
444                .dll_path(&dll_path)
445                .maybe_apply(apply)
446                .call()
447            {
448                Ok(injection) => {
449                    self.injections.push(injection);
450                }
451                Err(e) => {
452                    error!(%pid, ?e, "Failed to inject");
453                    if let Some(cb) = on_error.as_mut() {
454                        cb(pid, e);
455                    }
456                }
457            }
458        }
459
460        Ok(self)
461    }
462
463    /// Inject the DLL into all processes with the given name.
464    ///
465    /// Before [`DllInjectionVec::eject()`], the DLL file will be locked and can't be deleted.
466    ///
467    /// # Returns
468    /// - `Ok(DllInjectionVec)`: Successfully injected processes
469    /// - `Err(InjectError)`: Error during injection
470    #[builder]
471    pub fn inject_with_process_name(
472        &mut self,
473        /// Name of the process to inject into.
474        #[builder(start_fn)]
475        process_name: &str,
476        /// Path to the DLL
477        dll_path: &Path,
478        /// Optionally apply with the given input after injection.
479        apply: Option<&D::Input>,
480        /// Optional callback for errors during injection (called in the middle of the loop).
481        ///
482        /// Errors will always be logged.
483        on_error: Option<impl FnMut(Pid, InjectError)>,
484    ) -> Result<&mut Self, InjectError> {
485        // Find all processes with the given name
486        let processes = OwnedProcess::find_all_by_name(process_name);
487        if processes.is_empty() {
488            return Err(InjectError::ProcessNotFound(process_name.to_string()));
489        }
490        info!("Found {} {} processes", processes.len(), process_name);
491
492        self.inject(processes.into_iter())
493            .dll_path(dll_path)
494            .maybe_apply(apply)
495            .maybe_on_error(on_error)
496            .call()
497    }
498
499    /// Call [`apply`](DllInjection::apply) on all injections.
500    ///
501    /// Errors are reported via the `on_error` callback.
502    #[builder]
503    pub fn apply(
504        &self,
505        #[builder(start_fn)] input: &D::Input,
506        mut on_error: Option<impl FnMut(Pid, &dll_syringe::rpc::PayloadRpcError)>,
507    ) {
508        for injection in &self.injections {
509            if let Err(e) = injection.apply(input) {
510                if let Some(on_error) = on_error.as_mut() {
511                    on_error(injection.pid(), &e);
512                }
513            }
514        }
515    }
516
517    /// Call [`unapply`](DllInjection::unapply) on all injections.
518    ///
519    /// Errors are reported via the `on_error` callback.
520    #[builder]
521    pub fn unapply(
522        &self,
523        mut on_error: Option<impl FnMut(Pid, &dll_syringe::rpc::PayloadRpcError)>,
524    ) {
525        for injection in &self.injections {
526            if let Err(e) = injection.unapply() {
527                if let Some(on_error) = on_error.as_mut() {
528                    on_error(injection.pid(), &e);
529                }
530            }
531        }
532    }
533
534    /// Eject all DLL injections.
535    #[builder]
536    pub fn eject(
537        &mut self,
538        /// Optional callback for errors during ejection (called in the middle of the loop).
539        ///
540        /// Errors will always be logged.
541        mut on_error: Option<impl FnMut(Pid, InjectError)>,
542    ) {
543        for mut injection in self.injections.drain(..) {
544            let pid = injection.pid;
545            if let Err(e) = injection.eject() {
546                warn!(%pid, ?e, "Failed to eject");
547                if let Some(cb) = on_error.as_mut() {
548                    cb(pid, e);
549                }
550            }
551        }
552
553        info!("Successfully ejected");
554    }
555
556    /// Leak existing and new injections, preventing automatic cleanup on drop.
557    /// (Unlike [`DllInjection::leak`], this doesn't prevent manual `eject()`)
558    pub fn leak(&mut self) {
559        self.ejected = true;
560    }
561}
562
563impl<D: DllApp> Drop for DllInjectionVec<D> {
564    fn drop(&mut self) {
565        if self.ejected {
566            for injection in &mut self.injections {
567                injection.leak();
568            }
569            return;
570        }
571        self.eject().on_error(|_, _| ()).call();
572    }
573}
574
575/// A collection of injected processes that can be ejected together, with a shared input for apply.
576///
577/// Unlike [`DllInjectionVec`], this stores the input and applies it automatically during inject.
578#[derive(Deref, DerefMut)]
579pub struct DllInjectionVecWithInput<D: DllApp> {
580    dll_path: PathBuf,
581    input: Option<D::Input>,
582    #[deref]
583    #[deref_mut]
584    inner: DllInjectionVec<D>,
585}
586
587impl<D: DllApp> DllInjectionVecWithInput<D> {
588    /// Creates a new `DllInjectionVecWithInput` with the given DLL path.
589    ///
590    /// The DLL path is checked to ensure it exists.
591    pub fn new(dll_path: PathBuf) -> Result<Self, InjectError> {
592        Self::with_input(dll_path, None)
593    }
594
595    pub fn with_input(dll_path: PathBuf, input: Option<D::Input>) -> Result<Self, InjectError> {
596        if !dll_path.exists() {
597            return Err(InjectError::DllNotFound(dll_path));
598        }
599        Ok(Self {
600            dll_path,
601            input,
602            inner: Default::default(),
603        })
604    }
605
606    pub fn dll_path(&self) -> &PathBuf {
607        &self.dll_path
608    }
609
610    pub fn input(&self) -> Option<&D::Input> {
611        self.input.as_ref()
612    }
613}
614
615#[bon]
616impl<D: DllApp> DllInjectionVecWithInput<D> {
617    /// Inject the DLL into the given processes with the stored input and dll_path.
618    ///
619    /// Before [`DllInjectionVecWithInput::eject()`], the DLL file will be locked and can't be deleted.
620    ///
621    /// # Returns
622    /// - `Ok(DllInjectionVecWithInput)`: Successfully injected processes
623    /// - `Err(InjectError)`: Error during injection
624    #[builder]
625    pub fn inject(
626        &mut self,
627        /// Processes to inject into.
628        #[builder(start_fn)]
629        processes: impl Iterator<Item = OwnedProcess>,
630        /// Optional callback for errors during injection (called in the middle of the loop).
631        ///
632        /// Errors will always be logged.
633        on_error: Option<impl FnMut(Pid, InjectError)>,
634    ) -> Result<&mut Self, InjectError> {
635        self.inner
636            .inject(processes)
637            .dll_path(&self.dll_path)
638            .maybe_apply(self.input.as_ref())
639            .maybe_on_error(on_error)
640            .call()?;
641        Ok(self)
642    }
643
644    /// Inject the DLL into all processes with the given name.
645    ///
646    /// Before [`DllInjectionVecWithInput::eject()`], the DLL file will be locked and can't be deleted.
647    ///
648    /// # Returns
649    /// - `Ok(DllInjectionVecWithInput)`: Successfully injected processes
650    /// - `Err(InjectError)`: Error during injection
651    #[builder]
652    pub fn inject_with_process_name(
653        &mut self,
654        /// Name of the process to inject into.
655        #[builder(start_fn)]
656        process_name: &str,
657        /// Optional callback for errors during injection (called in the middle of the loop).
658        ///
659        /// Errors will always be logged.
660        on_error: Option<impl FnMut(Pid, InjectError)>,
661    ) -> Result<&mut Self, InjectError> {
662        self.inner
663            .inject_with_process_name(process_name)
664            .dll_path(&self.dll_path)
665            .maybe_apply(self.input.as_ref())
666            .maybe_on_error(on_error)
667            .call()?;
668        Ok(self)
669    }
670
671    /// Call [`apply`](DllInjection::apply) on all injections with the stored input.
672    ///
673    /// Updates the stored input to the new value.
674    #[builder]
675    pub fn apply(
676        &mut self,
677        #[builder(start_fn)] input: D::Input,
678        on_error: Option<impl FnMut(Pid, &dll_syringe::rpc::PayloadRpcError)>,
679    ) {
680        let input = self.input.insert(input);
681        self.inner.apply(input).maybe_on_error(on_error).call();
682    }
683
684    /// Call [`unapply`](DllInjection::unapply) on all injections.
685    ///
686    /// Updates the stored input to `None`.
687    #[builder]
688    pub fn unapply(
689        &mut self,
690        on_error: Option<impl FnMut(Pid, &dll_syringe::rpc::PayloadRpcError)>,
691    ) {
692        self.input = None;
693        self.inner.unapply().maybe_on_error(on_error).call()
694    }
695}
ib_hook/inject/dll/app.rs

ib_hook/inject/dll/
app.rs
