1
use std::{
2
    collections::HashMap,
3
    os::fd::{AsRawFd, FromRawFd, OwnedFd},
4
    path::PathBuf,
5
    sync::Arc,
6
};
7

            
8
use anyhow::Context;
9
use clap::Parser;
10
use tokio::{net::UdpSocket, sync::RwLock};
11
use tokio_util::sync::CancellationToken;
12
use tracing::level_filters::LevelFilter;
13
use tracing_subscriber::layer::SubscriberExt;
14

            
15
use roowho2_lib::server::{
16
    config::{DEFAULT_CONFIG_PATH, LogLevel},
17
    ignore_list::IgnoreList,
18
    rwhod::{RwhodStatusStore, rwhod_packet_receiver_task, rwhod_packet_sender_task},
19
    varlink_api::varlink_client_server_task,
20
};
21

            
22
#[derive(Parser)]
23
#[command(
24
    author = "Programvareverkstedet <projects@pvv.ntnu.no>",
25
    about,
26
    version
27
)]
28
struct Args {
29
    /// Path to configuration file
30
    #[arg(
31
        short = 'c',
32
        long = "config",
33
        default_value = DEFAULT_CONFIG_PATH,
34
        value_name = "PATH"
35
    )]
36
    config_path: PathBuf,
37
}
38

            
39
#[tokio::main]
40
async fn main() -> anyhow::Result<()> {
41
    let args = Args::parse();
42

            
43
    let config = toml::from_str::<roowho2_lib::server::config::Config>(
44
        &std::fs::read_to_string(&args.config_path).context(format!(
45
            "Failed to read configuration file {:?}",
46
            args.config_path
47
        ))?,
48
    )?;
49

            
50
    let log_filter = match config.log_level.unwrap_or(LogLevel::Info) {
51
        LogLevel::Info => LevelFilter::INFO,
52
        LogLevel::Debug => LevelFilter::DEBUG,
53
        LogLevel::Trace => LevelFilter::TRACE,
54
    };
55

            
56
    let subscriber = tracing_subscriber::registry()
57
        .with(log_filter)
58
        .with(tracing_journald::layer()?);
59

            
60
    tracing::subscriber::set_global_default(subscriber)
61
        .context("Failed to set global default tracing subscriber")?;
62

            
63
    let rwhod_ignore_list = IgnoreList::load_optional(config.rwhod.ignore_list_path.as_deref())?;
64
    let finger_ignore_list = IgnoreList::load_optional(config.fingerd.ignore_list_path.as_deref())?;
65

            
66
    let fd_map: HashMap<String, OwnedFd> =
67
        HashMap::from_iter(sd_notify::listen_fds_with_names()?.map(|(fd_num, name)| {
68
            (
69
                name.clone(),
70
                // SAFETY: please don't mess around with file descriptors in random places
71
                //         around the codebase lol
72
                unsafe { std::os::fd::OwnedFd::from_raw_fd(fd_num) },
73
            )
74
        }));
75

            
76
    let mut join_set = tokio::task::JoinSet::new();
77

            
78
    let whod_status_store = Arc::new(RwLock::new(HashMap::new()));
79

            
80
    let client_server_token = CancellationToken::new();
81
    let client_server_token_ = client_server_token.clone();
82
    tokio::spawn(async move {
83
        client_server_token_.cancelled().await;
84
        tracing::info!("RWHOD client-server is now accepting connections");
85
        #[cfg(feature = "systemd")]
86
        sd_notify::notify(&[sd_notify::NotifyState::Ready]).ok();
87
        Ok::<(), anyhow::Error>(())
88
    });
89

            
90
    if config.rwhod.enable {
91
        tracing::info!("Starting RWHOD server");
92

            
93
        let socket = fd_map
94
            .get("rwhod_socket")
95
            .map(|fd| {
96
                // SAFETY: see above
97
                let std_socket = unsafe { std::net::UdpSocket::from_raw_fd(fd.as_raw_fd()) };
98
                std_socket.set_nonblocking(true)?;
99
                UdpSocket::from_std(std_socket)
100
            })
101
            .context("RWHOD server is enabled, but socket fd not provided by systemd")??;
102

            
103
        join_set.spawn(rwhod_server(
104
            socket,
105
            whod_status_store.clone(),
106
            rwhod_ignore_list.clone(),
107
        ));
108
    } else {
109
        tracing::debug!("RWHOD server is disabled in configuration");
110
    }
111

            
112
    join_set.spawn(client_server(
113
        fd_map
114
            .get("client_socket")
115
            .context("RWHOD client-server socket fd not provided by systemd")?
116
            .try_clone()
117
            .context("Failed to clone RWHOD client-server socket fd")?,
118
        whod_status_store.clone(),
119
        config.rwhod.enable,
120
        config.fingerd.enable,
121
        finger_ignore_list,
122
        client_server_token,
123
    ));
124

            
125
    join_set.spawn(ctrl_c_handler());
126

            
127
    join_set.join_next().await.unwrap()??;
128

            
129
    Ok(())
130
}
131

            
132
async fn ctrl_c_handler() -> anyhow::Result<()> {
133
    tokio::signal::ctrl_c()
134
        .await
135
        .map_err(|e| anyhow::anyhow!("Failed to listen for Ctrl-C: {}", e))
136
}
137

            
138
async fn rwhod_server(
139
    socket: UdpSocket,
140
    whod_status_store: RwhodStatusStore,
141
    ignore_list: Option<IgnoreList>,
142
) -> anyhow::Result<()> {
143
    let socket = Arc::new(socket);
144

            
145
    let interfaces = roowho2_lib::server::rwhod::determine_relevant_interfaces()?;
146
    let sender_task = rwhod_packet_sender_task(socket.clone(), interfaces, ignore_list);
147
    let receiver_task = rwhod_packet_receiver_task(socket.clone(), whod_status_store);
148

            
149
    tokio::select! {
150
        res = sender_task => res?,
151
        res = receiver_task => res?,
152
    }
153

            
154
    Ok(())
155
}
156

            
157
async fn client_server(
158
    socket_fd: OwnedFd,
159
    whod_status_store: RwhodStatusStore,
160
    rwhod_enabled: bool,
161
    fingerd_enabled: bool,
162
    finger_ignore_list: Option<IgnoreList>,
163
    startup_token: CancellationToken,
164
) -> anyhow::Result<()> {
165
    // SAFETY: see above
166
    let std_socket =
167
        unsafe { std::os::unix::net::UnixListener::from_raw_fd(socket_fd.as_raw_fd()) };
168
    std_socket.set_nonblocking(true)?;
169
    let zlink_listener = zlink::unix::Listener::try_from(OwnedFd::from(std_socket))?;
170
    let client_server_task = varlink_client_server_task(
171
        zlink_listener,
172
        whod_status_store,
173
        rwhod_enabled,
174
        fingerd_enabled,
175
        finger_ignore_list,
176
    );
177

            
178
    startup_token.cancel();
179

            
180
    client_server_task.await?;
181

            
182
    Ok(())
183
}