fud/
util.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2025 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use smol::{
20    fs::{self, File},
21    stream::StreamExt,
22};
23use std::{
24    collections::HashSet,
25    path::{Path, PathBuf},
26    sync::Arc,
27    time::Instant,
28};
29
30pub use darkfi::geode::hash_to_string;
31use darkfi::{
32    net::{Message, MessageSubscription},
33    Error, Result,
34};
35
36use crate::proto::ResourceMessage;
37
38pub async fn get_all_files(dir: &Path) -> Result<Vec<(PathBuf, u64)>> {
39    let mut files = Vec::new();
40
41    let mut entries = fs::read_dir(dir).await.unwrap();
42
43    while let Some(entry) = entries.try_next().await.unwrap() {
44        let path = entry.path();
45
46        if path.is_dir() {
47            files.append(&mut Box::pin(get_all_files(&path)).await?);
48        } else {
49            let metadata = fs::metadata(&path).await?;
50            let file_size = metadata.len();
51            files.push((path, file_size));
52        }
53    }
54
55    Ok(files)
56}
57
58pub async fn create_all_files(files: &[PathBuf]) -> Result<()> {
59    for file_path in files.iter() {
60        if !file_path.exists() {
61            if let Some(dir) = file_path.parent() {
62                fs::create_dir_all(dir).await?;
63            }
64            File::create(&file_path).await?;
65        }
66    }
67
68    Ok(())
69}
70
71/// An enum to represent a set of files, where you can use `All` if you want
72/// all files without having to specify all of them.
73/// We could use an `Option<HashSet<PathBuf>>`, but this is more explicit.
74#[derive(Clone, Debug)]
75pub enum FileSelection {
76    All,
77    Set(HashSet<PathBuf>),
78}
79
80impl FromIterator<PathBuf> for FileSelection {
81    fn from_iter<I: IntoIterator<Item = PathBuf>>(iter: I) -> Self {
82        let paths: HashSet<PathBuf> = iter.into_iter().collect();
83        FileSelection::Set(paths)
84    }
85}
86
87/// Wait for a [`crate::proto::ResourceMessage`] on `msg_subscriber` with a timeout.
88/// If we receive a message with the wrong resource hash, it's skipped.
89pub async fn receive_resource_msg<M: Message + ResourceMessage + std::fmt::Debug>(
90    msg_subscriber: &MessageSubscription<M>,
91    resource_hash: blake3::Hash,
92    timeout_seconds: u64,
93) -> Result<Arc<M>> {
94    let start = Instant::now();
95    loop {
96        let elapsed = start.elapsed().as_secs();
97        if elapsed >= timeout_seconds {
98            return Err(Error::ConnectTimeout);
99        }
100        let remaining_timeout = timeout_seconds - elapsed;
101
102        let reply = msg_subscriber.receive_with_timeout(remaining_timeout).await?;
103        // Done if it's the right resource hash
104        if reply.resource_hash() == resource_hash {
105            return Ok(reply)
106        }
107    }
108}