1
use std::path::PathBuf;
2

            
3
use anyhow::Context;
4
use clap::{Parser, Subcommand};
5
use clap_verbosity_flag::{InfoLevel, Verbosity};
6
use tracing_subscriber::layer::SubscriberExt;
7

            
8
use muscl_lib::{
9
    core::common::{ASCII_BANNER, DEFAULT_CONFIG_PATH, KIND_REGARDS},
10
    server::{landlock::landlock_restrict_server, supervisor::Supervisor},
11
};
12

            
13
#[derive(Parser, Debug, Clone)]
14
pub struct ServerArgs {
15
    #[command(subcommand)]
16
    pub subcmd: ServerCommand,
17

            
18
    /// Enable systemd mode
19
    #[cfg(target_os = "linux")]
20
    #[arg(long)]
21
    pub systemd: bool,
22

            
23
    /// Disable Landlock sandboxing.
24
    ///
25
    /// This is useful if you are planning to reload the server's configuration.
26
    #[arg(long)]
27
    pub disable_landlock: bool,
28

            
29
    // NOTE: be careful not to add short options that collide with the `edit-privs` privilege
30
    //       characters. It should in theory be possible for `edit-privs` to ignore any options
31
    //       specified here, but in practice clap is being difficult to work with.
32
    /// Path to where the server's unix socket should be created. This is only relevant when
33
    /// not using systemd socket activation.
34
    #[arg(
35
        long = "socket",
36
        value_name = "PATH",
37
        value_hint = clap::ValueHint::FilePath,
38
    )]
39
    socket_path: Option<PathBuf>,
40

            
41
    /// Config file to use for the server.
42
    #[arg(
43
        long = "config",
44
        value_name = "PATH",
45
        value_hint = clap::ValueHint::FilePath,
46
    )]
47
    config_path: Option<PathBuf>,
48

            
49
    #[command(flatten)]
50
    verbosity: Verbosity<InfoLevel>,
51
}
52

            
53
#[derive(Subcommand, Debug, Clone)]
54
pub enum ServerCommand {
55
    /// Start the server and listen for incoming connections on the unix socket
56
    /// specified in the configuration file.
57
    Listen,
58

            
59
    /// Start the server using systemd socket activation.
60
    SocketActivate,
61
}
62

            
63
const LOG_LEVEL_WARNING: &str = r#"
64
===================================================
65
== WARNING: LOG LEVEL IS SET TO 'TRACE'!         ==
66
== THIS WILL CAUSE THE SERVER TO LOG SQL QUERIES ==
67
== THAT MAY CONTAIN SENSITIVE INFORMATION LIKE   ==
68
== PASSWORDS AND AUTHENTICATION TOKENS.          ==
69
== THIS IS INTENDED FOR DEBUGGING PURPOSES ONLY  ==
70
== AND SHOULD *NEVER* BE USED IN PRODUCTION.     ==
71
===================================================
72
"#;
73

            
74
const MIN_TOKIO_WORKER_THREADS: usize = 4;
75

            
76
fn main() -> anyhow::Result<()> {
77
    let args = ServerArgs::parse();
78

            
79
    if !args.disable_landlock {
80
        landlock_restrict_server(args.config_path.as_deref())
81
            .context("Failed to apply Landlock restrictions to the server process")?;
82
    }
83

            
84
    let worker_thread_count = std::cmp::max(num_cpus::get(), MIN_TOKIO_WORKER_THREADS);
85

            
86
    tokio::runtime::Builder::new_multi_thread()
87
        .worker_threads(worker_thread_count)
88
        .enable_all()
89
        .build()
90
        .context("Failed to start Tokio runtime")?
91
        .block_on(handle_command(args))?;
92

            
93
    Ok(())
94
}
95

            
96
fn trace_server_prelude() {
97
    let message = [ASCII_BANNER, "", KIND_REGARDS, ""].join("\n");
98
    tracing::info!(message);
99
}
100

            
101
async fn handle_command(args: ServerArgs) -> anyhow::Result<()> {
102
    let mut auto_detected_systemd_mode = false;
103

            
104
    #[cfg(target_os = "linux")]
105
    let systemd_mode = args.systemd || {
106
        if let Ok(true) = sd_notify::booted() {
107
            auto_detected_systemd_mode = true;
108
            true
109
        } else {
110
            false
111
        }
112
    };
113

            
114
    #[cfg(not(target_os = "linux"))]
115
    let systemd_mode = false;
116

            
117
    if systemd_mode {
118
        #[cfg(target_os = "linux")]
119
        {
120
            let subscriber = tracing_subscriber::Registry::default()
121
                .with(args.verbosity.tracing_level_filter())
122
                .with(tracing_journald::layer()?);
123

            
124
            tracing::subscriber::set_global_default(subscriber)
125
                .context("Failed to set global default tracing subscriber")?;
126

            
127
            trace_server_prelude();
128

            
129
            if args.verbosity.tracing_level_filter() >= tracing::Level::TRACE {
130
                tracing::warn!("{}", LOG_LEVEL_WARNING.trim());
131
            }
132

            
133
            if auto_detected_systemd_mode {
134
                tracing::debug!("Running in systemd mode, auto-detected");
135
            } else {
136
                tracing::debug!("Running in systemd mode");
137
            }
138
        }
139
    } else {
140
        let subscriber = tracing_subscriber::Registry::default()
141
            .with(args.verbosity.tracing_level_filter())
142
            .with(
143
                tracing_subscriber::fmt::layer()
144
                    .with_line_number(cfg!(debug_assertions))
145
                    .with_target(cfg!(debug_assertions))
146
                    .with_thread_ids(false)
147
                    .with_thread_names(false),
148
            );
149

            
150
        tracing::subscriber::set_global_default(subscriber)
151
            .context("Failed to set global default tracing subscriber")?;
152

            
153
        trace_server_prelude();
154

            
155
        tracing::debug!("Running in standalone mode");
156
    }
157

            
158
    let config_path = args
159
        .config_path
160
        .unwrap_or_else(|| PathBuf::from(DEFAULT_CONFIG_PATH));
161

            
162
    match args.subcmd {
163
        ServerCommand::Listen => {
164
            Supervisor::new(config_path, systemd_mode)
165
                .await?
166
                .run()
167
                .await
168
        }
169
        ServerCommand::SocketActivate => {
170
            if !args.systemd {
171
                anyhow::bail!(concat!(
172
                    "The `--systemd` flag must be used with the `socket-activate` command.\n",
173
                    "This command currently only supports socket activation under systemd."
174
                ));
175
            }
176

            
177
            Supervisor::new(config_path, systemd_mode)
178
                .await?
179
                .run()
180
                .await
181
        }
182
    }
183
}