Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 11 additions & 4 deletions crates/terraphim_rlm/src/executor/firecracker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -447,8 +447,13 @@ impl super::ExecutionEnvironment for FirecrackerExecutor {
) -> Result<SnapshotId, Self::Error> {
log::info!("Creating snapshot '{}' for session {}", name, session_id);

// Check snapshot limit for this session
let count = *self.snapshot_counts.read().get(session_id).unwrap_or(&0);
// Validate snapshot name for security (path traversal prevention)
crate::validation::validate_snapshot_name(name)?;

// Check snapshot limit for this session - use write lock for atomic check-and-increment
// to prevent race condition where multiple concurrent snapshots could exceed the limit
let mut snapshot_counts = self.snapshot_counts.write();
let count = *snapshot_counts.get(session_id).unwrap_or(&0);
if count >= self.config.max_snapshots_per_session {
return Err(RlmError::MaxSnapshotsReached {
max: self.config.max_snapshots_per_session,
Expand Down Expand Up @@ -498,8 +503,10 @@ impl super::ExecutionEnvironment for FirecrackerExecutor {
}
};

// Update tracking
*self.snapshot_counts.write().entry(*session_id).or_insert(0) += 1;
// Update tracking - use the existing write lock for atomic increment
*snapshot_counts.entry(*session_id).or_insert(0) += 1;
// Release the write lock by dropping it explicitly before await boundary
drop(snapshot_counts);

let result = SnapshotId::new(name, *session_id);

Expand Down
1 change: 1 addition & 0 deletions crates/terraphim_rlm/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ pub mod logger;
// Knowledge graph validation (Phase 5)
#[cfg(feature = "kg-validation")]
pub mod validator;
pub mod validation;

// MCP tools (Phase 6)
#[cfg(feature = "mcp")]
Expand Down
16 changes: 16 additions & 0 deletions crates/terraphim_rlm/src/mcp_tools.rs
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,14 @@ impl RlmMcpService {
.and_then(|v| v.as_str())
.ok_or_else(|| ErrorData::invalid_params("Missing 'code' parameter", None))?;

// Validate code size to prevent DoS via memory exhaustion
if let Err(e) = crate::validation::validate_code_input(code) {
return Err(ErrorData::invalid_params(
format!("Code validation failed: {}", e),
None,
));
}

let session_id = self.resolve_session_id(&args).await?;
// timeout_ms is available for future use when execution context supports it
let _timeout_ms = args.get("timeout_ms").and_then(|v| v.as_u64());
Expand Down Expand Up @@ -371,6 +379,14 @@ impl RlmMcpService {
.and_then(|v| v.as_str())
.ok_or_else(|| ErrorData::invalid_params("Missing 'command' parameter", None))?;

// Validate command size to prevent DoS via memory exhaustion
if let Err(e) = crate::validation::validate_code_input(command) {
return Err(ErrorData::invalid_params(
format!("Command validation failed: {}", e),
None,
));
}

let session_id = self.resolve_session_id(&args).await?;
// These are available for future use when execution context supports them
let _timeout_ms = args.get("timeout_ms").and_then(|v| v.as_u64());
Expand Down
Loading
Loading