swarm_nl/
util.rs

1// Copyright 2024 Algorealm, Inc.
2// Apache 2.0 License
3
4//! Utility helper functions.
5
6use crate::{
7	core::{replication::ReplBufferData, ByteVector, Core, StringVector},
8	prelude::*,
9	setup::BootstrapConfig,
10};
11use base58::FromBase58;
12use ini::Ini;
13use libp2p_identity::PeerId;
14use rand::{distributions::Alphanumeric, Rng};
15use std::{
16	collections::{HashMap, HashSet},
17	path::Path,
18	str::FromStr,
19	time::{SystemTime, UNIX_EPOCH},
20};
21
22/// Read an INI file containing bootstrap config information.
23pub fn read_ini_file(file_path: &str) -> SwarmNlResult<BootstrapConfig> {
24	// Read the file from disk
25	if let Ok(config) = Ini::load_from_file(file_path) {
26		// Get TCP port & UDP port
27		let (tcp_port, udp_port) = if let Some(section) = config.section(Some("ports")) {
28			(
29				section
30					.get("tcp")
31					.unwrap_or_default()
32					.parse::<Port>()
33					.unwrap_or_default(),
34				section
35					.get("udp")
36					.unwrap_or_default()
37					.parse::<Port>()
38					.unwrap_or_default(),
39			)
40		} else {
41			// Fallback to default ports
42			(MIN_PORT, MAX_PORT)
43		};
44
45		// Try to read the serialized keypair
46		// auth section
47		let (key_type, mut serialized_keypair) = if let Some(section) = config.section(Some("auth"))
48		{
49			(
50				// Get the preferred key type
51				section.get("crypto").unwrap_or_default(),
52				// Get serialized keypair
53				string_to_vec::<u8>(section.get("protobuf_keypair").unwrap_or_default()),
54			)
55		} else {
56			Default::default()
57		};
58
59		// Now, move onto reading the bootnodes if any
60		let section = config
61			.section(Some("bootstrap"))
62			.ok_or(SwarmNlError::BoostrapFileReadError(file_path.to_owned()))?;
63
64		// Get the provided bootnodes
65		let boot_nodes = string_to_hashmap(section.get("boot_nodes").unwrap_or_default());
66
67		// Now, move onto reading the blacklist if any
68		let blacklist = if let Some(section) = config.section(Some("blacklist")) {
69			string_to_vec(section.get("blacklist").unwrap_or_default())
70		} else {
71			Default::default()
72		};
73
74		Ok(BootstrapConfig::new()
75			.generate_keypair_from_protobuf(key_type, &mut serialized_keypair)
76			.with_bootnodes(boot_nodes)
77			.with_blacklist(blacklist)
78			.with_tcp(tcp_port)
79			.with_udp(udp_port))
80	} else {
81		// Return error
82		Err(SwarmNlError::BoostrapFileReadError(file_path.to_owned()))
83	}
84}
85
86/// Write value into config file.
87pub fn write_config<T: AsRef<Path> + ?Sized>(
88	section: &str,
89	key: &str,
90	new_value: &str,
91	file_path: &T,
92) -> bool {
93	if let Ok(mut conf) = Ini::load_from_file(file_path) {
94		// Set a value:
95		conf.set_to(Some(section), key.into(), new_value.into());
96		if let Ok(_) = conf.write_to_file(file_path) {
97			return true;
98		}
99	}
100	false
101}
102
103/// Parse string into a vector.
104fn string_to_vec<T: FromStr>(input: &str) -> Vec<T> {
105	input
106		.trim_matches(|c| c == '[' || c == ']')
107		.split(',')
108		.filter_map(|s| s.trim().parse::<T>().ok())
109		.fold(Vec::new(), |mut acc, item| {
110			acc.push(item);
111			acc
112		})
113}
114
115/// Parse string into a hashmap.
116fn string_to_hashmap(input: &str) -> HashMap<String, String> {
117	input
118		.trim_matches(|c| c == '[' || c == ']')
119		.split(',')
120		.filter(|s| s.contains(':'))
121		.fold(HashMap::new(), |mut acc, s| {
122			let mut parts = s.trim().splitn(2, ':');
123			if let (Some(key), Some(value)) = (parts.next(), parts.next()) {
124				if key.len() > 1 {
125					acc.insert(key.trim().to_owned(), value.trim().to_owned());
126				}
127			}
128			acc
129		})
130}
131
132/// Convert a peer ID string to [`PeerId`].
133pub fn string_to_peer_id(peer_id_string: &str) -> Option<PeerId> {
134	PeerId::from_bytes(&peer_id_string.from_base58().unwrap_or_default()).ok()
135}
136
137/// Generate a random string of variable length.
138pub fn generate_random_string(length: usize) -> String {
139	let mut rng = rand::thread_rng();
140	(0..length)
141		.map(|_| rng.sample(Alphanumeric) as char)
142		.collect()
143}
144
145/// Unmarshall data recieved as RPC during the execution of the eventual consistency algorithm to
146/// fill in missing messages in the node's buffer.
147pub fn unmarshal_messages(data: Vec<Vec<u8>>) -> Vec<ReplBufferData> {
148	let mut result = Vec::new();
149
150	for entry in data {
151		let serialized = String::from_utf8_lossy(&entry).to_string();
152		let entries: Vec<&str> = serialized.split(Core::ENTRY_DELIMITER).collect();
153
154		for entry in entries {
155			let fields: Vec<&str> = entry.split(Core::FIELD_DELIMITER).collect();
156			if fields.len() < 6 {
157				continue; // Skip malformed entries
158			}
159
160			let data_field: Vec<String> = fields[0]
161				.split(Core::DATA_DELIMITER)
162				.map(|s| s.to_string())
163				.collect();
164			let lamport_clock = fields[1].parse().unwrap_or(0);
165			let outgoing_timestamp = fields[2].parse().unwrap_or(0);
166			let incoming_timestamp = fields[3].parse().unwrap_or(0);
167			let message_id = fields[4].to_string();
168			let sender = fields[5];
169
170			// Parse peerId
171			if let Ok(peer_id) = sender.parse::<PeerId>() {
172				result.push(ReplBufferData {
173					data: data_field,
174					lamport_clock,
175					outgoing_timestamp,
176					incoming_timestamp,
177					message_id,
178					sender: peer_id,
179					confirmations: None, // Since eventual consistency
180				});
181			}
182		}
183	}
184
185	result
186}
187
188/// Get unix timestamp as string.
189pub fn get_unix_timestamp() -> Seconds {
190	// Get the current system time
191	let now = SystemTime::now();
192	// Calculate the duration since the Unix epoch
193	let duration_since_epoch = now.duration_since(UNIX_EPOCH).expect("Time went backwards");
194	// Return the Unix timestamp in seconds as a string
195	duration_since_epoch.as_secs()
196}
197
198/// Convert a [ByteVector] to a [StringVector].
199pub fn byte_vec_to_string_vec(input: ByteVector) -> StringVector {
200	input
201		.into_iter()
202		.map(|vec| String::from_utf8(vec).unwrap_or_else(|_| String::from("Invalid UTF-8")))
203		.collect()
204}
205
206/// Convert a [StringVector] to a [ByteVector].
207pub fn string_vec_to_byte_vec(input: StringVector) -> ByteVector {
208	input.into_iter().map(|s| s.into_bytes()).collect()
209}
210
211/// Marshall the shard network image into a [ByteVector].
212pub fn shard_image_to_bytes(
213	shard_image: HashMap<String, HashSet<PeerId>>,
214) -> Result<Vec<u8>, serde_json::Error> {
215	// Convert the PeerIds into serializable form
216	let serializable_image: HashMap<String, HashSet<PeerIdString>> = shard_image
217		.into_iter()
218		.map(|(shard_id, peers)| {
219			let string_ids = peers
220				.iter()
221				.map(|id| id.to_string())
222				.collect::<HashSet<_>>();
223			(shard_id, string_ids)
224		})
225		.collect();
226
227	// Serialize using JSON
228	serde_json::to_vec(&serializable_image)
229}
230
231/// Merge two shard network states.
232pub fn merge_shard_states(
233	local_state: &mut HashMap<String, HashSet<PeerId>>,
234	incoming_state: HashMap<String, HashSet<PeerId>>,
235) {
236	for (shard_id, incoming_peers) in incoming_state.iter() {
237		local_state
238			.entry(shard_id.to_owned())
239			.and_modify(|local_peers| {
240				// Use HashSet's `extend` method to add all unique peers from incoming_peers
241				local_peers.extend(incoming_peers);
242			})
243			.or_insert(incoming_peers.clone()); // Insert directly if the shard_id doesn't exist
244	}
245}
246
247/// Deserialize bytes into shard image structure.
248pub fn bytes_to_shard_image(bytes: Vec<u8>) -> HashMap<String, HashSet<PeerId>> {
249	// Deserialize into the serializable form
250	if let Ok(serializable_image) =
251		serde_json::from_slice::<HashMap<String, HashSet<PeerIdString>>>(&bytes)
252	{
253		// Convert back to PeerId form
254		let shard_image: HashMap<String, HashSet<PeerId>> = serializable_image
255			.into_iter()
256			.map(|(shard_id, peers)| {
257				let peer_ids: HashSet<PeerId> = peers
258					.into_iter()
259					.filter_map(|peer| peer.parse::<PeerId>().ok())
260					.collect();
261				(shard_id, peer_ids)
262			})
263			.collect();
264
265		return shard_image;
266	}
267
268	Default::default()
269}
270
271#[cfg(test)]
272mod tests {
273
274	use libp2p_identity::{KeyType, Keypair};
275
276	use super::*;
277	use std::fs;
278
279	// Define custom ports for testing
280	const CUSTOM_TCP_PORT: Port = 49666;
281	const CUSTOM_UDP_PORT: Port = 49852;
282
283	// Helper to create an INI file without a keypair and a valid range for ports.
284	fn create_test_ini_file_without_keypair(file_path: &str) {
285		let mut config = Ini::new();
286		config
287			.with_section(Some("ports"))
288			.set("tcp", CUSTOM_TCP_PORT.to_string())
289			.set("udp", CUSTOM_UDP_PORT.to_string());
290
291		config.with_section(Some("bootstrap")).set(
292			"boot_nodes",
293			"[12D3KooWGfbL6ZNGWqS11MoptH2A7DB1DG6u85FhXBUPXPVkVVRq:/ip4/192.168.1.205/tcp/1509]",
294		);
295		config
296			.with_section(Some("blacklist"))
297			.set("blacklist", "[]");
298		// Write config to a new INI file
299		config.write_to_file(file_path).unwrap_or_default();
300	}
301
302	// Helper to create an INI file with keypair
303	fn create_test_ini_file_with_keypair(file_path: &str, key_type: KeyType) {
304		let mut config = Ini::new();
305
306		match key_type {
307			KeyType::Ed25519 => {
308				let keypair_ed25519 = Keypair::generate_ed25519().to_protobuf_encoding().unwrap();
309				config
310					.with_section(Some("auth"))
311					.set("crypto", "ed25519")
312					.set("protobuf_keypair", &format!("{:?}", keypair_ed25519));
313			},
314			KeyType::Secp256k1 => {
315				let keypair_secp256k1 = Keypair::generate_secp256k1()
316					.to_protobuf_encoding()
317					.unwrap();
318				config
319					.with_section(Some("auth"))
320					.set("crypto", "secp256k1")
321					.set("protobuf_keypair", &format!("{:?}", keypair_secp256k1));
322			},
323			KeyType::Ecdsa => {
324				let keypair_ecdsa = Keypair::generate_ecdsa().to_protobuf_encoding().unwrap();
325				config
326					.with_section(Some("auth"))
327					.set("crypto", "ecdsa")
328					.set("protobuf_keypair", &format!("{:?}", keypair_ecdsa));
329			},
330			_ => {},
331		}
332
333		config.with_section(Some("bootstrap")).set(
334			"boot_nodes",
335			"[12D3KooWGfbL6ZNGWqS11MoptH2A7DB1DG6u85FhXBUPXPVkVVRq:/ip4/192.168.1.205/tcp/1509]",
336		);
337
338		config
339			.with_section(Some("blacklist"))
340			.set("blacklist", "[]");
341
342		// Write config to the new INI file
343		config.write_to_file(file_path).unwrap_or_default();
344	}
345
346	// Helper to clean up temp file
347	fn clean_up_temp_file(file_path: &str) {
348		fs::remove_file(file_path).unwrap_or_default();
349	}
350
351	#[test]
352	fn file_does_not_exist() {
353		// Try to read a non-existent file should panic
354		assert_eq!(read_ini_file("non_existent_file.ini").is_err(), true);
355	}
356
357	#[test]
358	fn write_config_works() {
359		// Create temp INI file
360		let file_path = "temp_test_write_ini_file.ini";
361
362		// Create INI file without keypair for simplicity
363		create_test_ini_file_without_keypair(file_path);
364
365		// Try to write some keypair to the INI file
366		let add_keypair = write_config(
367			"auth",
368			"serialized_keypair",
369			&format!(
370				"{:?}",
371				vec![
372					8, 1, 18, 64, 116, 193, 199, 84, 83, 25, 220, 116, 119, 194, 155, 173, 2, 241,
373					82, 0, 130, 225, 121, 9, 232, 244, 8, 253, 170, 13, 100, 24, 195, 179, 60, 133,
374					128, 221, 43, 214, 180, 33, 61, 73, 124, 161, 127, 119, 40, 146, 226, 50, 65,
375					35, 97, 188, 159, 169, 250, 241, 98, 36, 146, 9, 139, 98, 114, 224
376				]
377			),
378			file_path,
379		);
380
381		assert_eq!(add_keypair, true);
382
383		// Delete temp file
384		clean_up_temp_file(file_path);
385	}
386
387	// Read without keypair file
388	#[test]
389	fn read_ini_file_with_custom_setup_works() {
390		// Create temp INI file
391		let file_path = "temp_test_read_ini_file_custom.ini";
392
393		// We've set our ports to tcp=49666 and upd=49852
394		create_test_ini_file_without_keypair(file_path);
395
396		let ini_file_result: BootstrapConfig = read_ini_file(file_path).unwrap();
397
398		assert_eq!(ini_file_result.ports().0, CUSTOM_TCP_PORT);
399		assert_eq!(ini_file_result.ports().1, CUSTOM_UDP_PORT);
400
401		// Checking for the default keypair that's generated (ED25519) if none are provided
402		assert_eq!(ini_file_result.keypair().key_type(), KeyType::Ed25519);
403
404		// Delete temp file
405		clean_up_temp_file(file_path);
406	}
407
408	#[test]
409	fn read_ini_file_with_default_setup_works() {
410		// Create INI file
411		let file_path = "temp_test_ini_file_default.ini";
412		create_test_ini_file_with_keypair(file_path, KeyType::Ecdsa);
413
414		// Assert that the content has no [port] section
415		// let ini_file_content = fs::read_to_string(file_path).unwrap();
416		// assert!(!ini_file_content.contains("[port]"));
417
418		// But when we call read_ini_file it generates a BootstrapConfig with random ports
419		let ini_file_result = read_ini_file(file_path).unwrap();
420
421		// assert_eq!(ini_file_result.ports().0, MIN_PORT);
422		// assert_eq!(ini_file_result.ports().1, MAX_PORT);
423
424		// Checking that the default keypair matches the configured keytype
425		assert_eq!(ini_file_result.keypair().key_type(), KeyType::Ecdsa);
426
427		// Delete temp file
428		clean_up_temp_file(file_path);
429	}
430
431	#[test]
432	fn string_to_vec_works() {
433		// Define test input
434		let test_input_1 = "[1, 2, 3]";
435
436		// Call the function
437		let result: Vec<i32> = string_to_vec(test_input_1);
438
439		// Assert that the result is as expected
440		assert_eq!(result, vec![1, 2, 3]);
441	}
442
443	#[test]
444	fn string_to_hashmap_works() {
445		// Define test input
446		let input =
447			"[12D3KooWGfbL6ZNGWqS11MoptH2A7DB1DG6u85FhXBUPXPVkVVRq:/ip4/192.168.1.205/tcp/1509]";
448
449		// Call the function
450		let result = string_to_hashmap(input);
451
452		// Assert that the result is as expected
453		let mut expected = HashMap::new();
454		expected.insert(
455			"12D3KooWGfbL6ZNGWqS11MoptH2A7DB1DG6u85FhXBUPXPVkVVRq".to_string(),
456			"/ip4/192.168.1.205/tcp/1509".to_string(),
457		);
458
459		assert_eq!(result, expected);
460	}
461
462	#[test]
463	fn bootstrap_config_blacklist_works() {
464		let file_path = "bootstrap_config_blacklist_test.ini";
465
466		// Create a new INI file with a blacklist
467		let mut config = Ini::new();
468		config
469			.with_section(Some("ports"))
470			.set("tcp", CUSTOM_TCP_PORT.to_string())
471			.set("udp", CUSTOM_UDP_PORT.to_string());
472
473		config.with_section(Some("bootstrap")).set(
474			"boot_nodes",
475			"[12D3KooWGfbL6ZNGWqS11MoptH2A7DB1DG6u85FhXBUPXPVkVVRq:/ip4/192.168.1.205/tcp/1509]",
476		);
477
478		let blacklist_peer_id: PeerId = PeerId::random();
479		let black_list_peer_id_string = format!("[{}]", blacklist_peer_id.to_base58());
480
481		config
482			.with_section(Some("blacklist"))
483			.set("blacklist", black_list_peer_id_string);
484
485		// Write config to a new INI file
486		config.write_to_file(file_path).unwrap_or_default();
487
488		// Read the new file
489		let ini_file_result: BootstrapConfig = read_ini_file(file_path).unwrap();
490
491		assert_eq!(ini_file_result.blacklist().list.len(), 1);
492		assert!(ini_file_result
493			.blacklist()
494			.list
495			.contains(&blacklist_peer_id));
496
497		fs::remove_file(file_path).unwrap_or_default();
498	}
499}
swarm_nl/util.rs

swarm_nl/
util.rs
