Housekeeping and QOL
non functional update: - refactor stream of consciousness into separate files. - combine equivelent functions used in networking and wasm features. - cleanups and version bumps.
This commit is contained in:
@@ -0,0 +1,185 @@
|
||||
//! WebSocket connection handling for native (non-WASM) environments.
|
||||
|
||||
use crate::config::ConnectorConfig;
|
||||
use crate::error::{ConnectorError, Result};
|
||||
|
||||
use std::io::BufReader;
|
||||
use std::sync::Arc;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
|
||||
use url::Url;
|
||||
|
||||
#[cfg(feature = "tls")]
|
||||
use {
|
||||
rustls::{self, ClientConfig},
|
||||
rustls::{
|
||||
DigitallySignedStruct, RootCertStore, SignatureScheme,
|
||||
client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
|
||||
crypto::ring,
|
||||
pki_types::{CertificateDer, ServerName, UnixTime},
|
||||
},
|
||||
rustls_pemfile::Item,
|
||||
std::fs::File,
|
||||
tokio_tungstenite::Connector,
|
||||
};
|
||||
|
||||
pub type WsStream = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
|
||||
|
||||
/// Connect to the agent and return the WS stream
|
||||
pub async fn connect_to_agent(config: &ConnectorConfig) -> Result<WsStream> {
|
||||
#[cfg(feature = "tls")]
|
||||
ensure_crypto_provider();
|
||||
|
||||
let mut u = Url::parse(&config.url)?;
|
||||
if let Some(ca_path) = &config.tls_ca_path {
|
||||
if u.scheme() == "ws" {
|
||||
let _ = u.set_scheme("wss");
|
||||
}
|
||||
return connect_with_ca_and_config(u.as_str(), ca_path, config).await;
|
||||
}
|
||||
// No TLS - hostname verification is not applicable
|
||||
connect_without_ca_and_config(u.as_str(), config).await
|
||||
}
|
||||
|
||||
async fn connect_without_ca_and_config(url: &str, config: &ConnectorConfig) -> Result<WsStream> {
|
||||
let mut req = url.into_client_request()?;
|
||||
|
||||
// Apply WebSocket protocol configuration
|
||||
if let Some(version) = &config.ws_version {
|
||||
req.headers_mut().insert(
|
||||
"Sec-WebSocket-Version",
|
||||
version
|
||||
.parse()
|
||||
.map_err(|_| ConnectorError::protocol_error("Invalid WebSocket version"))?,
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(protocols) = &config.ws_protocols {
|
||||
let protocols_str = protocols.join(", ");
|
||||
req.headers_mut().insert(
|
||||
"Sec-WebSocket-Protocol",
|
||||
protocols_str
|
||||
.parse()
|
||||
.map_err(|_| ConnectorError::protocol_error("Invalid WebSocket protocols"))?,
|
||||
);
|
||||
}
|
||||
|
||||
let (ws, _) = connect_async(req).await?;
|
||||
Ok(ws)
|
||||
}
|
||||
|
||||
#[cfg(feature = "tls")]
|
||||
async fn connect_with_ca_and_config(
|
||||
url: &str,
|
||||
ca_path: &str,
|
||||
config: &ConnectorConfig,
|
||||
) -> Result<WsStream> {
|
||||
// Initialize the crypto provider for rustls
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
|
||||
let mut root = RootCertStore::empty();
|
||||
let mut reader = BufReader::new(File::open(ca_path)?);
|
||||
let mut der_certs = Vec::new();
|
||||
while let Ok(Some(item)) = rustls_pemfile::read_one(&mut reader) {
|
||||
if let Item::X509Certificate(der) = item {
|
||||
der_certs.push(der);
|
||||
}
|
||||
}
|
||||
root.add_parsable_certificates(der_certs);
|
||||
|
||||
let mut cfg = ClientConfig::builder()
|
||||
.with_root_certificates(root)
|
||||
.with_no_client_auth();
|
||||
|
||||
let mut req = url.into_client_request()?;
|
||||
|
||||
// Apply WebSocket protocol configuration
|
||||
if let Some(version) = &config.ws_version {
|
||||
req.headers_mut().insert(
|
||||
"Sec-WebSocket-Version",
|
||||
version
|
||||
.parse()
|
||||
.map_err(|_| ConnectorError::protocol_error("Invalid WebSocket version"))?,
|
||||
);
|
||||
}
|
||||
|
||||
if let Some(protocols) = &config.ws_protocols {
|
||||
let protocols_str = protocols.join(", ");
|
||||
req.headers_mut().insert(
|
||||
"Sec-WebSocket-Protocol",
|
||||
protocols_str
|
||||
.parse()
|
||||
.map_err(|_| ConnectorError::protocol_error("Invalid WebSocket protocols"))?,
|
||||
);
|
||||
}
|
||||
|
||||
if !config.verify_hostname {
|
||||
#[derive(Debug)]
|
||||
struct NoVerify;
|
||||
impl ServerCertVerifier for NoVerify {
|
||||
fn verify_server_cert(
|
||||
&self,
|
||||
_end_entity: &CertificateDer<'_>,
|
||||
_intermediates: &[CertificateDer<'_>],
|
||||
_server_name: &ServerName,
|
||||
_ocsp_response: &[u8],
|
||||
_now: UnixTime,
|
||||
) -> std::result::Result<ServerCertVerified, rustls::Error> {
|
||||
Ok(ServerCertVerified::assertion())
|
||||
}
|
||||
fn verify_tls12_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &CertificateDer<'_>,
|
||||
_dss: &DigitallySignedStruct,
|
||||
) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(HandshakeSignatureValid::assertion())
|
||||
}
|
||||
fn verify_tls13_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &CertificateDer<'_>,
|
||||
_dss: &DigitallySignedStruct,
|
||||
) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(HandshakeSignatureValid::assertion())
|
||||
}
|
||||
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
|
||||
vec![
|
||||
SignatureScheme::ECDSA_NISTP256_SHA256,
|
||||
SignatureScheme::ED25519,
|
||||
SignatureScheme::RSA_PSS_SHA256,
|
||||
]
|
||||
}
|
||||
}
|
||||
cfg.dangerous().set_certificate_verifier(Arc::new(NoVerify));
|
||||
eprintln!(
|
||||
"socktop_connector: hostname verification disabled (default). Set SOCKTOP_VERIFY_NAME=1 to enable strict SAN checking."
|
||||
);
|
||||
}
|
||||
let cfg = Arc::new(cfg);
|
||||
let (ws, _) = tokio_tungstenite::connect_async_tls_with_config(
|
||||
req,
|
||||
None,
|
||||
config.verify_hostname,
|
||||
Some(Connector::Rustls(cfg)),
|
||||
)
|
||||
.await?;
|
||||
Ok(ws)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "tls"))]
|
||||
async fn connect_with_ca_and_config(
|
||||
_url: &str,
|
||||
_ca_path: &str,
|
||||
_config: &ConnectorConfig,
|
||||
) -> Result<WsStream> {
|
||||
Err(ConnectorError::tls_error(
|
||||
"TLS support not compiled in",
|
||||
std::io::Error::new(std::io::ErrorKind::Unsupported, "TLS not available"),
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(feature = "tls")]
|
||||
fn ensure_crypto_provider() {
|
||||
let _ = ring::default_provider().install_default();
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! Networking module for native WebSocket connections.
|
||||
|
||||
pub mod connection;
|
||||
pub mod requests;
|
||||
|
||||
pub use connection::*;
|
||||
pub use requests::*;
|
||||
@@ -0,0 +1,84 @@
|
||||
//! WebSocket request handlers for native (non-WASM) environments.
|
||||
|
||||
use crate::networking::WsStream;
|
||||
use crate::utils::{gunzip_to_string, gunzip_to_vec, is_gzip};
|
||||
use crate::{DiskInfo, Metrics, ProcessInfo, ProcessesPayload, pb};
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use prost::Message as ProstMessage;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
/// Send a "get_metrics" request and await a single JSON reply
|
||||
pub async fn request_metrics(ws: &mut WsStream) -> Option<Metrics> {
|
||||
if ws.send(Message::Text("get_metrics".into())).await.is_err() {
|
||||
return None;
|
||||
}
|
||||
match ws.next().await {
|
||||
Some(Ok(Message::Binary(b))) => gunzip_to_string(&b)
|
||||
.ok()
|
||||
.and_then(|s| serde_json::from_str::<Metrics>(&s).ok()),
|
||||
Some(Ok(Message::Text(json))) => serde_json::from_str::<Metrics>(&json).ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a "get_disks" request and await a JSON Vec<DiskInfo>
|
||||
pub async fn request_disks(ws: &mut WsStream) -> Option<Vec<DiskInfo>> {
|
||||
if ws.send(Message::Text("get_disks".into())).await.is_err() {
|
||||
return None;
|
||||
}
|
||||
match ws.next().await {
|
||||
Some(Ok(Message::Binary(b))) => gunzip_to_string(&b)
|
||||
.ok()
|
||||
.and_then(|s| serde_json::from_str::<Vec<DiskInfo>>(&s).ok()),
|
||||
Some(Ok(Message::Text(json))) => serde_json::from_str::<Vec<DiskInfo>>(&json).ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a "get_processes" request and await a ProcessesPayload decoded from protobuf (binary, may be gzipped)
|
||||
pub async fn request_processes(ws: &mut WsStream) -> Option<ProcessesPayload> {
|
||||
if ws
|
||||
.send(Message::Text("get_processes".into()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
match ws.next().await {
|
||||
Some(Ok(Message::Binary(b))) => {
|
||||
let gz = is_gzip(&b);
|
||||
let data = if gz { gunzip_to_vec(&b).ok()? } else { b };
|
||||
match pb::Processes::decode(data.as_slice()) {
|
||||
Ok(pb) => {
|
||||
let rows: Vec<ProcessInfo> = pb
|
||||
.rows
|
||||
.into_iter()
|
||||
.map(|p: pb::Process| ProcessInfo {
|
||||
pid: p.pid,
|
||||
name: p.name,
|
||||
cpu_usage: p.cpu_usage,
|
||||
mem_bytes: p.mem_bytes,
|
||||
})
|
||||
.collect();
|
||||
Some(ProcessesPayload {
|
||||
process_count: pb.process_count as usize,
|
||||
top_processes: rows,
|
||||
})
|
||||
}
|
||||
Err(e) => {
|
||||
if std::env::var("SOCKTOP_DEBUG").ok().as_deref() == Some("1") {
|
||||
eprintln!("protobuf decode failed: {e}");
|
||||
}
|
||||
// Fallback: maybe it's JSON (bytes already decompressed if gz)
|
||||
match String::from_utf8(data) {
|
||||
Ok(s) => serde_json::from_str::<ProcessesPayload>(&s).ok(),
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Text(json))) => serde_json::from_str::<ProcessesPayload>(&json).ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user