jason/socktop
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reference: Usage as a lib #8

Browse Source 
- Implement protocol versioning
- migrate to thisError
- general error handling improvements in socktop_connector lib
- improve documentation
- increment  version
This commit is contained in:
jasonwitty
2025-09-07 18:55:23 -07:00
parent ffc246b705
commit 06cd6d0c82
9 changed files with 462 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
socktop_connector/src/connector.rs
+142 -43
View File
@@ -17,6 +17,7 @@ use tokio_tungstenite::{
};
use url::Url;
use crate::error::{ConnectorError, Result};
use crate::types::{AgentRequest, AgentResponse, DiskInfo, Metrics, ProcessInfo, ProcessesPayload};
#[cfg(feature = "tls")]
@@ -41,6 +42,8 @@ pub struct ConnectorConfig {
pub url: String,
pub tls_ca_path: Option<String>,
pub verify_hostname: bool,
pub ws_protocols: Option<Vec<String>>,
pub ws_version: Option<String>,
}
impl ConnectorConfig {
@@ -49,6 +52,8 @@ impl ConnectorConfig {
url: url.into(),
tls_ca_path: None,
verify_hostname: false,
ws_protocols: None,
ws_version: None,
}
}
@@ -61,6 +66,18 @@ impl ConnectorConfig {
self.verify_hostname = verify;
self
}
/// Set WebSocket sub-protocols to negotiate
pub fn with_protocols(mut self, protocols: Vec<String>) -> Self {
self.ws_protocols = Some(protocols);
self
}
/// Set WebSocket protocol version (default is "13")
pub fn with_version(mut self, version: impl Into<String>) -> Self {
self.ws_version = Some(version.into());
self
}
}
/// A WebSocket connector for communicating with socktop agents
@@ -79,39 +96,33 @@ impl SocktopConnector {
}
/// Connect to the agent
pub async fn connect(&mut self) -> Result<(), Box<dyn std::error::Error>> {
let stream = connect_to_agent(
&self.config.url,
self.config.tls_ca_path.as_deref(),
self.config.verify_hostname,
)
.await?;
pub async fn connect(&mut self) -> Result<()> {
let stream = connect_to_agent(&self.config).await?;
self.stream = Some(stream);
Ok(())
}
/// Send a request to the agent and get the response
pub async fn request(
&mut self,
request: AgentRequest,
) -> Result<AgentResponse, Box<dyn std::error::Error>> {
let stream = self.stream.as_mut().ok_or("Not connected")?;
pub async fn request(&mut self, request: AgentRequest) -> Result<AgentResponse> {
let stream = self.stream.as_mut().ok_or(ConnectorError::NotConnected)?;
match request {
AgentRequest::Metrics => {
let metrics = request_metrics(stream)
.await
.ok_or("Failed to get metrics")?;
.ok_or_else(|| ConnectorError::invalid_response("Failed to get metrics"))?;
Ok(AgentResponse::Metrics(metrics))
}
AgentRequest::Disks => {
let disks = request_disks(stream).await.ok_or("Failed to get disks")?;
let disks = request_disks(stream)
.await
.ok_or_else(|| ConnectorError::invalid_response("Failed to get disks"))?;
Ok(AgentResponse::Disks(disks))
}
AgentRequest::Processes => {
let processes = request_processes(stream)
.await
.ok_or("Failed to get processes")?;
.ok_or_else(|| ConnectorError::invalid_response("Failed to get processes"))?;
Ok(AgentResponse::Processes(processes))
}
}
@@ -123,7 +134,7 @@ impl SocktopConnector {
}
/// Disconnect from the agent
pub async fn disconnect(&mut self) -> Result<(), Box<dyn std::error::Error>> {
pub async fn disconnect(&mut self) -> Result<()> {
if let Some(mut stream) = self.stream.take() {
let _ = stream.close(None).await;
}
@@ -132,32 +143,54 @@ impl SocktopConnector {
}
// Connect to the agent and return the WS stream
async fn connect_to_agent(
url: &str,
tls_ca: Option<&str>,
verify_hostname: bool,
) -> Result<WsStream, Box<dyn std::error::Error>> {
async fn connect_to_agent(config: &ConnectorConfig) -> Result<WsStream> {
#[cfg(feature = "tls")]
ensure_crypto_provider();
let mut u = Url::parse(url)?;
if let Some(ca_path) = tls_ca {
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(u.as_str(), ca_path, verify_hostname).await;
return connect_with_ca_and_config(u.as_str(), ca_path, config).await;
}
// No TLS - hostname verification is not applicable
let (ws, _) = connect_async(u.as_str()).await?;
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(
async fn connect_with_ca_and_config(
url: &str,
ca_path: &str,
verify_hostname: bool,
) -> Result<WsStream, Box<dyn std::error::Error>> {
config: &ConnectorConfig,
) -> Result<WsStream> {
// Initialize the crypto provider for rustls
let _ = rustls::crypto::ring::default_provider().install_default();
@@ -175,8 +208,29 @@ async fn connect_with_ca(
.with_root_certificates(root)
.with_no_client_auth();
let req = url.into_client_request()?;
if !verify_hostname {
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 {
@@ -187,7 +241,7 @@ async fn connect_with_ca(
_server_name: &ServerName,
_ocsp_response: &[u8],
_now: UnixTime,
) -> Result<ServerCertVerified, rustls::Error> {
) -> std::result::Result<ServerCertVerified, rustls::Error> {
Ok(ServerCertVerified::assertion())
}
fn verify_tls12_signature(
@@ -195,7 +249,7 @@ async fn connect_with_ca(
_message: &[u8],
_cert: &CertificateDer<'_>,
_dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn verify_tls13_signature(
@@ -203,7 +257,7 @@ async fn connect_with_ca(
_message: &[u8],
_cert: &CertificateDer<'_>,
_dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
@@ -220,18 +274,26 @@ async fn connect_with_ca(
);
}
let cfg = Arc::new(cfg);
let (ws, _) =
connect_async_tls_with_config(req, None, verify_hostname, Some(Connector::Rustls(cfg)))
.await?;
let (ws, _) = 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(
async fn connect_with_ca_and_config(
_url: &str,
_ca_path: &str,
) -> Result<WsStream, Box<dyn std::error::Error>> {
Err("TLS support not compiled in".into())
_config: &ConnectorConfig,
) -> Result<WsStream> {
Err(ConnectorError::tls_error(
"TLS support not compiled in",
std::io::Error::new(std::io::ErrorKind::Unsupported, "TLS not available"),
))
}
// Send a "get_metrics" request and await a single JSON reply
@@ -334,9 +396,7 @@ fn is_gzip(bytes: &[u8]) -> bool {
/// certificate involved, hostname verification is not applicable.
///
/// For TLS connections with certificate pinning, use `connect_to_socktop_agent_with_tls()`.
pub async fn connect_to_socktop_agent(
url: impl Into<String>,
) -> Result<SocktopConnector, Box<dyn std::error::Error>> {
pub async fn connect_to_socktop_agent(url: impl Into<String>) -> Result<SocktopConnector> {
let config = ConnectorConfig::new(url);
let mut connector = SocktopConnector::new(config);
connector.connect().await?;
@@ -354,7 +414,7 @@ pub async fn connect_to_socktop_agent_with_tls(
url: impl Into<String>,
ca_path: impl Into<String>,
verify_hostname: bool,
) -> Result<SocktopConnector, Box<dyn std::error::Error>> {
) -> Result<SocktopConnector> {
let config = ConnectorConfig::new(url)
.with_tls_ca(ca_path)
.with_hostname_verification(verify_hostname);
@@ -362,3 +422,42 @@ pub async fn connect_to_socktop_agent_with_tls(
connector.connect().await?;
Ok(connector)
}
/// Convenience function to create a connector with custom WebSocket protocol configuration.
///
/// This function allows you to specify WebSocket protocol version and sub-protocols.
/// Most users should use the simpler `connect_to_socktop_agent()` function instead.
///
/// # Example
/// ```no_run
/// use socktop_connector::connect_to_socktop_agent_with_config;
///
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let connector = connect_to_socktop_agent_with_config(
/// "ws://localhost:3000/ws",
/// Some(vec!["socktop".to_string()]), // WebSocket sub-protocols
/// Some("13".to_string()), // WebSocket version (13 is standard)
/// ).await?;
/// # Ok(())
/// # }
/// ```
pub async fn connect_to_socktop_agent_with_config(
url: impl Into<String>,
protocols: Option<Vec<String>>,
version: Option<String>,
) -> Result<SocktopConnector> {
let mut config = ConnectorConfig::new(url);
if let Some(protocols) = protocols {
config = config.with_protocols(protocols);
}
if let Some(version) = version {
config = config.with_version(version);
}
let mut connector = SocktopConnector::new(config);
connector.connect().await?;
Ok(connector)
}
socktop_connector/src/error.rs
+136
View File
@@ -0,0 +1,136 @@
//! Error types for socktop_connector
use thiserror::Error;
/// Errors that can occur when using socktop_connector
#[derive(Error, Debug)]
pub enum ConnectorError {
/// WebSocket connection failed
#[error("WebSocket connection failed: {source}")]
ConnectionFailed {
#[from]
source: tokio_tungstenite::tungstenite::Error,
},
/// URL parsing error
#[error("Invalid URL: {url}")]
InvalidUrl {
url: String,
#[source]
source: url::ParseError,
},
/// TLS certificate error
#[error("TLS certificate error: {message}")]
TlsError {
message: String,
#[source]
source: Box<dyn std::error::Error + Send + Sync>,
},
/// Certificate file not found or invalid
#[error("Certificate file error at '{path}': {message}")]
CertificateError { path: String, message: String },
/// Invalid server response format
#[error("Invalid response from server: {message}")]
InvalidResponse { message: String },
/// JSON parsing error
#[error("JSON parsing error: {source}")]
JsonError {
#[from]
source: serde_json::Error,
},
/// Request/response protocol error
#[error("Protocol error: {message}")]
ProtocolError { message: String },
/// Connection is not established
#[error("Not connected to server")]
NotConnected,
/// Connection was closed unexpectedly
#[error("Connection closed: {reason}")]
ConnectionClosed { reason: String },
/// IO error (network, file system, etc.)
#[error("IO error: {source}")]
IoError {
#[from]
source: std::io::Error,
},
/// Compression/decompression error
#[error("Compression error: {message}")]
CompressionError { message: String },
/// Protocol Buffer parsing error
#[error("Protocol buffer error: {source}")]
ProtobufError {
#[from]
source: prost::DecodeError,
},
}
/// Result type alias for connector operations
pub type Result<T> = std::result::Result<T, ConnectorError>;
impl ConnectorError {
/// Create a TLS error with context
pub fn tls_error(
message: impl Into<String>,
source: impl std::error::Error + Send + Sync + 'static,
) -> Self {
Self::TlsError {
message: message.into(),
source: Box::new(source),
}
}
/// Create a certificate error
pub fn certificate_error(path: impl Into<String>, message: impl Into<String>) -> Self {
Self::CertificateError {
path: path.into(),
message: message.into(),
}
}
/// Create a protocol error
pub fn protocol_error(message: impl Into<String>) -> Self {
Self::ProtocolError {
message: message.into(),
}
}
/// Create an invalid response error
pub fn invalid_response(message: impl Into<String>) -> Self {
Self::InvalidResponse {
message: message.into(),
}
}
/// Create a connection closed error
pub fn connection_closed(reason: impl Into<String>) -> Self {
Self::ConnectionClosed {
reason: reason.into(),
}
}
/// Create a compression error
pub fn compression_error(message: impl Into<String>) -> Self {
Self::CompressionError {
message: message.into(),
}
}
}
impl From<url::ParseError> for ConnectorError {
fn from(source: url::ParseError) -> Self {
Self::InvalidUrl {
url: "unknown".to_string(), // We don't have the URL in the error context
source,
}
}
}
socktop_connector/src/lib.rs
+3 -4
View File
@@ -141,16 +141,15 @@
#![cfg_attr(docsrs, feature(doc_cfg))]
pub mod connector;
pub mod error;
pub mod types;
pub use connector::{
ConnectorConfig, SocktopConnector, WsStream, connect_to_socktop_agent,
connect_to_socktop_agent_with_tls,
connect_to_socktop_agent_with_config, connect_to_socktop_agent_with_tls,
};
pub use error::{ConnectorError, Result};
pub use types::{
AgentRequest, AgentResponse, DiskInfo, GpuInfo, Metrics, NetworkInfo, ProcessInfo,
ProcessesPayload,
};
/// Re-export commonly used error type
pub use anyhow::Error;