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:
@@ -13,6 +13,7 @@ pub mod clients;
|
||||
pub mod commands;
|
||||
pub mod metrics;
|
||||
pub mod sites;
|
||||
pub mod tunnel;
|
||||
|
||||
use axum::{
|
||||
routing::{delete, get, post, put},
|
||||
@@ -63,4 +64,8 @@ pub fn routes() -> Router<AppState> {
|
||||
.route("/agent/register-legacy", post(agents::register_legacy))
|
||||
.route("/agent/heartbeat", post(agents::heartbeat))
|
||||
.route("/agent/command-result", post(agents::command_result))
|
||||
// Tunnel management
|
||||
.route("/v1/tunnel/open", post(tunnel::open_tunnel))
|
||||
.route("/v1/tunnel/close", post(tunnel::close_tunnel))
|
||||
.route("/v1/tunnel/status/:session_id", get(tunnel::get_tunnel_status))
|
||||
}
|
||||
|
||||
231
projects/msp-tools/guru-rmm/server/src/api/tunnel.rs
Normal file
231
projects/msp-tools/guru-rmm/server/src/api/tunnel.rs
Normal file
@@ -0,0 +1,231 @@
|
||||
//! Tunnel session management endpoints
|
||||
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{error, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::auth::AuthUser;
|
||||
use crate::db;
|
||||
use crate::ws::ServerMessage;
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct OpenTunnelRequest {
|
||||
pub agent_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct OpenTunnelResponse {
|
||||
pub session_id: String,
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CloseTunnelRequest {
|
||||
pub session_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct CloseTunnelResponse {
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct TunnelStatusResponse {
|
||||
pub session_id: String,
|
||||
pub agent_id: String,
|
||||
pub status: String,
|
||||
pub opened_at: String,
|
||||
pub last_activity: String,
|
||||
}
|
||||
|
||||
/// POST /api/v1/tunnel/open
|
||||
/// Open a new tunnel session to an agent
|
||||
pub async fn open_tunnel(
|
||||
State(state): State<AppState>,
|
||||
user: AuthUser,
|
||||
Json(req): Json<OpenTunnelRequest>,
|
||||
) -> Result<Json<OpenTunnelResponse>, (StatusCode, String)> {
|
||||
// Parse agent_id
|
||||
let agent_id = Uuid::parse_str(&req.agent_id)
|
||||
.map_err(|_| (StatusCode::BAD_REQUEST, "Invalid agent_id format".to_string()))?;
|
||||
|
||||
// Check if agent exists and is online
|
||||
let agent_connected = state.agents.read().await.is_connected(&agent_id);
|
||||
if !agent_connected {
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
"Agent not connected".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Check for existing active session
|
||||
let has_session = db::has_active_session(&state.db, user.user_id, agent_id)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
|
||||
|
||||
if has_session {
|
||||
return Err((
|
||||
StatusCode::CONFLICT,
|
||||
"Active session already exists for this agent".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Generate session ID
|
||||
let session_id = Uuid::new_v4().to_string();
|
||||
|
||||
// Create session in database
|
||||
let _session = db::create_tech_session(&state.db, &session_id, user.user_id, agent_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
// Handle unique constraint violation (PostgreSQL error code 23505)
|
||||
if let Some(db_err) = e.as_database_error() {
|
||||
if db_err.code().as_deref() == Some("23505") {
|
||||
return (
|
||||
StatusCode::CONFLICT,
|
||||
"Active session already exists for this agent".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
error!("Failed to create tunnel session: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
|
||||
})?;
|
||||
|
||||
// Send TunnelOpen message to agent via WebSocket
|
||||
let tunnel_open_msg = ServerMessage::TunnelOpen {
|
||||
session_id: session_id.clone(),
|
||||
tech_id: user.user_id,
|
||||
};
|
||||
|
||||
let sent = state.agents.read().await.send_to(&agent_id, tunnel_open_msg).await;
|
||||
if !sent {
|
||||
// Clean up database session if send failed
|
||||
if let Err(e) = db::close_tech_session(&state.db, &session_id).await {
|
||||
error!("Failed to cleanup session {} after send failure: {}", session_id, e);
|
||||
}
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"Failed to send tunnel open message to agent".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Json(OpenTunnelResponse {
|
||||
session_id,
|
||||
status: "active".to_string(),
|
||||
}))
|
||||
}
|
||||
|
||||
/// POST /api/v1/tunnel/close
|
||||
/// Close an existing tunnel session
|
||||
pub async fn close_tunnel(
|
||||
State(state): State<AppState>,
|
||||
user: AuthUser,
|
||||
Json(req): Json<CloseTunnelRequest>,
|
||||
) -> Result<Json<CloseTunnelResponse>, (StatusCode, String)> {
|
||||
// Validate session_id format
|
||||
if Uuid::parse_str(&req.session_id).is_err() {
|
||||
return Err((StatusCode::BAD_REQUEST, "Invalid session_id format".to_string()));
|
||||
}
|
||||
|
||||
// Verify session ownership
|
||||
let is_owner = db::verify_session_ownership(&state.db, &req.session_id, user.user_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("Failed to verify session ownership: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
|
||||
})?;
|
||||
|
||||
if !is_owner {
|
||||
return Err((
|
||||
StatusCode::FORBIDDEN,
|
||||
"Session not found or not owned by user".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Get session to find agent_id
|
||||
let session = db::get_tech_session(&state.db, &req.session_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("Failed to get session: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
|
||||
})?
|
||||
.ok_or((StatusCode::NOT_FOUND, "Session not found".to_string()))?;
|
||||
|
||||
// Send TunnelClose message to agent
|
||||
let tunnel_close_msg = ServerMessage::TunnelClose {
|
||||
session_id: req.session_id.clone(),
|
||||
};
|
||||
|
||||
if !state.agents.read().await.send_to(&session.agent_id, tunnel_close_msg).await {
|
||||
warn!(
|
||||
"Failed to send TunnelClose message to agent {} for session {}",
|
||||
session.agent_id, req.session_id
|
||||
);
|
||||
}
|
||||
|
||||
// Close session in database
|
||||
match db::close_tech_session(&state.db, &req.session_id).await {
|
||||
Ok(rows) if rows == 0 => {
|
||||
warn!("No rows updated when closing session {}", req.session_id);
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
error!("Failed to close session in database: {}", e);
|
||||
return Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Json(CloseTunnelResponse {
|
||||
status: "closed".to_string(),
|
||||
}))
|
||||
}
|
||||
|
||||
/// GET /api/v1/tunnel/status/:session_id
|
||||
/// Get tunnel session status
|
||||
pub async fn get_tunnel_status(
|
||||
State(state): State<AppState>,
|
||||
user: AuthUser,
|
||||
Path(session_id): Path<String>,
|
||||
) -> Result<Json<TunnelStatusResponse>, (StatusCode, String)> {
|
||||
// Validate session_id format
|
||||
if Uuid::parse_str(&session_id).is_err() {
|
||||
return Err((StatusCode::BAD_REQUEST, "Invalid session_id format".to_string()));
|
||||
}
|
||||
|
||||
// Verify session ownership
|
||||
let is_owner = db::verify_session_ownership(&state.db, &session_id, user.user_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("Failed to verify session ownership: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
|
||||
})?;
|
||||
|
||||
if !is_owner {
|
||||
return Err((
|
||||
StatusCode::FORBIDDEN,
|
||||
"Session not found or not owned by user".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Get session
|
||||
let session = db::get_tech_session(&state.db, &session_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
error!("Failed to get session: {}", e);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, e.to_string())
|
||||
})?
|
||||
.ok_or((StatusCode::NOT_FOUND, "Session not found".to_string()))?;
|
||||
|
||||
Ok(Json(TunnelStatusResponse {
|
||||
session_id: session.session_id,
|
||||
agent_id: session.agent_id.to_string(),
|
||||
status: session.status,
|
||||
opened_at: session.opened_at.to_rfc3339(),
|
||||
last_activity: session.last_activity.to_rfc3339(),
|
||||
}))
|
||||
}
|
||||
@@ -7,6 +7,7 @@ pub mod clients;
|
||||
pub mod commands;
|
||||
pub mod metrics;
|
||||
pub mod sites;
|
||||
pub mod tunnel;
|
||||
pub mod updates;
|
||||
pub mod users;
|
||||
|
||||
@@ -15,5 +16,6 @@ pub use clients::*;
|
||||
pub use commands::*;
|
||||
pub use metrics::*;
|
||||
pub use sites::*;
|
||||
pub use tunnel::*;
|
||||
pub use updates::*;
|
||||
pub use users::*;
|
||||
|
||||
151
projects/msp-tools/guru-rmm/server/src/db/tunnel.rs
Normal file
151
projects/msp-tools/guru-rmm/server/src/db/tunnel.rs
Normal file
@@ -0,0 +1,151 @@
|
||||
//! Database operations for tunnel sessions
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlx::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct TechSession {
|
||||
pub id: i32,
|
||||
pub session_id: String,
|
||||
pub tech_id: Uuid,
|
||||
pub agent_id: Uuid,
|
||||
pub opened_at: DateTime<Utc>,
|
||||
pub last_activity: DateTime<Utc>,
|
||||
pub closed_at: Option<DateTime<Utc>>,
|
||||
pub status: String,
|
||||
}
|
||||
|
||||
/// Create a new tech session
|
||||
pub async fn create_tech_session(
|
||||
pool: &PgPool,
|
||||
session_id: &str,
|
||||
tech_id: Uuid,
|
||||
agent_id: Uuid,
|
||||
) -> Result<TechSession, sqlx::Error> {
|
||||
sqlx::query_as::<_, TechSession>(
|
||||
r#"
|
||||
INSERT INTO tech_sessions (session_id, tech_id, agent_id, status)
|
||||
VALUES ($1, $2, $3, 'active')
|
||||
RETURNING *
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(tech_id)
|
||||
.bind(agent_id)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Get tech session by session_id
|
||||
pub async fn get_tech_session(
|
||||
pool: &PgPool,
|
||||
session_id: &str,
|
||||
) -> Result<Option<TechSession>, sqlx::Error> {
|
||||
sqlx::query_as::<_, TechSession>(
|
||||
r#"
|
||||
SELECT * FROM tech_sessions
|
||||
WHERE session_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Update last_activity timestamp
|
||||
pub async fn update_session_activity(
|
||||
pool: &PgPool,
|
||||
session_id: &str,
|
||||
) -> Result<u64, sqlx::Error> {
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE tech_sessions
|
||||
SET last_activity = NOW()
|
||||
WHERE session_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
/// Close a session
|
||||
pub async fn close_tech_session(
|
||||
pool: &PgPool,
|
||||
session_id: &str,
|
||||
) -> Result<u64, sqlx::Error> {
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE tech_sessions
|
||||
SET status = 'closed', closed_at = NOW()
|
||||
WHERE session_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
/// Check if tech owns session (for authorization)
|
||||
pub async fn verify_session_ownership(
|
||||
pool: &PgPool,
|
||||
session_id: &str,
|
||||
tech_id: Uuid,
|
||||
) -> Result<bool, sqlx::Error> {
|
||||
let result = sqlx::query_scalar::<_, bool>(
|
||||
r#"
|
||||
SELECT EXISTS(
|
||||
SELECT 1 FROM tech_sessions
|
||||
WHERE session_id = $1 AND tech_id = $2 AND status = 'active'
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind(tech_id)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Check if there's an active session for tech + agent pair
|
||||
pub async fn has_active_session(
|
||||
pool: &PgPool,
|
||||
tech_id: Uuid,
|
||||
agent_id: Uuid,
|
||||
) -> Result<bool, sqlx::Error> {
|
||||
let result = sqlx::query_scalar::<_, bool>(
|
||||
r#"
|
||||
SELECT EXISTS(
|
||||
SELECT 1 FROM tech_sessions
|
||||
WHERE tech_id = $1 AND agent_id = $2 AND status = 'active'
|
||||
)
|
||||
"#,
|
||||
)
|
||||
.bind(tech_id)
|
||||
.bind(agent_id)
|
||||
.fetch_one(pool)
|
||||
.await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Close all active sessions for an agent (when agent disconnects)
|
||||
pub async fn close_agent_tunnel_sessions(
|
||||
pool: &PgPool,
|
||||
agent_id: Uuid,
|
||||
) -> Result<u64, sqlx::Error> {
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE tech_sessions
|
||||
SET status = 'closed', closed_at = NOW()
|
||||
WHERE agent_id = $1 AND status = 'active'
|
||||
"#,
|
||||
)
|
||||
.bind(agent_id)
|
||||
.execute(pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
@@ -85,6 +85,9 @@ pub enum AgentMessage {
|
||||
WatchdogEvent(WatchdogEventPayload),
|
||||
UpdateResult(UpdateResultPayload),
|
||||
Heartbeat,
|
||||
TunnelReady { session_id: String },
|
||||
TunnelData { channel_id: String, data: TunnelDataPayload },
|
||||
TunnelError { channel_id: String, error: String },
|
||||
}
|
||||
|
||||
/// Messages from server to agent
|
||||
@@ -98,6 +101,9 @@ pub enum ServerMessage {
|
||||
Update(UpdatePayload),
|
||||
Ack { message_id: Option<String> },
|
||||
Error { code: String, message: String },
|
||||
TunnelOpen { session_id: String, tech_id: Uuid },
|
||||
TunnelClose { session_id: String },
|
||||
TunnelData { channel_id: String, data: TunnelDataPayload },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -297,6 +303,21 @@ pub struct UpdateResultPayload {
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Tunnel data payload types
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", content = "payload")]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TunnelDataPayload {
|
||||
/// Terminal command execution (Phase 1)
|
||||
Terminal { command: String },
|
||||
/// Terminal output response
|
||||
TerminalOutput {
|
||||
stdout: String,
|
||||
stderr: String,
|
||||
exit_code: Option<i32>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Result of successful agent authentication
|
||||
struct AuthResult {
|
||||
agent_id: Uuid,
|
||||
@@ -479,7 +500,25 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
|
||||
|
||||
// Cleanup
|
||||
state.agents.write().await.remove(&agent_id);
|
||||
let _ = db::update_agent_status(&state.db, agent_id, "offline").await;
|
||||
|
||||
// Update agent status
|
||||
if let Err(e) = db::update_agent_status(&state.db, agent_id, "offline").await {
|
||||
error!("Failed to update agent status for {}: {}", agent_id, e);
|
||||
}
|
||||
|
||||
// Close all active tunnel sessions for this agent
|
||||
match db::close_agent_tunnel_sessions(&state.db, agent_id).await {
|
||||
Ok(count) if count > 0 => {
|
||||
info!("Closed {} active tunnel session(s) for agent {}", count, agent_id);
|
||||
}
|
||||
Ok(_) => {
|
||||
debug!("No active tunnel sessions to close for agent {}", agent_id);
|
||||
}
|
||||
Err(e) => {
|
||||
error!("Failed to close tunnel sessions for agent {}: {}", agent_id, e);
|
||||
}
|
||||
}
|
||||
|
||||
send_task.abort();
|
||||
|
||||
info!("Agent {} connection closed", agent_id);
|
||||
@@ -745,6 +784,57 @@ async fn handle_agent_message(
|
||||
}
|
||||
}
|
||||
|
||||
AgentMessage::TunnelReady { session_id } => {
|
||||
info!(
|
||||
"Agent {} tunnel ready: session_id={}",
|
||||
agent_id, session_id
|
||||
);
|
||||
|
||||
// Update session activity timestamp
|
||||
if let Err(e) = db::update_session_activity(&state.db, &session_id).await {
|
||||
error!(
|
||||
"Failed to update session activity for {}: {}",
|
||||
session_id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
AgentMessage::TunnelData { channel_id, data } => {
|
||||
debug!(
|
||||
"Received tunnel data from agent {}: channel_id={}, type={:?}",
|
||||
agent_id, channel_id, data
|
||||
);
|
||||
|
||||
// Phase 2: Forward data to connected clients via WebSocket or REST API
|
||||
// For now, just log the data
|
||||
match data {
|
||||
TunnelDataPayload::TerminalOutput { stdout, stderr, exit_code } => {
|
||||
if !stdout.is_empty() {
|
||||
debug!("Terminal stdout: {}", stdout.trim());
|
||||
}
|
||||
if !stderr.is_empty() {
|
||||
debug!("Terminal stderr: {}", stderr.trim());
|
||||
}
|
||||
if let Some(code) = exit_code {
|
||||
debug!("Terminal exit code: {}", code);
|
||||
}
|
||||
}
|
||||
TunnelDataPayload::Terminal { command } => {
|
||||
debug!("Terminal command echo: {}", command);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
AgentMessage::TunnelError { channel_id, error } => {
|
||||
error!(
|
||||
"Tunnel error from agent {}: channel_id={}, error={}",
|
||||
agent_id, channel_id, error
|
||||
);
|
||||
|
||||
// Phase 2: Forward error to connected clients
|
||||
// For now, just log the error
|
||||
}
|
||||
|
||||
AgentMessage::Auth(_) => {
|
||||
warn!("Received unexpected auth message from already authenticated agent");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user