Stalwart/tests/src/utils/server.rs

/*
 * SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
 *
 * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
 */

use crate::{
    AssertConfig,
    store::TempDir,
    utils::{
        account::Account,
        cleanup::{search_store_destroy, store_destroy},
        registry::{RegistryEnvStores, UnwrapRegistryId, build_data_store},
        storage::assert_is_empty,
    },
};
use ahash::AHashMap;
use common::{
    BuildServer, Caches, Core, Data, Inner, Server,
    config::{
        server::{Listeners, ServerProtocol},
        storage::Storage,
        telemetry::Telemetry,
    },
    manager::{boot::build_ipc, defaults::BootstrapDefaults},
};
use http::HttpSessionManager;
use imap::core::ImapSessionManager;
use managesieve::core::ManageSieveSessionManager;
use pop3::Pop3SessionManager;
use registry::{
    schema::{
        enums::{DataStoreType, EventPolicy, NetworkListenerProtocol, TracingLevel},
        prelude::{Object, SocketAddr},
        structs::{NetworkListener, Tracer, TracerStdout},
    },
    types::{EnumImpl, map::Map},
};
use services::{SpawnServices, broadcast::subscriber::spawn_broadcast_subscriber};
use smtp::{SpawnQueueManager, core::SmtpSessionManager};
use std::{str::FromStr, sync::Arc};
use store::{
    RegistryStore, Store,
    registry::{bootstrap::Bootstrap, write::RegistryWrite},
};
use tokio::sync::watch;
use trc::EventType;
use types::id::Id;

pub struct TestServer {
    pub server: Server,
    accounts: AHashMap<&'static str, Account>,
    pub temp_dir: TempDir,
    shutdown_tx: watch::Sender<bool>,
}

pub struct TestServerBuilder {
    bootstrap: Bootstrap,
    temp_dir: TempDir,
    reset: bool,
}

impl TestServerBuilder {
    pub async fn new(test_name: &str, reset: bool) -> Self {
        let temp_dir = TempDir::new(test_name, reset);
        let path = temp_dir.path.to_string_lossy().to_string();
        let data_store = build_data_store(
            std::env::var("STORE")
                .map(|store| DataStoreType::parse(&store).expect("Invalid store type"))
                .expect(concat!(
                    "Missing or invalid store type. Try ",
                    "running `STORE=<store_type> cargo test`"
                )),
            &path,
        );
        let store = Store::build(data_store).await.unwrap();

        store.create_tables().await.unwrap();

        // Delete old store if requested
        if reset {
            store_destroy(&store).await;
        }

        Self {
            bootstrap: Bootstrap::new(
                RegistryStore::new(&path, store, "mail.example.org".to_string(), 1, None).await,
            )
            .await,
            temp_dir,
            reset,
        }
    }

    pub async fn with_listener(
        self,
        protocol: NetworkListenerProtocol,
        name: &str,
        port: u16,
        use_tls: bool,
    ) -> Self {
        self.insert_object(NetworkListener {
            bind: Map::new(vec![SocketAddr::from_str(&format!("[::]:{port}")).unwrap()]),
            name: name.to_string(),
            protocol,
            use_tls,
            ..Default::default()
        })
        .await;
        self
    }

    pub async fn with_object(self, object: impl Into<Object>) -> Self {
        self.insert_object(object).await;
        self
    }

    pub async fn insert_object(&self, object: impl Into<Object>) -> Id {
        self.bootstrap
            .registry
            .write(RegistryWrite::insert(&object.into()))
            .await
            .unwrap()
            .unwrap_id(trc::location!())
    }

    pub async fn build(mut self) -> TestServer {
        // Register stores from environment
        self.bootstrap.registry.insert_stores_from_env().await;

        // Enable logging if requested
        let level = std::env::var("LOG")
            .map(|log| TracingLevel::parse(&log).expect("Invalid log level"))
            .ok();
        self.bootstrap
            .registry
            .write(RegistryWrite::insert(
                &Tracer::Stdout(TracerStdout {
                    enable: level.is_some(),
                    level: level.unwrap_or(TracingLevel::Info),
                    ansi: true,
                    multiline: false,
                    events: Map::new(
                        EventType::variants()
                            .iter()
                            .filter(|ev| {
                                let ev = ev.as_str();
                                ev.starts_with("network.")
                                    || ev == "telemetry.webhook-error"
                                    || ev == "http.request-body"
                            })
                            .copied()
                            .collect(),
                    ),
                    events_policy: EventPolicy::Exclude,
                    ..Default::default()
                })
                .into(),
            ))
            .await
            .unwrap()
            .unwrap_id(trc::location!());

        // Start listeners
        let mut servers = Listeners::parse(&mut self.bootstrap).await;
        servers.bind_and_drop_priv(&mut self.bootstrap);

        // Parse storage
        let storage = Storage::parse(&mut self.bootstrap).await;

        // Reset search store
        if self.reset {
            search_store_destroy(&storage.search).await;
        }

        // Parse telemetry
        let telemetry = Telemetry::parse(&mut self.bootstrap, &storage).await;

        // Add safe defaults if missing
        self.bootstrap.insert_safe_defaults().await;

        // Parse components
        let core = Box::pin(Core::parse(&mut self.bootstrap, storage)).await;
        let data = Data::parse(&mut self.bootstrap).await;
        let cache = Caches::parse(&mut self.bootstrap).await;

        // Enable telemetry
        telemetry.enable(true);

        // Build inner
        let (ipc, mut ipc_rxs) = build_ipc(!core.storage.coordinator.is_none());
        let inner = Arc::new(Inner {
            shared_core: core.into_shared(),
            data,
            ipc,
            cache,
        });

        // Parse TCP acceptors
        servers
            .parse_tcp_acceptors(&mut self.bootstrap, inner.clone())
            .await;

        // Start services
        self.bootstrap.assert_no_errors();
        ipc_rxs.spawn_queue_manager(inner.clone());
        ipc_rxs.spawn_services(inner.clone());

        // Spawn servers
        let (shutdown_tx, shutdown_rx) = servers.spawn(|server, acceptor, shutdown_rx| {
            match &server.protocol {
                ServerProtocol::Smtp | ServerProtocol::Lmtp => server.spawn(
                    SmtpSessionManager::new(inner.clone()),
                    inner.clone(),
                    acceptor,
                    shutdown_rx,
                ),
                ServerProtocol::Http => server.spawn(
                    HttpSessionManager::new(inner.clone()),
                    inner.clone(),
                    acceptor,
                    shutdown_rx,
                ),
                ServerProtocol::Imap => server.spawn(
                    ImapSessionManager::new(inner.clone()),
                    inner.clone(),
                    acceptor,
                    shutdown_rx,
                ),
                ServerProtocol::Pop3 => server.spawn(
                    Pop3SessionManager::new(inner.clone()),
                    inner.clone(),
                    acceptor,
                    shutdown_rx,
                ),
                ServerProtocol::ManageSieve => server.spawn(
                    ManageSieveSessionManager::new(inner.clone()),
                    inner.clone(),
                    acceptor,
                    shutdown_rx,
                ),
            };
        });

        // Start broadcast subscriber
        spawn_broadcast_subscriber(inner.clone(), shutdown_rx);

        TestServer {
            server: inner.build_server(),
            temp_dir: self.temp_dir,
            accounts: Default::default(),
            shutdown_tx,
        }
    }
}

impl TestServer {
    pub fn account(&self, name: &str) -> &Account {
        self.accounts.get(name).unwrap()
    }

    pub async fn assert_is_empty(&self) {
        assert_is_empty(&self.server).await;
    }

    pub async fn destroy_store(&self) {
        store_destroy(self.server.store()).await;
    }

    pub fn shutdown(&self) {
        let _ = self.shutdown_tx.send(true);
    }
}