/ src / util / cli.rs
cli.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  
 19  use std::{
 20      fs,
 21      io::Write,
 22      path::Path,
 23      str,
 24      sync::{Arc, Mutex},
 25      time::Instant,
 26  };
 27  
 28  use crate::Result;
 29  
 30  /*
 31  #[derive(Clone, Default)]
 32  pub struct Config<T> {
 33      config: PhantomData<T>,
 34  }
 35  
 36  impl<T: Serialize + DeserializeOwned> Config<T> {
 37      pub fn load(path: PathBuf) -> Result<T> {
 38          if Path::new(&path).exists() {
 39              let toml = fs::read(&path)?;
 40              let str_buff = str::from_utf8(&toml)?;
 41              let config: T = toml::from_str(str_buff)?;
 42              Ok(config)
 43          } else {
 44              let path = path.to_str();
 45              if path.is_some() {
 46                  println!("Could not find/parse configuration file in: {}", path.unwrap());
 47              } else {
 48                  println!("Could not find/parse configuration file");
 49              }
 50              println!("Please follow the instructions in the README");
 51              Err(Error::ConfigNotFound)
 52          }
 53      }
 54  }
 55  */
 56  
 57  pub fn spawn_config(path: &Path, contents: &[u8]) -> Result<()> {
 58      if !path.exists() {
 59          if let Some(parent) = path.parent() {
 60              fs::create_dir_all(parent)?;
 61          }
 62  
 63          let mut file = fs::File::create(path)?;
 64          file.write_all(contents)?;
 65          println!("Config file created in {path:?}. Please review it and try again.");
 66          std::process::exit(2);
 67      }
 68  
 69      Ok(())
 70  }
 71  
 72  /// This macro is used for a standard way of daemonizing darkfi binaries
 73  /// with TOML config file configuration, and argument parsing.
 74  ///
 75  /// It also spawns a multithreaded async executor and passes it into the
 76  /// given function.
 77  ///
 78  /// The Cargo.toml dependencies needed for this are:
 79  /// ```text
 80  /// darkfi = { path = "../../", features = ["util"] }
 81  /// easy-parallel = "3.2.0"
 82  /// signal-hook-async-std = "0.2.2"
 83  /// signal-hook = "0.3.15"
 84  /// tracing-subscriber = "0.3.19"
 85  /// tracing-appender = "0.2.3"
 86  /// smol = "1.2.5"
 87  ///
 88  /// # Argument parsing
 89  /// serde = {version = "1.0.135", features = ["derive"]}
 90  /// structopt = "0.3.26"
 91  /// structopt-toml = "0.5.1"
 92  /// ```
 93  ///
 94  /// Example usage:
 95  /// ```
 96  /// use darkfi::{async_daemonize, cli_desc, Result};
 97  /// use smol::stream::StreamExt;
 98  /// use structopt_toml::{serde::Deserialize, structopt::StructOpt, StructOptToml};
 99  ///
100  /// const CONFIG_FILE: &str = "daemond_config.toml";
101  /// const CONFIG_FILE_CONTENTS: &str = include_str!("../daemond_config.toml");
102  ///
103  /// #[derive(Clone, Debug, Deserialize, StructOpt, StructOptToml)]
104  /// #[serde(default)]
105  /// #[structopt(name = "daemond", about = cli_desc!())]
106  /// struct Args {
107  ///     #[structopt(short, long)]
108  ///     /// Configuration file to use
109  ///     config: Option<String>,
110  ///
111  ///     #[structopt(short, long)]
112  ///     /// Set log file to ouput into
113  ///     log: Option<String>,
114  ///
115  ///     #[structopt(short, parse(from_occurrences))]
116  ///     /// Increase verbosity (-vvv supported)
117  ///     verbose: u8,
118  /// }
119  ///
120  /// async_daemonize!(realmain);
121  /// async fn realmain(args: Args, ex: Arc<smol::Executor<'static>>) -> Result<()> {
122  ///     println!("Hello, world!");
123  ///     Ok(())
124  /// }
125  /// ```
126  #[cfg(feature = "async-daemonize")]
127  #[macro_export]
128  macro_rules! async_daemonize {
129      ($realmain:ident) => {
130          fn main() -> Result<()> {
131              let args = match Args::from_args_with_toml("") {
132                  Ok(v) => v,
133                  Err(e) => {
134                      eprintln!("Unable to get args: {e}");
135                      return Err(Error::ConfigInvalid)
136                  }
137              };
138              let cfg_path =
139                  match darkfi::util::path::get_config_path(args.config.clone(), CONFIG_FILE) {
140                      Ok(v) => v,
141                      Err(e) => {
142                          eprintln!("Unable to get config path `{:?}`: {e}", args.config);
143                          return Err(e)
144                      }
145                  };
146              if let Err(e) =
147                  darkfi::util::cli::spawn_config(&cfg_path, CONFIG_FILE_CONTENTS.as_bytes())
148              {
149                  eprintln!("Spawn config failed `{cfg_path:?}`: {e}");
150                  return Err(e)
151              }
152              let cfg_text = match std::fs::read_to_string(&cfg_path) {
153                  Ok(c) => c,
154                  Err(e) => {
155                      eprintln!("Read config failed `{cfg_path:?}`: {e}");
156                      return Err(e.into())
157                  }
158              };
159              let args = match Args::from_args_with_toml(&cfg_text) {
160                  Ok(v) => v,
161                  Err(e) => {
162                      eprintln!("Parsing config failed `{cfg_path:?}`: {e}");
163                      return Err(Error::ConfigInvalid)
164                  }
165              };
166  
167              // If a log file has been configured, create a terminal and file logger.
168              // Otherwise, output to terminal logger only.
169              match args.log {
170                  Some(ref log_path) => {
171                      let log_path = match darkfi::util::path::expand_path(log_path) {
172                          Ok(v) => v,
173                          Err(e) => {
174                              eprintln!("Expanding log path failed `{log_path:?}`: {e}");
175                              return Err(e)
176                          }
177                      };
178                      let log_file = match std::fs::File::create(&log_path) {
179                          Ok(v) => v,
180                          Err(e) => {
181                              eprintln!("Creating log file failed `{log_path:?}`: {e}");
182                              return Err(e.into())
183                          }
184                      };
185  
186                      // hold guard until process stops to ensure buffer logs are flushed to file
187                      let (non_blocking, _guard) = tracing_appender::non_blocking(log_file);
188                      if let Err(e) =
189                          darkfi::util::logger::setup_logging(args.verbose, Some(non_blocking))
190                      {
191                          eprintln!("Unable to init logger with term + logfile combo: {e}");
192                          return Err(e.into())
193                      }
194                  }
195                  None => {
196                      if let Err(e) = darkfi::util::logger::setup_logging(args.verbose, None) {
197                          eprintln!("Unable to init term logger: {e}");
198                          return Err(e.into())
199                      }
200                  }
201              }
202  
203              // https://docs.rs/smol/latest/smol/struct.Executor.html#examples
204              let n_threads = std::thread::available_parallelism().unwrap().get();
205              let ex = std::sync::Arc::new(smol::Executor::new());
206              let (signal, shutdown) = smol::channel::unbounded::<()>();
207              let (_, result) = easy_parallel::Parallel::new()
208                  // Run four executor threads
209                  .each(0..n_threads, |_| smol::future::block_on(ex.run(shutdown.recv())))
210                  // Run the main future on the current thread.
211                  .finish(|| {
212                      smol::future::block_on(async {
213                          $realmain(args, ex.clone()).await?;
214                          drop(signal);
215                          Ok::<(), darkfi::Error>(())
216                      })
217                  });
218  
219              result
220          }
221  
222          /// Auxiliary structure used to keep track of signals
223          struct SignalHandler {
224              /// Termination signal channel receiver
225              term_rx: smol::channel::Receiver<()>,
226              /// Signals handle
227              handle: signal_hook_async_std::Handle,
228              /// SIGHUP publisher to retrieve new configuration,
229              sighup_pub: darkfi::system::PublisherPtr<Args>,
230          }
231  
232          impl SignalHandler {
233              fn new(
234                  ex: std::sync::Arc<smol::Executor<'static>>,
235              ) -> Result<(Self, smol::Task<Result<()>>)> {
236                  let (term_tx, term_rx) = smol::channel::bounded::<()>(1);
237                  let signals = signal_hook_async_std::Signals::new([
238                      signal_hook::consts::SIGHUP,
239                      signal_hook::consts::SIGTERM,
240                      signal_hook::consts::SIGINT,
241                      signal_hook::consts::SIGQUIT,
242                  ])?;
243                  let handle = signals.handle();
244                  let sighup_pub = darkfi::system::Publisher::new();
245                  let signals_task =
246                      ex.spawn(handle_signals(signals, term_tx, sighup_pub.clone(), ex.clone()));
247  
248                  Ok((Self { term_rx, handle, sighup_pub }, signals_task))
249              }
250  
251              /// Handler waits for termination signal
252              async fn wait_termination(&self, signals_task: smol::Task<Result<()>>) -> Result<()> {
253                  self.term_rx.recv().await?;
254                  print!("\r");
255                  self.handle.close();
256                  signals_task.await?;
257  
258                  Ok(())
259              }
260          }
261  
262          /// Auxiliary task to handle SIGINT for forceful process abort
263          async fn handle_abort(mut signals: signal_hook_async_std::Signals) {
264              let mut n_sigint = 0;
265              while let Some(signal) = signals.next().await {
266                  n_sigint += 1;
267                  if n_sigint == 2 {
268                      print!("\r");
269                      info!("Aborting. Good luck.");
270                      std::process::abort();
271                  }
272              }
273          }
274  
275          /// Auxiliary task to handle SIGHUP, SIGTERM, SIGINT and SIGQUIT signals
276          async fn handle_signals(
277              mut signals: signal_hook_async_std::Signals,
278              term_tx: smol::channel::Sender<()>,
279              publisher: darkfi::system::PublisherPtr<Args>,
280              ex: std::sync::Arc<smol::Executor<'static>>,
281          ) -> Result<()> {
282              while let Some(signal) = signals.next().await {
283                  match signal {
284                      signal_hook::consts::SIGHUP => {
285                          let args = Args::from_args_with_toml("").unwrap();
286                          let cfg_path =
287                              darkfi::util::path::get_config_path(args.config, CONFIG_FILE)?;
288                          darkfi::util::cli::spawn_config(
289                              &cfg_path,
290                              CONFIG_FILE_CONTENTS.as_bytes(),
291                          )?;
292                          let args = Args::from_args_with_toml(&std::fs::read_to_string(cfg_path)?);
293                          if args.is_err() {
294                              println!("handle_signals():: Error parsing the config file");
295                              continue
296                          }
297                          publisher.notify(args.unwrap()).await;
298                      }
299                      signal_hook::consts::SIGINT => {
300                          // Spawn a new background task to listen for more SIGINT.
301                          // This lets us forcefully abort the process if necessary.
302                          let signals =
303                              signal_hook_async_std::Signals::new([signal_hook::consts::SIGINT])?;
304                          let handle = signals.handle();
305                          ex.spawn(handle_abort(signals)).detach();
306  
307                          term_tx.send(()).await?;
308                      }
309                      signal_hook::consts::SIGTERM | signal_hook::consts::SIGQUIT => {
310                          term_tx.send(()).await?;
311                      }
312  
313                      _ => println!("handle_signals():: Unsupported signal"),
314                  }
315              }
316              Ok(())
317          }
318      };
319  }
320  
321  pub fn fg_red(message: &str) -> String {
322      format!("\x1b[31m{message}\x1b[0m")
323  }
324  
325  pub fn fg_green(message: &str) -> String {
326      format!("\x1b[32m{message}\x1b[0m")
327  }
328  
329  pub fn fg_reset() -> String {
330      "\x1b[0m".to_string()
331  }
332  
333  pub struct ProgressInc {
334      position: Arc<Mutex<u64>>,
335      timer: Arc<Mutex<Option<Instant>>>,
336  }
337  
338  impl Default for ProgressInc {
339      fn default() -> Self {
340          Self::new()
341      }
342  }
343  
344  impl ProgressInc {
345      pub fn new() -> Self {
346          eprint!("\x1b[?25l");
347          Self { position: Arc::new(Mutex::new(0)), timer: Arc::new(Mutex::new(None)) }
348      }
349  
350      pub fn inc(&self, n: u64) {
351          let mut position = self.position.lock().unwrap();
352  
353          if *position == 0 {
354              *self.timer.lock().unwrap() = Some(Instant::now());
355          }
356  
357          *position += n;
358  
359          let binding = self.timer.lock().unwrap();
360          let Some(elapsed) = binding.as_ref() else { return };
361          let elapsed = elapsed.elapsed();
362          let pos = *position;
363  
364          eprint!("\r[{elapsed:?}] {pos} attempts");
365      }
366  
367      pub fn position(&self) -> u64 {
368          *self.position.lock().unwrap()
369      }
370  
371      pub fn finish_and_clear(&self) {
372          *self.timer.lock().unwrap() = None;
373          eprint!("\r\x1b[2K\x1b[?25h");
374      }
375  }