1
use std::path::PathBuf;
2

            
3
use anyhow::Context;
4
use clap::Parser;
5
use clap_verbosity_flag::Verbosity;
6
use systemd_journal_logger::JournalLog;
7

            
8
use crate::server::{
9
    config::{ServerConfigArgs, read_config_from_path_with_arg_overrides},
10
    server_loop::{
11
        listen_for_incoming_connections_with_socket_path,
12
        listen_for_incoming_connections_with_systemd_socket,
13
    },
14
};
15

            
16
#[derive(Parser, Debug, Clone)]
17
pub struct ServerArgs {
18
    #[command(subcommand)]
19
    subcmd: ServerCommand,
20

            
21
    #[command(flatten)]
22
    config_overrides: ServerConfigArgs,
23

            
24
    #[arg(long)]
25
    systemd: bool,
26
}
27

            
28
#[derive(Parser, Debug, Clone)]
29
pub enum ServerCommand {
30
    #[command()]
31
    Listen,
32

            
33
    #[command()]
34
    SocketActivate,
35
}
36

            
37
const LOG_LEVEL_WARNING: &str = r#"
38
===================================================
39
== WARNING: LOG LEVEL IS SET TO 'TRACE'!         ==
40
== THIS WILL CAUSE THE SERVER TO LOG SQL QUERIES ==
41
== THAT MAY CONTAIN SENSITIVE INFORMATION LIKE   ==
42
== PASSWORDS AND AUTHENTICATION TOKENS.          ==
43
== THIS IS INTENDED FOR DEBUGGING PURPOSES ONLY  ==
44
== AND SHOULD *NEVER* BE USED IN PRODUCTION.     ==
45
===================================================
46
"#;
47

            
48
pub async fn handle_command(
49
    socket_path: Option<PathBuf>,
50
    config_path: Option<PathBuf>,
51
    verbosity: Verbosity,
52
    args: ServerArgs,
53
) -> anyhow::Result<()> {
54
    let mut auto_detected_systemd_mode = false;
55
    let systemd_mode = args.systemd || {
56
        if let Ok(true) = sd_notify::booted() {
57
            auto_detected_systemd_mode = true;
58
            true
59
        } else {
60
            false
61
        }
62
    };
63
    if systemd_mode {
64
        JournalLog::new()
65
            .context("Failed to initialize journald logging")?
66
            .install()
67
            .context("Failed to install journald logger")?;
68

            
69
        log::set_max_level(verbosity.log_level_filter());
70

            
71
        if verbosity.log_level_filter() >= log::LevelFilter::Trace {
72
            log::warn!("{}", LOG_LEVEL_WARNING.trim());
73
        }
74

            
75
        if auto_detected_systemd_mode {
76
            log::info!("Running in systemd mode, auto-detected");
77
        } else {
78
            log::info!("Running in systemd mode");
79
        }
80

            
81
        start_watchdog_thread_if_enabled();
82
    } else {
83
        env_logger::Builder::new()
84
            .filter_level(verbosity.log_level_filter())
85
            .init();
86

            
87
        log::info!("Running in standalone mode");
88
    }
89

            
90
    let config = read_config_from_path_with_arg_overrides(config_path, args.config_overrides)?;
91

            
92
    match args.subcmd {
93
        ServerCommand::Listen => {
94
            listen_for_incoming_connections_with_socket_path(socket_path, config).await
95
        }
96
        ServerCommand::SocketActivate => {
97
            if !args.systemd {
98
                anyhow::bail!(concat!(
99
                    "The `--systemd` flag must be used with the `socket-activate` command.\n",
100
                    "This command currently only supports socket activation under systemd."
101
                ));
102
            }
103

            
104
            listen_for_incoming_connections_with_systemd_socket(config).await
105
        }
106
    }
107
}
108

            
109
fn start_watchdog_thread_if_enabled() {
110
    let mut micro_seconds: u64 = 0;
111
    let watchdog_enabled = sd_notify::watchdog_enabled(false, &mut micro_seconds);
112

            
113
    if watchdog_enabled {
114
        micro_seconds = micro_seconds.max(2_000_000).div_ceil(2);
115

            
116
        tokio::spawn(async move {
117
            log::debug!(
118
                "Starting systemd watchdog thread with {} millisecond interval",
119
                micro_seconds.div_ceil(1000)
120
            );
121
            loop {
122
                tokio::time::sleep(tokio::time::Duration::from_micros(micro_seconds)).await;
123
                if let Err(err) = sd_notify::notify(false, &[sd_notify::NotifyState::Watchdog]) {
124
                    log::warn!("Failed to notify systemd watchdog: {}", err);
125
                } else {
126
                    log::trace!("Ping sent to systemd watchdog");
127
                }
128
            }
129
        });
130
    } else {
131
        log::debug!("Systemd watchdog not enabled, skipping watchdog thread");
132
    }
133
}