jasonwitty
e51054811c
- socktop connector allows you to communicate with socktop agent directly from you code without needing to implement the agent API directly. - will also be used for non tui implementation of "socktop collector" in the future. - moved to rust 2024 to take advantage of some new features that helped with refactor. - fixed everything that exploded with update. - added rust docs for lib.
92 lines
2.7 KiB
Rust
92 lines
2.7 KiB
Rust
//! socktop agent entrypoint: sets up sysinfo handles and serves a WebSocket endpoint at /ws.
|
|
|
|
mod gpu;
|
|
mod metrics;
|
|
mod proto;
|
|
// sampler module removed (metrics now purely request-driven)
|
|
mod state;
|
|
mod types;
|
|
mod ws;
|
|
|
|
use axum::{Router, http::StatusCode, routing::get};
|
|
use std::net::SocketAddr;
|
|
use std::str::FromStr;
|
|
|
|
mod tls;
|
|
|
|
use state::AppState;
|
|
|
|
fn arg_flag(name: &str) -> bool {
|
|
std::env::args().any(|a| a == name)
|
|
}
|
|
fn arg_value(name: &str) -> Option<String> {
|
|
let mut it = std::env::args();
|
|
while let Some(a) = it.next() {
|
|
if a == name {
|
|
return it.next();
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {
|
|
tracing_subscriber::fmt::init();
|
|
|
|
// Version flag (print and exit). Keep before heavy initialization.
|
|
if arg_flag("--version") || arg_flag("-V") {
|
|
println!("socktop_agent {}", env!("CARGO_PKG_VERSION"));
|
|
return Ok(());
|
|
}
|
|
|
|
let state = AppState::new();
|
|
|
|
// No background samplers: metrics collected on-demand per websocket request.
|
|
|
|
// Web app: route /ws to the websocket handler
|
|
async fn healthz() -> StatusCode {
|
|
println!("/healthz request");
|
|
StatusCode::OK
|
|
}
|
|
let app = Router::new()
|
|
.route("/ws", get(ws::ws_handler))
|
|
.route("/healthz", get(healthz))
|
|
.with_state(state.clone());
|
|
|
|
let enable_ssl =
|
|
arg_flag("--enableSSL") || std::env::var("SOCKTOP_ENABLE_SSL").ok().as_deref() == Some("1");
|
|
if enable_ssl {
|
|
// Port can be overridden by --port or SOCKTOP_PORT; default to 8443 when SSL
|
|
let port = arg_value("--port")
|
|
.or_else(|| arg_value("-p"))
|
|
.or_else(|| std::env::var("SOCKTOP_PORT").ok())
|
|
.and_then(|s| s.parse::<u16>().ok())
|
|
.unwrap_or(8443);
|
|
|
|
let (cert_path, key_path) = tls::ensure_self_signed_cert()?;
|
|
let cfg = axum_server::tls_rustls::RustlsConfig::from_pem_file(cert_path, key_path).await?;
|
|
|
|
let addr = SocketAddr::from_str(&format!("0.0.0.0:{port}"))?;
|
|
println!("socktop_agent: TLS enabled. Listening on wss://{addr}/ws");
|
|
axum_server::bind_rustls(addr, cfg)
|
|
.serve(app.into_make_service())
|
|
.await?;
|
|
return Ok(());
|
|
}
|
|
|
|
// Non-TLS HTTP/WS path
|
|
let port = arg_value("--port")
|
|
.or_else(|| arg_value("-p"))
|
|
.or_else(|| std::env::var("SOCKTOP_PORT").ok())
|
|
.and_then(|s| s.parse::<u16>().ok())
|
|
.unwrap_or(3000);
|
|
let addr = SocketAddr::from(([0, 0, 0, 0], port));
|
|
println!("socktop_agent: Listening on ws://{addr}/ws");
|
|
axum_server::bind(addr)
|
|
.serve(app.into_make_service())
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
|
|
// Unit tests for CLI parsing moved to `tests/port_parse.rs`.
|