malcontent-dns-parental-controls/tests/common.rs

// SPDX-FileCopyrightText: 2022 Matteo Settenvini <matteo.settenvini@montecristosoftware.eu>
// SPDX-License-Identifier: GPL-3.0-or-later

#![allow(non_upper_case_globals)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
#![allow(dead_code)]

include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
include!(concat!(env!("CARGO_MANIFEST_DIR"), "/src/constants.rs"));

use {
    anyhow::{anyhow, bail, ensure, Result},
    futures_util::TryStreamExt,
    libc::{freeaddrinfo, gai_strerror, getaddrinfo},
    nix::sys::socket::{SockaddrLike as _, SockaddrStorage},
    nix::unistd::Uid,
    std::collections::HashMap,
    std::env,
    std::ffi::{CStr, CString},
    std::net::{IpAddr, Ipv4Addr, Ipv6Addr},
    std::path::PathBuf,
    std::process::Command,
    std::str::FromStr,
    std::sync::Once,
    tokio::task,
    tokio::task::JoinHandle,
};

// Adapted from rusty_forkfork (which inherits it from rusty_fork)
// to allow a custom pre-fork function
#[macro_export]
macro_rules! fork_test {
    (#![rusty_fork(timeout_ms = $timeout:expr)]
     $(
         $(#[$meta:meta])*
         fn $test_name:ident() $(-> $ret:ty)? $body:block
    )*) => { $(
        $(#[$meta])*
        fn $test_name() {
            // Eagerly convert everything to function pointers so that all
            // tests use the same instantiation of `fork`.
            fn body_fn() $(-> $ret)? $body
            let body: fn () $(-> $ret)? = body_fn;

            fn supervise_fn(child: &mut rusty_forkfork::ChildWrapper,
                            _file: &mut ::std::fs::File) {
                rusty_forkfork::fork_test::supervise_child(child, $timeout)
            }
            let supervise:
                fn (&mut rusty_forkfork::ChildWrapper, &mut ::std::fs::File) =
                supervise_fn;

            rusty_forkfork::fork(
                rusty_forkfork::rusty_fork_test_name!($test_name),
                rusty_forkfork::rusty_fork_id!(),
                $crate::common::prefork_setup,
                supervise, body).expect("forking test failed");
        }
    )* };

    ($(
         $(#[$meta:meta])*
         fn $test_name:ident() $(-> $ret:ty)? $body:block
    )*) => {
        fork_test! {
            #![rusty_fork(timeout_ms = 0)]

            $($(#[$meta])* fn $test_name() $(-> $ret)? $body)*
        }
    };
}

pub fn prefork_setup(_child: &mut Command) {
    let out_dir = PathBuf::from(env!("OUT_DIR"));

    PREFORK.call_once(|| {
        std::env::set_var("LD_LIBRARY_PATH", &out_dir);

        let library_path = test_cdylib::build_current_project();
        let mut library_filename = library_path.file_name().unwrap().to_owned();
        library_filename.push(".2"); // required for NSS modules

        let dest = out_dir.join(library_filename);
        std::fs::copy(library_path, &dest).unwrap();
    });
}

pub type Eai = EaiRetcode;

static PREFORK: Once = Once::new();
static SETUP: Once = Once::new();

pub fn setup() -> Result<()> {
    SETUP.call_once(|| {
        let nss_config_status = unsafe {
            let db = CString::new("hosts").unwrap();
            let resolvers = CString::new("malcontent dns").unwrap();
            __nss_configure_lookup(db.as_ptr(), resolvers.as_ptr())
        };

        if nss_config_status != 0 {
            panic!(
                "Unable to configure NSS to load module: __nss_configure_lookup() returned {}",
                nss_config_status
            );
        }
    });

    Ok(())
}

pub fn resolve_with_system(family: libc::c_int, host: &str) -> Result<IpAddr> {
    let process = Command::new("getent")
        .arg(match family {
            libc::AF_INET => "ahostsv4",
            libc::AF_INET6 => "ahostsv6",
            _ => panic!("INET family must be either IPv4 or IPv6"),
        })
        .arg(host)
        .output()?;
    ensure!(
        process.status.success(),
        "Failed to run getent to check host IP"
    );
    let output = String::from_utf8(process.stdout)?;
    let addr_string = output
        .as_str()
        .split(' ')
        .next()
        .ok_or(anyhow!("Unparseable output from getent"))?;
    Ok(IpAddr::from_str(&addr_string)?)
}

pub fn resolve_with_module(family: libc::c_int, hostname: &str) -> Result<IpAddr> {
    assert!(
        SETUP.is_completed(),
        "Forgot to call common::setup() at beginning of test?"
    );
    let c_hostname = CString::new(hostname).unwrap();
    unsafe {
        let mut addr = std::ptr::null_mut();
        let hints = libc::addrinfo {
            ai_family: family,
            ai_flags: 0,
            ai_socktype: 0,
            ai_protocol: 0,
            ai_addrlen: 0,
            ai_addr: std::ptr::null_mut(),
            ai_canonname: std::ptr::null_mut(),
            ai_next: std::ptr::null_mut(),
        };
        let getaddrinfo_status =
            getaddrinfo(c_hostname.as_ptr(), std::ptr::null(), &hints, &mut addr);

        let error = CStr::from_ptr(gai_strerror(getaddrinfo_status));
        assert_eq!(
            getaddrinfo_status,
            Eai::Success.0,
            "Should have gotten hostname for {}, instead got {}",
            hostname,
            error.to_str().unwrap()
        );
        let addr_storage = SockaddrStorage::from_raw((*addr).ai_addr, Some((*addr).ai_addrlen))
            .ok_or(anyhow!("Garbled addrinfo from getaddrinfo()"))?;
        let ip = convert_addrinfo(&addr_storage)?;
        freeaddrinfo(addr);
        Ok(ip)
    }
}

fn convert_addrinfo(sa: &SockaddrStorage) -> Result<IpAddr> {
    if let Some(addr) = sa.as_sockaddr_in() {
        Ok(IpAddr::V4(Ipv4Addr::from(addr.ip())))
    } else if let Some(addr) = sa.as_sockaddr_in6() {
        Ok(IpAddr::V6(Ipv6Addr::from(addr.ip())))
    } else {
        bail!("addrinfo is not either an IPv4 or IPv6 address")
    }
}

pub type Restrictions = Vec<IpAddr>;

pub fn mock_dbus(responses: HashMap<Uid, Vec<Restrictions>>) -> JoinHandle<Result<()>> {
    async fn dbus_loop(mut responses: HashMap<Uid, Vec<Restrictions>>) -> Result<()> {
        use zbus::MessageType::*;

        let mut responses_size: usize = responses.values().map(|v| v.len()).sum();
        for r in responses.values_mut() {
            r.reverse(); // we pop responses from the back, so...
        }

        let connection = zbus::ConnectionBuilder::session()?.build().await?;
        let mut stream = zbus::MessageStream::from(&connection);
        while responses_size > 0 {
            let msg = stream
                .try_next()
                .await?
                .ok_or(anyhow!("Unparseable DBus message"))?;

            if msg.header()?.message_type()? == MethodCall
                && msg
                    .interface()
                    .ok_or(anyhow!("Invoked method has no interface"))?
                    == DBUS_INTERFACE
                && msg
                    .member()
                    .ok_or(anyhow!("Invoked method has no member"))?
                    == DBUS_GET_RESTRICTIONS_METHOD
            {
                let user_id: u32 = msg.body()?;
                match responses.get_mut(&Uid::from(user_id)) {
                    Some(answers) => match answers.pop() {
                        Some(a) => {
                            responses_size = responses_size - 1;
                            let ips: Vec<String> = a.into_iter().map(|ip| ip.to_string()).collect();
                            connection.reply(&msg, &ips).await
                        }
                        None => {
                            connection
                                .reply_error(
                                    &msg,
                                    "MockExhausted",
                                    &format!("DBus mock is saturated for user with id {}", user_id),
                                )
                                .await
                        }
                    },
                    None => {
                        connection
                            .reply_error(
                                &msg,
                                "MockExhausted",
                                &format!(
                                    "No mocked invocations available for user with id {}",
                                    user_id
                                ),
                            )
                            .await
                    }
                }?;
            }
        }

        Ok(())
    }

    task::spawn(async { dbus_loop(responses).await })
}