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
9pub struct VmiSession<'a, Os>
15where
16 Os: VmiOs,
17{
18 core: &'a VmiCore<Os::Driver>,
20
21 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 pub fn new(core: &'a VmiCore<Os::Driver>, os: &'a Os) -> Self {
53 Self { core, os }
54 }
55
56 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 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 pub fn core(&self) -> &'a VmiCore<Os::Driver> {
74 self.core
75 }
76
77 pub fn underlying_os(&self) -> &'a Os {
79 self.os
80 }
81}
82
83impl<'a, Os> VmiSession<'a, Os>
88where
89 Os: VmiOs,
90 Os::Driver: VmiEventControl,
91{
92 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
109impl<'a, Os> VmiSession<'a, Os>
114where
115 Os: VmiOs,
116 Os::Driver: VmiEventControl + VmiVmControl,
117{
118 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 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}