mirror of
https://github.com/mofeng-git/One-KVM.git
synced 2026-10-28 08:20:49 +08:00
refactor: 简化运行时与服务协调
提取运行时生命周期管理,集中远程访问与 USB 协调,并缩小 Web 路由状态依赖。
This commit is contained in:
832
src/main.rs
832
src/main.rs
@@ -3,43 +3,20 @@ use std::future::Future;
|
||||
use std::io::Write;
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum_server::tls_rustls::RustlsConfig;
|
||||
use clap::{Args, Parser, Subcommand, ValueEnum};
|
||||
use futures::{stream::FuturesUnordered, StreamExt};
|
||||
use rustls::crypto::{ring, CryptoProvider};
|
||||
use tokio::sync::{broadcast, mpsc};
|
||||
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
|
||||
|
||||
use one_kvm::atx::AtxController;
|
||||
use one_kvm::audio::{AudioController, AudioControllerConfig, AudioQuality};
|
||||
use one_kvm::auth::{SessionStore, TwoFactorService, UserStore};
|
||||
use one_kvm::computer_use::ComputerUseManager;
|
||||
use one_kvm::config::{self, AppConfig, ConfigStore};
|
||||
use one_kvm::config;
|
||||
use one_kvm::db::DatabasePool;
|
||||
use one_kvm::events::EventBus;
|
||||
use one_kvm::extensions::ExtensionManager;
|
||||
use one_kvm::hid::{HidBackendType, HidController};
|
||||
#[cfg(unix)]
|
||||
use one_kvm::msd::MsdController;
|
||||
#[cfg(unix)]
|
||||
use one_kvm::otg::OtgService;
|
||||
use one_kvm::platform::PlatformCapabilities;
|
||||
use one_kvm::rtsp::RtspService;
|
||||
use one_kvm::rustdesk::RustDeskService;
|
||||
use one_kvm::state::{AppState, ShutdownAction};
|
||||
use one_kvm::update::UpdateService;
|
||||
use one_kvm::runtime::{RuntimeBuilder, WebConfigOverrides};
|
||||
use one_kvm::state::ShutdownAction;
|
||||
use one_kvm::utils::bind_tcp_listener;
|
||||
use one_kvm::video::codec_constraints::{
|
||||
enforce_constraints_with_stream_manager, validate_third_party_codec_compatibility,
|
||||
StreamCodecConstraints,
|
||||
};
|
||||
use one_kvm::video::format::{PixelFormat, Resolution};
|
||||
use one_kvm::video::{Streamer, VideoStreamManager};
|
||||
use one_kvm::vnc::VncService;
|
||||
use one_kvm::web;
|
||||
use one_kvm::webrtc::{WebRtcStreamer, WebRtcStreamerConfig};
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, ValueEnum)]
|
||||
enum LogLevel {
|
||||
@@ -147,28 +124,20 @@ async fn main() -> anyhow::Result<()> {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let (db, config_store, mut config) = load_runtime_config(&data_dir).await?;
|
||||
|
||||
if let Some(addr) = args.address {
|
||||
config.web.bind_address = addr.clone();
|
||||
config.web.bind_addresses = vec![addr];
|
||||
}
|
||||
if let Some(port) = args.http_port {
|
||||
config.web.http_port = port;
|
||||
}
|
||||
if let Some(port) = args.https_port {
|
||||
config.web.https_port = port;
|
||||
}
|
||||
if args.enable_https {
|
||||
config.web.https_enabled = true;
|
||||
}
|
||||
|
||||
if let Some(cert_path) = args.ssl_cert {
|
||||
config.web.ssl_cert_path = Some(cert_path.to_string_lossy().to_string());
|
||||
}
|
||||
if let Some(key_path) = args.ssl_key {
|
||||
config.web.ssl_key_path = Some(key_path.to_string_lossy().to_string());
|
||||
}
|
||||
let overrides = WebConfigOverrides {
|
||||
address: args.address,
|
||||
http_port: args.http_port,
|
||||
https_port: args.https_port,
|
||||
enable_https: args.enable_https,
|
||||
ssl_cert: args.ssl_cert,
|
||||
ssl_key: args.ssl_key,
|
||||
};
|
||||
let mut runtime = RuntimeBuilder::new(data_dir.clone())
|
||||
.with_web_overrides(overrides)
|
||||
.build()
|
||||
.await?;
|
||||
let config = runtime.config();
|
||||
let state = runtime.state().clone();
|
||||
|
||||
let bind_ips = resolve_bind_addresses(&config.web)?;
|
||||
let scheme = if config.web.https_enabled {
|
||||
@@ -187,500 +156,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::info!("Server will listen on: {}://{}", scheme, addr);
|
||||
}
|
||||
|
||||
let session_store = SessionStore::new(config.auth.session_timeout_secs as i64);
|
||||
|
||||
let user_store = UserStore::new(db.clone_pool());
|
||||
let two_factor = TwoFactorService::new(db.clone_pool());
|
||||
|
||||
let (shutdown_tx, _) = broadcast::channel::<ShutdownAction>(1);
|
||||
|
||||
let events = Arc::new(EventBus::new());
|
||||
tracing::info!("Event bus initialized");
|
||||
|
||||
let (video_format, video_resolution) = parse_video_config(&config);
|
||||
tracing::debug!(
|
||||
"Parsed video config: {} @ {}x{}",
|
||||
video_format,
|
||||
video_resolution.width,
|
||||
video_resolution.height
|
||||
);
|
||||
|
||||
let streamer = Streamer::new();
|
||||
streamer.set_event_bus(events.clone()).await;
|
||||
if let Some(ref device_path) = config.video.device {
|
||||
if let Err(e) = streamer
|
||||
.apply_video_config(
|
||||
device_path,
|
||||
video_format,
|
||||
video_resolution,
|
||||
config.video.fps,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"Failed to initialize video with config: {}, will auto-detect",
|
||||
e
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
"Video configured: {} @ {}x{} {}",
|
||||
device_path,
|
||||
video_resolution.width,
|
||||
video_resolution.height,
|
||||
video_format
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let webrtc_streamer = {
|
||||
let webrtc_config = WebRtcStreamerConfig {
|
||||
resolution: video_resolution,
|
||||
input_format: video_format,
|
||||
fps: config.video.fps,
|
||||
bitrate_preset: config.stream.bitrate_preset,
|
||||
encoder_backend: one_kvm::stream_encoder::encoder_type_to_backend(
|
||||
config.stream.encoder.clone(),
|
||||
),
|
||||
webrtc: {
|
||||
let mut stun_servers = vec![];
|
||||
let mut turn_servers = vec![];
|
||||
|
||||
let has_custom_stun = config
|
||||
.stream
|
||||
.stun_server
|
||||
.as_ref()
|
||||
.map(|s| !s.is_empty())
|
||||
.unwrap_or(false);
|
||||
let has_custom_turn = config
|
||||
.stream
|
||||
.turn_server
|
||||
.as_ref()
|
||||
.map(|s| !s.is_empty())
|
||||
.unwrap_or(false);
|
||||
|
||||
if !has_custom_stun && !has_custom_turn {
|
||||
use one_kvm::webrtc::config::public_ice;
|
||||
let stun = public_ice::stun_server().to_string();
|
||||
tracing::info!("Using public STUN server: {}", stun);
|
||||
stun_servers.push(stun);
|
||||
} else {
|
||||
if let Some(ref stun) = config.stream.stun_server {
|
||||
if !stun.is_empty() {
|
||||
stun_servers.push(stun.clone());
|
||||
tracing::info!("Using custom STUN server: {}", stun);
|
||||
}
|
||||
}
|
||||
if let Some(ref turn) = config.stream.turn_server {
|
||||
if !turn.is_empty() {
|
||||
let username = config.stream.turn_username.clone().unwrap_or_default();
|
||||
let credential =
|
||||
config.stream.turn_password.clone().unwrap_or_default();
|
||||
turn_servers.push(one_kvm::webrtc::config::TurnServer::new(
|
||||
turn.clone(),
|
||||
username.clone(),
|
||||
credential,
|
||||
));
|
||||
tracing::info!(
|
||||
"Using custom TURN server: {} (user: {})",
|
||||
turn,
|
||||
username
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
one_kvm::webrtc::config::WebRtcConfig {
|
||||
stun_servers,
|
||||
turn_servers,
|
||||
..Default::default()
|
||||
}
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
WebRtcStreamer::with_config(webrtc_config)
|
||||
};
|
||||
tracing::info!("WebRTC streamer created");
|
||||
|
||||
#[cfg(unix)]
|
||||
let otg_service = Arc::new(OtgService::new());
|
||||
#[cfg(unix)]
|
||||
tracing::info!("OTG Service created");
|
||||
|
||||
#[cfg(unix)]
|
||||
if let Err(e) = otg_service
|
||||
.apply_config(&config.hid, &config.msd, &config.otg_network, &config.uac)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to apply OTG config: {}", e);
|
||||
}
|
||||
|
||||
let hid_backend = match config.hid.backend {
|
||||
config::HidBackend::Otg => HidBackendType::Otg,
|
||||
config::HidBackend::Ch9329 => HidBackendType::Ch9329 {
|
||||
port: config.hid.ch9329_port.clone(),
|
||||
baud_rate: config.hid.ch9329_baudrate,
|
||||
hybrid_mouse: config.hid.ch9329_hybrid_mouse,
|
||||
},
|
||||
config::HidBackend::None => HidBackendType::None,
|
||||
};
|
||||
#[cfg(unix)]
|
||||
let hid = Arc::new(HidController::new(hid_backend, Some(otg_service.clone())));
|
||||
#[cfg(not(unix))]
|
||||
let hid = Arc::new(HidController::new(hid_backend));
|
||||
hid.set_event_bus(events.clone()).await;
|
||||
if let Err(e) = hid.init().await {
|
||||
tracing::warn!("Failed to initialize HID backend: {}", e);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
let msd = if config.msd.enabled {
|
||||
let ventoy_resource_dir = data_dir.join("ventoy");
|
||||
let controller = MsdController::new(otg_service.clone(), config.msd.msd_dir_path());
|
||||
if let Err(e) = controller.init(&ventoy_resource_dir).await {
|
||||
tracing::warn!("Failed to initialize MSD controller: {}", e);
|
||||
None
|
||||
} else {
|
||||
controller.set_event_bus(events.clone()).await;
|
||||
Some(controller)
|
||||
}
|
||||
} else {
|
||||
tracing::info!("MSD disabled in configuration");
|
||||
None
|
||||
};
|
||||
|
||||
let atx = if config.atx.enabled {
|
||||
let controller_config = config.atx.to_controller_config();
|
||||
let controller = AtxController::new(controller_config);
|
||||
|
||||
if let Err(e) = controller.init().await {
|
||||
tracing::warn!("Failed to initialize ATX controller: {}", e);
|
||||
None
|
||||
} else {
|
||||
Some(controller)
|
||||
}
|
||||
} else {
|
||||
tracing::info!("ATX disabled in configuration");
|
||||
None
|
||||
};
|
||||
|
||||
let audio = {
|
||||
let audio_config = AudioControllerConfig {
|
||||
enabled: config.audio.enabled,
|
||||
device: config.audio.device.clone(),
|
||||
quality: match config.audio.quality.parse::<AudioQuality>() {
|
||||
Ok(q) => q,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Invalid audio quality in config (value={:?}): {}, using balanced",
|
||||
config.audio.quality,
|
||||
e
|
||||
);
|
||||
AudioQuality::Balanced
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
let controller = AudioController::new(audio_config);
|
||||
controller.set_event_bus(events.clone()).await;
|
||||
|
||||
if config.audio.enabled {
|
||||
tracing::info!(
|
||||
"Audio enabled: {}, quality={}",
|
||||
config.audio.device,
|
||||
config.audio.quality
|
||||
);
|
||||
if let Err(e) = controller.start_streaming().await {
|
||||
tracing::warn!("Failed to start audio streaming: {}", e);
|
||||
}
|
||||
} else {
|
||||
tracing::info!("Audio disabled in configuration");
|
||||
}
|
||||
|
||||
Arc::new(controller)
|
||||
};
|
||||
|
||||
let extensions = Arc::new(ExtensionManager::new());
|
||||
tracing::info!("Extension manager initialized");
|
||||
|
||||
webrtc_streamer.set_hid_controller(hid.clone()).await;
|
||||
|
||||
webrtc_streamer.set_audio_controller(audio.clone()).await;
|
||||
if config.audio.enabled {
|
||||
if let Err(e) = webrtc_streamer.set_audio_enabled(true).await {
|
||||
tracing::warn!("Failed to enable WebRTC audio: {}", e);
|
||||
} else {
|
||||
tracing::debug!("WebRTC audio enabled");
|
||||
}
|
||||
}
|
||||
|
||||
let (device_path, actual_resolution, actual_format, actual_fps, jpeg_quality) =
|
||||
streamer.current_capture_config().await;
|
||||
tracing::debug!(
|
||||
"Initial video config: {}x{} {:?} @ {}fps",
|
||||
actual_resolution.width,
|
||||
actual_resolution.height,
|
||||
actual_format,
|
||||
actual_fps
|
||||
);
|
||||
webrtc_streamer
|
||||
.update_video_config(actual_resolution, actual_format, actual_fps)
|
||||
.await;
|
||||
if let Some(device_path) = device_path {
|
||||
let device_info = streamer.current_device().await;
|
||||
webrtc_streamer
|
||||
.set_capture_device(device_path, jpeg_quality, device_info)
|
||||
.await;
|
||||
tracing::debug!("WebRTC streamer configured for direct capture");
|
||||
} else {
|
||||
tracing::warn!("No capture device configured for WebRTC");
|
||||
}
|
||||
|
||||
let stream_manager = VideoStreamManager::with_webrtc_streamer(
|
||||
streamer.clone(),
|
||||
webrtc_streamer.clone() as std::sync::Arc<dyn one_kvm::video::traits::VideoOutput>,
|
||||
);
|
||||
stream_manager.set_event_bus(events.clone()).await;
|
||||
stream_manager.set_config_store(config_store.clone()).await;
|
||||
{
|
||||
let stream_manager_weak = Arc::downgrade(&stream_manager);
|
||||
audio
|
||||
.set_recovered_callback(Arc::new(move || {
|
||||
if let Some(stream_manager) = stream_manager_weak.upgrade() {
|
||||
tokio::spawn(async move {
|
||||
stream_manager.reconnect_webrtc_audio_sources().await;
|
||||
});
|
||||
}
|
||||
}))
|
||||
.await;
|
||||
}
|
||||
|
||||
let initial_mode = config.stream.mode.clone();
|
||||
if let Err(e) = stream_manager.init_with_mode(initial_mode.clone()).await {
|
||||
tracing::warn!(
|
||||
"Failed to initialize stream manager with mode {:?}: {}",
|
||||
initial_mode,
|
||||
e
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
"Video stream manager initialized with mode: {:?}",
|
||||
initial_mode
|
||||
);
|
||||
}
|
||||
|
||||
let third_party_codec_config_valid = match validate_third_party_codec_compatibility(&config) {
|
||||
Ok(()) => true,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"Third-party access codec configuration is invalid; RustDesk/VNC/RTSP will not start: {}",
|
||||
e
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
let rustdesk = if third_party_codec_config_valid && config.rustdesk.is_valid() {
|
||||
tracing::info!(
|
||||
"Initializing RustDesk service: ID={} -> {}",
|
||||
config.rustdesk.device_id,
|
||||
config.rustdesk.rendezvous_addr()
|
||||
);
|
||||
let service = RustDeskService::new(
|
||||
config.rustdesk.clone(),
|
||||
stream_manager.clone(),
|
||||
hid.clone(),
|
||||
audio.clone(),
|
||||
);
|
||||
Some(Arc::new(service))
|
||||
} else {
|
||||
if config.rustdesk.enabled {
|
||||
tracing::warn!(
|
||||
"RustDesk enabled but configuration is incomplete (missing server or credentials)"
|
||||
);
|
||||
} else {
|
||||
tracing::info!("RustDesk disabled in configuration");
|
||||
}
|
||||
None
|
||||
};
|
||||
|
||||
let rtsp = if third_party_codec_config_valid && config.rtsp.enabled {
|
||||
tracing::info!(
|
||||
"Initializing RTSP service: rtsp://{}:{}/{}",
|
||||
config.rtsp.bind,
|
||||
config.rtsp.port,
|
||||
config.rtsp.path
|
||||
);
|
||||
let service = RtspService::new(config.rtsp.clone(), stream_manager.clone());
|
||||
Some(Arc::new(service))
|
||||
} else {
|
||||
tracing::info!("RTSP disabled in configuration");
|
||||
None
|
||||
};
|
||||
|
||||
let vnc = if third_party_codec_config_valid && config.vnc.enabled {
|
||||
tracing::info!(
|
||||
"Initializing VNC service: {}:{} ({:?})",
|
||||
config.vnc.bind,
|
||||
config.vnc.port,
|
||||
config.vnc.encoding
|
||||
);
|
||||
Some(Arc::new(VncService::new(
|
||||
config.vnc.clone(),
|
||||
stream_manager.clone(),
|
||||
hid.clone(),
|
||||
)))
|
||||
} else {
|
||||
tracing::info!("VNC disabled in configuration");
|
||||
None
|
||||
};
|
||||
|
||||
let update_service = Arc::new(UpdateService::new());
|
||||
let computer_use = ComputerUseManager::new(config_store.clone(), hid.clone());
|
||||
|
||||
let state = AppState::new(
|
||||
db.clone(),
|
||||
config_store.clone(),
|
||||
session_store,
|
||||
user_store,
|
||||
two_factor,
|
||||
#[cfg(unix)]
|
||||
otg_service,
|
||||
stream_manager,
|
||||
webrtc_streamer.clone(),
|
||||
hid,
|
||||
computer_use,
|
||||
#[cfg(unix)]
|
||||
msd,
|
||||
atx,
|
||||
audio,
|
||||
rustdesk.clone(),
|
||||
vnc.clone(),
|
||||
rtsp.clone(),
|
||||
extensions.clone(),
|
||||
events.clone(),
|
||||
update_service,
|
||||
shutdown_tx.clone(),
|
||||
data_dir.clone(),
|
||||
);
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
// Initialize UAC playback writer if UAC is enabled.
|
||||
if config.uac.enabled {
|
||||
let uac_cfg = one_kvm::audio::uac::UacPlaybackConfig {
|
||||
sample_rate: config.uac.sample_rate,
|
||||
channels: config.uac.channels as u16,
|
||||
..Default::default()
|
||||
};
|
||||
match one_kvm::audio::uac::UacPlayback::start(uac_cfg) {
|
||||
Ok(writer) => {
|
||||
*state.uac_playback.write().await = Some(writer);
|
||||
tracing::info!("UAC playback writer started");
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to start UAC playback writer: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if config.watchdog.enabled {
|
||||
if let Err(error) = state.watchdog.enable().await {
|
||||
tracing::error!(
|
||||
"Configured hardware watchdog failed to start; web service will continue: {}",
|
||||
error
|
||||
);
|
||||
} else {
|
||||
tracing::info!("Hardware watchdog started");
|
||||
}
|
||||
}
|
||||
|
||||
extensions.set_event_bus(events.clone()).await;
|
||||
|
||||
if let Some(ref service) = rustdesk {
|
||||
if let Err(e) = service.start().await {
|
||||
tracing::error!("Failed to start RustDesk service: {}", e);
|
||||
} else {
|
||||
if let Some(updated_config) = service.save_credentials() {
|
||||
if let Err(e) = config_store
|
||||
.update(|cfg| {
|
||||
cfg.rustdesk.public_key = updated_config.public_key.clone();
|
||||
cfg.rustdesk.private_key = updated_config.private_key.clone();
|
||||
cfg.rustdesk.signing_public_key = updated_config.signing_public_key.clone();
|
||||
cfg.rustdesk.signing_private_key =
|
||||
updated_config.signing_private_key.clone();
|
||||
cfg.rustdesk.uuid = updated_config.uuid.clone();
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to save RustDesk credentials: {}", e);
|
||||
}
|
||||
}
|
||||
tracing::info!("RustDesk service started");
|
||||
}
|
||||
}
|
||||
if let Some(ref service) = vnc {
|
||||
if let Err(e) = service.start().await {
|
||||
tracing::error!("Failed to start VNC service: {}", e);
|
||||
} else {
|
||||
tracing::info!("VNC service started");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref service) = rtsp {
|
||||
if let Err(e) = service.start().await {
|
||||
tracing::error!("Failed to start RTSP service: {}", e);
|
||||
} else {
|
||||
tracing::info!("RTSP service started");
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let runtime_config = state.runtime_third_party_config().await;
|
||||
let constraints = StreamCodecConstraints::from_config(&runtime_config);
|
||||
state
|
||||
.stream_manager
|
||||
.set_runtime_codec_constraints(constraints.clone())
|
||||
.await;
|
||||
match enforce_constraints_with_stream_manager(&state.stream_manager, &constraints).await {
|
||||
Ok(result) if result.changed => {
|
||||
if let Some(message) = result.message {
|
||||
tracing::info!("{}", message);
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => tracing::warn!("Failed to enforce startup codec constraints: {}", e),
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let ext_config = config_store.get();
|
||||
extensions.start_enabled(&ext_config.extensions).await;
|
||||
}
|
||||
|
||||
{
|
||||
let extensions_clone = extensions.clone();
|
||||
let config_store_clone = config_store.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let config = config_store_clone.get();
|
||||
extensions_clone.health_check(&config.extensions).await;
|
||||
}
|
||||
});
|
||||
tracing::info!("Extension health check task started");
|
||||
}
|
||||
|
||||
state.publish_device_info().await;
|
||||
|
||||
spawn_device_info_broadcaster(state.clone(), events);
|
||||
|
||||
let app = web::create_router(state.clone());
|
||||
let app = runtime.router();
|
||||
|
||||
let listeners = bind_tcp_listeners(&bind_ips, bind_port)?;
|
||||
|
||||
let shutdown_signal = {
|
||||
let shutdown_tx = state.shutdown_tx.clone();
|
||||
let mut shutdown_rx = shutdown_tx.subscribe();
|
||||
async move {
|
||||
tokio::select! {
|
||||
@@ -740,7 +221,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
servers.push(server);
|
||||
}
|
||||
|
||||
run_servers_until_shutdown(servers, shutdown_signal, &state, "HTTPS").await
|
||||
run_servers_until_shutdown(servers, shutdown_signal, "HTTPS").await
|
||||
} else {
|
||||
let servers = FuturesUnordered::new();
|
||||
for listener in listeners {
|
||||
@@ -752,9 +233,10 @@ async fn main() -> anyhow::Result<()> {
|
||||
servers.push(async move { server.await });
|
||||
}
|
||||
|
||||
run_servers_until_shutdown(servers, shutdown_signal, &state, "HTTP").await
|
||||
run_servers_until_shutdown(servers, shutdown_signal, "HTTP").await
|
||||
};
|
||||
|
||||
runtime.shutdown().await;
|
||||
tracing::info!("Server shutdown complete");
|
||||
if let ShutdownAction::Restart { exe_path } = shutdown_action {
|
||||
restart_current_process(exe_path)?;
|
||||
@@ -840,14 +322,13 @@ async fn open_database_pool(data_dir: &Path) -> anyhow::Result<DatabasePool> {
|
||||
async fn run_servers_until_shutdown<F, E>(
|
||||
mut servers: FuturesUnordered<F>,
|
||||
shutdown_signal: impl Future<Output = ShutdownAction>,
|
||||
state: &Arc<AppState>,
|
||||
protocol: &'static str,
|
||||
) -> ShutdownAction
|
||||
where
|
||||
F: Future<Output = Result<(), E>> + Send,
|
||||
E: std::fmt::Display,
|
||||
{
|
||||
let action = tokio::select! {
|
||||
tokio::select! {
|
||||
action = shutdown_signal => {
|
||||
action
|
||||
}
|
||||
@@ -857,9 +338,7 @@ where
|
||||
}
|
||||
ShutdownAction::Exit
|
||||
}
|
||||
};
|
||||
cleanup(state).await;
|
||||
action
|
||||
}
|
||||
}
|
||||
|
||||
fn restart_current_process(exe_path: Option<PathBuf>) -> anyhow::Result<()> {
|
||||
@@ -896,64 +375,6 @@ async fn run_cli_command(command: CliCommand, data_dir: PathBuf) -> anyhow::Resu
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_runtime_config(
|
||||
data_dir: &Path,
|
||||
) -> anyhow::Result<(DatabasePool, ConfigStore, AppConfig)> {
|
||||
tokio::fs::create_dir_all(data_dir).await?;
|
||||
|
||||
let db = open_database_pool(data_dir).await?;
|
||||
let config_store = ConfigStore::new(db.clone_pool())?;
|
||||
config_store.load().await?;
|
||||
let mut config = (*config_store.get()).clone();
|
||||
config.apply_platform_defaults();
|
||||
|
||||
prepare_linux_runtime_dirs(data_dir, &config_store, &mut config).await?;
|
||||
|
||||
Ok((db, config_store, config))
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn prepare_linux_runtime_dirs(
|
||||
data_dir: &Path,
|
||||
config_store: &ConfigStore,
|
||||
config: &mut AppConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut msd_dir_updated = false;
|
||||
if config.msd.msd_dir.trim().is_empty() {
|
||||
let msd_dir = data_dir.join("msd");
|
||||
config.msd.msd_dir = msd_dir.to_string_lossy().to_string();
|
||||
msd_dir_updated = true;
|
||||
} else if !PathBuf::from(&config.msd.msd_dir).is_absolute() {
|
||||
let msd_dir = data_dir.join(&config.msd.msd_dir);
|
||||
tracing::warn!(
|
||||
"MSD directory is relative, rebasing to {}",
|
||||
msd_dir.display()
|
||||
);
|
||||
config.msd.msd_dir = msd_dir.to_string_lossy().to_string();
|
||||
msd_dir_updated = true;
|
||||
}
|
||||
if msd_dir_updated {
|
||||
config_store.set(config.clone()).await?;
|
||||
}
|
||||
let msd_dir = PathBuf::from(&config.msd.msd_dir);
|
||||
if let Err(e) = tokio::fs::create_dir_all(msd_dir.join("images")).await {
|
||||
tracing::warn!("Failed to create MSD images directory: {}", e);
|
||||
}
|
||||
if let Err(e) = tokio::fs::create_dir_all(msd_dir.join("ventoy")).await {
|
||||
tracing::warn!("Failed to create MSD ventoy directory: {}", e);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
async fn prepare_linux_runtime_dirs(
|
||||
_data_dir: &Path,
|
||||
_config_store: &ConfigStore,
|
||||
_config: &mut AppConfig,
|
||||
) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_user_action(
|
||||
action: UserAction,
|
||||
users: &UserStore,
|
||||
@@ -1065,17 +486,6 @@ fn bind_tcp_listeners(addrs: &[IpAddr], port: u16) -> anyhow::Result<Vec<std::ne
|
||||
Ok(listeners)
|
||||
}
|
||||
|
||||
fn parse_video_config(config: &AppConfig) -> (PixelFormat, Resolution) {
|
||||
let format = config
|
||||
.video
|
||||
.format
|
||||
.as_ref()
|
||||
.and_then(|f: &String| f.parse::<PixelFormat>().ok())
|
||||
.unwrap_or(PixelFormat::Mjpeg);
|
||||
let resolution = Resolution::new(config.video.width, config.video.height);
|
||||
(format, resolution)
|
||||
}
|
||||
|
||||
fn generate_self_signed_cert() -> anyhow::Result<rcgen::CertifiedKey<rcgen::KeyPair>> {
|
||||
use rcgen::generate_simple_self_signed;
|
||||
|
||||
@@ -1088,197 +498,3 @@ fn generate_self_signed_cert() -> anyhow::Result<rcgen::CertifiedKey<rcgen::KeyP
|
||||
let certified_key = generate_simple_self_signed(subject_alt_names)?;
|
||||
Ok(certified_key)
|
||||
}
|
||||
|
||||
fn spawn_device_info_broadcaster(state: Arc<AppState>, events: Arc<EventBus>) {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
enum DeviceInfoTrigger {
|
||||
Event,
|
||||
Lagged { topic: &'static str, count: u64 },
|
||||
}
|
||||
|
||||
const DEVICE_INFO_TOPICS: &[&str] = &[
|
||||
"stream.state_changed",
|
||||
"stream.config_applied",
|
||||
"stream.mode_ready",
|
||||
];
|
||||
const DEBOUNCE_MS: u64 = 100;
|
||||
|
||||
let (trigger_tx, mut trigger_rx) = mpsc::unbounded_channel();
|
||||
|
||||
for topic in DEVICE_INFO_TOPICS {
|
||||
let Some(mut rx) = events.subscribe_topic(topic) else {
|
||||
tracing::warn!(
|
||||
"DeviceInfo broadcaster missing topic subscription: {}",
|
||||
topic
|
||||
);
|
||||
continue;
|
||||
};
|
||||
|
||||
let trigger_tx = trigger_tx.clone();
|
||||
let topic_name = *topic;
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match rx.recv().await {
|
||||
Ok(_) => {
|
||||
if trigger_tx.send(DeviceInfoTrigger::Event).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => {
|
||||
if trigger_tx
|
||||
.send(DeviceInfoTrigger::Lagged {
|
||||
topic: topic_name,
|
||||
count,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
{
|
||||
let mut dirty_rx = events.subscribe_device_info_dirty();
|
||||
let trigger_tx = trigger_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
match dirty_rx.recv().await {
|
||||
Ok(()) => {
|
||||
if trigger_tx.send(DeviceInfoTrigger::Event).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(count)) => {
|
||||
if trigger_tx
|
||||
.send(DeviceInfoTrigger::Lagged {
|
||||
topic: "device_info_dirty",
|
||||
count,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut last_broadcast = Instant::now() - Duration::from_millis(DEBOUNCE_MS);
|
||||
let mut pending_broadcast = false;
|
||||
|
||||
loop {
|
||||
let recv_result = if pending_broadcast {
|
||||
let remaining =
|
||||
DEBOUNCE_MS.saturating_sub(last_broadcast.elapsed().as_millis() as u64);
|
||||
tokio::time::timeout(Duration::from_millis(remaining), trigger_rx.recv()).await
|
||||
} else {
|
||||
Ok(trigger_rx.recv().await)
|
||||
};
|
||||
|
||||
match recv_result {
|
||||
Ok(Some(DeviceInfoTrigger::Event)) => {
|
||||
pending_broadcast = true;
|
||||
}
|
||||
Ok(Some(DeviceInfoTrigger::Lagged { topic, count })) => {
|
||||
tracing::warn!(
|
||||
"DeviceInfo broadcaster lagged by {} events on topic {}",
|
||||
count,
|
||||
topic
|
||||
);
|
||||
pending_broadcast = true;
|
||||
}
|
||||
Ok(None) => {
|
||||
tracing::info!("Event bus closed, stopping DeviceInfo broadcaster");
|
||||
break;
|
||||
}
|
||||
Err(_timeout) => {}
|
||||
}
|
||||
|
||||
if pending_broadcast && last_broadcast.elapsed() >= Duration::from_millis(DEBOUNCE_MS) {
|
||||
state.publish_device_info().await;
|
||||
tracing::trace!("Broadcasted DeviceInfo (debounced)");
|
||||
last_broadcast = Instant::now();
|
||||
pending_broadcast = false;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!(
|
||||
"DeviceInfo broadcaster task started (debounce: {}ms)",
|
||||
DEBOUNCE_MS
|
||||
);
|
||||
}
|
||||
|
||||
async fn cleanup(state: &Arc<AppState>) {
|
||||
state.extensions.stop_all().await;
|
||||
tracing::info!("Extensions stopped");
|
||||
|
||||
if let Some(ref service) = *state.rustdesk.read().await {
|
||||
if let Err(e) = service.stop().await {
|
||||
tracing::warn!("Failed to stop RustDesk service: {}", e);
|
||||
} else {
|
||||
tracing::info!("RustDesk service stopped");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref service) = *state.vnc.read().await {
|
||||
if let Err(e) = service.stop().await {
|
||||
tracing::warn!("Failed to stop VNC service: {}", e);
|
||||
} else {
|
||||
tracing::info!("VNC service stopped");
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref service) = *state.rtsp.read().await {
|
||||
if let Err(e) = service.stop().await {
|
||||
tracing::warn!("Failed to stop RTSP service: {}", e);
|
||||
} else {
|
||||
tracing::info!("RTSP service stopped");
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = state.stream_manager.stop().await {
|
||||
tracing::warn!("Failed to stop streamer: {}", e);
|
||||
}
|
||||
|
||||
if let Err(e) = state.hid.shutdown().await {
|
||||
tracing::warn!("Failed to shutdown HID: {}", e);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
if let Some(msd) = state.msd.write().await.as_mut() {
|
||||
if let Err(e) = msd.shutdown().await {
|
||||
tracing::warn!("Failed to shutdown MSD: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
if let Err(e) = state.otg_service.shutdown().await {
|
||||
tracing::warn!("Failed to shutdown OTG: {}", e);
|
||||
}
|
||||
|
||||
if let Some(atx) = state.atx.write().await.as_mut() {
|
||||
if let Err(e) = atx.shutdown().await {
|
||||
tracing::warn!("Failed to shutdown ATX: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = state.audio.shutdown().await {
|
||||
tracing::warn!("Failed to shutdown audio: {}", e);
|
||||
}
|
||||
|
||||
if let Err(error) = state.watchdog.disable().await {
|
||||
tracing::error!(
|
||||
"CRITICAL: failed to disable hardware watchdog during shutdown: {}",
|
||||
error
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user