swage_core/
swage.rs

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
19/// Main orchestrator for conducting end-to-end Rowhammer experiments.
20///
21/// The `Swage` struct combines an allocator, hammerer, and victim to execute
22/// complete Rowhammer attack experiments with profiling, reproducibility checks,
23/// and result collection. It manages the full lifecycle:
24///
25/// 1. Memory allocation using [`ConsecAllocator`]
26/// 2. Profiling to identify vulnerable memory locations
27/// 3. Reproducibility verification to filter transient bit flips
28/// 4. Attack execution with hammering and victim checking
29///
30/// # Type Parameters
31///
32/// * `PH` - Profile hammerer type (implements [`Hammering`])
33/// * `H` - Attack hammerer type (implements [`Hammering`])
34/// * `AE` - Allocator error type
35/// * `VE` - Victim error type
36///
37/// # Examples
38///
39/// Use [`Swage::builder()`] to construct a `Swage` instance with the required components.
40pub 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/// Profiling results from a series of hammering rounds.
52///
53/// Contains the bit flips that were consistently reproduced during profiling
54/// and the data pattern used to induce them.
55#[derive(Debug, Serialize, Clone)]
56pub struct RoundProfile {
57    /// Bit flips that met the reproducibility threshold
58    pub bit_flips: Vec<BitFlip>,
59    /// Data pattern used during profiling
60    pub pattern: DataPattern,
61}
62
63/// Configuration parameters for Swage experiments.
64///
65/// Controls profiling behavior, reproducibility requirements, and execution timeouts.
66pub struct SwageConfig {
67    /// Number of profiling rounds to identify vulnerable bit flips
68    pub profiling_rounds: u64,
69    /// Minimum fraction of rounds a bit flip must appear during profiling to be considered reproducible (0.0-1.0)
70    pub reproducibility_threshold: f64,
71
72    /// Timeout for total hammering operation (None = unlimited)
73    pub hammering_timeout: Option<Duration>,
74    /// Number of times to repeat the attack (None = unlimited)
75    pub repetitions: Option<u64>,
76    /// Overall experiment timeout (None = no timeout)
77    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/// Results from a complete Rowhammer experiment.
93///
94/// Contains all attack results, profiling data, timestamp, and optional metadata.
95///
96/// # Type Parameters
97///
98/// * `T` - Success result type
99/// * `E` - Error type
100#[derive(Serialize)]
101pub struct ExperimentData<T, E> {
102    /// ISO 8601 timestamp of when the experiment ran
103    date: String,
104    /// Results from each attack repetition
105    results: Vec<std::result::Result<T, E>>,
106    /// Profiling data from the experiment
107    profiling: RoundProfile,
108    /// Additional JSON metadata (implementation-specific)
109    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    /// Creates a new Swage builder.
129    ///
130    /// # Returns
131    ///
132    /// A builder for configuring and constructing a Swage instance
133    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        //unsafe { shm_unlink(CString::new("HAMMER_SHM").unwrap().as_ptr()) };
173        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); // TODO maybe remove this?
270            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    /// Start the attack.
305    ///
306    /// Returns a vector of ExperimentData with VictimResults and possible Error observed.
307    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
366/// Hammer a given `memory` region `num_rounds` times to profile for vulnerable addresses.
367fn 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/// Data pattern selection for configuration.
437///
438/// Used to specify which type of data pattern to use in the aggressors.
439#[derive(Clone, Copy)]
440pub enum DataPatternKind {
441    /// Random data pattern
442    Random,
443    /// All zeros (0x00)
444    Zero,
445    /// All ones (0xFF)
446    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}