1use crate::MemCheck;
2use crate::allocator::{ConsecAllocator, alloc_memory};
3use crate::hammerer::Hammering;
4use crate::memory::{BitFlip, BytePointer, ConsecBlocks, DataPattern, Initializable};
5use crate::util::{NamedProgress, PAGE_MASK, Rng, Size};
6use crate::victim::{HammerVictimError, VictimOrchestrator, VictimResult};
7use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
8use log::{debug, info, warn};
9use serde::{Serialize, Serializer};
10use std::collections::HashMap;
11use std::time::{Duration, Instant};
12use thiserror::Error;
13
14pub type ProfileHammererFactory<H> = Box<dyn Fn(ConsecBlocks) -> H>;
15pub type HammererFactory<H1, H2> = Box<dyn Fn(H1, ConsecBlocks, RoundProfile) -> H2>;
16pub type VictimFactory<E> =
17 Box<dyn Fn(ConsecBlocks, RoundProfile) -> Result<Box<dyn VictimOrchestrator>, E>>;
18
19pub struct Swage<PH: Hammering, H: Hammering, AE: std::error::Error, VE: std::error::Error> {
41 allocator: Box<dyn ConsecAllocator<Error = AE>>,
42 profile_hammerer_factory: ProfileHammererFactory<PH>,
43 profile_data_pattern: DataPatternKind,
44 hammerer_factory: HammererFactory<PH, H>,
45 victim_factory: VictimFactory<VE>,
46 pattern_size: usize,
47 progress: Option<MultiProgress>,
48 config: SwageConfig,
49}
50
51#[derive(Debug, Serialize, Clone)]
56pub struct RoundProfile {
57 pub bit_flips: Vec<BitFlip>,
59 pub pattern: DataPattern,
61}
62
63pub struct SwageConfig {
67 pub profiling_rounds: u64,
69 pub reproducibility_threshold: f64,
71
72 pub hammering_timeout: Option<Duration>,
74 pub repetitions: Option<u64>,
76 pub timeout: Option<Duration>,
78}
79
80impl Default for SwageConfig {
81 fn default() -> Self {
82 Self {
83 profiling_rounds: 10,
84 reproducibility_threshold: 0.8,
85 hammering_timeout: None,
86 repetitions: Some(1),
87 timeout: None,
88 }
89 }
90}
91
92#[derive(Serialize)]
101pub struct ExperimentData<T, E> {
102 date: String,
104 results: Vec<std::result::Result<T, E>>,
106 profiling: RoundProfile,
108 data: Option<serde_json::Value>,
110}
111
112impl<T, E> ExperimentData<T, E> {
113 fn new(
114 results: Vec<std::result::Result<T, E>>,
115 profiling: RoundProfile,
116 data: Option<serde_json::Value>,
117 ) -> Self {
118 Self {
119 date: chrono::Local::now().to_rfc3339(),
120 results,
121 profiling,
122 data,
123 }
124 }
125}
126
127impl<H: Hammering, AE: std::error::Error, VE: std::error::Error> Swage<H, H, AE, VE> {
128 pub fn builder() -> SwageBuilder<H, H, AE, VE> {
134 SwageBuilder::default()
135 }
136}
137
138#[derive(Debug, Error)]
139pub enum HammerError<AE: std::error::Error, HE: std::error::Error, VE: std::error::Error> {
140 #[error(transparent)]
141 AllocationFailed(AE),
142 #[error(transparent)]
143 HammeringFailed(HE),
144 #[error(transparent)]
145 VictimFailed(VE),
146 #[error("No vulnerable cells found during profiling")]
147 NoVulnerableCells,
148 #[error(transparent)]
149 VictimError(#[from] HammerVictimError),
150}
151
152impl<AE: std::error::Error, HE: std::error::Error, VE: std::error::Error> Serialize
153 for HammerError<AE, HE, VE>
154{
155 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
156 where
157 S: Serializer,
158 {
159 serializer.serialize_str(&self.to_string())
160 }
161}
162
163impl<PH: Hammering, H: Hammering, AE: std::error::Error, VE: std::error::Error>
164 Swage<PH, H, AE, VE>
165{
166 fn round(
167 &mut self,
168 start: Instant,
169 hammering_time: &mut Duration,
170 ) -> ExperimentData<VictimResult, HammerError<AE, H::Error, VE>> {
171 info!("Starting bait allocation");
172 let memory = match alloc_memory(self.allocator.as_mut(), Size::B(self.pattern_size)) {
174 Ok(memory) => memory,
175 Err(e) => {
176 warn!("Failed to allocate memory: {}", e);
177 return ExperimentData::new(
178 vec![Err(HammerError::AllocationFailed(e))],
179 RoundProfile {
180 bit_flips: vec![],
181 pattern: DataPattern::Random(Box::new(Rng::from_seed(rand::random()))),
182 },
183 None,
184 );
185 }
186 };
187 info!("Allocated {} bytes of memory", memory.len());
188
189 info!("Profiling memory for vulnerable addresses");
190
191 let hammerer = (self.profile_hammerer_factory)(memory.clone());
192
193 let profiling = hammer_profile(
194 &hammerer,
195 memory.clone(),
196 self.profile_data_pattern,
197 self.config.profiling_rounds,
198 self.config.reproducibility_threshold,
199 self.progress.clone(),
200 );
201 debug!("Profiling results: {:?}", profiling);
202 if profiling.bit_flips.is_empty() {
203 warn!("No vulnerable addresses found");
204 memory.dealloc();
205 return ExperimentData::new(
206 vec![Err(HammerError::NoVulnerableCells)],
207 profiling.clone(),
208 None,
209 );
210 }
211
212 let flips = profiling.bit_flips.clone();
213 let dpattern = profiling.pattern.clone();
214
215 let hammerer = (self.hammerer_factory)(hammerer, memory.clone(), profiling.clone());
216
217 let mut victim = match (self.victim_factory)(memory.clone(), profiling.clone()) {
218 Ok(v) => v,
219 Err(e) => {
220 return ExperimentData::new(
221 vec![Err(HammerError::VictimFailed(e))],
222 profiling,
223 None,
224 );
225 }
226 };
227
228 match victim.as_mut().start() {
229 Ok(_) => {}
230 Err(e) => {
231 warn!("Failed to start victim: {:?}", e);
232 victim.stop();
233 memory.dealloc();
234 return ExperimentData::new(
235 vec![Err(HammerError::VictimError(e))],
236 profiling.clone(),
237 victim.serialize(),
238 );
239 }
240 }
241 let flip_pages = flips
242 .iter()
243 .map(|f| (f.addr & !PAGE_MASK) as *const u8)
244 .collect::<Vec<_>>();
245
246 let hammer_progress = match (self.config.hammering_timeout, self.progress.as_mut()) {
247 (Some(hammering_timeout), Some(p)) => {
248 let p = p.add(ProgressBar::new(hammering_timeout.as_secs()));
249 p.set_style(ProgressStyle::named_bar("Total hammering time"));
250 p.set_position((hammering_timeout - *hammering_time).as_secs());
251 Some(p)
252 }
253 _ => None,
254 };
255
256 let mut results: Vec<Result<VictimResult, HammerError<AE, H::Error, VE>>> = vec![];
257 loop {
258 if check_timeout(self.config.timeout, Instant::now() - start) {
259 info!("Timeout reached. Stopping.");
260 break;
261 }
262 if check_timeout(self.config.hammering_timeout, *hammering_time) {
263 info!("Hammering timeout reached. Stopping.");
264 break;
265 }
266 if let Some(hammer_progress) = &hammer_progress {
267 hammer_progress.set_position(hammering_time.as_secs());
268 }
269 memory.initialize_excluding(dpattern.clone(), &flip_pages); victim.init();
271 let hammer_start = Instant::now();
272 let result = hammerer.hammer();
273 *hammering_time += Instant::now().duration_since(hammer_start);
274 match result {
275 Ok(_) => {}
276 Err(err) => results.push(Err(HammerError::HammeringFailed(err))),
277 };
278 let result = victim.check();
279 match result {
280 Ok(result) => {
281 info!("Hammering successful: {:?}", result);
282 results.push(Ok(result));
283 }
284 Err(HammerVictimError::NoFlips) => {
285 warn!("No flips detected");
286 results.push(Err(HammerError::VictimError(HammerVictimError::NoFlips)));
287 }
288 Err(e) => {
289 warn!("Hammering failed: {:?}", e);
290 results.push(Err(HammerError::VictimError(e)));
291 break;
292 }
293 }
294 if self.config.timeout.is_none() && self.config.hammering_timeout.is_none() {
295 info!("No timeout set, stopping after one round");
296 break;
297 }
298 }
299 victim.stop();
300 memory.dealloc();
301 ExperimentData::new(results, profiling.clone(), victim.serialize())
302 }
303
304 pub fn run(mut self) -> Vec<ExperimentData<VictimResult, HammerError<AE, H::Error, VE>>> {
308 let mut experiments = vec![];
309
310 let repetitions = self.config.repetitions;
311 let timeout = self.config.timeout;
312 let hammering_timeout = self.config.hammering_timeout;
313
314 let start = Instant::now();
315 let mut hammering_time = Duration::ZERO;
316 let timeout_progress = match (timeout, self.progress.as_mut()) {
317 (Some(timeout), Some(p)) => {
318 let p = p.add(ProgressBar::new(timeout.as_secs()));
319 p.set_style(ProgressStyle::named_bar("Global timeout"));
320 p.set_position(0);
321 p.tick();
322 p.enable_steady_tick(Duration::from_secs(1));
323 Some(p)
324 }
325 _ => None,
326 };
327 let rep_progress = match (repetitions, self.progress.as_mut()) {
328 (Some(repetitions), Some(p)) => {
329 let p = p.add(ProgressBar::new(repetitions));
330 p.set_style(ProgressStyle::named_bar("Repetitions"));
331 Some(p)
332 }
333 _ => None,
334 };
335 for rep in 0..repetitions.unwrap_or(u64::MAX) {
336 if let Some(rep_progress) = &rep_progress {
337 rep_progress.set_position(rep + 1);
338 }
339 if let Some(timeout_progress) = &timeout_progress {
340 timeout_progress.set_position((Instant::now() - start).as_secs());
341 }
342 if rep > 0 && check_timeout(timeout, Instant::now() - start) {
343 info!("Timeout reached. Stopping.");
344 break;
345 }
346 if let Some(hammering_timeout) = hammering_timeout {
347 if hammering_time >= hammering_timeout {
348 info!("Hammering timeout reached. Stopping.");
349 break;
350 }
351 info!(
352 "Hammering time left: {} minutes",
353 (hammering_timeout - hammering_time).as_secs() / 60,
354 );
355 }
356 experiments.push(self.round(start, &mut hammering_time));
357 }
358 experiments
359 }
360}
361
362fn check_timeout(timeout: Option<Duration>, duration: Duration) -> bool {
363 timeout.is_some_and(|timeout| duration > timeout)
364}
365
366fn hammer_profile<E: std::error::Error>(
368 hammerer: &dyn Hammering<Error = E>,
369 memory: ConsecBlocks,
370 pattern: DataPatternKind,
371 num_rounds: u64,
372 reproducibility_threshold: f64,
373 progress: Option<MultiProgress>,
374) -> RoundProfile {
375 let p = progress.as_ref().map(|p| {
376 let p = p.add(ProgressBar::new(num_rounds));
377 p.set_style(ProgressStyle::named_bar("Profiling round"));
378 p.enable_steady_tick(Duration::from_secs(1));
379 p
380 });
381
382 const _SHM_SEED: u64 = 9804201662804659191;
383 let mut candidates = HashMap::new();
384 let min_repro_count = (reproducibility_threshold * num_rounds as f64) as u64;
385 let pattern = match pattern {
386 DataPatternKind::Random => DataPattern::Random(Box::new(Rng::from_seed(rand::random()))),
387 DataPatternKind::One => DataPattern::One,
388 DataPatternKind::Zero => DataPattern::Zero,
389 };
390 for r in 1..=num_rounds {
391 if let Some(p) = p.as_ref() {
392 p.set_position(r);
393 }
394 if candidates.is_empty() && r > num_rounds - min_repro_count {
395 warn!(
396 "No candidates and only {} round(s) left. Stopping profiling, continuing with next pattern",
397 num_rounds - r
398 );
399 break;
400 }
401 let mut victim = MemCheck::new(memory.clone(), pattern.clone(), vec![].into());
402 victim.init();
403 let result = hammerer.hammer();
404 match result {
405 Ok(_) => {
406 let result = victim.check();
407 let bit_flips = match result {
408 Ok(result) => {
409 info!("Profiling hammering round successful: {:?}", result);
410 result.bit_flips()
411 }
412 Err(e) => {
413 warn!("Profiling hammering round not successful: {:?}", e);
414 vec![]
415 }
416 };
417 for flip in bit_flips {
418 let entry = candidates.entry(flip).or_insert(0);
419 *entry += 1;
420 }
421 }
422 Err(e) => {
423 warn!("Profiling hammering round not successful: {:?}", e);
424 }
425 }
426 let remaining_rounds = num_rounds - r;
427 candidates.retain(|_, v| *v + remaining_rounds >= min_repro_count);
428 info!("Profiling round {} candidates: {:?}", r, candidates);
429 }
430 RoundProfile {
431 bit_flips: candidates.keys().cloned().collect(),
432 pattern,
433 }
434}
435
436#[derive(Clone, Copy)]
440pub enum DataPatternKind {
441 Random,
443 Zero,
445 One,
447}
448
449pub struct SwageBuilder<PH: Hammering, H: Hammering, AE: std::error::Error, VE: std::error::Error> {
450 allocator: Option<Box<dyn ConsecAllocator<Error = AE>>>,
451 profile_hammerer_factory: Option<ProfileHammererFactory<PH>>,
452 profile_data_pattern: DataPatternKind,
453 hammerer_factory: HammererFactory<PH, H>,
454 victim_factory: Option<VictimFactory<VE>>,
455 pattern_size: Option<usize>,
456 progress: Option<MultiProgress>,
457 config: SwageConfig,
458}
459
460impl<H: Hammering, AE: std::error::Error, VE: std::error::Error> Default
461 for SwageBuilder<H, H, AE, VE>
462{
463 fn default() -> Self {
464 SwageBuilder {
465 allocator: None,
466 profile_hammerer_factory: None,
467 profile_data_pattern: DataPatternKind::Random,
468 hammerer_factory: Box::new(|h, _, _| h),
469 victim_factory: None,
470 pattern_size: None,
471 progress: None,
472 config: SwageConfig::default(),
473 }
474 }
475}
476
477impl<PH: Hammering, H: Hammering, AE: std::error::Error, VE: std::error::Error>
478 SwageBuilder<PH, H, AE, VE>
479{
480 pub fn allocator<A: ConsecAllocator + 'static>(
481 self,
482 allocator: A,
483 ) -> SwageBuilder<PH, H, A::Error, VE> {
484 SwageBuilder {
485 allocator: Some(Box::new(allocator)),
486 profile_hammerer_factory: self.profile_hammerer_factory,
487 profile_data_pattern: self.profile_data_pattern,
488 hammerer_factory: self.hammerer_factory,
489 victim_factory: self.victim_factory,
490 pattern_size: self.pattern_size,
491 progress: self.progress,
492 config: self.config,
493 }
494 }
495
496 pub fn profile_hammerer_factory(
497 mut self,
498 profile_hammerer_factory: impl Fn(ConsecBlocks) -> PH + 'static,
499 ) -> Self {
500 self.profile_hammerer_factory = Some(Box::new(profile_hammerer_factory));
501 self
502 }
503
504 pub fn profile_data_pattern(mut self, profile_data_pattern: DataPatternKind) -> Self {
505 self.profile_data_pattern = profile_data_pattern;
506 self
507 }
508
509 pub fn hammerer_factory<H1: Hammering>(
510 self,
511 hammerer_factory: impl Fn(PH, ConsecBlocks, RoundProfile) -> H1 + 'static,
512 ) -> SwageBuilder<PH, H1, AE, VE> {
513 SwageBuilder {
514 allocator: self.allocator,
515 profile_hammerer_factory: self.profile_hammerer_factory,
516 profile_data_pattern: self.profile_data_pattern,
517 hammerer_factory: Box::new(hammerer_factory),
518 victim_factory: self.victim_factory,
519 pattern_size: self.pattern_size,
520 progress: self.progress,
521 config: self.config,
522 }
523 }
524
525 pub fn victim_factory(
526 mut self,
527 victim_factory: impl Fn(ConsecBlocks, RoundProfile) -> Result<Box<dyn VictimOrchestrator>, VE>
528 + 'static,
529 ) -> Self {
530 self.victim_factory = Some(Box::new(victim_factory));
531 self
532 }
533
534 pub fn pattern_size(mut self, pattern_size: usize) -> Self {
535 self.pattern_size = Some(pattern_size);
536 self
537 }
538
539 pub fn progress(mut self, progress: MultiProgress) -> Self {
540 self.progress = Some(progress);
541 self
542 }
543
544 pub fn config(mut self, config: SwageConfig) -> Self {
545 self.config = config;
546 self
547 }
548
549 pub fn build(self) -> Result<Swage<PH, H, AE, VE>, Error> {
550 if !(self.config.timeout.is_some()
551 || self.config.repetitions.is_some()
552 || self.config.hammering_timeout.is_some())
553 {
554 return Err(Error::InvalidConfig(
555 "At least one of timeout, repetitions or hammering_timeout must be set".into(),
556 ));
557 }
558 Ok(Swage {
559 allocator: self.allocator.ok_or(Error::Allocator)?,
560 profile_hammerer_factory: self
561 .profile_hammerer_factory
562 .ok_or(Error::ProfileHammerer)?,
563 profile_data_pattern: self.profile_data_pattern,
564 hammerer_factory: self.hammerer_factory,
565 victim_factory: self.victim_factory.ok_or(Error::Victim)?,
566 progress: self.progress,
567 pattern_size: self.pattern_size.ok_or(Error::PatternSize)?,
568 config: self.config,
569 })
570 }
571}
572
573#[derive(Debug, Error)]
574pub enum Error {
575 #[error("No allocator specified")]
576 Allocator,
577 #[error("No profiling hammerer specified")]
578 ProfileHammerer,
579 #[error("No victim specified")]
580 Victim,
581 #[error("Pattern size not specified")]
582 PatternSize,
583 #[error("Invalid config: {0}")]
584 InvalidConfig(String),
585}