Skip to main content

vmi_core/ctx/
session.rs

1use std::{io::ErrorKind, time::Duration};
2
3use super::{VmiState, context::VmiContext};
4use crate::{
5    Architecture, VmiCore, VmiError, VmiEventControl, VmiHandler, VmiVmControl,
6    os::{NoOS, VmiOs},
7};
8
9/// A VMI session.
10///
11/// The session combines a [`VmiCore`] with an OS-specific [`VmiOs`]
12/// implementation to provide unified access to both low-level VMI operations
13/// and higher-level OS abstractions.
14pub struct VmiSession<'a, Os>
15where
16    Os: VmiOs,
17{
18    /// The VMI core providing low-level VM introspection capabilities.
19    core: &'a VmiCore<Os::Driver>,
20
21    /// The OS-specific operations and abstractions.
22    os: &'a Os,
23}
24
25impl<Os> Clone for VmiSession<'_, Os>
26where
27    Os: VmiOs,
28{
29    fn clone(&self) -> Self {
30        *self
31    }
32}
33
34impl<Os> Copy for VmiSession<'_, Os> where Os: VmiOs {}
35
36impl<Os> std::ops::Deref for VmiSession<'_, Os>
37where
38    Os: VmiOs,
39{
40    type Target = VmiCore<Os::Driver>;
41
42    fn deref(&self) -> &Self::Target {
43        self.core
44    }
45}
46
47impl<'a, Os> VmiSession<'a, Os>
48where
49    Os: VmiOs,
50{
51    /// Creates a new VMI session.
52    pub fn new(core: &'a VmiCore<Os::Driver>, os: &'a Os) -> Self {
53        Self { core, os }
54    }
55
56    /// Creates a new VMI state with the specified registers.
57    pub fn with_registers(
58        &'a self,
59        registers: &'a <Os::Architecture as Architecture>::Registers,
60    ) -> VmiState<'a, Os> {
61        VmiState::new(self, registers)
62    }
63
64    /// Creates a new VMI session without an OS-specific implementation.
65    pub fn without_os(&self) -> VmiSession<'a, NoOS<Os::Driver>> {
66        VmiSession {
67            core: self.core,
68            os: const { &NoOS(std::marker::PhantomData) },
69        }
70    }
71
72    /// Returns the VMI core.
73    pub fn core(&self) -> &'a VmiCore<Os::Driver> {
74        self.core
75    }
76
77    /// Returns the underlying OS-specific implementation.
78    pub fn underlying_os(&self) -> &'a Os {
79        self.os
80    }
81}
82
83///////////////////////////////////////////////////////////////////////////////
84// VmiEventControl
85///////////////////////////////////////////////////////////////////////////////
86
87impl<'a, Os> VmiSession<'a, Os>
88where
89    Os: VmiOs,
90    Os::Driver: VmiEventControl,
91{
92    /// Waits for an event to occur and processes it with the provided handler.
93    ///
94    /// This method blocks until an event occurs or the specified timeout is
95    /// reached. When an event occurs, it is passed to the provided callback
96    /// function for processing.
97    pub fn wait_for_event(
98        &self,
99        timeout: Duration,
100        handler: &mut impl VmiHandler<Os>,
101    ) -> Result<(), VmiError> {
102        self.core.wait_for_event(timeout, |event| {
103            let state = VmiState::new(self, event.registers());
104            handler.handle_event(VmiContext::new(&state, event))
105        })
106    }
107}
108
109///////////////////////////////////////////////////////////////////////////////
110// VmiEventControl + VmiVmControl
111///////////////////////////////////////////////////////////////////////////////
112
113impl<'a, Os> VmiSession<'a, Os>
114where
115    Os: VmiOs,
116    Os::Driver: VmiEventControl + VmiVmControl,
117{
118    /// Enters the main event handling loop that processes VMI events until
119    /// finished.
120    pub fn handle<Handler>(
121        &self,
122        handler_factory: impl FnOnce(&VmiSession<Os>) -> Result<Handler, VmiError>,
123    ) -> Result<Option<Handler::Output>, VmiError>
124    where
125        Handler: VmiHandler<Os>,
126    {
127        self.handle_with_timeout(Duration::from_millis(5000), handler_factory)
128    }
129
130    /// Enters the main event handling loop that processes VMI events until
131    /// finished, with a timeout for each event.
132    pub fn handle_with_timeout<Handler>(
133        &self,
134        timeout: Duration,
135        handler_factory: impl FnOnce(&VmiSession<Os>) -> Result<Handler, VmiError>,
136    ) -> Result<Option<Handler::Output>, VmiError>
137    where
138        Handler: VmiHandler<Os>,
139    {
140        let mut result;
141        let mut handler = handler_factory(self)?;
142
143        loop {
144            result = handler.poll();
145
146            if result.is_some() {
147                break;
148            }
149
150            match self.wait_for_event(timeout, &mut handler) {
151                Err(VmiError::Timeout) => {
152                    tracing::trace!("timeout");
153                    handler.handle_timeout(self);
154                }
155                Err(VmiError::Io(err)) if err.kind() == ErrorKind::Interrupted => {
156                    tracing::trace!("interrupted");
157                    handler.handle_interrupted(self);
158                    break;
159                }
160                Err(err) => return Err(err),
161                Ok(_) => {}
162            }
163        }
164
165        tracing::trace!("disabling monitor");
166        handler.cleanup(self);
167        self.core.reset_state()?;
168        tracing::trace!(pending_events = self.events_pending());
169
170        let _pause_guard = self.pause_guard()?;
171        if self.events_pending() > 0 {
172            match self.wait_for_event(Duration::from_millis(0), &mut handler) {
173                Err(VmiError::Timeout) => {
174                    tracing::trace!("timeout");
175                }
176                Err(err) => return Err(err),
177                Ok(_) => {}
178            }
179        }
180
181        Ok(result)
182    }
183}