Implement GuruRMM Phase 1: Real-time tunnel infrastructure

Complete bidirectional tunnel communication between server and agents,
enabling persistent secure channels for future command execution and
file operations. Agents transition from heartbeat mode to tunnel mode
on-demand while maintaining WebSocket connection.

Server Implementation:
- Database layer (db/tunnel.rs): Session CRUD, ownership validation,
  cleanup on disconnect (prevents orphaned sessions)
- API endpoints (api/tunnel.rs): POST /open, POST /close, GET /status
  with JWT auth, UUID validation, proper HTTP status codes
- Protocol extension (ws/mod.rs): TunnelOpen/Close/Data messages,
  agent response handlers (TunnelReady/Data/Error)
- Migration (006_tunnel_sessions.sql): tech_sessions table with
  partial unique constraint, foreign keys with CASCADE, audit table

Agent Implementation:
- State machine (tunnel/mod.rs): AgentMode (Heartbeat ↔ Tunnel),
  channel multiplexing, concurrent session prevention
- WebSocket handlers (transport/websocket.rs): Open/close tunnel,
  mode switching without dropping connection, cleanup on disconnect
- Protocol extension (transport/mod.rs): TunnelReady/Data/Error
  messages matching server definitions
- Unit tests: Lifecycle and channel management coverage

Key Features:
- Security: JWT auth, session ownership verification, SQL injection
  prevention, constraint-based duplicate session blocking
- Cleanup: Automatic session closure on agent disconnect (both sides),
  channel cleanup, graceful state transitions
- Error handling: Proper HTTP status codes (400/403/404/409/500),
  comprehensive Result types, detailed logging
- Extensibility: Channel types ready (Terminal/File/Registry/Service),
  TunnelDataPayload enum for Phase 2+ expansion

Phase 1 Scope (Implemented):
- Tunnel session lifecycle management
- Mode switching (heartbeat ↔ tunnel)
- Protocol message routing
- Database session tracking

Phase 2 Next Steps:
- Terminal command execution (tokio::process::Command)
- Client WebSocket connections for output streaming
- Command audit logging
- File transfer operations

Verification:
- Server compiles successfully (0 errors)
- Agent unit tests pass (tunnel lifecycle, channel management)
- Code review approved (protocol alignment verified)
- Database constraints enforce referential integrity
- Cleanup tested (session closure on disconnect)

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
This commit is contained in:
2026-04-14 07:10:09 -07:00
parent 9940faf34a
commit 2e6d1a67dd
14 changed files with 2293 additions and 4 deletions

View File

@@ -38,6 +38,18 @@ pub enum AgentMessage {
/// Heartbeat to keep connection alive
Heartbeat,
/// Tunnel ready confirmation (agent → server)
TunnelReady { session_id: String },
/// Tunnel data (bidirectional)
TunnelData {
channel_id: String,
data: TunnelDataPayload,
},
/// Tunnel error (agent → server)
TunnelError { channel_id: String, error: String },
}
/// Authentication payload
@@ -157,6 +169,18 @@ pub enum ServerMessage {
/// Error message
Error { code: String, message: String },
/// Tunnel open request (server → agent)
TunnelOpen { session_id: String, tech_id: Uuid },
/// Tunnel close request (server → agent)
TunnelClose { session_id: String },
/// Tunnel data (bidirectional)
TunnelData {
channel_id: String,
data: TunnelDataPayload,
},
}
/// Authentication acknowledgment payload
@@ -311,3 +335,19 @@ pub enum UpdateStatus {
/// Rolled back to previous version
RolledBack,
}
/// Tunnel data payload types (Phase 1: Terminal only)
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
#[serde(rename_all = "snake_case")]
pub enum TunnelDataPayload {
/// Terminal command execution request (server → agent)
Terminal { command: String },
/// Terminal output response (agent → server)
TerminalOutput {
stdout: String,
stderr: String,
exit_code: Option<i32>,
},
}

View File

@@ -18,9 +18,10 @@ use tokio::time::{interval, timeout};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use tracing::{debug, error, info, warn};
use super::{AgentMessage, AuthPayload, CommandPayload, ServerMessage, UpdatePayload, UpdateResultPayload, UpdateStatus};
use super::{AgentMessage, AuthPayload, CommandPayload, ServerMessage, TunnelDataPayload, UpdatePayload, UpdateResultPayload, UpdateStatus};
use crate::claude::{ClaudeExecutor, ClaudeTaskCommand};
use crate::metrics::NetworkState;
use crate::tunnel::TunnelManager;
use crate::updater::{AgentUpdater, UpdaterConfig};
use crate::AppState;
@@ -203,6 +204,9 @@ impl WebSocketClient {
}
});
// Create tunnel manager for mode switching
let mut tunnel_manager = TunnelManager::new();
// Main message loop
let result: Result<()> = loop {
tokio::select! {
@@ -224,6 +228,15 @@ impl WebSocketClient {
AgentMessage::Heartbeat => {
debug!("Sent heartbeat");
}
AgentMessage::TunnelReady { session_id } => {
info!("Sent TunnelReady for session: {}", session_id);
}
AgentMessage::TunnelData { channel_id, .. } => {
debug!("Sent TunnelData on channel: {}", channel_id);
}
AgentMessage::TunnelError { channel_id, error } => {
warn!("Sent TunnelError on channel {}: {}", channel_id, error);
}
_ => {
debug!("Sent message: {:?}", std::mem::discriminant(&msg));
}
@@ -234,7 +247,7 @@ impl WebSocketClient {
Some(msg_result) = read.next() => {
match msg_result {
Ok(Message::Text(text)) => {
if let Err(e) = Self::handle_server_message(&text, &tx).await {
if let Err(e) = Self::handle_server_message(&text, &tx, &mut tunnel_manager).await {
error!("Error handling message: {}", e);
}
}
@@ -277,6 +290,9 @@ impl WebSocketClient {
heartbeat_task.abort();
*state.connected.write().await = false;
// Force close tunnel if active
tunnel_manager.force_close();
result
}
@@ -284,6 +300,7 @@ impl WebSocketClient {
async fn handle_server_message(
text: &str,
tx: &mpsc::Sender<AgentMessage>,
tunnel_manager: &mut TunnelManager,
) -> Result<()> {
let msg: ServerMessage =
serde_json::from_str(text).context("Failed to parse server message")?;
@@ -315,11 +332,107 @@ impl WebSocketClient {
);
Self::handle_update(payload, tx.clone()).await;
}
ServerMessage::TunnelOpen { session_id, tech_id } => {
info!(
"Received tunnel open request: session={}, tech={}",
session_id, tech_id
);
Self::handle_tunnel_open(session_id, tech_id, tunnel_manager, tx.clone()).await;
}
ServerMessage::TunnelClose { session_id } => {
info!("Received tunnel close request: session={}", session_id);
Self::handle_tunnel_close(session_id, tunnel_manager, tx.clone()).await;
}
ServerMessage::TunnelData { channel_id, data } => {
debug!("Received tunnel data on channel: {}", channel_id);
Self::handle_tunnel_data(channel_id, data, tunnel_manager, tx.clone()).await;
}
}
Ok(())
}
/// Handle tunnel open request
async fn handle_tunnel_open(
session_id: String,
tech_id: uuid::Uuid,
tunnel_manager: &mut TunnelManager,
tx: mpsc::Sender<AgentMessage>,
) {
match tunnel_manager.open_tunnel(session_id.clone(), tech_id) {
Ok(_) => {
info!("Tunnel opened successfully: {}", session_id);
// Send TunnelReady confirmation
let ready_msg = AgentMessage::TunnelReady {
session_id: session_id.clone(),
};
if let Err(e) = tx.send(ready_msg).await {
error!("Failed to send TunnelReady message: {}", e);
}
}
Err(e) => {
error!("Failed to open tunnel: {}", e);
// Send error back to server
let error_msg = AgentMessage::TunnelError {
channel_id: "system".to_string(),
error: format!("Failed to open tunnel: {}", e),
};
let _ = tx.send(error_msg).await;
}
}
}
/// Handle tunnel close request
async fn handle_tunnel_close(
session_id: String,
tunnel_manager: &mut TunnelManager,
tx: mpsc::Sender<AgentMessage>,
) {
match tunnel_manager.close_tunnel(&session_id) {
Ok(_) => {
info!("Tunnel closed successfully: {}", session_id);
}
Err(e) => {
warn!("Error closing tunnel: {}", e);
// Send error back to server
let error_msg = AgentMessage::TunnelError {
channel_id: "system".to_string(),
error: format!("Failed to close tunnel: {}", e),
};
let _ = tx.send(error_msg).await;
}
}
}
/// Handle tunnel data (Phase 1: Terminal commands only)
async fn handle_tunnel_data(
channel_id: String,
data: TunnelDataPayload,
_tunnel_manager: &TunnelManager,
tx: mpsc::Sender<AgentMessage>,
) {
match data {
TunnelDataPayload::Terminal { command } => {
info!("Terminal command on channel {}: {}", channel_id, command);
// Phase 1: Just log and respond with placeholder
// Phase 2 will implement actual command execution
let response = AgentMessage::TunnelData {
channel_id,
data: TunnelDataPayload::TerminalOutput {
stdout: String::new(),
stderr: "Terminal execution not yet implemented (Phase 2)".to_string(),
exit_code: Some(-1),
},
};
let _ = tx.send(response).await;
}
TunnelDataPayload::TerminalOutput { .. } => {
// This shouldn't be sent to the agent, it's agent → server only
warn!("Received TerminalOutput on agent (unexpected)");
}
}
}
/// Handle an update command from the server
async fn handle_update(payload: UpdatePayload, tx: mpsc::Sender<AgentMessage>) {
// Send starting status