rpz

main.rs (20244B)

---

 //! # `rpz`
 //!
 //! Consult [`README.md`](https://crates.io/crates/rpz).
 #![expect(
     unstable_features,
     reason = "we already require nightly, so we elect to use never"
 )]
 #![expect(
     clippy::multiple_crate_versions,
     reason = "dependencies haven't updated to newest crates"
 )]
 #![feature(never_type)]
 #![cfg_attr(doc, feature(doc_auto_cfg))]
 /// Contains a wrapper of block and unblock `RpzDomain`s
 /// which can be used to write to a `File` or `stdout`.
 mod app;
 /// Module for reading and parsing passed arguments.
 mod args;
 /// Module for the TOML config file.
 mod config;
 /// Contains functions for `pledge(2)` and `unveil(2)` on OpenBSD platforms when compiled
 /// with the `priv_sep` feature; otherwise almost all functions are no-ops.
 mod priv_sep;
 use crate::{
     app::Domains,
     args::{ArgsErr, ConfigPath, Opts},
     config::Config,
 };
 #[cfg(all(feature = "priv_sep", not(target_os = "openbsd")))]
 use ::priv_sep as _;
 use ascii_domain as _;
 use core::{
     error::Error,
     fmt::{self, Display, Formatter},
     time::Duration,
 };
 use num_bigint as _;
 #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
 use priv_sep::NulOrIoErr;
 use reqwest::Client;
 use rpz::{
     dom::FirefoxDomainErr,
     file::{AbsFilePath, ExtFileErr, ExternalFiles, Files, HttpUrl, LocalFiles, Summary},
 };
 use std::{
     collections::HashSet,
     fs,
     io::{self, Read as _, Write as _},
     sync::OnceLock,
 };
 use tokio::runtime::Builder;
 use toml::de;
 use url as _;
 use zfc as _;
 /// The HTTP(S) client that is used to download all files.
 /// It is initialized exactly once in `main` before being used.
 static CLIENT: OnceLock<Client> = OnceLock::new();
 /// The output printed to `stdout` when `-h`/`--help` are passed
 /// to the program.
 const HELP: &str = "Response policy zone (RPZ) file generator
 
 Usage: rpz [OPTIONS]
 
 Options:
   -V, --version               Print version info and exit
   -h, --help                  Print help
   -q, --quiet                 Do not print messages
   -v, --verbose               Print summary parsing information about each file
   -f, --file <CONFIG_FILE>    Required option to pass '-' for stdin or the absolute path to the config file";
 /// The output printed to `stdout` when `-v`/`--version` are passed
 /// to the program.
 const VERSION: &str = concat!("rpz ", env!("CARGO_PKG_VERSION"));
 /// The User-Agent header value sent to HTTP(S) servers.
 const USER_AGENT: &str = concat!("rpz/", env!("CARGO_PKG_VERSION"));
 /// Error returned from the program.
 enum E {
     /// Variant for errors due to incorrect arguments being passed.
     Args(ArgsErr),
     /// Variant for errors due to issues with the TOML config file.
     Config(de::Error),
     /// Variant for errors due to calls to `unveil`.
     #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
     Unveil(NulOrIoErr),
     /// Variant for IO errors.
     Io(io::Error),
     /// Variant for errors due to downloading external HTTP(S) block files.
     ExtFile(ExtFileErr),
     /// Variant when there are no block entries to be written.
     NoBlockEntries,
 }
 impl fmt::Debug for E {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match *self {
             Self::Args(ref e) => write!(f, "{e}.\nFor more information, try '--help'."),
             Self::Config(_) | Self::Io(_) | Self::ExtFile(_) | Self::NoBlockEntries => {
                 <Self as Display>::fmt(self, f)
             }
             #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
             Self::Unveil(_) => <Self as Display>::fmt(self, f),
         }
     }
 }
 impl Display for E {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match *self {
             Self::Args(ref e) => e.fmt(f),
             Self::Config(ref e) => e.fmt(f),
             #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
             Self::Unveil(ref err) => write!(f, "unveil(2) failed with {err}"),
             Self::Io(ref e) => e.fmt(f),
             Self::ExtFile(ref e) => e.fmt(f),
             Self::NoBlockEntries => f.write_str("there are no domains to block"),
         }
     }
 }
 impl Error for E {}
 impl From<ArgsErr> for E {
     fn from(value: ArgsErr) -> Self {
         Self::Args(value)
     }
 }
 impl From<de::Error> for E {
     fn from(value: de::Error) -> Self {
         Self::Config(value)
     }
 }
 impl From<io::Error> for E {
     fn from(value: io::Error) -> Self {
         Self::Io(value)
     }
 }
 impl From<ExtFileErr> for E {
     fn from(value: ExtFileErr) -> Self {
         Self::ExtFile(value)
     }
 }
 #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
 impl From<NulOrIoErr> for E {
     fn from(value: NulOrIoErr) -> Self {
         Self::Unveil(value)
     }
 }
 #[cfg(not(all(feature = "priv_sep", target_os = "openbsd")))]
 impl From<!> for E {
     fn from(value: !) -> Self {
         value
     }
 }
 /// Reads `Config` from `conf`.
 fn get_config(conf: ConfigPath) -> Result<Config, E> {
     toml::from_str::<Config>(
         match conf {
             ConfigPath::Stdin => {
                 let mut file = String::new();
                 _ = io::stdin().lock().read_to_string(&mut file)?;
                 file
             }
             ConfigPath::Path(path) => {
                 priv_sep::unveil_read_file(path.as_path())?;
                 let file = fs::read_to_string(path.as_path())?;
                 priv_sep::unveil_none(path)?;
                 file
             }
         }
         .as_str(),
     )
     .map_err(E::Config)
 }
 /// Gets `LocalFiles` from `local_dir`.
 fn get_local_files(local_dir: Option<AbsFilePath<true>>) -> Result<Option<LocalFiles>, E> {
     local_dir.map_or_else(
         || Ok(None),
         |dir| {
             priv_sep::unveil_read_dir(dir.as_path())
                 .map_err(E::from)
                 .and_then(|exists| {
                     if exists {
                         LocalFiles::from_path(dir.clone())
                             .map_err(E::Io)
                             .and_then(|files| {
                                 priv_sep::unveil_none(dir).map_err(E::from).map(|()| files)
                             })
                     } else {
                         Ok(None)
                     }
                 })
         },
     )
 }
 /// Downloads block files from HTTP(S) servers.
 #[expect(clippy::unreachable, reason = "there is a bug and we want to crash")]
 fn get_external_files(
     timeout: Duration,
     adblock: HashSet<HttpUrl>,
     domain: HashSet<HttpUrl>,
     hosts: HashSet<HttpUrl>,
     wildcard: HashSet<HttpUrl>,
 ) -> Result<Files, E> {
     Builder::new_current_thread()
         .enable_all()
         .build()
         .map_or_else(
             |e| Err(E::Io(e)),
             |runtime| {
                 runtime.block_on(async {
                     Files::from_external(
                         {
                             let mut files = ExternalFiles::new();
                             CLIENT
                                 .set(
                                     Client::builder()
                                         .user_agent(USER_AGENT)
                                         .use_rustls_tls()
                                         .build()
                                         .map_err(ExtFileErr::Http)?,
                                 )
                                 .unwrap_or_else(|_e| {
                                     unreachable!("there is a bug in OnceLock::set")
                                 });
                             let client = CLIENT
                                 .get()
                                 .unwrap_or_else(|| unreachable!("there is a bug in OnceLock::get"));
                             files.add_adblock(client, adblock);
                             files.add_domain(client, domain);
                             files.add_hosts(client, hosts);
                             files.add_wildcard(client, wildcard);
                             files
                         },
                         timeout,
                     )
                     .await
                     .map_err(E::ExtFile)
                 })
             },
         )
 }
 /// Verbosity of what is written to `stdout`.
 #[derive(Clone, Copy)]
 enum Verbosity {
     /// Suppress all summary info.
     None,
     /// Write normal amount of info.
     Normal,
     /// Write verbose information.
     High,
 }
 /// Writes to `stdout` the summary information in the event the quiet
 /// option was not passed.
 fn write_summary(
     summaries: Vec<Summary<'_, FirefoxDomainErr>>,
     verbose: bool,
     unblock_count: usize,
     block_count: usize,
 ) -> Result<(), io::Error> {
     let mut stdout = io::stdout().lock();
     let mut domain_count = 0usize;
     let mut comment_count = 0usize;
     let mut blank_count = 0usize;
     let mut error_count = 0usize;
     if verbose {
         summaries.into_iter().try_fold((), |(), summary| {
             domain_count = domain_count.saturating_add(summary.domain_count);
             comment_count = comment_count.saturating_add(summary.comment_count);
             blank_count = blank_count.saturating_add(summary.blank_count);
             error_count = error_count.saturating_add(summary.errors.values().sum());
             writeln!(&mut stdout, "{summary}")
         })?;
     } else {
         summaries.into_iter().fold((), |(), summary| {
             domain_count = domain_count.saturating_add(summary.domain_count);
             comment_count = comment_count.saturating_add(summary.comment_count);
             blank_count = blank_count.saturating_add(summary.blank_count);
             error_count = error_count.saturating_add(summary.errors.values().sum());
         });
     }
     writeln!(
         stdout,
         "unblock count written: {}\nblock count written: {}\ntotal lines written: {}\ndomains parsed: {}\ncomments parsed: {}\nblanks parsed: {}\nparsing errors: {}",
         unblock_count,
         block_count,
         unblock_count.saturating_add(block_count),
         domain_count,
         comment_count,
         blank_count,
         error_count,
     )
 }
 #[expect(
     clippy::arithmetic_side_effects,
     clippy::unreachable,
     reason = "math is correct and we want to crash if there is a bug"
 )]
 fn main() -> Result<(), E> {
     let mut promises = priv_sep::pledge_init()?;
     priv_sep::veil_all()?;
     let (conf, verbosity) = match Opts::from_args()? {
         Opts::Help => return writeln!(io::stdout().lock(), "{HELP}").map_err(E::Io),
         Opts::Version => return writeln!(io::stdout().lock(), "{VERSION}").map_err(E::Io),
         Opts::Config(path) => (path, Verbosity::Normal),
         Opts::ConfigQuiet(path) => (path, Verbosity::None),
         Opts::ConfigVerbose(path) => (path, Verbosity::High),
         Opts::Verbose | Opts::Quiet => return Err(E::Args(ArgsErr::ConfigPathNotPassed)),
         Opts::None => return Err(E::Args(ArgsErr::NoArgs)),
     };
     let config = get_config(conf)?;
     // We use a temp file to write to that way we can avoid overwriting
     // a file just for the process to error. Once the temp file is written to,
     // we rename it to the desired name.
     let tmp_rpz = config.rpz.as_ref().map_or_else(
         || {
             priv_sep::pledge_away_create_write(&mut promises)
                 .map_err(E::from)
                 .map(|()| None)
         },
         |file| {
             priv_sep::unveil_create(file.as_path())
                 .and_then(|()| {
                     let mut rpz = file.clone();
                     _ = rpz.append("tmp");
                     priv_sep::unveil_create_read_write(rpz.as_path()).map(|()| Some(rpz))
                 })
                 .map_err(E::from)
         },
     )?;
     let local_files = get_local_files(config.local_dir)?;
     let (mut domains, mut summaries) = if let Some(files) = local_files.as_ref() {
         Domains::new_with(files)
     } else {
         (
             Domains::new(),
             Vec::with_capacity(
                 config.adblock.len()
                     + config.domain.len()
                     + config.hosts.len()
                     + config.wildcard.len(),
             ),
         )
     };
     let ext_files = if config.adblock.is_empty()
         && config.domain.is_empty()
         && config.hosts.is_empty()
         && config.wildcard.is_empty()
     {
         if domains.block().is_empty() {
             return Err(E::NoBlockEntries);
         }
         priv_sep::pledge_away_net(&mut promises)
             .and_then(|()| priv_sep::pledge_away_unveil(&mut promises).map(|()| Files::new()))
             .map_err(E::from)
     } else {
         priv_sep::unveil_https().map_err(E::from).and_then(|()| {
             priv_sep::pledge_away_unveil(&mut promises)
                 .map_err(E::from)
                 .and_then(|()| {
                     get_external_files(
                         config.timeout.unwrap_or(Duration::from_secs(3600)),
                         config.adblock,
                         config.domain,
                         config.hosts,
                         config.wildcard,
                     )
                     .and_then(|files| {
                         priv_sep::pledge_away_net(&mut promises)
                             .map_err(E::from)
                             .map(|()| files)
                     })
                 })
         })
     }?;
     domains.add_block_files(&ext_files, &mut summaries);
     if domains.block().is_empty() {
         return Err(E::NoBlockEntries);
     }
     let (unblock_count, block_count) = domains.write(config.rpz.map(|file| {
         (
             file,
             tmp_rpz.unwrap_or_else(|| unreachable!("there is a bug in main")),
         )
     }))?;
     if matches!(verbosity, Verbosity::None) {
         Ok(())
     } else {
         priv_sep::pledge_away_all_but_stdio(&mut promises)?;
         write_summary(
             summaries,
             matches!(verbosity, Verbosity::High),
             unblock_count,
             block_count,
         )
         .map_err(E::Io)
     }
 }
 #[cfg(test)]
 pub(crate) mod test_prog {
     use std::fs;
     use std::process::Command;
     /// The path to the program dir with no block subdirectories.
     pub(crate) const PROG_DIR_NO_SUB: &str = "/home/zack/projects/rpz/target/";
     /// The path to the program dir with block subdirectories.
     pub(crate) const PROG_DIR: &str = "/home/zack/projects/rpz/target/release/";
     /// The path to the program.
     const PROG: &str = "/home/zack/projects/rpz/target/release/rpz";
     /// Message to append to the appropriate variable above if an issue occurs.
     pub(crate) const ERR_MSG: &str = " does not exist, so program testing cannot occur";
     /// Verify the correct directories and files exist.
     pub(crate) fn verify_files() {
         if !fs::metadata(PROG)
             .expect(format!("{PROG}{ERR_MSG}").as_str())
             .is_file()
         {
             panic!("{PROG} is not an executable file")
         } else if !fs::metadata(PROG_DIR)
             .expect(format!("{PROG_DIR}{ERR_MSG}").as_str())
             .is_dir()
         {
             panic!("{PROG_DIR} is not a directory")
         } else if !fs::metadata(PROG_DIR_NO_SUB)
             .expect(format!("{PROG_DIR_NO_SUB}{ERR_MSG}").as_str())
             .is_dir()
         {
             panic!("{PROG_DIR_NO_SUB} is not a directory")
         } else {
             let mut files = String::from(PROG_DIR);
             let len = files.len();
             files.push_str("block/adblock/foo");
             if fs::metadata(files.as_str())
                 .expect(format!("{files}{ERR_MSG}. it must only contain '||bar.com'").as_str())
                 .is_file()
             {
                 files.truncate(len);
                 files.push_str("block/domain/foo");
                 if fs::metadata(files.as_str())
                     .expect(
                         format!("{files}{ERR_MSG}. it must only contain 'www.example.com'")
                             .as_str(),
                     )
                     .is_file()
                 {
                     files.truncate(len);
                     files.push_str("unblock/hosts/foo");
                     if fs::metadata(files.as_str())
                         .expect(
                             format!("{files}{ERR_MSG}. it must only contain '0.0.0.0 www.bar.com'")
                                 .as_str(),
                         )
                         .is_file()
                     {
                         files.truncate(len);
                         files.push_str("unblock/wildcard/foo");
                         if !fs::metadata(files.as_str())
                             .expect(
                                 format!("{files}{ERR_MSG}. it must only contain '*.foo.com'")
                                     .as_str(),
                             )
                             .is_file()
                         {
                             panic!("{files} is not a file");
                         }
                     } else {
                         panic!("{files} is not a file");
                     }
                 } else {
                     panic!("{files} is not a file");
                 }
             } else {
                 panic!("{files} is not a file");
             }
         }
     }
     pub(crate) fn get_command() -> Command {
         Command::new(PROG)
     }
 }
 #[cfg(test)]
 mod tests {
     use crate::{E, test_prog};
     use std::io::Write;
     use std::process::Stdio;
     use std::thread;
     #[test]
     #[ignore]
     fn test_app() {
         test_prog::verify_files();
         assert!(
             test_prog::get_command()
                 .args(["-f", "-"])
                 .stderr(Stdio::piped())
                 .stdin(Stdio::piped())
                 .stdout(Stdio::null())
                 .spawn()
                 .map_or(false, |mut cmd| {
                     cmd.stdin.take().map_or(false, |mut stdin| {
                         thread::spawn(move || {
                             stdin
                                 .write_all(
                                     format!("local_dir=\"{}\"", test_prog::PROG_DIR_NO_SUB)
                                         .as_bytes(),
                                 )
                                 .map_or(false, |_| true)
                         })
                         .join()
                         .map_or(false, |v| v)
                     }) && cmd.wait_with_output().map_or(false, |output| {
                         !output.status.success()
                             && output.stderr
                                 == format!("Error: {:?}\n", E::NoBlockEntries).into_bytes()
                     })
                 })
         );
         assert!(test_prog::get_command()
             .args(["-vf", "-"])
             .stderr(Stdio::null())
             .stdin(Stdio::piped())
             .stdout(Stdio::piped())
             .spawn()
             .map_or(false, |mut cmd| {
                 cmd.stdin.take().map_or(false, |mut stdin| {
                     thread::spawn(move || {
                         stdin
                             .write_all(format!("local_dir=\"{}\"", test_prog::PROG_DIR).as_bytes())
                             .map_or(false, |_| true)
                     })
                     .join()
                     .map_or(false, |v| v)
                 }) && cmd.wait_with_output().map_or(false, |output| {
                     output.status.success()
                         && output.stdout
                             == b"www.bar.com CNAME rpz-passthru.
 bar.com CNAME .
 *.bar.com CNAME .
 www.example.com CNAME .
 (Hosts) /home/zack/projects/rpz/target/release/unblock/hosts/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
 (Wildcard) /home/zack/projects/rpz/target/release/unblock/wildcard/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
 (Adblock) /home/zack/projects/rpz/target/release/block/adblock/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
 (Domain-only) /home/zack/projects/rpz/target/release/block/domain/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
 unblock count written: 1
 block count written: 3
 total lines written: 4
 domains parsed: 4
 comments parsed: 0
 blanks parsed: 0
 parsing errors: 0\n"
                 })
             }));
     }
 }