1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
use std::{
    collections::HashMap,
    fmt::{Debug, Formatter},
};

use log::debug;

use crate::{
    agent::AgentDescriptor,
    error::Error,
    protocol::ProtocolBehavior,
    put::{Put, PutName},
    trace::TraceContext,
};

pub const DUMMY_PUT: PutName = PutName(['D', 'U', 'M', 'Y', 'Y', 'D', 'U', 'M', 'M', 'Y']);

/// Registry for [Factories](Factory). An instance of this is usually defined statically and then
/// used throughout the fuzzer.
pub struct PutRegistry<PB> {
    factories: HashMap<String, Box<dyn Factory<PB>>>,
    default_put: String,
}

impl<PB: ProtocolBehavior> PartialEq for PutRegistry<PB> {
    fn eq(&self, other: &Self) -> bool {
        self.default_put == other.default_put
            && self.factories.len() == other.factories.len()
            && self
                .factories
                .keys()
                .all(|id| other.factories.contains_key(id))
    }
}

impl<PB: ProtocolBehavior> Debug for PutRegistry<PB> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PutRegistry (default only)")
            .field("default", &self.default().name())
            .finish()
    }
}

impl<PB: ProtocolBehavior> PutRegistry<PB> {
    pub fn new<SI, I, S>(puts: I, default: S) -> Self
    where
        SI: Into<String>,
        I: IntoIterator<Item = (SI, Box<dyn Factory<PB>>)>,
        S: Into<String>,
    {
        let result = Self {
            factories: puts
                .into_iter()
                .map(|(id, f)| (Into::<String>::into(id), f))
                .collect(),
            default_put: default.into(),
        };

        // check that the default PUT is actually in the registry
        let _ = result.find_by_id(&result.default_put);

        result
    }

    pub fn default(&self) -> &dyn Factory<PB> {
        self.find_by_id(&self.default_put)
            .unwrap_or_else(|| panic!("default PUT {} is not in registry", &self.default_put))
    }

    pub fn puts(&self) -> impl Iterator<Item = (&str, &dyn Factory<PB>)> {
        self.factories
            .iter()
            .map(|(n, f)| (n.as_str(), f.to_owned().as_ref()))
    }

    pub fn find_by_id<S: AsRef<str>>(&self, id: S) -> Option<&dyn Factory<PB>> {
        self.factories
            .get(id.as_ref())
            .map(|f| f.to_owned().as_ref())
    }

    /// To be called at the beginning of all fuzzing campaigns!
    pub fn determinism_set_reseed_all_factories(&self) {
        debug!("== Set and reseed all ({}):", self.factories.len());
        for (_, factory) in self.factories.iter() {
            factory.determinism_set_reseed();
        }
    }

    pub fn determinism_reseed_all_factories(&self) {
        debug!("== Reseed all ({}):", self.factories.len());
        for (_, factory) in self.factories.iter() {
            factory.determinism_reseed();
        }
    }
}

impl<PB: ProtocolBehavior> Clone for PutRegistry<PB> {
    fn clone(&self) -> Self {
        Self::new(
            self.factories
                .iter()
                .map(|(n, f)| (n.clone(), f.clone_factory())),
            self.default_put.clone(),
        )
    }
}

#[derive(Debug)]
pub enum PutKind {
    CPUT,
    Rust,
}

/// Factory for instantiating programs-under-test.
pub trait Factory<PB: ProtocolBehavior> {
    fn create(
        &self,
        context: &TraceContext<PB>,
        agent_descriptor: &AgentDescriptor,
    ) -> Result<Box<dyn Put<PB>>, Error>;

    fn kind(&self) -> PutKind;
    fn name(&self) -> PutName;
    fn versions(&self) -> Vec<(String, String)>;

    fn determinism_set_reseed(&self);
    fn determinism_reseed(&self);

    fn clone_factory(&self) -> Box<dyn Factory<PB>>;
}