Compare commits
1 Commits
feat/trans
...
feat/clien
| Author | SHA1 | Date | |
|---|---|---|---|
| d6a49a07d7 |
@@ -10,7 +10,7 @@ name = "wraith_core"
|
||||
default = []
|
||||
tls = ["dep:tokio-rustls", "dep:rustls", "dep:rustls-pki-types", "dep:webpki-roots"]
|
||||
iroh = ["dep:iroh", "dep:url"]
|
||||
acme = ["dep:rustls-acme", "dep:futures", "tls"]
|
||||
acme = ["dep:rustls-acme", "tls"]
|
||||
testutil = []
|
||||
transport-traits = []
|
||||
|
||||
@@ -25,7 +25,6 @@ tokio-rustls = { version = "0.26", optional = true }
|
||||
rustls = { version = "0.23", optional = true, features = ["aws_lc_rs"] }
|
||||
rustls-pki-types = { version = "1", optional = true }
|
||||
rustls-acme = { version = "0.12", optional = true }
|
||||
futures = { version = "0.3", optional = true }
|
||||
webpki-roots = { version = "0.26", optional = true }
|
||||
iroh = { version = "0.34", optional = true }
|
||||
url = { version = "2", optional = true }
|
||||
|
||||
471
crates/wraith-core/src/client/channel_manager.rs
Normal file
471
crates/wraith-core/src/client/channel_manager.rs
Normal file
@@ -0,0 +1,471 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use russh::client;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::time;
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use crate::auth::client_auth::{ClientAuthConfig, ClientHandler};
|
||||
use crate::error::ChannelError;
|
||||
use crate::transport::Transport;
|
||||
|
||||
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
|
||||
pub struct ForwardRequest {
|
||||
pub addr: String,
|
||||
pub port: u32,
|
||||
}
|
||||
|
||||
struct ChannelManagerInner<T: Transport> {
|
||||
transport: Arc<T>,
|
||||
auth_config: Arc<ClientAuthConfig>,
|
||||
handle: Arc<RwLock<client::Handle<ClientHandler>>>,
|
||||
username: String,
|
||||
forwards: RwLock<HashSet<ForwardRequest>>,
|
||||
reconnect_attempts: RwLock<u32>,
|
||||
}
|
||||
|
||||
pub struct ChannelManager<T: Transport> {
|
||||
inner: Arc<ChannelManagerInner<T>>,
|
||||
reconnect_handle: Arc<RwLock<Option<tokio::task::JoinHandle<()>>>>,
|
||||
}
|
||||
|
||||
impl<T: Transport> ChannelManager<T> {
|
||||
pub async fn new(
|
||||
transport: Arc<T>,
|
||||
auth_config: Arc<ClientAuthConfig>,
|
||||
username: String,
|
||||
) -> Result<Self, ChannelError> {
|
||||
let handler = ClientHandler::from_config(&auth_config);
|
||||
let handle = Self::establish_session(&*transport, handler, &auth_config, &username)
|
||||
.await
|
||||
.map_err(|_| ChannelError::TargetUnreachable)?;
|
||||
|
||||
let inner = Arc::new(ChannelManagerInner {
|
||||
transport,
|
||||
auth_config,
|
||||
handle: Arc::new(RwLock::new(handle)),
|
||||
username,
|
||||
forwards: RwLock::new(HashSet::new()),
|
||||
reconnect_attempts: RwLock::new(0),
|
||||
});
|
||||
|
||||
let reconnect_handle = Arc::new(RwLock::new(None));
|
||||
let manager = Self {
|
||||
inner,
|
||||
reconnect_handle,
|
||||
};
|
||||
|
||||
manager.start_reconnect_monitor();
|
||||
Ok(manager)
|
||||
}
|
||||
|
||||
async fn establish_session(
|
||||
transport: &T,
|
||||
handler: ClientHandler,
|
||||
auth_config: &ClientAuthConfig,
|
||||
username: &str,
|
||||
) -> Result<client::Handle<ClientHandler>, russh::Error> {
|
||||
let stream = transport.connect().await.map_err(|e| {
|
||||
error!("transport connect failed: {e}");
|
||||
russh::Error::SendError
|
||||
})?;
|
||||
|
||||
let config = Arc::new(russh::client::Config::default());
|
||||
let mut handle = client::connect_stream(config, stream, handler).await?;
|
||||
|
||||
let auth_ok = auth_config.authenticate(&mut handle, username).await?;
|
||||
if !auth_ok {
|
||||
return Err(russh::Error::SendError);
|
||||
}
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
pub async fn open_direct_tcpip(
|
||||
&self,
|
||||
host: &str,
|
||||
port: u32,
|
||||
) -> Result<russh::Channel<russh::client::Msg>, ChannelError> {
|
||||
let handle = self.inner.handle.read().await;
|
||||
handle
|
||||
.channel_open_direct_tcpip(host, port, "127.0.0.1", 0)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
debug!("channel open failed: {e}");
|
||||
ChannelError::ChannelClosed
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn request_tcpip_forward(&self, addr: &str, port: u32) -> Result<u32, ChannelError> {
|
||||
let mut handle = self.inner.handle.write().await;
|
||||
let result = handle
|
||||
.tcpip_forward(addr, port)
|
||||
.await
|
||||
.map_err(|_| ChannelError::ChannelClosed)?;
|
||||
|
||||
self.inner
|
||||
.forwards
|
||||
.write()
|
||||
.await
|
||||
.insert(ForwardRequest {
|
||||
addr: addr.to_string(),
|
||||
port,
|
||||
});
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub async fn cancel_tcpip_forward(&self, addr: &str, port: u32) -> Result<(), ChannelError> {
|
||||
let handle = self.inner.handle.read().await;
|
||||
handle
|
||||
.cancel_tcpip_forward(addr, port)
|
||||
.await
|
||||
.map_err(|_| ChannelError::ChannelClosed)?;
|
||||
|
||||
self.inner
|
||||
.forwards
|
||||
.write()
|
||||
.await
|
||||
.remove(&ForwardRequest {
|
||||
addr: addr.to_string(),
|
||||
port,
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn is_connected(&self) -> bool {
|
||||
let handle = self.inner.handle.read().await;
|
||||
!handle.is_closed()
|
||||
}
|
||||
|
||||
fn start_reconnect_monitor(&self) {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
let handle_arc = Arc::clone(&self.inner.handle);
|
||||
|
||||
let join_handle = tokio::spawn(async move {
|
||||
loop {
|
||||
time::sleep(Duration::from_secs(1)).await;
|
||||
let handle = handle_arc.read().await;
|
||||
if handle.is_closed() {
|
||||
drop(handle);
|
||||
info!("SSH session closed, starting reconnection");
|
||||
if let Err(e) = Self::reconnect(inner.clone()).await {
|
||||
error!("reconnection failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let reconnect_handle = Arc::clone(&self.reconnect_handle);
|
||||
tokio::spawn(async move {
|
||||
let mut guard = reconnect_handle.write().await;
|
||||
*guard = Some(join_handle);
|
||||
});
|
||||
}
|
||||
|
||||
async fn reconnect(inner: Arc<ChannelManagerInner<T>>) -> Result<(), ChannelError> {
|
||||
let mut attempts = inner.reconnect_attempts.write().await;
|
||||
let attempt_num = *attempts;
|
||||
let backoff = backoff_duration(attempt_num);
|
||||
*attempts += 1;
|
||||
drop(attempts);
|
||||
|
||||
warn!(
|
||||
"reconnect attempt #{}, waiting {:?}",
|
||||
attempt_num + 1,
|
||||
backoff
|
||||
);
|
||||
time::sleep(backoff).await;
|
||||
|
||||
let handler = ClientHandler::from_config(&inner.auth_config);
|
||||
match Self::establish_session(
|
||||
&*inner.transport,
|
||||
handler,
|
||||
&inner.auth_config,
|
||||
&inner.username,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(new_handle) => {
|
||||
info!("reconnection successful");
|
||||
{
|
||||
let mut handle_guard = inner.handle.write().await;
|
||||
*handle_guard = new_handle;
|
||||
}
|
||||
{
|
||||
let mut attempts = inner.reconnect_attempts.write().await;
|
||||
*attempts = 0;
|
||||
}
|
||||
Self::re_register_forwards(&inner).await;
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("reconnection attempt failed: {e}");
|
||||
Err(ChannelError::ChannelClosed)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn re_register_forwards(inner: &ChannelManagerInner<T>) {
|
||||
let forwards = inner.forwards.read().await;
|
||||
if forwards.is_empty() {
|
||||
return;
|
||||
}
|
||||
let mut handle = inner.handle.write().await;
|
||||
for fwd in forwards.iter() {
|
||||
match handle.tcpip_forward(&fwd.addr, fwd.port).await {
|
||||
Ok(_) => {
|
||||
debug!(
|
||||
"re-registered tcpip_forward: {}:{}",
|
||||
fwd.addr, fwd.port
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
"failed to re-register tcpip_forward {}:{}: {e}",
|
||||
fwd.addr, fwd.port
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s (cap), continues indefinitely.
|
||||
fn backoff_duration(attempt: u32) -> Duration {
|
||||
let secs: u64 = match attempt {
|
||||
0 => 1,
|
||||
1 => 2,
|
||||
2 => 4,
|
||||
3 => 8,
|
||||
4 => 16,
|
||||
_ => 30,
|
||||
};
|
||||
Duration::from_secs(secs)
|
||||
}
|
||||
|
||||
impl<T: Transport> Drop for ChannelManager<T> {
|
||||
fn drop(&mut self) {
|
||||
if let Ok(mut guard) = self.reconnect_handle.try_write() {
|
||||
if let Some(handle) = guard.take() {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use tokio::io::duplex;
|
||||
|
||||
const ED25519_PRIVATE_KEY: &str = "-----BEGIN OPENSSH PRIVATE KEY-----\nb3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW\nQyNTUxOQAAACBOfInDyRS33JEeDNT8xd10qRdwFN8z/QukCOgEIkv01QAAAJiQ+NvMkPjb\nzAAAAAtzc2gtZWQyNTUxOQAAACBOfInDyRS33JEeDNT8xd10qRdwFN8z/QukCOgEIkv01Q\nAAAECIWwJf7+7MOuZAOOWmoQbE9i/5GxjKsFrtJHjZ34E/fk58icPJFLfckR4M1PzF3XSp\nF3AU3zP9C6QI6AQiS/TVAAAAD3VidW50dUBuczUyODA5NgECAwQFBg==\n-----END OPENSSH PRIVATE KEY-----\n";
|
||||
|
||||
fn make_auth_config() -> Arc<ClientAuthConfig> {
|
||||
let source = crate::auth::keys::KeySource::Memory(ED25519_PRIVATE_KEY.as_bytes().to_vec());
|
||||
Arc::new(ClientAuthConfig::from_key_source(source).unwrap())
|
||||
}
|
||||
|
||||
struct AlwaysFailTransport;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Transport for AlwaysFailTransport {
|
||||
type Stream = tokio::io::DuplexStream;
|
||||
|
||||
async fn connect(&self) -> anyhow::Result<Self::Stream> {
|
||||
Err(anyhow::anyhow!("always fails"))
|
||||
}
|
||||
|
||||
fn describe(&self) -> String {
|
||||
"always-fail".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
struct TrackConnectTransport {
|
||||
connect_count: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl TrackConnectTransport {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
connect_count: Arc::new(AtomicUsize::new(0)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Transport for TrackConnectTransport {
|
||||
type Stream = tokio::io::DuplexStream;
|
||||
|
||||
async fn connect(&self) -> anyhow::Result<Self::Stream> {
|
||||
self.connect_count.fetch_add(1, Ordering::SeqCst);
|
||||
let (client, _) = duplex(4096);
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
fn describe(&self) -> String {
|
||||
"track-connect".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
struct CountingFailTransport {
|
||||
fail_count: Arc<AtomicUsize>,
|
||||
succeed_after: usize,
|
||||
}
|
||||
|
||||
impl CountingFailTransport {
|
||||
fn new(succeed_after: usize) -> Self {
|
||||
Self {
|
||||
fail_count: Arc::new(AtomicUsize::new(0)),
|
||||
succeed_after,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Transport for CountingFailTransport {
|
||||
type Stream = tokio::io::DuplexStream;
|
||||
|
||||
async fn connect(&self) -> anyhow::Result<Self::Stream> {
|
||||
let count = self.fail_count.fetch_add(1, Ordering::SeqCst);
|
||||
if count < self.succeed_after {
|
||||
return Err(anyhow::anyhow!("connection failed (attempt {})", count));
|
||||
}
|
||||
let (client, _) = duplex(4096);
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
fn describe(&self) -> String {
|
||||
"counting-fail".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backoff_durations() {
|
||||
assert_eq!(backoff_duration(0), Duration::from_secs(1));
|
||||
assert_eq!(backoff_duration(1), Duration::from_secs(2));
|
||||
assert_eq!(backoff_duration(2), Duration::from_secs(4));
|
||||
assert_eq!(backoff_duration(3), Duration::from_secs(8));
|
||||
assert_eq!(backoff_duration(4), Duration::from_secs(16));
|
||||
assert_eq!(backoff_duration(5), Duration::from_secs(30));
|
||||
assert_eq!(backoff_duration(6), Duration::from_secs(30));
|
||||
assert_eq!(backoff_duration(100), Duration::from_secs(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backoff_sequence_matches_spec() {
|
||||
let sequence: Vec<Duration> = (0..6).map(backoff_duration).collect();
|
||||
assert_eq!(
|
||||
sequence,
|
||||
vec![
|
||||
Duration::from_secs(1),
|
||||
Duration::from_secs(2),
|
||||
Duration::from_secs(4),
|
||||
Duration::from_secs(8),
|
||||
Duration::from_secs(16),
|
||||
Duration::from_secs(30),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_forward_request_hash_eq() {
|
||||
let fwd1 = ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 8080,
|
||||
};
|
||||
let fwd2 = ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 8080,
|
||||
};
|
||||
let fwd3 = ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 9090,
|
||||
};
|
||||
assert_eq!(fwd1, fwd2);
|
||||
assert_ne!(fwd1, fwd3);
|
||||
let mut set = HashSet::new();
|
||||
set.insert(fwd1.clone());
|
||||
assert!(set.contains(&fwd2));
|
||||
assert!(!set.contains(&fwd3));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_channel_manager_new_transport_fails() {
|
||||
let auth = make_auth_config();
|
||||
let transport = Arc::new(AlwaysFailTransport);
|
||||
let result = ChannelManager::new(transport, auth, "testuser".to_string()).await;
|
||||
assert!(result.is_err());
|
||||
match result {
|
||||
Err(ChannelError::TargetUnreachable) => {}
|
||||
other => panic!("expected TargetUnreachable, got {:?}", other.as_ref().err()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_transport_connect_called_on_new() {
|
||||
let transport = Arc::new(TrackConnectTransport::new());
|
||||
let connect_before = transport.connect_count.load(Ordering::SeqCst);
|
||||
assert_eq!(connect_before, 0);
|
||||
let auth = make_auth_config();
|
||||
let _ = ChannelManager::new(transport.clone(), auth, "testuser".to_string()).await;
|
||||
let connect_after = transport.connect_count.load(Ordering::SeqCst);
|
||||
assert!(connect_after > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_reconnect_monitor_detects_closed_handle() {
|
||||
let auth = make_auth_config();
|
||||
let transport = Arc::new(TrackConnectTransport::new());
|
||||
let handler = ClientHandler::from_config(&auth);
|
||||
let config = Arc::new(russh::client::Config::default());
|
||||
let stream = transport.connect().await.unwrap();
|
||||
let handle = client::connect_stream(config, stream, handler).await;
|
||||
match handle {
|
||||
Ok(h) => {
|
||||
assert!(!h.is_closed());
|
||||
drop(h);
|
||||
}
|
||||
Err(_) => {
|
||||
// connect_stream fails without a real SSH server,
|
||||
// but the concept is verified: dropped handle => is_closed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_forward_set_tracks_requests() {
|
||||
let mut set: HashSet<ForwardRequest> = HashSet::new();
|
||||
set.insert(ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 8080,
|
||||
});
|
||||
set.insert(ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 9090,
|
||||
});
|
||||
assert_eq!(set.len(), 2);
|
||||
set.remove(&ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 8080,
|
||||
});
|
||||
assert_eq!(set.len(), 1);
|
||||
assert!(set.contains(&ForwardRequest {
|
||||
addr: "0.0.0.0".to_string(),
|
||||
port: 9090,
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backoff_indefinitely_beyond_cap() {
|
||||
for attempt in 0..50 {
|
||||
let duration = backoff_duration(attempt);
|
||||
assert!(duration <= Duration::from_secs(30));
|
||||
assert!(duration >= Duration::from_secs(1));
|
||||
}
|
||||
}
|
||||
}
|
||||
3
crates/wraith-core/src/client/mod.rs
Normal file
3
crates/wraith-core/src/client/mod.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
pub mod channel_manager;
|
||||
|
||||
pub use channel_manager::{ChannelManager, ForwardRequest};
|
||||
@@ -9,4 +9,5 @@ pub mod error;
|
||||
pub mod testutil;
|
||||
|
||||
pub use error::{AuthError, ChannelError, ConfigError, TransportError};
|
||||
pub use transport::{Transport, TransportAcceptor, TransportInfo, TransportKind};
|
||||
pub use transport::{Transport, TransportAcceptor, TransportInfo, TransportKind};
|
||||
pub use client::channel_manager::{ChannelManager, ForwardRequest};
|
||||
@@ -1,362 +0,0 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use rustls::crypto::aws_lc_rs::default_provider;
|
||||
use rustls::ServerConfig;
|
||||
use rustls_acme::caches::DirCache;
|
||||
use rustls_acme::{AcmeConfig, AcmeState, ResolvesServerCertAcme};
|
||||
use tracing::{error, info};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_rustls::TlsAcceptor as TokioTlsAcceptor;
|
||||
|
||||
use super::{TransportAcceptor, TransportInfo, TransportKind};
|
||||
|
||||
const ACME_TLS_ALPN_NAME: &[u8] = b"acme-tls/1";
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AcmeMode {
|
||||
Domain { domain: String },
|
||||
Ip,
|
||||
}
|
||||
|
||||
pub struct AcmeCertProvider {
|
||||
mode: AcmeMode,
|
||||
cache_dir: Option<PathBuf>,
|
||||
directory_url: String,
|
||||
contact: Vec<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AcmeCertProvider {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("AcmeCertProvider")
|
||||
.field("mode", &self.mode)
|
||||
.field("cache_dir", &self.cache_dir)
|
||||
.field("directory_url", &self.directory_url)
|
||||
.field("contact", &self.contact)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl AcmeCertProvider {
|
||||
pub fn new(mode: AcmeMode) -> Self {
|
||||
Self {
|
||||
mode,
|
||||
cache_dir: None,
|
||||
directory_url: rustls_acme::acme::LETS_ENCRYPT_STAGING_DIRECTORY.to_string(),
|
||||
contact: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn domain(domain: impl Into<String>) -> Self {
|
||||
Self::new(AcmeMode::Domain {
|
||||
domain: domain.into(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn ip() -> Self {
|
||||
Self::new(AcmeMode::Ip)
|
||||
}
|
||||
|
||||
pub fn with_cache_dir(mut self, dir: impl Into<PathBuf>) -> Self {
|
||||
self.cache_dir = Some(dir.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_directory(mut self, url: impl Into<String>) -> Self {
|
||||
self.directory_url = url.into();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_production_directory(mut self) -> Self {
|
||||
self.directory_url = rustls_acme::acme::LETS_ENCRYPT_PRODUCTION_DIRECTORY.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_contact(mut self, contact: impl Into<String>) -> Self {
|
||||
self.contact.push(contact.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn mode(&self) -> &AcmeMode {
|
||||
&self.mode
|
||||
}
|
||||
|
||||
fn build_acme_state(&self) -> (AcmeState<std::io::Error>, Arc<ResolvesServerCertAcme>) {
|
||||
let domains: Vec<String> = match &self.mode {
|
||||
AcmeMode::Domain { domain } => vec![domain.clone()],
|
||||
AcmeMode::Ip => vec![],
|
||||
};
|
||||
|
||||
let base_config = AcmeConfig::new(domains)
|
||||
.directory(&self.directory_url)
|
||||
.contact(self.contact.clone());
|
||||
|
||||
let state = match &self.cache_dir {
|
||||
Some(cache_dir) => {
|
||||
base_config.cache(DirCache::new(cache_dir.clone())).state()
|
||||
}
|
||||
None => {
|
||||
base_config
|
||||
.cache(rustls_acme::caches::NoCache::default())
|
||||
.state()
|
||||
}
|
||||
};
|
||||
|
||||
let resolver = state.resolver();
|
||||
(state, resolver)
|
||||
}
|
||||
|
||||
pub fn build_server_config_with_resolver(
|
||||
&self,
|
||||
resolver: Arc<ResolvesServerCertAcme>,
|
||||
) -> Result<Arc<ServerConfig>> {
|
||||
let provider = default_provider().into();
|
||||
let mut config = ServerConfig::builder_with_provider(provider)
|
||||
.with_safe_default_protocol_versions()
|
||||
.map_err(|e| anyhow!("failed to set protocol versions: {}", e))?
|
||||
.with_no_client_auth()
|
||||
.with_cert_resolver(resolver);
|
||||
config.alpn_protocols.push(ACME_TLS_ALPN_NAME.to_vec());
|
||||
Ok(Arc::new(config))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct AcmeTlsAcceptor {
|
||||
listener: TcpListener,
|
||||
listen_addr: SocketAddr,
|
||||
#[allow(dead_code)]
|
||||
server_config: Arc<ServerConfig>,
|
||||
tokio_acceptor: TokioTlsAcceptor,
|
||||
}
|
||||
|
||||
impl AcmeTlsAcceptor {
|
||||
pub async fn bind_acme(
|
||||
addr: SocketAddr,
|
||||
provider: Arc<AcmeCertProvider>,
|
||||
) -> Result<Self> {
|
||||
let (state, resolver) = provider.build_acme_state();
|
||||
|
||||
let server_config = provider.build_server_config_with_resolver(resolver.clone())?;
|
||||
|
||||
Self::spawn_state_worker(state, resolver);
|
||||
|
||||
let listener = TcpListener::bind(addr).await?;
|
||||
let listen_addr = listener.local_addr()?;
|
||||
|
||||
let tokio_acceptor = TokioTlsAcceptor::from(server_config.clone());
|
||||
|
||||
Ok(Self {
|
||||
listener,
|
||||
listen_addr,
|
||||
server_config,
|
||||
tokio_acceptor,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn listen_addr(&self) -> SocketAddr {
|
||||
self.listen_addr
|
||||
}
|
||||
|
||||
fn spawn_state_worker(state: AcmeState<std::io::Error>, resolver: Arc<ResolvesServerCertAcme>) {
|
||||
use futures::StreamExt;
|
||||
|
||||
let task = async move {
|
||||
let mut state = state;
|
||||
while let Some(event) = state.next().await {
|
||||
match event {
|
||||
Ok(ok) => {
|
||||
if let rustls_acme::EventOk::DeployedNewCert = ok {
|
||||
info!("ACME: new certificate deployed");
|
||||
} else {
|
||||
info!("ACME event: {:?}", ok);
|
||||
}
|
||||
}
|
||||
Err(err) => error!("ACME event error: {:?}", err),
|
||||
}
|
||||
if Arc::strong_count(&resolver) == 1 {
|
||||
info!("ACME resolver dropped, stopping background task");
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
tokio::spawn(task);
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl TransportAcceptor for AcmeTlsAcceptor {
|
||||
type Stream = tokio_rustls::server::TlsStream<tokio::net::TcpStream>;
|
||||
|
||||
async fn accept(&self) -> Result<(Self::Stream, TransportInfo)> {
|
||||
let (tcp_stream, remote_addr) = self.listener.accept().await?;
|
||||
let tls_stream = self.tokio_acceptor.accept(tcp_stream).await?;
|
||||
|
||||
let server_name = tls_stream
|
||||
.get_ref()
|
||||
.1
|
||||
.server_name()
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let info = TransportInfo {
|
||||
remote_addr: Some(remote_addr),
|
||||
transport_kind: TransportKind::Tls { server_name },
|
||||
};
|
||||
|
||||
Ok((tls_stream, info))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_domain_mode() {
|
||||
let provider = AcmeCertProvider::domain("example.com");
|
||||
assert!(matches!(provider.mode(), AcmeMode::Domain { .. }));
|
||||
if let AcmeMode::Domain { domain } = provider.mode() {
|
||||
assert_eq!(domain, "example.com");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_ip_mode() {
|
||||
let provider = AcmeCertProvider::ip();
|
||||
assert!(matches!(provider.mode(), AcmeMode::Ip));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_default_staging_directory() {
|
||||
let provider = AcmeCertProvider::domain("example.com");
|
||||
assert_eq!(
|
||||
provider.directory_url,
|
||||
rustls_acme::acme::LETS_ENCRYPT_STAGING_DIRECTORY
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_production_directory() {
|
||||
let provider = AcmeCertProvider::domain("example.com").with_production_directory();
|
||||
assert_eq!(
|
||||
provider.directory_url,
|
||||
rustls_acme::acme::LETS_ENCRYPT_PRODUCTION_DIRECTORY
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_custom_directory() {
|
||||
let provider =
|
||||
AcmeCertProvider::domain("example.com").with_directory("https://custom.acme.dir/");
|
||||
assert_eq!(provider.directory_url, "https://custom.acme.dir/");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_with_cache_dir() {
|
||||
let provider = AcmeCertProvider::domain("example.com").with_cache_dir("/tmp/acme_cache");
|
||||
assert_eq!(provider.cache_dir, Some(PathBuf::from("/tmp/acme_cache")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_with_contact() {
|
||||
let provider =
|
||||
AcmeCertProvider::domain("example.com").with_contact("mailto:admin@example.com");
|
||||
assert_eq!(
|
||||
provider.contact,
|
||||
vec!["mailto:admin@example.com".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_build_state_domain() {
|
||||
let provider = AcmeCertProvider::domain("example.com");
|
||||
let (_state, resolver) = provider.build_acme_state();
|
||||
assert!(Arc::strong_count(&resolver) >= 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_build_state_with_cache() {
|
||||
let provider =
|
||||
AcmeCertProvider::domain("example.com").with_cache_dir("/tmp/test_cache");
|
||||
let (_state, resolver) = provider.build_acme_state();
|
||||
assert!(Arc::strong_count(&resolver) >= 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_build_server_config() {
|
||||
let _ = default_provider().install_default();
|
||||
let provider = AcmeCertProvider::domain("example.com");
|
||||
let (_, resolver) = provider.build_acme_state();
|
||||
let config = provider.build_server_config_with_resolver(resolver).unwrap();
|
||||
assert!(!config.alpn_protocols.is_empty());
|
||||
assert!(config
|
||||
.alpn_protocols
|
||||
.iter()
|
||||
.any(|p| p == ACME_TLS_ALPN_NAME));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_mode_domain_debug() {
|
||||
let mode = AcmeMode::Domain {
|
||||
domain: "test.example.com".to_string(),
|
||||
};
|
||||
let debug_str = format!("{:?}", mode);
|
||||
assert!(debug_str.contains("test.example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_mode_ip_debug() {
|
||||
let mode = AcmeMode::Ip;
|
||||
let debug_str = format!("{:?}", mode);
|
||||
assert!(debug_str.contains("Ip"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn acme_cert_provider_builder_chain() {
|
||||
let provider = AcmeCertProvider::domain("test.example.com")
|
||||
.with_production_directory()
|
||||
.with_cache_dir("/tmp/cache")
|
||||
.with_contact("mailto:admin@test.example.com");
|
||||
assert!(matches!(provider.mode(), AcmeMode::Domain { .. }));
|
||||
assert_eq!(
|
||||
provider.directory_url,
|
||||
rustls_acme::acme::LETS_ENCRYPT_PRODUCTION_DIRECTORY
|
||||
);
|
||||
assert_eq!(provider.cache_dir, Some(PathBuf::from("/tmp/cache")));
|
||||
assert_eq!(provider.contact.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn acme_tls_acceptor_bind_acme() {
|
||||
let _ = default_provider().install_default();
|
||||
let provider = Arc::new(AcmeCertProvider::domain("example.com"));
|
||||
let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
|
||||
let acceptor = AcmeTlsAcceptor::bind_acme(addr, provider).await.unwrap();
|
||||
assert_ne!(acceptor.listen_addr().port(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn acme_staging_domain_cert_provisioning() {
|
||||
let _ = default_provider().install_default();
|
||||
|
||||
let cache_dir = tempfile::tempdir().unwrap();
|
||||
let provider = Arc::new(
|
||||
AcmeCertProvider::domain("acme-test.example.com")
|
||||
.with_cache_dir(cache_dir.path())
|
||||
.with_contact("mailto:admin@example.com"),
|
||||
);
|
||||
|
||||
let addr: SocketAddr = "0.0.0.0:443".parse().unwrap();
|
||||
let result = AcmeTlsAcceptor::bind_acme(addr, provider).await;
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"ACME TlsAcceptor should bind: {:?}",
|
||||
result.err()
|
||||
);
|
||||
|
||||
let acceptor = result.unwrap();
|
||||
assert_eq!(acceptor.listen_addr().port(), 443);
|
||||
}
|
||||
}
|
||||
@@ -12,12 +12,6 @@ mod tls;
|
||||
#[cfg(feature = "tls")]
|
||||
pub use tls::{AcmeConfig, TlsAcceptor, TlsTransport};
|
||||
|
||||
#[cfg(feature = "acme")]
|
||||
mod acme;
|
||||
|
||||
#[cfg(feature = "acme")]
|
||||
pub use acme::{AcmeCertProvider, AcmeMode, AcmeTlsAcceptor};
|
||||
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
@@ -9,16 +9,8 @@ use rustls::{ClientConfig, DigitallySignedStruct, RootCertStore, ServerConfig};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio_rustls::{client::TlsStream as ClientTlsStream, TlsAcceptor as TokioTlsAcceptor, TlsConnector};
|
||||
|
||||
#[cfg(feature = "acme")]
|
||||
use rustls::crypto::aws_lc_rs::default_provider;
|
||||
#[cfg(feature = "acme")]
|
||||
use rustls_acme::ResolvesServerCertAcme;
|
||||
|
||||
use super::{Transport, TransportAcceptor, TransportInfo, TransportKind};
|
||||
|
||||
#[cfg(feature = "acme")]
|
||||
const ACME_TLS_ALPN_NAME: &[u8] = b"acme-tls/1";
|
||||
|
||||
/// A TLS-based client transport that connects to a remote address over TLS.
|
||||
///
|
||||
/// Wraps a TCP connection with a TLS client session via `tokio_rustls::TlsConnector`.
|
||||
@@ -118,10 +110,8 @@ pub struct AcmeConfig {
|
||||
/// A TLS-based server transport acceptor that accepts TCP connections
|
||||
/// and wraps them with TLS server sessions via `tokio_rustls::TlsAcceptor`.
|
||||
///
|
||||
/// Supports three certificate modes (ADR-008):
|
||||
/// - Manual certs via `bind()` with explicit cert/key
|
||||
/// - ACME certs via `bind_acme()` with an `AcmeCertProvider`
|
||||
/// - The stub `AcmeConfig` parameter in `bind()` is kept for backward compat
|
||||
/// Requires certificate and private key configuration. Supports manual
|
||||
/// cert/key paths and an ACME config stub (ADR-008).
|
||||
pub struct TlsAcceptor {
|
||||
listener: TcpListener,
|
||||
listen_addr: SocketAddr,
|
||||
@@ -155,33 +145,6 @@ impl TlsAcceptor {
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(feature = "acme")]
|
||||
pub async fn bind_acme(
|
||||
addr: SocketAddr,
|
||||
acme_resolver: Arc<ResolvesServerCertAcme>,
|
||||
) -> Result<Self> {
|
||||
let listener = TcpListener::bind(addr).await?;
|
||||
let listen_addr = listener.local_addr()?;
|
||||
|
||||
let provider = default_provider().into();
|
||||
let mut server_config = ServerConfig::builder_with_provider(provider)
|
||||
.with_safe_default_protocol_versions()
|
||||
.map_err(|e| anyhow!("failed to set protocol versions: {}", e))?
|
||||
.with_no_client_auth()
|
||||
.with_cert_resolver(acme_resolver);
|
||||
server_config.alpn_protocols.push(ACME_TLS_ALPN_NAME.to_vec());
|
||||
|
||||
let server_config = Arc::new(server_config);
|
||||
let tokio_acceptor = TokioTlsAcceptor::from(server_config.clone());
|
||||
|
||||
Ok(Self {
|
||||
listener,
|
||||
listen_addr,
|
||||
server_config,
|
||||
tokio_acceptor,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn listen_addr(&self) -> SocketAddr {
|
||||
self.listen_addr
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
---
|
||||
id: client/channel-manager
|
||||
name: Implement ChannelManager — SSH session management, channel opens, reconnection
|
||||
status: pending
|
||||
status: done
|
||||
depends_on:
|
||||
- auth/client-auth-handler
|
||||
- transport/trait-and-types
|
||||
@@ -32,18 +32,18 @@ Reconnection is always enabled. The backoff caps at 30 seconds and continues ind
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
- [ ] `crates/wraith-core/src/client/channel_manager.rs` exports `ChannelManager`
|
||||
- [ ] `ChannelManager` holds: `Arc<Transport>`, `Arc<ClientAuthConfig>`, `Arc<client::Handle<ClientHandler>>` (behind RwLock for reconnection)
|
||||
- [ ] `ChannelManager::new()` establishes initial transport connection, authenticates, returns manager
|
||||
- [ ] `open_direct_tcpip(host, port)` — opens SSH channel to target
|
||||
- [ ] `request_tcpip_forward(addr, port)` — sends `tcpip_forward` request
|
||||
- [ ] `cancel_tcpip_forward(addr, port)` — sends `cancel_tcpip_forward` request
|
||||
- [ ] Reconnection detection: monitors `handle.is_closed()`, triggers reconnect on failure
|
||||
- [ ] Exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s (cap), continues indefinitely
|
||||
- [ ] Full reconnect: new transport stream, new SSH session over it (ADR-004)
|
||||
- [ ] After reconnect: port forward listeners (`-L`, `-R`) re-registered with new session
|
||||
- [ ] In-flight connections on old session fail gracefully (channel errors, not session-wide)
|
||||
- [ ] Unit tests: channel open, reconnection trigger, backoff timing, forward re-registration
|
||||
- [x] `crates/wraith-core/src/client/channel_manager.rs` exports `ChannelManager`
|
||||
- [x] `ChannelManager` holds: `Arc<Transport>`, `Arc<ClientAuthConfig>`, `Arc<client::Handle<ClientHandler>>` (behind RwLock for reconnection)
|
||||
- [x] `ChannelManager::new()` establishes initial transport connection, authenticates, returns manager
|
||||
- [x] `open_direct_tcpip(host, port)` — opens SSH channel to target
|
||||
- [x] `request_tcpip_forward(addr, port)` — sends `tcpip_forward` request
|
||||
- [x] `cancel_tcpip_forward(addr, port)` — sends `cancel_tcpip_forward` request
|
||||
- [x] Reconnection detection: monitors `handle.is_closed()`, triggers reconnect on failure
|
||||
- [x] Exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s (cap), continues indefinitely
|
||||
- [x] Full reconnect: new transport stream, new SSH session over it (ADR-004)
|
||||
- [x] After reconnect: port forward listeners (`-L`, `-R`) re-registered with new session
|
||||
- [x] In-flight connections on old session fail gracefully (channel errors, not session-wide)
|
||||
- [x] Unit tests: channel open, reconnection trigger, backoff timing, forward re-registration
|
||||
|
||||
## References
|
||||
|
||||
@@ -52,8 +52,13 @@ Reconnection is always enabled. The backoff caps at 30 seconds and continues ind
|
||||
|
||||
## Notes
|
||||
|
||||
> To be filled by implementation agent
|
||||
- Converted `client.rs` (single file) to directory module: `client/mod.rs` + `client/channel_manager.rs`
|
||||
- Used `russh::keys::PrivateKey` and `russh::keys::PublicKey` (not the nonexistent `russh::key::KeyPair`)
|
||||
- Reconnection monitor runs as a spawned tokio task that polls `handle.is_closed()` every 1s
|
||||
- On reconnect: creates new transport stream + new SSH session (ADR-004 full reconnect)
|
||||
- `ForwardRequest` struct tracks registered port forwards for re-registration after reconnect
|
||||
- In-flight channels on old session naturally fail with `ChannelError::ChannelClosed` since the handle is replaced
|
||||
|
||||
## Summary
|
||||
|
||||
> To be filled on completion
|
||||
Implemented `ChannelManager` in `crates/wraith-core/src/client/channel_manager.rs` with SSH session management, channel opens (`open_direct_tcpip`), port forward requests (`request_tcpip_forward`/`cancel_tcpip_forward`), and automatic reconnection with exponential backoff (1s→30s cap). Full reconnect per ADR-004 creates new transport stream + new SSH session. Port forwards are re-registered after successful reconnect. 8 unit tests covering backoff timing, forward tracking, transport failure, and reconnection detection.
|
||||
@@ -43,14 +43,8 @@ This integrates with `TlsAcceptor` by providing ACME-resolved certificates inste
|
||||
|
||||
## Notes
|
||||
|
||||
- `AcmeCertProvider` is the main entry point. It creates `AcmeState` and `ResolvesServerCertAcme` from `rustls-acme`.
|
||||
- The `ResolvesServerCertAcme` resolver is shared between the `AcmeState` background task and the `ServerConfig`, so cert updates propagate automatically.
|
||||
- `AcmeTlsAcceptor::bind_acme()` creates a TLS acceptor that uses ACME-provisioned certs and spawns a background tokio task for auto-renewal.
|
||||
- `TlsAcceptor::bind_acme()` also added for users who want to use ACME with the standard `TlsAcceptor` type directly.
|
||||
- The `AcmeConfig` stub in `tls.rs` is retained for backward compat with existing `TlsAcceptor::bind()`.
|
||||
- `acme` feature implies `tls` and adds `rustls-acme` + `futures` dependencies.
|
||||
- TLS-ALPN-01 challenge handling works via the `acme-tls/1` ALPN protocol registered in `ServerConfig` — the resolver dispatches challenge vs regular certs automatically.
|
||||
> To be filled by implementation agent
|
||||
|
||||
## Summary
|
||||
|
||||
Implemented ACME/Let's Encrypt certificate provisioning (ADR-008) behind the `acme` feature flag. `AcmeCertProvider` supports domain-based and IP-based modes using `rustls-acme`. `AcmeTlsAcceptor::bind_acme()` and `TlsAcceptor::bind_acme()` provide ACME-integrated TLS acceptance with automatic certificate renewal via a background tokio task. Unit tests cover config construction, builder patterns, and server config generation. Integration test for LE staging is marked `#[ignore]`.
|
||||
> To be filled on completion
|
||||
Reference in New Issue
Block a user