mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-11-09 09:15:23 +08:00
542 lines
18 KiB
Rust
542 lines
18 KiB
Rust
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
//! SSH tunnel handler for the multiplexed gateway.
|
|
|
|
use axum::{Router, extract::State, http::Method, response::IntoResponse, routing::any};
|
|
use http::StatusCode;
|
|
use hyper::Request;
|
|
use hyper_util::rt::TokioIo;
|
|
use openshell_core::proto::{Sandbox, SandboxPhase, SshSession};
|
|
use prost::Message;
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
use tokio::io::AsyncWriteExt;
|
|
use tracing::{info, warn};
|
|
|
|
use crate::ServerState;
|
|
use crate::persistence::{ObjectType, Store};
|
|
|
|
const HEADER_SANDBOX_ID: &str = "x-sandbox-id";
|
|
const HEADER_TOKEN: &str = "x-sandbox-token";
|
|
|
|
/// Maximum concurrent SSH tunnel connections per session token.
|
|
const MAX_CONNECTIONS_PER_TOKEN: u32 = 3;
|
|
|
|
/// Redact a bearer token for safe logging — show only the last 4 characters.
|
|
fn redact_token(token: &str) -> String {
|
|
if token.len() <= 4 {
|
|
"****".to_string()
|
|
} else {
|
|
format!("****{}", &token[token.len() - 4..])
|
|
}
|
|
}
|
|
|
|
/// Maximum concurrent SSH tunnel connections per sandbox.
|
|
const MAX_CONNECTIONS_PER_SANDBOX: u32 = 20;
|
|
|
|
fn acquire_connection_slots(
|
|
token_counts: &std::sync::Mutex<std::collections::HashMap<String, u32>>,
|
|
sandbox_counts: &std::sync::Mutex<std::collections::HashMap<String, u32>>,
|
|
token: &str,
|
|
sandbox_id: &str,
|
|
) -> Result<(), ConnectionLimit> {
|
|
{
|
|
let mut counts = token_counts.lock().unwrap();
|
|
let count = counts.entry(token.to_string()).or_insert(0);
|
|
if *count >= MAX_CONNECTIONS_PER_TOKEN {
|
|
return Err(ConnectionLimit::PerToken);
|
|
}
|
|
*count += 1;
|
|
}
|
|
|
|
{
|
|
let mut counts = sandbox_counts.lock().unwrap();
|
|
let count = counts.entry(sandbox_id.to_string()).or_insert(0);
|
|
if *count >= MAX_CONNECTIONS_PER_SANDBOX {
|
|
decrement_connection_count(token_counts, token);
|
|
return Err(ConnectionLimit::PerSandbox);
|
|
}
|
|
*count += 1;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
enum ConnectionLimit {
|
|
PerToken,
|
|
PerSandbox,
|
|
}
|
|
|
|
pub fn router(state: Arc<ServerState>) -> Router {
|
|
Router::new()
|
|
.route("/connect/ssh", any(ssh_connect))
|
|
.with_state(state)
|
|
}
|
|
|
|
async fn ssh_connect(
|
|
State(state): State<Arc<ServerState>>,
|
|
req: Request<axum::body::Body>,
|
|
) -> impl IntoResponse {
|
|
if req.method() != Method::CONNECT {
|
|
return StatusCode::METHOD_NOT_ALLOWED.into_response();
|
|
}
|
|
|
|
let sandbox_id = match header_value(req.headers(), HEADER_SANDBOX_ID) {
|
|
Ok(value) => value,
|
|
Err(status) => return status.into_response(),
|
|
};
|
|
let token = match header_value(req.headers(), HEADER_TOKEN) {
|
|
Ok(value) => value,
|
|
Err(status) => return status.into_response(),
|
|
};
|
|
|
|
let session = match state.store.get_message::<SshSession>(&token).await {
|
|
Ok(Some(session)) => session,
|
|
Ok(None) => return StatusCode::UNAUTHORIZED.into_response(),
|
|
Err(err) => {
|
|
warn!(error = %err, "Failed to fetch SSH session");
|
|
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
|
|
}
|
|
};
|
|
|
|
if session.revoked || session.sandbox_id != sandbox_id {
|
|
return StatusCode::UNAUTHORIZED.into_response();
|
|
}
|
|
|
|
// Check token expiry (0 means no expiry for backward compatibility).
|
|
if session.expires_at_ms > 0 {
|
|
let now_ms = i64::try_from(
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_millis(),
|
|
)
|
|
.unwrap_or(i64::MAX);
|
|
if now_ms > session.expires_at_ms {
|
|
return StatusCode::UNAUTHORIZED.into_response();
|
|
}
|
|
}
|
|
|
|
let sandbox = match state.store.get_message::<Sandbox>(&sandbox_id).await {
|
|
Ok(Some(sandbox)) => sandbox,
|
|
Ok(None) => return StatusCode::NOT_FOUND.into_response(),
|
|
Err(err) => {
|
|
warn!(error = %err, "Failed to fetch sandbox");
|
|
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
|
|
}
|
|
};
|
|
|
|
if SandboxPhase::try_from(sandbox.phase).ok() != Some(SandboxPhase::Ready) {
|
|
return StatusCode::PRECONDITION_FAILED.into_response();
|
|
}
|
|
|
|
// Enforce connection caps *before* opening a relay — otherwise denied
|
|
// calls churn pending relay slots and wake the supervisor until the relay
|
|
// timeout elapses.
|
|
if let Err(limit) = acquire_connection_slots(
|
|
&state.ssh_connections_by_token,
|
|
&state.ssh_connections_by_sandbox,
|
|
&token,
|
|
&sandbox_id,
|
|
) {
|
|
match limit {
|
|
ConnectionLimit::PerToken => {
|
|
warn!(token = %redact_token(&token), "SSH tunnel: per-token connection limit reached");
|
|
}
|
|
ConnectionLimit::PerSandbox => {
|
|
warn!(sandbox_id = %sandbox_id, "SSH tunnel: per-sandbox connection limit reached");
|
|
}
|
|
}
|
|
return StatusCode::TOO_MANY_REQUESTS.into_response();
|
|
}
|
|
|
|
// Open a relay channel through the supervisor session. Use a generous
|
|
// 30s session-wait timeout because `/connect/ssh` is typically called
|
|
// immediately after `sandbox create`, so we need to cover the supervisor's
|
|
// initial TLS + gRPC handshake on a cold-started pod. The old
|
|
// direct-connect path tolerated ~34s here for similar reasons.
|
|
let (channel_id, relay_rx) = match state
|
|
.supervisor_sessions
|
|
.open_relay(&sandbox_id, Duration::from_secs(30))
|
|
.await
|
|
{
|
|
Ok(pair) => pair,
|
|
Err(status) => {
|
|
warn!(sandbox_id = %sandbox_id, error = %status.message(), "SSH tunnel: supervisor session not available");
|
|
decrement_connection_count(&state.ssh_connections_by_token, &token);
|
|
decrement_connection_count(&state.ssh_connections_by_sandbox, &sandbox_id);
|
|
return StatusCode::BAD_GATEWAY.into_response();
|
|
}
|
|
};
|
|
|
|
let sandbox_id_clone = sandbox_id.clone();
|
|
let token_clone = token.clone();
|
|
let state_clone = state.clone();
|
|
|
|
let upgrade = hyper::upgrade::on(req);
|
|
tokio::spawn(async move {
|
|
// Wait for the supervisor to open its `RelayStream` and deliver the
|
|
// bridge half of the relay.
|
|
let mut relay = match tokio::time::timeout(Duration::from_secs(10), relay_rx).await {
|
|
Ok(Ok(stream)) => stream,
|
|
Ok(Err(_)) => {
|
|
warn!(sandbox_id = %sandbox_id_clone, channel_id = %channel_id, "SSH tunnel: relay channel dropped");
|
|
decrement_connection_count(&state_clone.ssh_connections_by_token, &token_clone);
|
|
decrement_connection_count(
|
|
&state_clone.ssh_connections_by_sandbox,
|
|
&sandbox_id_clone,
|
|
);
|
|
return;
|
|
}
|
|
Err(_) => {
|
|
warn!(sandbox_id = %sandbox_id_clone, channel_id = %channel_id, "SSH tunnel: relay open timed out");
|
|
decrement_connection_count(&state_clone.ssh_connections_by_token, &token_clone);
|
|
decrement_connection_count(
|
|
&state_clone.ssh_connections_by_sandbox,
|
|
&sandbox_id_clone,
|
|
);
|
|
return;
|
|
}
|
|
};
|
|
|
|
info!(sandbox_id = %sandbox_id_clone, channel_id = %channel_id, "SSH tunnel: relay established, bridging client");
|
|
|
|
match upgrade.await {
|
|
Ok(upgraded) => {
|
|
let mut upgraded = TokioIo::new(upgraded);
|
|
let _ = tokio::io::copy_bidirectional(&mut upgraded, &mut relay).await;
|
|
let _ = AsyncWriteExt::shutdown(&mut upgraded).await;
|
|
}
|
|
Err(err) => {
|
|
warn!(error = %err, "SSH upgrade failed");
|
|
}
|
|
}
|
|
|
|
// Decrement connection counts on tunnel completion.
|
|
decrement_connection_count(&state_clone.ssh_connections_by_token, &token_clone);
|
|
decrement_connection_count(&state_clone.ssh_connections_by_sandbox, &sandbox_id_clone);
|
|
});
|
|
|
|
StatusCode::OK.into_response()
|
|
}
|
|
|
|
fn header_value(headers: &http::HeaderMap, name: &str) -> Result<String, StatusCode> {
|
|
let value = headers
|
|
.get(name)
|
|
.ok_or(StatusCode::UNAUTHORIZED)?
|
|
.to_str()
|
|
.map_err(|_| StatusCode::BAD_REQUEST)?
|
|
.trim()
|
|
.to_string();
|
|
if value.is_empty() {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
}
|
|
Ok(value)
|
|
}
|
|
|
|
impl ObjectType for SshSession {
|
|
fn object_type() -> &'static str {
|
|
"ssh_session"
|
|
}
|
|
}
|
|
|
|
/// Decrement a connection count entry, removing it if it reaches zero.
|
|
fn decrement_connection_count(
|
|
counts: &std::sync::Mutex<std::collections::HashMap<String, u32>>,
|
|
key: &str,
|
|
) {
|
|
let mut map = counts.lock().unwrap();
|
|
if let Some(count) = map.get_mut(key) {
|
|
*count = count.saturating_sub(1);
|
|
if *count == 0 {
|
|
map.remove(key);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Spawn a background task that periodically reaps expired and revoked SSH sessions.
|
|
pub fn spawn_session_reaper(store: Arc<Store>, interval: Duration) {
|
|
tokio::spawn(async move {
|
|
// Initial delay to let startup settle.
|
|
tokio::time::sleep(interval).await;
|
|
|
|
loop {
|
|
if let Err(e) = reap_expired_sessions(&store).await {
|
|
warn!(error = %e, "SSH session reaper sweep failed");
|
|
}
|
|
tokio::time::sleep(interval).await;
|
|
}
|
|
});
|
|
}
|
|
|
|
async fn reap_expired_sessions(store: &Store) -> Result<(), String> {
|
|
let now_ms = i64::try_from(
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_millis(),
|
|
)
|
|
.unwrap_or(i64::MAX);
|
|
|
|
let records = store
|
|
.list(SshSession::object_type(), 1000, 0)
|
|
.await
|
|
.map_err(|e| e.to_string())?;
|
|
|
|
let mut reaped = 0u32;
|
|
for record in records {
|
|
let session: SshSession = match Message::decode(record.payload.as_slice()) {
|
|
Ok(s) => s,
|
|
Err(_) => continue,
|
|
};
|
|
|
|
let should_delete =
|
|
// Expired sessions (expires_at_ms > 0 means expiry is set).
|
|
(session.expires_at_ms > 0 && now_ms > session.expires_at_ms)
|
|
// Revoked sessions — already invalidated, just cleaning up storage.
|
|
|| session.revoked;
|
|
|
|
if should_delete {
|
|
use openshell_core::ObjectId;
|
|
if let Err(e) = store
|
|
.delete(SshSession::object_type(), session.object_id())
|
|
.await
|
|
{
|
|
warn!(session_id = %session.object_id(), error = %e, "Failed to reap SSH session");
|
|
} else {
|
|
reaped += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
if reaped > 0 {
|
|
info!(count = reaped, "SSH session reaper: cleaned up sessions");
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::persistence::Store;
|
|
use std::collections::HashMap;
|
|
use std::sync::Mutex;
|
|
|
|
fn make_session(id: &str, sandbox_id: &str, expires_at_ms: i64, revoked: bool) -> SshSession {
|
|
SshSession {
|
|
metadata: Some(openshell_core::proto::datamodel::v1::ObjectMeta {
|
|
id: id.to_string(),
|
|
name: format!("session-{id}"),
|
|
created_at_ms: 1000,
|
|
labels: HashMap::new(),
|
|
}),
|
|
sandbox_id: sandbox_id.to_string(),
|
|
token: id.to_string(),
|
|
expires_at_ms,
|
|
revoked,
|
|
}
|
|
}
|
|
|
|
fn now_ms() -> i64 {
|
|
i64::try_from(
|
|
std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_millis(),
|
|
)
|
|
.unwrap_or(i64::MAX)
|
|
}
|
|
|
|
// ---- Connection limit tests ----
|
|
|
|
#[test]
|
|
fn decrement_removes_entry_at_zero() {
|
|
let counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
counts.lock().unwrap().insert("tok1".to_string(), 1);
|
|
decrement_connection_count(&counts, "tok1");
|
|
assert!(counts.lock().unwrap().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn decrement_reduces_count() {
|
|
let counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
counts.lock().unwrap().insert("tok1".to_string(), 5);
|
|
decrement_connection_count(&counts, "tok1");
|
|
assert_eq!(*counts.lock().unwrap().get("tok1").unwrap(), 4);
|
|
}
|
|
|
|
#[test]
|
|
fn decrement_missing_key_is_noop() {
|
|
let counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
decrement_connection_count(&counts, "nonexistent");
|
|
assert!(counts.lock().unwrap().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn per_token_connection_limit_enforced() {
|
|
let counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
counts
|
|
.lock()
|
|
.unwrap()
|
|
.insert("tok1".to_string(), MAX_CONNECTIONS_PER_TOKEN);
|
|
let current = *counts.lock().unwrap().get("tok1").unwrap();
|
|
assert!(current >= MAX_CONNECTIONS_PER_TOKEN);
|
|
}
|
|
|
|
#[test]
|
|
fn per_sandbox_connection_limit_enforced() {
|
|
let counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
counts
|
|
.lock()
|
|
.unwrap()
|
|
.insert("sbx1".to_string(), MAX_CONNECTIONS_PER_SANDBOX);
|
|
let current = *counts.lock().unwrap().get("sbx1").unwrap();
|
|
assert!(current >= MAX_CONNECTIONS_PER_SANDBOX);
|
|
}
|
|
|
|
#[test]
|
|
fn acquire_connection_slots_rejects_per_token_limit_without_touching_sandbox() {
|
|
let token_counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
let sandbox_counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
token_counts
|
|
.lock()
|
|
.unwrap()
|
|
.insert("tok1".to_string(), MAX_CONNECTIONS_PER_TOKEN);
|
|
|
|
let result = acquire_connection_slots(&token_counts, &sandbox_counts, "tok1", "sbx1");
|
|
|
|
assert!(matches!(result, Err(ConnectionLimit::PerToken)));
|
|
assert!(sandbox_counts.lock().unwrap().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn acquire_connection_slots_rolls_back_token_increment_on_sandbox_limit() {
|
|
let token_counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
let sandbox_counts: Mutex<HashMap<String, u32>> = Mutex::new(HashMap::new());
|
|
sandbox_counts
|
|
.lock()
|
|
.unwrap()
|
|
.insert("sbx1".to_string(), MAX_CONNECTIONS_PER_SANDBOX);
|
|
|
|
let result = acquire_connection_slots(&token_counts, &sandbox_counts, "tok1", "sbx1");
|
|
|
|
assert!(matches!(result, Err(ConnectionLimit::PerSandbox)));
|
|
assert!(token_counts.lock().unwrap().is_empty());
|
|
}
|
|
|
|
// ---- Session reaper tests ----
|
|
|
|
#[tokio::test]
|
|
async fn reaper_deletes_expired_sessions() {
|
|
let store = Store::connect("sqlite::memory:?cache=shared")
|
|
.await
|
|
.unwrap();
|
|
|
|
let expired = make_session("expired1", "sbx1", now_ms() - 60_000, false);
|
|
store.put_message(&expired).await.unwrap();
|
|
|
|
let valid = make_session("valid1", "sbx1", now_ms() + 3_600_000, false);
|
|
store.put_message(&valid).await.unwrap();
|
|
|
|
reap_expired_sessions(&store).await.unwrap();
|
|
|
|
assert!(
|
|
store
|
|
.get_message::<SshSession>("expired1")
|
|
.await
|
|
.unwrap()
|
|
.is_none(),
|
|
"expired session should be reaped"
|
|
);
|
|
assert!(
|
|
store
|
|
.get_message::<SshSession>("valid1")
|
|
.await
|
|
.unwrap()
|
|
.is_some(),
|
|
"valid session should be kept"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn reaper_deletes_revoked_sessions() {
|
|
let store = Store::connect("sqlite::memory:?cache=shared")
|
|
.await
|
|
.unwrap();
|
|
|
|
let revoked = make_session("revoked1", "sbx1", 0, true);
|
|
store.put_message(&revoked).await.unwrap();
|
|
|
|
let active = make_session("active1", "sbx1", 0, false);
|
|
store.put_message(&active).await.unwrap();
|
|
|
|
reap_expired_sessions(&store).await.unwrap();
|
|
|
|
assert!(
|
|
store
|
|
.get_message::<SshSession>("revoked1")
|
|
.await
|
|
.unwrap()
|
|
.is_none(),
|
|
"revoked session should be reaped"
|
|
);
|
|
assert!(
|
|
store
|
|
.get_message::<SshSession>("active1")
|
|
.await
|
|
.unwrap()
|
|
.is_some(),
|
|
"active session should be kept"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn reaper_preserves_zero_expiry_sessions() {
|
|
let store = Store::connect("sqlite::memory:?cache=shared")
|
|
.await
|
|
.unwrap();
|
|
|
|
// expires_at_ms = 0 means no expiry (backward compatible).
|
|
let no_expiry = make_session("noexpiry1", "sbx1", 0, false);
|
|
store.put_message(&no_expiry).await.unwrap();
|
|
|
|
reap_expired_sessions(&store).await.unwrap();
|
|
|
|
assert!(
|
|
store
|
|
.get_message::<SshSession>("noexpiry1")
|
|
.await
|
|
.unwrap()
|
|
.is_some(),
|
|
"session with no expiry should be preserved"
|
|
);
|
|
}
|
|
|
|
// ---- Expiry validation logic tests ----
|
|
|
|
#[test]
|
|
fn expired_session_is_detected() {
|
|
let session = make_session("tok1", "sbx1", now_ms() - 1000, false);
|
|
let is_expired = session.expires_at_ms > 0 && now_ms() > session.expires_at_ms;
|
|
assert!(is_expired, "session in the past should be expired");
|
|
}
|
|
|
|
#[test]
|
|
fn future_session_is_not_expired() {
|
|
let session = make_session("tok1", "sbx1", now_ms() + 3_600_000, false);
|
|
let is_expired = session.expires_at_ms > 0 && now_ms() > session.expires_at_ms;
|
|
assert!(!is_expired, "session in the future should not be expired");
|
|
}
|
|
|
|
#[test]
|
|
fn zero_expiry_is_not_expired() {
|
|
let session = make_session("tok1", "sbx1", 0, false);
|
|
let is_expired = session.expires_at_ms > 0 && now_ms() > session.expires_at_ms;
|
|
assert!(
|
|
!is_expired,
|
|
"session with zero expiry should never be expired"
|
|
);
|
|
}
|
|
}
|