feat: 初始化项目基础架构和核心功能
- 添加项目基础结构：Cargo.toml、.gitignore、设备UID和密钥文件 - 实现前端Vue3项目结构：路由、登录页面、设备管理页面 - 添加核心协议定义(crates/protocol)：设备状态、资产、USB事件等 - 实现客户端监控模块：系统状态收集、资产收集 - 实现服务端基础API和插件系统 - 添加数据库迁移脚本：设备管理、资产跟踪、告警系统等 - 实现前端设备状态展示和基本交互 - 添加使用时长统计和水印功能插件
This commit is contained in:
iven
380
crates/client/src/network/mod.rs
Normal file
380
crates/client/src/network/mod.rs
Normal file
@@ -0,0 +1,380 @@
|
||||
use anyhow::Result;
|
||||
use std::time::Duration;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tracing::{info, debug, warn};
|
||||
use csm_protocol::{Frame, MessageType, RegisterRequest, RegisterResponse, HeartbeatPayload, WatermarkConfigPayload, UsbPolicyPayload};
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
|
||||
use crate::ClientState;
|
||||
|
||||
/// Holds senders for all plugin config channels
|
||||
pub struct PluginChannels {
|
||||
pub watermark_tx: tokio::sync::watch::Sender<Option<WatermarkConfigPayload>>,
|
||||
pub web_filter_tx: tokio::sync::watch::Sender<crate::web_filter::WebFilterConfig>,
|
||||
pub software_blocker_tx: tokio::sync::watch::Sender<crate::software_blocker::SoftwareBlockerConfig>,
|
||||
pub popup_blocker_tx: tokio::sync::watch::Sender<crate::popup_blocker::PopupBlockerConfig>,
|
||||
pub usb_audit_tx: tokio::sync::watch::Sender<crate::usb_audit::UsbAuditConfig>,
|
||||
pub usage_timer_tx: tokio::sync::watch::Sender<crate::usage_timer::UsageConfig>,
|
||||
pub usb_policy_tx: tokio::sync::watch::Sender<Option<UsbPolicyPayload>>,
|
||||
}
|
||||
|
||||
/// Connect to server and run the main communication loop
|
||||
pub async fn connect_and_run(
|
||||
state: &ClientState,
|
||||
data_rx: &mut tokio::sync::mpsc::Receiver<Frame>,
|
||||
plugins: &PluginChannels,
|
||||
) -> Result<()> {
|
||||
let tcp_stream = TcpStream::connect(&state.server_addr).await?;
|
||||
info!("TCP connected to {}", state.server_addr);
|
||||
|
||||
if state.use_tls {
|
||||
let tls_stream = wrap_tls(tcp_stream, &state.server_addr).await?;
|
||||
run_comm_loop(tls_stream, state, data_rx, plugins).await
|
||||
} else {
|
||||
run_comm_loop(tcp_stream, state, data_rx, plugins).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrap a TCP stream with TLS.
|
||||
/// Supports custom CA certificate via CSM_TLS_CA_CERT env var (path to PEM file).
|
||||
/// Supports skipping verification via CSM_TLS_SKIP_VERIFY=true (development only).
|
||||
async fn wrap_tls(stream: TcpStream, server_addr: &str) -> Result<tokio_rustls::client::TlsStream<TcpStream>> {
|
||||
let mut root_store = rustls::RootCertStore::empty();
|
||||
|
||||
// Load custom CA certificate if specified
|
||||
if let Ok(ca_path) = std::env::var("CSM_TLS_CA_CERT") {
|
||||
let ca_pem = std::fs::read(&ca_path)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to read CA cert {}: {}", ca_path, e))?;
|
||||
let certs = rustls_pemfile::certs(&mut &ca_pem[..])
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to parse CA cert: {:?}", e))?;
|
||||
for cert in certs {
|
||||
root_store.add(cert)?;
|
||||
}
|
||||
info!("Loaded custom CA certificates from {}", ca_path);
|
||||
}
|
||||
|
||||
// Always include system roots as fallback
|
||||
root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
|
||||
let config = if std::env::var("CSM_TLS_SKIP_VERIFY").as_deref() == Ok("true") {
|
||||
warn!("TLS certificate verification DISABLED — do not use in production!");
|
||||
rustls::ClientConfig::builder()
|
||||
.dangerous()
|
||||
.with_custom_certificate_verifier(std::sync::Arc::new(NoVerifier))
|
||||
.with_no_client_auth()
|
||||
} else {
|
||||
rustls::ClientConfig::builder()
|
||||
.with_root_certificates(root_store)
|
||||
.with_no_client_auth()
|
||||
};
|
||||
|
||||
let connector = tokio_rustls::TlsConnector::from(std::sync::Arc::new(config));
|
||||
|
||||
// Extract hostname from server_addr (strip port)
|
||||
let domain = server_addr.split(':').next().unwrap_or("localhost").to_string();
|
||||
let server_name = rustls_pki_types::ServerName::try_from(domain.clone())
|
||||
.map_err(|e| anyhow::anyhow!("Invalid TLS server name '{}': {:?}", domain, e))?;
|
||||
|
||||
let tls_stream = connector.connect(server_name, stream).await?;
|
||||
info!("TLS handshake completed with {}", domain);
|
||||
Ok(tls_stream)
|
||||
}
|
||||
|
||||
/// A no-op certificate verifier for development use (CSM_TLS_SKIP_VERIFY=true).
|
||||
#[derive(Debug)]
|
||||
struct NoVerifier;
|
||||
|
||||
impl rustls::client::danger::ServerCertVerifier for NoVerifier {
|
||||
fn verify_server_cert(
|
||||
&self,
|
||||
_end_entity: &rustls_pki_types::CertificateDer,
|
||||
_intermediates: &[rustls_pki_types::CertificateDer],
|
||||
_server_name: &rustls_pki_types::ServerName,
|
||||
_ocsp_response: &[u8],
|
||||
_now: rustls_pki_types::UnixTime,
|
||||
) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
|
||||
Ok(rustls::client::danger::ServerCertVerified::assertion())
|
||||
}
|
||||
|
||||
fn verify_tls12_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &rustls_pki_types::CertificateDer,
|
||||
_dss: &rustls::DigitallySignedStruct,
|
||||
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
|
||||
}
|
||||
|
||||
fn verify_tls13_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &rustls_pki_types::CertificateDer,
|
||||
_dss: &rustls::DigitallySignedStruct,
|
||||
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
|
||||
}
|
||||
|
||||
fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
|
||||
vec![
|
||||
rustls::SignatureScheme::RSA_PKCS1_SHA256,
|
||||
rustls::SignatureScheme::RSA_PKCS1_SHA384,
|
||||
rustls::SignatureScheme::RSA_PKCS1_SHA512,
|
||||
rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
|
||||
rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
|
||||
rustls::SignatureScheme::RSA_PSS_SHA256,
|
||||
rustls::SignatureScheme::RSA_PSS_SHA384,
|
||||
rustls::SignatureScheme::RSA_PSS_SHA512,
|
||||
rustls::SignatureScheme::ED25519,
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
/// Main communication loop over any read+write stream
|
||||
async fn run_comm_loop<S>(
|
||||
mut stream: S,
|
||||
state: &ClientState,
|
||||
data_rx: &mut tokio::sync::mpsc::Receiver<Frame>,
|
||||
plugins: &PluginChannels,
|
||||
) -> Result<()>
|
||||
where
|
||||
S: AsyncReadExt + AsyncWriteExt + Unpin,
|
||||
{
|
||||
// Send registration
|
||||
let register = RegisterRequest {
|
||||
device_uid: state.device_uid.clone(),
|
||||
hostname: hostname::get()
|
||||
.map(|h| h.to_string_lossy().to_string())
|
||||
.unwrap_or_else(|_| "unknown".to_string()),
|
||||
registration_token: state.registration_token.clone(),
|
||||
os_version: get_os_info(),
|
||||
mac_address: None,
|
||||
};
|
||||
|
||||
let frame = Frame::new_json(MessageType::Register, ®ister)?;
|
||||
stream.write_all(&frame.encode()).await?;
|
||||
info!("Registration request sent");
|
||||
|
||||
let mut buffer = vec![0u8; 65536];
|
||||
let mut read_buf = Vec::with_capacity(65536);
|
||||
// Clamp heartbeat interval to sane range [5, 3600] to prevent CPU spin or effective disable
|
||||
let heartbeat_secs = state.config.heartbeat_interval_secs.clamp(5, 3600);
|
||||
let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(heartbeat_secs));
|
||||
heartbeat_interval.tick().await; // Skip first tick
|
||||
|
||||
// HMAC key — set after receiving RegisterResponse
|
||||
let mut device_secret: Option<String> = state.device_secret.clone();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// Read from server
|
||||
result = stream.read(&mut buffer) => {
|
||||
let n = result?;
|
||||
if n == 0 {
|
||||
return Err(anyhow::anyhow!("Server closed connection"));
|
||||
}
|
||||
read_buf.extend_from_slice(&buffer[..n]);
|
||||
|
||||
// Guard against unbounded buffer growth from a malicious server
|
||||
if read_buf.len() > 1_048_576 {
|
||||
return Err(anyhow::anyhow!("Read buffer exceeded 1MB, server may be malicious"));
|
||||
}
|
||||
|
||||
// Process complete frames
|
||||
loop {
|
||||
match Frame::decode(&read_buf)? {
|
||||
Some(frame) => {
|
||||
let consumed = frame.encoded_size();
|
||||
read_buf.drain(..consumed);
|
||||
// Capture device_secret from registration response
|
||||
if frame.msg_type == MessageType::RegisterResponse {
|
||||
if let Ok(resp) = frame.decode_payload::<RegisterResponse>() {
|
||||
device_secret = Some(resp.device_secret.clone());
|
||||
crate::save_device_secret(&resp.device_secret);
|
||||
info!("Device secret received and persisted, HMAC enabled for heartbeats");
|
||||
}
|
||||
}
|
||||
handle_server_message(frame, plugins)?;
|
||||
}
|
||||
None => break, // Incomplete frame, wait for more data
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Send queued data
|
||||
frame = data_rx.recv() => {
|
||||
let frame = frame.ok_or_else(|| anyhow::anyhow!("Channel closed"))?;
|
||||
stream.write_all(&frame.encode()).await?;
|
||||
}
|
||||
|
||||
// Heartbeat
|
||||
_ = heartbeat_interval.tick() => {
|
||||
let timestamp = chrono::Utc::now().to_rfc3339();
|
||||
let hmac_value = compute_hmac(device_secret.as_deref(), &state.device_uid, ×tamp);
|
||||
let heartbeat = HeartbeatPayload {
|
||||
device_uid: state.device_uid.clone(),
|
||||
timestamp,
|
||||
hmac: hmac_value,
|
||||
};
|
||||
let frame = Frame::new_json(MessageType::Heartbeat, &heartbeat)?;
|
||||
stream.write_all(&frame.encode()).await?;
|
||||
debug!("Heartbeat sent (hmac={})", !heartbeat.hmac.is_empty());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_server_message(frame: Frame, plugins: &PluginChannels) -> Result<()> {
|
||||
match frame.msg_type {
|
||||
MessageType::RegisterResponse => {
|
||||
let resp: RegisterResponse = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid registration response: {}", e))?;
|
||||
info!("Registration accepted by server (server version: {})", resp.config.server_version);
|
||||
}
|
||||
MessageType::PolicyUpdate => {
|
||||
let policy: serde_json::Value = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid policy update: {}", e))?;
|
||||
info!("Received policy update: {}", policy);
|
||||
}
|
||||
MessageType::ConfigUpdate => {
|
||||
info!("Received config update");
|
||||
}
|
||||
MessageType::TaskExecute => {
|
||||
warn!("Task execution requested (not yet implemented)");
|
||||
}
|
||||
MessageType::WatermarkConfig => {
|
||||
let config: WatermarkConfigPayload = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid watermark config: {}", e))?;
|
||||
info!("Received watermark config: enabled={}", config.enabled);
|
||||
plugins.watermark_tx.send(Some(config))?;
|
||||
}
|
||||
MessageType::UsbPolicyUpdate => {
|
||||
let policy: UsbPolicyPayload = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid USB policy: {}", e))?;
|
||||
info!("Received USB policy: type={}, enabled={}", policy.policy_type, policy.enabled);
|
||||
plugins.usb_policy_tx.send(Some(policy))?;
|
||||
}
|
||||
MessageType::WebFilterRuleUpdate => {
|
||||
let payload: serde_json::Value = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid web filter update: {}", e))?;
|
||||
info!("Received web filter rules update");
|
||||
let rules: Vec<crate::web_filter::WebFilterRule> = payload.get("rules")
|
||||
.and_then(|r| serde_json::from_value(r.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let config = crate::web_filter::WebFilterConfig { enabled: true, rules };
|
||||
plugins.web_filter_tx.send(config)?;
|
||||
}
|
||||
MessageType::SoftwareBlacklist => {
|
||||
let payload: serde_json::Value = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid software blacklist: {}", e))?;
|
||||
info!("Received software blacklist update");
|
||||
let blacklist: Vec<crate::software_blocker::BlacklistEntry> = payload.get("blacklist")
|
||||
.and_then(|r| serde_json::from_value(r.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let config = crate::software_blocker::SoftwareBlockerConfig { enabled: true, blacklist };
|
||||
plugins.software_blocker_tx.send(config)?;
|
||||
}
|
||||
MessageType::PopupRules => {
|
||||
let payload: serde_json::Value = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid popup rules: {}", e))?;
|
||||
info!("Received popup blocker rules update");
|
||||
let rules: Vec<crate::popup_blocker::PopupRule> = payload.get("rules")
|
||||
.and_then(|r| serde_json::from_value(r.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
let config = crate::popup_blocker::PopupBlockerConfig { enabled: true, rules };
|
||||
plugins.popup_blocker_tx.send(config)?;
|
||||
}
|
||||
MessageType::PluginEnable => {
|
||||
let payload: csm_protocol::PluginControlPayload = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid plugin enable: {}", e))?;
|
||||
info!("Plugin enabled: {}", payload.plugin_name);
|
||||
// Route to appropriate plugin channel based on plugin_name
|
||||
handle_plugin_control(&payload, plugins, true)?;
|
||||
}
|
||||
MessageType::PluginDisable => {
|
||||
let payload: csm_protocol::PluginControlPayload = frame.decode_payload()
|
||||
.map_err(|e| anyhow::anyhow!("Invalid plugin disable: {}", e))?;
|
||||
info!("Plugin disabled: {}", payload.plugin_name);
|
||||
handle_plugin_control(&payload, plugins, false)?;
|
||||
}
|
||||
_ => {
|
||||
debug!("Unhandled message type: {:?}", frame.msg_type);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_plugin_control(
|
||||
payload: &csm_protocol::PluginControlPayload,
|
||||
plugins: &PluginChannels,
|
||||
enabled: bool,
|
||||
) -> Result<()> {
|
||||
match payload.plugin_name.as_str() {
|
||||
"watermark" => {
|
||||
if !enabled {
|
||||
// Send disabled config to remove overlay
|
||||
plugins.watermark_tx.send(None)?;
|
||||
}
|
||||
// When enabling, server will push the actual config next
|
||||
}
|
||||
"web_filter" => {
|
||||
if !enabled {
|
||||
// Clear hosts rules on disable
|
||||
plugins.web_filter_tx.send(crate::web_filter::WebFilterConfig { enabled: false, rules: vec![] })?;
|
||||
}
|
||||
// When enabling, server will push rules
|
||||
}
|
||||
"software_blocker" => {
|
||||
if !enabled {
|
||||
plugins.software_blocker_tx.send(crate::software_blocker::SoftwareBlockerConfig { enabled: false, blacklist: vec![] })?;
|
||||
}
|
||||
}
|
||||
"popup_blocker" => {
|
||||
if !enabled {
|
||||
plugins.popup_blocker_tx.send(crate::popup_blocker::PopupBlockerConfig { enabled: false, rules: vec![] })?;
|
||||
}
|
||||
}
|
||||
"usb_audit" => {
|
||||
if !enabled {
|
||||
plugins.usb_audit_tx.send(crate::usb_audit::UsbAuditConfig { enabled: false, monitored_extensions: vec![] })?;
|
||||
}
|
||||
}
|
||||
"usage_timer" => {
|
||||
if !enabled {
|
||||
plugins.usage_timer_tx.send(crate::usage_timer::UsageConfig { enabled: false, ..Default::default() })?;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
warn!("Unknown plugin: {}", payload.plugin_name);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Compute HMAC-SHA256 for heartbeat verification.
|
||||
/// Format: HMAC-SHA256(device_secret, "{device_uid}\n{timestamp}")
|
||||
fn compute_hmac(secret: Option<&str>, device_uid: &str, timestamp: &str) -> String {
|
||||
let secret = match secret {
|
||||
Some(s) if !s.is_empty() => s,
|
||||
_ => return String::new(),
|
||||
};
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
let message = format!("{}\n{}", device_uid, timestamp);
|
||||
let mut mac = match HmacSha256::new_from_slice(secret.as_bytes()) {
|
||||
Ok(m) => m,
|
||||
Err(_) => return String::new(),
|
||||
};
|
||||
mac.update(message.as_bytes());
|
||||
hex::encode(mac.finalize().into_bytes())
|
||||
}
|
||||
|
||||
fn get_os_info() -> String {
|
||||
use sysinfo::System;
|
||||
let name = System::name().unwrap_or_else(|| "Unknown".to_string());
|
||||
let version = System::os_version().unwrap_or_else(|| "Unknown".to_string());
|
||||
format!("{} {}", name, version)
|
||||
}
|
||||
Reference in New Issue
Block a user