rpz

Response policy zone (RPZ) file generator.
git clone https://git.philomathiclife.com/repos/rpz
Log | Files | Refs | README

main.rs (20942B)


      1 //! # `rpz`
      2 //!
      3 //! Consult [`README.md`](https://crates.io/crates/rpz).
      4 #![feature(never_type)]
      5 #![cfg_attr(doc, feature(doc_auto_cfg))]
      6 #![deny(
      7     future_incompatible,
      8     let_underscore,
      9     missing_docs,
     10     nonstandard_style,
     11     rust_2018_compatibility,
     12     rust_2018_idioms,
     13     rust_2021_compatibility,
     14     rust_2024_compatibility,
     15     unsafe_code,
     16     unused,
     17     warnings,
     18     clippy::all,
     19     clippy::cargo,
     20     clippy::complexity,
     21     clippy::correctness,
     22     clippy::nursery,
     23     clippy::pedantic,
     24     clippy::perf,
     25     clippy::restriction,
     26     clippy::style,
     27     clippy::suspicious
     28 )]
     29 #![allow(
     30     unknown_lints,
     31     reason = "OpenBSD does not recognize clippy::return_and_then"
     32 )]
     33 #![allow(clippy::pub_use, reason = "when priv_sep is used, we export UnveilErr")]
     34 #![expect(
     35     clippy::arbitrary_source_item_ordering,
     36     clippy::blanket_clippy_restriction_lints,
     37     clippy::implicit_return,
     38     clippy::min_ident_chars,
     39     clippy::missing_trait_methods,
     40     clippy::multiple_crate_versions,
     41     clippy::question_mark_used,
     42     clippy::ref_patterns,
     43     clippy::return_and_then,
     44     clippy::single_call_fn,
     45     clippy::single_char_lifetime_names,
     46     clippy::unseparated_literal_suffix,
     47     reason = "never want to use these lints"
     48 )]
     49 /// Contains a wrapper of block and unblock `RpzDomain`s
     50 /// which can be used to write to a `File` or `stdout`.
     51 mod app;
     52 /// Module for reading and parsing passed arguments.
     53 mod args;
     54 /// Module for the TOML config file.
     55 mod config;
     56 /// Contains functions for `pledge(2)` and `unveil(2)` on OpenBSD platforms when compiled
     57 /// with the `priv_sep` feature; otherwise almost all functions are no-ops.
     58 mod priv_sep;
     59 use crate::{
     60     app::Domains,
     61     args::{ArgsErr, ConfigPath, Opts},
     62     config::Config,
     63 };
     64 use core::{
     65     error::Error,
     66     fmt::{self, Display, Formatter},
     67     time::Duration,
     68 };
     69 #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
     70 use priv_sep::UnveilErr;
     71 use reqwest::Client;
     72 use rpz::{
     73     dom::FirefoxDomainErr,
     74     file::{AbsFilePath, ExtFileErr, ExternalFiles, Files, HttpUrl, LocalFiles, Summary},
     75 };
     76 use std::{
     77     collections::HashSet,
     78     fs,
     79     io::{self, Read as _, Write as _},
     80     sync::OnceLock,
     81 };
     82 use tokio::runtime::Builder;
     83 use toml::de;
     84 /// The HTTP(S) client that is used to download all files.
     85 /// It is initialized exactly once in `main` before being used.
     86 static CLIENT: OnceLock<Client> = OnceLock::new();
     87 /// The output printed to `stdout` when `-h`/`--help` are passed
     88 /// to the program.
     89 const HELP: &str = "Response policy zone (RPZ) file generator
     90 
     91 Usage: rpz [OPTIONS]
     92 
     93 Options:
     94   -V, --version               Print version info and exit
     95   -h, --help                  Print help
     96   -q, --quiet                 Do not print messages
     97   -v, --verbose               Print summary parsing information about each file
     98   -f, --file <CONFIG_FILE>    Required option to pass '-' for stdin or the absolute path to the config file";
     99 /// The output printed to `stdout` when `-v`/`--version` are passed
    100 /// to the program.
    101 const VERSION: &str = concat!("rpz ", env!("CARGO_PKG_VERSION"));
    102 /// The User-Agent header value sent to HTTP(S) servers.
    103 const USER_AGENT: &str = concat!("rpz/", env!("CARGO_PKG_VERSION"));
    104 /// Error returned from the program.
    105 enum E {
    106     /// Variant for errors due to incorrect arguments being passed.
    107     Args(ArgsErr),
    108     /// Variant for errors due to issues with the TOML config file.
    109     Config(de::Error),
    110     /// Variant for errors due to calls to `unveil`.
    111     #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
    112     Unveil(UnveilErr),
    113     /// Variant for IO errors.
    114     Io(io::Error),
    115     /// Variant for errors due to downloading external HTTP(S) block files.
    116     ExtFile(ExtFileErr),
    117     /// Variant when there are no block entries to be written.
    118     NoBlockEntries,
    119 }
    120 impl fmt::Debug for E {
    121     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
    122         match *self {
    123             Self::Args(ref e) => write!(f, "{e}.\nFor more information, try '--help'."),
    124             Self::Config(_) | Self::Io(_) | Self::ExtFile(_) | Self::NoBlockEntries => {
    125                 <Self as Display>::fmt(self, f)
    126             }
    127             #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
    128             Self::Unveil(_) => <Self as Display>::fmt(self, f),
    129         }
    130     }
    131 }
    132 impl Display for E {
    133     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
    134         match *self {
    135             Self::Args(ref e) => e.fmt(f),
    136             Self::Config(ref e) => e.fmt(f),
    137             #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
    138             Self::Unveil(ref err) => write!(f, "unveil(2) failed with {err}"),
    139             Self::Io(ref e) => e.fmt(f),
    140             Self::ExtFile(ref e) => e.fmt(f),
    141             Self::NoBlockEntries => f.write_str("there are no domains to block"),
    142         }
    143     }
    144 }
    145 impl Error for E {}
    146 impl From<ArgsErr> for E {
    147     fn from(value: ArgsErr) -> Self {
    148         Self::Args(value)
    149     }
    150 }
    151 impl From<de::Error> for E {
    152     fn from(value: de::Error) -> Self {
    153         Self::Config(value)
    154     }
    155 }
    156 impl From<io::Error> for E {
    157     fn from(value: io::Error) -> Self {
    158         Self::Io(value)
    159     }
    160 }
    161 impl From<ExtFileErr> for E {
    162     fn from(value: ExtFileErr) -> Self {
    163         Self::ExtFile(value)
    164     }
    165 }
    166 #[cfg(all(feature = "priv_sep", target_os = "openbsd"))]
    167 impl From<UnveilErr> for E {
    168     fn from(value: UnveilErr) -> Self {
    169         Self::Unveil(value)
    170     }
    171 }
    172 #[cfg(not(all(feature = "priv_sep", target_os = "openbsd")))]
    173 impl From<!> for E {
    174     fn from(value: !) -> Self {
    175         value
    176     }
    177 }
    178 /// Reads `Config` from `conf`.
    179 fn get_config(conf: ConfigPath) -> Result<Config, E> {
    180     toml::from_str::<Config>(
    181         match conf {
    182             ConfigPath::Stdin => {
    183                 let mut file = String::new();
    184                 io::stdin().lock().read_to_string(&mut file)?;
    185                 file
    186             }
    187             ConfigPath::Path(path) => {
    188                 priv_sep::unveil_read_file(path.as_path())?;
    189                 let file = fs::read_to_string(path.as_path())?;
    190                 priv_sep::unveil_none(path)?;
    191                 file
    192             }
    193         }
    194         .as_str(),
    195     )
    196     .map_err(E::Config)
    197 }
    198 /// Gets `LocalFiles` from `local_dir`.
    199 fn get_local_files(local_dir: Option<AbsFilePath<true>>) -> Result<Option<LocalFiles>, E> {
    200     local_dir.map_or_else(
    201         || Ok(None),
    202         |dir| {
    203             priv_sep::unveil_read_dir(dir.as_path())
    204                 .map_err(E::from)
    205                 .and_then(|exists| {
    206                     if exists {
    207                         LocalFiles::from_path(dir.clone())
    208                             .map_err(E::Io)
    209                             .and_then(|files| {
    210                                 priv_sep::unveil_none(dir).map_err(E::from).map(|()| files)
    211                             })
    212                     } else {
    213                         Ok(None)
    214                     }
    215                 })
    216         },
    217     )
    218 }
    219 /// Downloads block files from HTTP(S) servers.
    220 #[expect(clippy::unreachable, reason = "there is a bug and we want to crash")]
    221 fn get_external_files(
    222     timeout: Duration,
    223     adblock: HashSet<HttpUrl>,
    224     domain: HashSet<HttpUrl>,
    225     hosts: HashSet<HttpUrl>,
    226     wildcard: HashSet<HttpUrl>,
    227 ) -> Result<Files, E> {
    228     Builder::new_current_thread()
    229         .enable_all()
    230         .build()
    231         .map_or_else(
    232             |e| Err(E::Io(e)),
    233             |runtime| {
    234                 runtime.block_on(async {
    235                     Files::from_external(
    236                         {
    237                             let mut files = ExternalFiles::new();
    238                             CLIENT
    239                                 .set(
    240                                     Client::builder()
    241                                         .user_agent(USER_AGENT)
    242                                         .use_rustls_tls()
    243                                         .build()
    244                                         .map_err(ExtFileErr::Http)?,
    245                                 )
    246                                 .unwrap_or_else(|_e| {
    247                                     unreachable!("there is a bug in OnceLock::set")
    248                                 });
    249                             let client = CLIENT
    250                                 .get()
    251                                 .unwrap_or_else(|| unreachable!("there is a bug in OnceLock::get"));
    252                             files.add_adblock(client, adblock);
    253                             files.add_domain(client, domain);
    254                             files.add_hosts(client, hosts);
    255                             files.add_wildcard(client, wildcard);
    256                             files
    257                         },
    258                         timeout,
    259                     )
    260                     .await
    261                     .map_err(E::ExtFile)
    262                 })
    263             },
    264         )
    265 }
    266 /// Verbosity of what is written to `stdout`.
    267 #[derive(Clone, Copy)]
    268 enum Verbosity {
    269     /// Suppress all summary info.
    270     None,
    271     /// Write normal amount of info.
    272     Normal,
    273     /// Write verbose information.
    274     High,
    275 }
    276 /// Writes to `stdout` the summary information in the event the quiet
    277 /// option was not passed.
    278 fn write_summary(
    279     summaries: Vec<Summary<'_, FirefoxDomainErr>>,
    280     verbose: bool,
    281     unblock_count: usize,
    282     block_count: usize,
    283 ) -> Result<(), io::Error> {
    284     let mut stdout = io::stdout().lock();
    285     let mut domain_count = 0usize;
    286     let mut comment_count = 0usize;
    287     let mut blank_count = 0usize;
    288     let mut error_count = 0usize;
    289     if verbose {
    290         summaries.into_iter().try_fold((), |(), summary| {
    291             domain_count = domain_count.saturating_add(summary.domain_count);
    292             comment_count = comment_count.saturating_add(summary.comment_count);
    293             blank_count = blank_count.saturating_add(summary.blank_count);
    294             error_count = error_count.saturating_add(summary.errors.values().sum());
    295             writeln!(&mut stdout, "{summary}")
    296         })?;
    297     } else {
    298         summaries.into_iter().fold((), |(), summary| {
    299             domain_count = domain_count.saturating_add(summary.domain_count);
    300             comment_count = comment_count.saturating_add(summary.comment_count);
    301             blank_count = blank_count.saturating_add(summary.blank_count);
    302             error_count = error_count.saturating_add(summary.errors.values().sum());
    303         });
    304     }
    305     writeln!(
    306         stdout,
    307         "unblock count written: {}\nblock count written: {}\ntotal lines written: {}\ndomains parsed: {}\ncomments parsed: {}\nblanks parsed: {}\nparsing errors: {}",
    308         unblock_count,
    309         block_count,
    310         unblock_count.saturating_add(block_count),
    311         domain_count,
    312         comment_count,
    313         blank_count,
    314         error_count,
    315     )
    316 }
    317 #[expect(
    318     clippy::arithmetic_side_effects,
    319     clippy::unreachable,
    320     reason = "math is correct and we want to crash if there is a bug"
    321 )]
    322 fn main() -> Result<(), E> {
    323     let mut promises = priv_sep::pledge_init()?;
    324     priv_sep::veil_all()?;
    325     let (conf, verbosity) = match Opts::from_args()? {
    326         Opts::Help => return writeln!(io::stdout().lock(), "{HELP}").map_err(E::Io),
    327         Opts::Version => return writeln!(io::stdout().lock(), "{VERSION}").map_err(E::Io),
    328         Opts::Config(path) => (path, Verbosity::Normal),
    329         Opts::ConfigQuiet(path) => (path, Verbosity::None),
    330         Opts::ConfigVerbose(path) => (path, Verbosity::High),
    331         Opts::Verbose | Opts::Quiet => return Err(E::Args(ArgsErr::ConfigPathNotPassed)),
    332         Opts::None => return Err(E::Args(ArgsErr::NoArgs)),
    333     };
    334     let config = get_config(conf)?;
    335     // We use a temp file to write to that way we can avoid overwriting
    336     // a file just for the process to error. Once the temp file is written to,
    337     // we rename it to the desired name.
    338     let tmp_rpz = config.rpz.as_ref().map_or_else(
    339         || {
    340             priv_sep::pledge_away_create_write(&mut promises)
    341                 .map_err(E::from)
    342                 .map(|()| None)
    343         },
    344         |file| {
    345             priv_sep::unveil_create(file.as_path())
    346                 .and_then(|()| {
    347                     let mut rpz = file.clone();
    348                     rpz.append("tmp");
    349                     priv_sep::unveil_create_read_write(rpz.as_path()).map(|()| Some(rpz))
    350                 })
    351                 .map_err(E::from)
    352         },
    353     )?;
    354     let local_files = get_local_files(config.local_dir)?;
    355     let (mut domains, mut summaries) = if let Some(files) = local_files.as_ref() {
    356         Domains::new_with(files)
    357     } else {
    358         (
    359             Domains::new(),
    360             Vec::with_capacity(
    361                 config.adblock.len()
    362                     + config.domain.len()
    363                     + config.hosts.len()
    364                     + config.wildcard.len(),
    365             ),
    366         )
    367     };
    368     let ext_files = if config.adblock.is_empty()
    369         && config.domain.is_empty()
    370         && config.hosts.is_empty()
    371         && config.wildcard.is_empty()
    372     {
    373         if domains.block().is_empty() {
    374             return Err(E::NoBlockEntries);
    375         }
    376         priv_sep::pledge_away_net(&mut promises)
    377             .and_then(|()| priv_sep::pledge_away_unveil(&mut promises).map(|()| Files::new()))
    378             .map_err(E::from)
    379     } else {
    380         priv_sep::unveil_https().map_err(E::from).and_then(|()| {
    381             priv_sep::pledge_away_unveil(&mut promises)
    382                 .map_err(E::from)
    383                 .and_then(|()| {
    384                     get_external_files(
    385                         config.timeout.unwrap_or(Duration::from_secs(3600)),
    386                         config.adblock,
    387                         config.domain,
    388                         config.hosts,
    389                         config.wildcard,
    390                     )
    391                     .and_then(|files| {
    392                         priv_sep::pledge_away_net(&mut promises)
    393                             .map_err(E::from)
    394                             .map(|()| files)
    395                     })
    396                 })
    397         })
    398     }?;
    399     domains.add_block_files(&ext_files, &mut summaries);
    400     if domains.block().is_empty() {
    401         return Err(E::NoBlockEntries);
    402     }
    403     let (unblock_count, block_count) = domains.write(config.rpz.map(|file| {
    404         (
    405             file,
    406             tmp_rpz.unwrap_or_else(|| unreachable!("there is a bug in main")),
    407         )
    408     }))?;
    409     if matches!(verbosity, Verbosity::None) {
    410         Ok(())
    411     } else {
    412         priv_sep::pledge_away_all_but_stdio(&mut promises)?;
    413         write_summary(
    414             summaries,
    415             matches!(verbosity, Verbosity::High),
    416             unblock_count,
    417             block_count,
    418         )
    419         .map_err(E::Io)
    420     }
    421 }
    422 #[cfg(test)]
    423 pub(crate) mod test_prog {
    424     use std::fs;
    425     use std::process::Command;
    426     /// The path to the program dir with no block subdirectories.
    427     pub const PROG_DIR_NO_SUB: &str = "/home/zack/projects/rpz/target/";
    428     /// The path to the program dir with block subdirectories.
    429     pub const PROG_DIR: &str = "/home/zack/projects/rpz/target/release/";
    430     /// The path to the program.
    431     const PROG: &str = "/home/zack/projects/rpz/target/release/rpz";
    432     /// Message to append to the appropriate variable above if an issue occurs.
    433     pub const ERR_MSG: &str = " does not exist, so program testing cannot occur";
    434     /// Verify the correct directories and files exist.
    435     pub fn verify_files() {
    436         if !fs::metadata(PROG)
    437             .expect(format!("{PROG}{ERR_MSG}").as_str())
    438             .is_file()
    439         {
    440             panic!("{PROG} is not an executable file")
    441         } else if !fs::metadata(PROG_DIR)
    442             .expect(format!("{PROG_DIR}{ERR_MSG}").as_str())
    443             .is_dir()
    444         {
    445             panic!("{PROG_DIR} is not a directory")
    446         } else if !fs::metadata(PROG_DIR_NO_SUB)
    447             .expect(format!("{PROG_DIR_NO_SUB}{ERR_MSG}").as_str())
    448             .is_dir()
    449         {
    450             panic!("{PROG_DIR_NO_SUB} is not a directory")
    451         } else {
    452             let mut files = String::from(PROG_DIR);
    453             let len = files.len();
    454             files.push_str("block/adblock/foo");
    455             if fs::metadata(files.as_str())
    456                 .expect(format!("{files}{ERR_MSG}. it must only contain '||bar.com'").as_str())
    457                 .is_file()
    458             {
    459                 files.truncate(len);
    460                 files.push_str("block/domain/foo");
    461                 if fs::metadata(files.as_str())
    462                     .expect(
    463                         format!("{files}{ERR_MSG}. it must only contain 'www.example.com'")
    464                             .as_str(),
    465                     )
    466                     .is_file()
    467                 {
    468                     files.truncate(len);
    469                     files.push_str("unblock/hosts/foo");
    470                     if fs::metadata(files.as_str())
    471                         .expect(
    472                             format!("{files}{ERR_MSG}. it must only contain '0.0.0.0 www.bar.com'")
    473                                 .as_str(),
    474                         )
    475                         .is_file()
    476                     {
    477                         files.truncate(len);
    478                         files.push_str("unblock/wildcard/foo");
    479                         if !fs::metadata(files.as_str())
    480                             .expect(
    481                                 format!("{files}{ERR_MSG}. it must only contain '*.foo.com'")
    482                                     .as_str(),
    483                             )
    484                             .is_file()
    485                         {
    486                             panic!("{files} is not a file");
    487                         }
    488                     } else {
    489                         panic!("{files} is not a file");
    490                     }
    491                 } else {
    492                     panic!("{files} is not a file");
    493                 }
    494             } else {
    495                 panic!("{files} is not a file");
    496             }
    497         }
    498     }
    499     pub fn get_command() -> Command {
    500         Command::new(PROG)
    501     }
    502 }
    503 #[cfg(test)]
    504 mod tests {
    505     use crate::{E, test_prog};
    506     use std::io::Write;
    507     use std::process::Stdio;
    508     use std::thread;
    509     #[test]
    510     #[ignore]
    511     fn test_app() {
    512         test_prog::verify_files();
    513         assert!(
    514             test_prog::get_command()
    515                 .args(["-f", "-"])
    516                 .stderr(Stdio::piped())
    517                 .stdin(Stdio::piped())
    518                 .stdout(Stdio::null())
    519                 .spawn()
    520                 .map_or(false, |mut cmd| {
    521                     cmd.stdin.take().map_or(false, |mut stdin| {
    522                         thread::spawn(move || {
    523                             stdin
    524                                 .write_all(
    525                                     format!("local_dir=\"{}\"", test_prog::PROG_DIR_NO_SUB)
    526                                         .as_bytes(),
    527                                 )
    528                                 .map_or(false, |_| true)
    529                         })
    530                         .join()
    531                         .map_or(false, |v| v)
    532                     }) && cmd.wait_with_output().map_or(false, |output| {
    533                         !output.status.success()
    534                             && output.stderr
    535                                 == format!("Error: {:?}\n", E::NoBlockEntries).into_bytes()
    536                     })
    537                 })
    538         );
    539         assert!(test_prog::get_command()
    540             .args(["-vf", "-"])
    541             .stderr(Stdio::null())
    542             .stdin(Stdio::piped())
    543             .stdout(Stdio::piped())
    544             .spawn()
    545             .map_or(false, |mut cmd| {
    546                 cmd.stdin.take().map_or(false, |mut stdin| {
    547                     thread::spawn(move || {
    548                         stdin
    549                             .write_all(format!("local_dir=\"{}\"", test_prog::PROG_DIR).as_bytes())
    550                             .map_or(false, |_| true)
    551                     })
    552                     .join()
    553                     .map_or(false, |v| v)
    554                 }) && cmd.wait_with_output().map_or(false, |output| {
    555                     output.status.success()
    556                         && output.stdout
    557                             == b"www.bar.com CNAME rpz-passthru.
    558 bar.com CNAME .
    559 *.bar.com CNAME .
    560 www.example.com CNAME .
    561 (Hosts) /home/zack/projects/rpz/target/release/unblock/hosts/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
    562 (Wildcard) /home/zack/projects/rpz/target/release/unblock/wildcard/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
    563 (Adblock) /home/zack/projects/rpz/target/release/block/adblock/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
    564 (Domain-only) /home/zack/projects/rpz/target/release/block/domain/foo - domains parsed: 1, comments parsed: 0, blanks parsed: 0, parsing errors: 0
    565 unblock count written: 1
    566 block count written: 3
    567 total lines written: 4
    568 domains parsed: 4
    569 comments parsed: 0
    570 blanks parsed: 0
    571 parsing errors: 0\n"
    572                 })
    573             }));
    574     }
    575 }