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

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

include!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/src/policy_checker/dbus.rs"
));

use {
    std::collections::HashMap,
    std::sync::atomic::{AtomicUsize, Ordering},
    zbus::dbus_interface,
};

#[derive(Debug)]
pub struct MalcontentDBusMock {
    responses: HashMap<Uid, Vec<Restrictions>>,
    invocations_left: AtomicUsize,
}

#[dbus_interface(name = "com.endlessm.ParentalControls.Dns")]
impl MalcontentDBusMock {
    fn get_restrictions(&mut self, user_id: u32) -> Restrictions {
        let uid = Uid::from_raw(user_id);
        let answers = self.responses.get_mut(&uid).expect(&format!(
            "MockError: No mocked invocations available for user with id {}",
            uid
        ));

        let restrictions = answers.pop().expect(&format!(
            "MockError: DBus mock is saturated for user with id {}",
            uid
        ));

        self.invocations_left.fetch_sub(1, Ordering::SeqCst);
        restrictions
    }
}

impl MalcontentDBusMock {
    pub fn new(mut responses: HashMap<Uid, Vec<Restrictions>>) -> Self {
        let responses_size: usize = responses.values().map(Vec::len).sum();

        for r in responses.values_mut() {
            r.reverse(); // we pop responses from the back, so...
        }

        let ret = Self {
            responses,
            invocations_left: AtomicUsize::new(responses_size),
        };

        ret
    }
}

impl Drop for MalcontentDBusMock {
    fn drop(&mut self) {
        let invocations_left = self.invocations_left.load(Ordering::Acquire);
        assert_eq!(
            invocations_left, 0,
            "MockError: During teardown, {} invocations are still left on the mock object",
            invocations_left
        );
    }
}

pub struct DBusServerGuard {
    handle: tokio::task::JoinHandle<()>,
}

impl Drop for DBusServerGuard {
    fn drop(&mut self) {
        self.handle.abort();
    }
}

pub async fn mock_dbus(responses: HashMap<Uid, Vec<Restrictions>>) -> Result<DBusServerGuard> {
    let guid = zbus::Guid::generate();
    let socket_path =
        std::path::PathBuf::from(std::env!("CARGO_TARGET_TMPDIR")).join(guid.as_str());

    if socket_path.exists() {
        std::fs::remove_file(&socket_path)?;
    }

    let socket = tokio::net::UnixListener::bind(&socket_path)?;
    std::env::set_var("TEST_DBUS_SOCKET", &socket_path);

    let mock = MalcontentDBusMock::new(responses);
    let handle = tokio::spawn(async move {
        let (stream, _) = socket
            .accept()
            .await
            .expect("Server socket closed unexpectedly");
        std::fs::remove_file(socket_path).unwrap(); // Once we accepted, we can already remove the socket

        let _ = zbus::ConnectionBuilder::unix_stream(stream)
            .server(&guid)
            .p2p()
            .name("com.endlessm.ParentalControls")
            .expect("Unable to serve given dbus name")
            .serve_at("/com/endlessm/ParentalControls/Dns", mock)
            .expect("Unable to server malcontent dbus mock object")
            .build()
            .await;

        std::future::pending::<()>().await;
    });

    Ok(DBusServerGuard { handle })
}