1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
|
// -*- coding: utf-8 -*-
//
// Copyright (C) 2024 Michael Büsch <m@bues.ch>
//
// Licensed under the Apache License version 2.0
// or the MIT license, at your option.
// SPDX-License-Identifier: Apache-2.0 OR MIT
#![forbid(unsafe_code)]
mod firewall;
mod processor;
mod server;
use crate::{
firewall::{nftables::NftFirewall, FirewallMaintain},
processor::Processor,
server::Server,
};
use anyhow::{self as ah, format_err as err, Context as _};
use clap::Parser;
use letmein_conf::{Config, ConfigVariant, INSTALL_PREFIX, SERVER_CONF_PATH};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::{
signal::unix::{signal, SignalKind},
sync::{self, Mutex, RwLock, RwLockReadGuard, Semaphore},
task, time,
};
const FW_MAINTAIN_PERIOD: Duration = Duration::from_millis(5000);
pub type ConfigRef<'a> = RwLockReadGuard<'a, Config>;
#[derive(Parser, Debug, Clone)]
struct Opts {
/// Override the default path to the configuration file.
#[arg(short, long)]
config: Option<PathBuf>,
/// Maximum number of simultaneous connections.
#[arg(short, long, default_value = "8")]
num_connections: usize,
/// Force-disable use of systemd socket.
///
/// Do not use systemd socket,
/// even if a systemd socket has been passed to the application.
#[arg(long, default_value = "false")]
no_systemd: bool,
}
impl Opts {
pub fn get_config(&self) -> PathBuf {
if let Some(config) = &self.config {
config.clone()
} else {
format!("{INSTALL_PREFIX}{SERVER_CONF_PATH}").into()
}
}
}
#[tokio::main(flavor = "current_thread")]
async fn main() -> ah::Result<()> {
let opts = Opts::parse();
let mut conf = Config::new(ConfigVariant::Server);
conf.load(&opts.get_config())
.context("Configuration file")?;
let conf = Arc::new(RwLock::new(conf));
let fw = Arc::new(Mutex::new(NftFirewall::new(&conf.read().await).await?));
let mut sigterm = signal(SignalKind::terminate()).unwrap();
let mut sigint = signal(SignalKind::interrupt()).unwrap();
let mut sighup = signal(SignalKind::hangup()).unwrap();
let (exit_sock_tx, mut exit_sock_rx) = sync::mpsc::channel(1);
let (exit_fw_tx, mut exit_fw_rx) = sync::mpsc::channel(1);
let srv = Server::new(&conf.read().await, opts.no_systemd)
.await
.context("Server init")?;
// Task: Socket handler.
let conf_clone = Arc::clone(&conf);
let fw_clone = Arc::clone(&fw);
task::spawn(async move {
let conn_semaphore = Semaphore::new(opts.num_connections);
loop {
let conf = Arc::clone(&conf_clone);
let fw = Arc::clone(&fw_clone);
match srv.accept().await {
Ok(conn) => {
// Socket connection handler.
if let Ok(_permit) = conn_semaphore.acquire().await {
task::spawn(async move {
let conf = conf.read().await;
let mut proc = Processor::new(conn, &conf, fw);
if let Err(e) = proc.run().await {
eprintln!("Client error: {e}");
}
});
}
}
Err(e) => {
let _ = exit_sock_tx.send(Err(e)).await;
break;
}
}
}
});
// Task: Firewall.
let conf_clone = Arc::clone(&conf);
let fw_clone = Arc::clone(&fw);
task::spawn(async move {
let mut interval = time::interval(FW_MAINTAIN_PERIOD);
loop {
interval.tick().await;
let conf = conf_clone.read().await;
let mut fw = fw_clone.lock().await;
if let Err(e) = fw.maintain(&conf).await {
let _ = exit_fw_tx.send(Err(e)).await;
break;
}
}
});
// Task: Main loop.
let mut exitcode;
loop {
tokio::select! {
_ = sigterm.recv() => {
eprintln!("SIGTERM: Terminating.");
exitcode = Ok(());
break;
}
_ = sigint.recv() => {
exitcode = Err(err!("Interrupted by SIGINT."));
break;
}
_ = sighup.recv() => {
println!("SIGHUP: Reloading.");
{
let mut conf = conf.write().await;
if let Err(e) = conf.load(&opts.get_config()) {
eprintln!("Failed to load configuration file: {e}");
}
}
{
let conf = conf.read().await;
let mut fw = fw.lock().await;
if let Err(e) = fw.reload(&conf).await {
eprintln!("Failed to reload filewall rules: {e}");
}
}
}
code = exit_sock_rx.recv() => {
exitcode = code.unwrap_or_else(|| Err(err!("Unknown error code.")));
break;
}
code = exit_fw_rx.recv() => {
exitcode = code.unwrap_or_else(|| Err(err!("Unknown error code.")));
break;
}
}
}
// Exiting...
// Try to remove all firewall rules.
{
let conf = conf.read().await;
let mut fw = fw.lock().await;
if let Err(e) = fw.clear(&conf).await {
eprintln!("WARNING: Failed to remove firewall rules: {e}");
if exitcode.is_ok() {
exitcode = Err(err!("Failed to remove firewall rules"));
}
}
}
exitcode
}
// vim: ts=4 sw=4 expandtab
|
