azcomputerguru/guru-connect
1
0
Fork 0
Code Issues 4 Pull Requests Actions Packages Projects Releases Wiki Activity

Add PostgreSQL database persistence

Browse Source 
- Add connect_machines, connect_sessions, connect_session_events, connect_support_codes tables
- Implement db module with connection pooling (sqlx)
- Add machine persistence across server restarts
- Add audit logging for session/viewer events
- Support codes now persisted to database

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Mike Swanson
2025-12-28 19:51:01 -07:00
parent 448d3b75ac
commit f6bf0cfd26
10 changed files with 788 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
server/src/relay/mod.rs
View File

@@ -18,6 +18,7 @@ use uuid::Uuid;
use crate::proto;
use crate::session::SessionManager;
use crate::db::{self, Database};
use crate::AppState;
#[derive(Debug, Deserialize)]
@@ -53,8 +54,9 @@ pub async fn agent_ws_handler(
let support_code = params.support_code;
let sessions = state.sessions.clone();
let support_codes = state.support_codes.clone();
let db = state.db.clone();
ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, agent_id, agent_name, support_code))
ws.on_upgrade(move |socket| handle_agent_connection(socket, sessions, support_codes, db, agent_id, agent_name, support_code))
}
/// WebSocket handler for viewer connections
@@ -66,8 +68,9 @@ pub async fn viewer_ws_handler(
let session_id = params.session_id;
let viewer_name = params.viewer_name;
let sessions = state.sessions.clone();
let db = state.db.clone();
ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, session_id, viewer_name))
ws.on_upgrade(move |socket| handle_viewer_connection(socket, sessions, db, session_id, viewer_name))
}
/// Handle an agent WebSocket connection
@@ -75,6 +78,7 @@ async fn handle_agent_connection(
socket: WebSocket,
sessions: SessionManager,
support_codes: crate::support_codes::SupportCodeManager,
db: Option<Database>,
agent_id: String,
agent_name: String,
support_code: Option<String>,
@@ -110,11 +114,54 @@ async fn handle_agent_connection(
info!("Session created: {} (agent in idle mode)", session_id);
// Database: upsert machine and create session record
let machine_id = if let Some(ref db) = db {
match db::machines::upsert_machine(db.pool(), &agent_id, &agent_name, is_persistent).await {
Ok(machine) => {
// Create session record
let _ = db::sessions::create_session(
db.pool(),
session_id,
machine.id,
support_code.is_some(),
support_code.as_deref(),
).await;
// Log session started event
let _ = db::events::log_event(
db.pool(),
session_id,
db::events::EventTypes::SESSION_STARTED,
None, None, None, None,
).await;
Some(machine.id)
}
Err(e) => {
warn!("Failed to upsert machine in database: {}", e);
None
}
}
} else {
None
};
// If a support code was provided, mark it as connected
if let Some(ref code) = support_code {
info!("Linking support code {} to session {}", code, session_id);
support_codes.mark_connected(code, Some(agent_name.clone()), Some(agent_id.clone())).await;
support_codes.link_session(code, session_id).await;
// Database: update support code
if let Some(ref db) = db {
let _ = db::support_codes::mark_code_connected(
db.pool(),
code,
Some(session_id),
Some(&agent_name),
Some(&agent_id),
).await;
}
}
// Use Arc<Mutex> for sender so we can use it from multiple places
@@ -233,10 +280,33 @@ async fn handle_agent_connection(
// Mark agent as disconnected (persistent agents stay in list as offline)
sessions_cleanup.mark_agent_disconnected(session_id).await;
// Database: end session and mark machine offline
if let Some(ref db) = db {
// End the session record
let _ = db::sessions::end_session(db.pool(), session_id, "ended").await;
// Mark machine as offline
let _ = db::machines::mark_machine_offline(db.pool(), &agent_id).await;
// Log session ended event
let _ = db::events::log_event(
db.pool(),
session_id,
db::events::EventTypes::SESSION_ENDED,
None, None, None, None,
).await;
}
// Mark support code as completed if one was used (unless cancelled)
if let Some(ref code) = support_code_cleanup {
if !support_codes_cleanup.is_cancelled(code).await {
support_codes_cleanup.mark_completed(code).await;
// Database: mark code as completed
if let Some(ref db) = db {
let _ = db::support_codes::mark_code_completed(db.pool(), code).await;
}
info!("Support code {} marked as completed", code);
}
}
@@ -248,6 +318,7 @@ async fn handle_agent_connection(
async fn handle_viewer_connection(
socket: WebSocket,
sessions: SessionManager,
db: Option<Database>,
session_id_str: String,
viewer_name: String,
) {
@@ -274,6 +345,18 @@ async fn handle_viewer_connection(
info!("Viewer {} ({}) joined session: {}", viewer_name, viewer_id, session_id);
// Database: log viewer joined event
if let Some(ref db) = db {
let _ = db::events::log_event(
db.pool(),
session_id,
db::events::EventTypes::VIEWER_JOINED,
Some(&viewer_id),
Some(&viewer_name),
None, None,
).await;
}
let (mut ws_sender, mut ws_receiver) = socket.split();
// Task to forward frames from agent to this viewer
@@ -287,6 +370,7 @@ async fn handle_viewer_connection(
let sessions_cleanup = sessions.clone();
let viewer_id_cleanup = viewer_id.clone();
let viewer_name_cleanup = viewer_name.clone();
// Main loop: receive input from viewer and forward to agent
while let Some(msg) = ws_receiver.next().await {
@@ -330,5 +414,18 @@ async fn handle_viewer_connection(
// Cleanup (this sends StopStream to agent if last viewer)
frame_forward.abort();
sessions_cleanup.leave_session(session_id, &viewer_id_cleanup).await;
// Database: log viewer left event
if let Some(ref db) = db {
let _ = db::events::log_event(
db.pool(),
session_id,
db::events::EventTypes::VIEWER_LEFT,
Some(&viewer_id_cleanup),
Some(&viewer_name_cleanup),
None, None,
).await;
}
info!("Viewer {} left session: {}", viewer_id_cleanup, session_id);
}