vmi_core/ctx/
session.rs

1use std::{io::ErrorKind, time::Duration};
2
3use super::{VmiState, context::VmiContext};
4use crate::{
5    Architecture, VcpuId, VmiCore, VmiError, VmiHandler,
6    driver::{VmiDriver, VmiEventControl, VmiQueryRegisters, VmiVmControl},
7    os::{NoOS, VmiOs},
8};
9
10/// A VMI session.
11///
12/// The session combines a [`VmiCore`] with an OS-specific [`VmiOs`]
13/// implementation to provide unified access to both low-level VMI operations
14/// and higher-level OS abstractions.
15pub struct VmiSession<'a, Os>
16where
17    Os: VmiOs,
18{
19    /// The VMI core providing low-level VM introspection capabilities.
20    core: &'a VmiCore<Os::Driver>,
21
22    /// The OS-specific operations and abstractions.
23    os: &'a Os,
24}
25
26impl<Os> Clone for VmiSession<'_, Os>
27where
28    Os: VmiOs,
29{
30    fn clone(&self) -> Self {
31        *self
32    }
33}
34
35impl<Os> Copy for VmiSession<'_, Os> where Os: VmiOs {}
36
37impl<Os> std::ops::Deref for VmiSession<'_, Os>
38where
39    Os: VmiOs,
40{
41    type Target = VmiCore<Os::Driver>;
42
43    fn deref(&self) -> &Self::Target {
44        self.core
45    }
46}
47
48impl<'a, Os> VmiSession<'a, Os>
49where
50    Os: VmiOs,
51{
52    /// Creates a new VMI session.
53    pub fn new(core: &'a VmiCore<Os::Driver>, os: &'a Os) -> Self {
54        Self { core, os }
55    }
56
57    /// Creates a new VMI state with the specified registers.
58    pub fn with_registers(
59        &'a self,
60        registers: &'a <Os::Architecture as Architecture>::Registers,
61    ) -> VmiState<'a, Os> {
62        VmiState::new(self, registers)
63    }
64
65    /// Creates a new VMI session without an OS-specific implementation.
66    pub fn without_os(&self) -> VmiSession<'a, NoOS<Os::Driver>> {
67        VmiSession {
68            core: self.core,
69            os: const { &NoOS(std::marker::PhantomData) },
70        }
71    }
72
73    /// Returns the VMI core.
74    pub fn core(&self) -> &'a VmiCore<Os::Driver> {
75        self.core
76    }
77
78    /// Returns the underlying OS-specific implementation.
79    pub fn underlying_os(&self) -> &'a Os {
80        self.os
81    }
82}
83
84///////////////////////////////////////////////////////////////////////////////
85// VmiEventControl
86///////////////////////////////////////////////////////////////////////////////
87
88impl<'a, Os> VmiSession<'a, Os>
89where
90    Os: VmiOs,
91    Os::Driver: VmiEventControl,
92{
93    /// Waits for an event to occur and processes it with the provided handler.
94    ///
95    /// This method blocks until an event occurs or the specified timeout is
96    /// reached. When an event occurs, it is passed to the provided callback
97    /// function for processing.
98    pub fn wait_for_event(
99        &self,
100        timeout: Duration,
101        handler: &mut impl VmiHandler<Os>,
102    ) -> Result<(), VmiError> {
103        self.core.wait_for_event(timeout, |event| {
104            let state = VmiState::new(self, event.registers());
105            handler.handle_event(VmiContext::new(&state, event))
106        })
107    }
108}
109
110///////////////////////////////////////////////////////////////////////////////
111// VmiEventControl + VmiVmControl
112///////////////////////////////////////////////////////////////////////////////
113
114impl<'a, Os> VmiSession<'a, Os>
115where
116    Os: VmiOs,
117    Os::Driver: VmiEventControl + VmiVmControl,
118{
119    /// Enters the main event handling loop that processes VMI events until
120    /// finished.
121    pub fn handle<Handler>(
122        &self,
123        handler_factory: impl FnOnce(&VmiSession<Os>) -> Result<Handler, VmiError>,
124    ) -> Result<Option<Handler::Output>, VmiError>
125    where
126        Handler: VmiHandler<Os>,
127    {
128        self.handle_with_timeout(Duration::from_millis(5000), handler_factory)
129    }
130
131    /// Enters the main event handling loop that processes VMI events until
132    /// finished, with a timeout for each event.
133    pub fn handle_with_timeout<Handler>(
134        &self,
135        timeout: Duration,
136        handler_factory: impl FnOnce(&VmiSession<Os>) -> Result<Handler, VmiError>,
137    ) -> Result<Option<Handler::Output>, VmiError>
138    where
139        Handler: VmiHandler<Os>,
140    {
141        let mut result;
142        let mut handler = handler_factory(self)?;
143
144        loop {
145            result = handler.poll();
146
147            if result.is_some() {
148                break;
149            }
150
151            match self.wait_for_event(timeout, &mut handler) {
152                Ok(_) => {}
153                Err(VmiError::Timeout) => {
154                    tracing::trace!("timeout");
155                    handler.handle_timeout(self);
156                }
157                Err(VmiError::Io(err)) if err.kind() == ErrorKind::Interrupted => {
158                    tracing::trace!("interrupted");
159                    handler.handle_interrupted(self);
160                    break;
161                }
162                Err(err) => return Err(err),
163            }
164        }
165
166        tracing::trace!("disabling monitor");
167        handler.cleanup(self);
168        self.core.reset_state()?;
169        tracing::trace!(pending_events = self.events_pending());
170
171        let _pause_guard = self.core().pause_guard()?;
172        if self.events_pending() > 0 {
173            match self.wait_for_event(Duration::from_millis(0), &mut handler) {
174                Err(VmiError::Timeout) => {
175                    tracing::trace!("timeout");
176                }
177                Err(err) => return Err(err),
178                Ok(_) => {}
179            }
180        }
181
182        Ok(result)
183    }
184}
185
186///////////////////////////////////////////////////////////////////////////////
187// VmiQueryRegisters + VmiVmControl
188///////////////////////////////////////////////////////////////////////////////
189
190impl<'a, Os> VmiSession<'a, Os>
191where
192    Os: VmiOs,
193    Os::Driver: VmiQueryRegisters + VmiVmControl,
194{
195    /// Pauses the virtual machine, snapshots the boot CPU registers, and
196    /// returns a guard that resumes the VM when dropped.
197    pub fn pause_guard(&self) -> Result<VmiSessionPauseGuard<'_, Os>, VmiError> {
198        VmiSessionPauseGuard::new(self)
199    }
200}
201
202/// A guard that pauses the virtual machine and snapshots the boot CPU
203/// registers on creation, then resumes the VM on drop.
204///
205/// Unlike [`VmiPauseGuard`], this guard carries enough register context to
206/// build a [`VmiState`] via [`state`], so callers can introspect the paused
207/// guest without separately querying registers.
208///
209/// [`VmiPauseGuard`]: crate::VmiPauseGuard
210/// [`state`]: Self::state
211pub struct VmiSessionPauseGuard<'a, Os>
212where
213    Os: VmiOs,
214    Os::Driver: VmiQueryRegisters + VmiVmControl,
215{
216    session: &'a VmiSession<'a, Os>,
217    registers: <<Os::Driver as VmiDriver>::Architecture as Architecture>::Registers,
218}
219
220impl<'a, Os> VmiSessionPauseGuard<'a, Os>
221where
222    Os: VmiOs,
223    Os::Driver: VmiQueryRegisters + VmiVmControl,
224{
225    /// Creates a new pause guard.
226    pub fn new(session: &'a VmiSession<'a, Os>) -> Result<Self, VmiError> {
227        session.driver().pause()?;
228
229        let registers = match session.registers(VcpuId(0)) {
230            Ok(registers) => registers,
231            Err(err) => {
232                if let Err(resume_err) = session.driver().resume() {
233                    tracing::error!(
234                        err = %resume_err,
235                        "failed to resume after register-query failure"
236                    );
237                }
238                return Err(err);
239            }
240        };
241
242        Ok(Self { session, registers })
243    }
244
245    /// Returns the VMI session.
246    pub fn session(&self) -> &VmiSession<'a, Os> {
247        self.session
248    }
249
250    /// Returns the boot CPU (`VcpuId(0)`) registers captured when
251    /// the guard was created.
252    pub fn registers(&self) -> &<Os::Architecture as Architecture>::Registers {
253        &self.registers
254    }
255
256    /// Returns the state captured when the guard was created, bound to
257    /// the boot CPU (`VcpuId(0)`) registers.
258    pub fn state(&self) -> VmiState<'_, Os> {
259        VmiState::new(self.session, &self.registers)
260    }
261}
262
263impl<Os> Drop for VmiSessionPauseGuard<'_, Os>
264where
265    Os: VmiOs,
266    Os::Driver: VmiQueryRegisters + VmiVmControl,
267{
268    fn drop(&mut self) {
269        if let Err(err) = self.session.driver().resume() {
270            tracing::error!(%err, "failed to resume the virtual machine");
271        }
272    }
273}
vmi_core/ctx/session.rs

vmi_core/ctx/
session.rs
