manycastr/orchestrator/
config.rs

1use crate::ALL_WORKERS;
2use crate::custom_module::manycastr::{Address, ProtocolType};
3use crate::custom_module::parse_src_address;
4use log::info;
5use std::collections::{HashMap, HashSet};
6use std::fmt;
7use std::fmt::Display;
8use std::fs;
9use std::path::Path;
10use std::sync::{Arc, Mutex};
11
12/// An origin allow-list rule: a source address and the protocols permitted for it.
13#[derive(Debug, Clone)]
14pub struct AllowedOrigin {
15    /// Anycast source address, or a `unicastv4`/`unicastv6` marker
16    pub src: Address,
17    /// Permitted protocols (`None` = any protocol)
18    pub protocols: Option<Vec<ProtocolType>>,
19}
20
21impl AllowedOrigin {
22    /// Whether this rule permits the given protocol.
23    pub fn allows(&self, p_type: ProtocolType) -> bool {
24        self.protocols
25            .as_ref()
26            .is_none_or(|protocols| protocols.contains(&p_type))
27    }
28}
29
30/// Print an allow-list rule, e.g. "10.0.0.1 (all)" or "unicastv4 (icmp|tcp)"
31impl Display for AllowedOrigin {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        match &self.protocols {
34            None => write!(f, "{} (all)", self.src),
35            Some(protocols) => {
36                let protocols: Vec<&str> = protocols.iter().map(|p| p.as_str()).collect();
37                write!(f, "{} ({})", self.src, protocols.join("|"))
38            }
39        }
40    }
41}
42
43/// Load the origin allow-list from a file (`--origins`).
44/// Each line defines one rule:
45/// src_addr, protocol[, protocol...]
46///
47/// `src_addr` is an anycast IP address or `unicastv4`/`unicastv6`;
48/// the protocol list may be `all` to allow all protocols.
49///
50/// # Arguments
51/// * `origins_path` - the path to the origins file
52///
53/// # Returns
54/// * The list of allowed origins
55///
56/// # Panics
57/// If the file does not exist, contains malformed entries, or defines no rules.
58pub fn load_allowed_origins(origins_path: &String) -> Vec<AllowedOrigin> {
59    if !Path::new(origins_path).exists() {
60        panic!("[Orchestrator] Origins file {origins_path} not found!");
61    }
62
63    let content =
64        fs::read_to_string(origins_path).expect("[Orchestrator] Could not read the origins file.");
65
66    let mut allowed_origins = Vec::new();
67
68    for (i, line) in content.lines().enumerate() {
69        let line_number = i + 1;
70
71        let trimmed_line = line.trim();
72
73        // Skip empty lines and comments
74        if trimmed_line.is_empty() || trimmed_line.starts_with('#') {
75            continue;
76        }
77
78        // Format: "src_addr, protocol[, protocol...]"
79        let mut parts = trimmed_line.split(',').map(str::trim);
80        let src = parse_src_address(parts.next().unwrap());
81
82        let protocol_tokens: Vec<&str> = parts.collect();
83        if protocol_tokens.is_empty() {
84            panic!(
85                "[Orchestrator] Error on line {line_number}: Malformed entry. Expected 'src_addr, protocol[, protocol...]' ('all' allows all protocols), found '{line}'"
86            );
87        }
88
89        // 'all' allows all protocols
90        let protocols = if protocol_tokens
91            .iter()
92            .any(|p| p.eq_ignore_ascii_case("all"))
93        {
94            None
95        } else {
96            Some(
97                protocol_tokens
98                    .iter()
99                    .map(|p| {
100                        ProtocolType::from_str(p).unwrap_or_else(|| {
101                            panic!(
102                                "[Orchestrator] Error on line {line_number}: Unknown protocol '{p}'. Expected icmp, dns, tcp, chaos, or all."
103                            )
104                        })
105                    })
106                    .collect(),
107            )
108        };
109
110        allowed_origins.push(AllowedOrigin { src, protocols });
111    }
112
113    if allowed_origins.is_empty() {
114        panic!("[Orchestrator] No origin rules found in {origins_path}");
115    }
116
117    info!(
118        "[Orchestrator] {} allowed origins loaded.",
119        allowed_origins.len()
120    );
121
122    allowed_origins
123}
124
125/// Load the worker configuration from a file.
126/// This provides a static mapping of hostnames to worker IDs.
127/// Formats the file as follows:
128/// hostname,id
129///
130/// # Arguments
131/// * `config_path` - the path to the configuration file
132///
133/// # Returns
134/// * The worker ID for any new hostname, which is the maximum ID + 1 in the configuration file
135/// * A mapping of hostnames to worker IDs
136///
137/// # Panics
138/// If the configuration file does not exist, or if there are malformed entries, duplicate hostnames, or duplicate IDs.
139pub fn load_worker_config(config_path: &String) -> (Arc<Mutex<u32>>, Option<HashMap<String, u32>>) {
140    if !Path::new(config_path).exists() {
141        panic!("[Orchestrator] Configuration file {config_path} not found!");
142    }
143
144    let config_content = fs::read_to_string(config_path)
145        .expect("[Orchestrator] Could not read the configuration file.");
146
147    let mut hosts = HashMap::new();
148    let mut used_ids = HashSet::new();
149
150    for (i, line) in config_content.lines().enumerate() {
151        let line_number = i + 1;
152
153        let trimmed_line = line.trim();
154
155        // Skip empty lines and comments
156        if trimmed_line.is_empty() || trimmed_line.starts_with('#') {
157            continue;
158        }
159
160        // Format: "hostname,id"
161        let parts: Vec<&str> = trimmed_line.split(',').collect();
162        if parts.len() != 2 {
163            panic!(
164                "[Orchestrator] Error on line {line_number}: Malformed entry. Expected 'hostname,id', found '{line}'"
165            );
166        }
167
168        let hostname = parts[0].trim().to_string();
169        let id = match parts[1].trim().parse::<u32>() {
170            Ok(val) => val,
171            Err(_) => {
172                panic!(
173                    "[Orchestrator] Error on line {line_number}: Invalid ID '{}'. ID must be an integer.",
174                    parts[1].trim()
175                );
176            }
177        };
178
179        // Check for duplicate hostname before inserting.
180        if hosts.contains_key(&hostname) {
181            panic!(
182                "[Orchestrator] Error on line {line_number}: Duplicate hostname '{hostname}' found. Hostnames must be unique."
183            );
184        }
185
186        // Insert the ID (if it is not already used)
187        if !used_ids.insert(id) {
188            panic!(
189                "[Orchestrator] Error on line {line_number}: Duplicate ID '{id}' found. IDs must be unique."
190            );
191        }
192
193        // Avoid special worker IDs
194        if id == ALL_WORKERS {
195            panic!(
196                "[Orchestrator] Error on line {line_number}: ID '{id}' is reserved for special purposes. Please use a different ID."
197            );
198        }
199
200        hosts.insert(hostname, id);
201    }
202
203    info!("[Orchestrator] {} hosts loaded.", hosts.len());
204
205    // Current worker ID is the maximum ID + 1 in the configuration file
206    let current_worker_id = hosts.values().max().map_or(1, |&max_id| max_id + 1);
207
208    (Arc::new(Mutex::new(current_worker_id)), Some(hosts))
209}
manycastr/orchestrator/config.rs

manycastr/orchestrator/
config.rs
