Files
alknet/crates/alknet-core/src/interface/ssh.rs

983 lines
30 KiB
Rust

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Instant;
use anyhow::Result;
use arc_swap::ArcSwap;
use async_trait::async_trait;
use russh::keys::ssh_key::HashAlg;
use russh::server::{self, Config};
use russh::Channel;
use russh::ChannelId;
use tokio::sync::mpsc;
use crate::auth::identity::{Identity, IdentityProvider};
use crate::call::frame::{FrameFramedReader, FrameFramedWriter};
use crate::call::EventEnvelope;
use crate::config::DynamicConfig;
use crate::interface::session::{InterfaceEvent, InterfaceSession};
use crate::interface::{StreamInterface, StreamInterfaceConfig, TransportStream};
use crate::server::control_channel::{
ControlChannelHandler, ControlChannelRouter, DuplexStream, ALKNET_CONTROL_DESTINATION,
ALKNET_PREFIX,
};
use crate::server::rate_limit::{AuthAttemptLimiter, ConnectionRateLimiter};
use crate::transport::TransportKind;
struct SshHandler {
dynamic: Arc<ArcSwap<DynamicConfig>>,
identity_provider: Arc<dyn IdentityProvider>,
outbound_proxy: Option<crate::server::handler::ProxyConfig>,
remote_addr: Option<SocketAddr>,
transport: TransportKind,
connection_limiter: Arc<ConnectionRateLimiter>,
connection_allowed: bool,
auth_limiter: AuthAttemptLimiter,
authenticated_identity: Option<Identity>,
control_channel_router: ControlChannelRouter,
bridge_event_tx: Option<mpsc::Sender<InterfaceEvent>>,
bridge_envelope_rx: Option<mpsc::Receiver<EventEnvelope>>,
connected_at: Instant,
}
impl SshHandler {
fn new(
dynamic: Arc<ArcSwap<DynamicConfig>>,
identity_provider: Arc<dyn IdentityProvider>,
outbound_proxy: Option<crate::server::handler::ProxyConfig>,
remote_addr: Option<SocketAddr>,
transport: TransportKind,
connection_limiter: Arc<ConnectionRateLimiter>,
max_auth_attempts: usize,
) -> Self {
let allowed = if let Some(addr) = remote_addr {
let ip = addr.ip();
if connection_limiter.check(ip) {
connection_limiter.on_connect(ip);
tracing::info!(
remote_addr = %addr,
transport = %transport,
"connection opened"
);
true
} else {
tracing::info!(
remote_addr = %addr,
transport = %transport,
"connection rejected"
);
false
}
} else {
true
};
Self {
dynamic,
identity_provider,
outbound_proxy,
remote_addr,
transport,
connection_limiter,
connection_allowed: allowed,
auth_limiter: AuthAttemptLimiter::new(max_auth_attempts),
authenticated_identity: None,
control_channel_router: ControlChannelRouter::without_handler(),
bridge_event_tx: None,
bridge_envelope_rx: None,
connected_at: Instant::now(),
}
}
#[allow(dead_code)]
fn with_control_channel_router(mut self, router: ControlChannelRouter) -> Self {
self.control_channel_router = router;
self
}
fn with_bridge_channels(
mut self,
event_tx: mpsc::Sender<InterfaceEvent>,
envelope_rx: mpsc::Receiver<EventEnvelope>,
) -> Self {
self.bridge_event_tx = Some(event_tx);
self.bridge_envelope_rx = Some(envelope_rx);
self
}
fn has_control_channel_bridge(&self) -> bool {
self.bridge_event_tx.is_some() && self.bridge_envelope_rx.is_some()
}
}
impl Drop for SshHandler {
fn drop(&mut self) {
if let Some(addr) = self.remote_addr {
if self.connection_allowed {
self.connection_limiter.on_disconnect(addr.ip());
let duration = self.connected_at.elapsed();
tracing::info!(
remote_addr = %addr,
duration_secs = duration.as_secs_f64(),
"connection closed"
);
}
}
}
}
#[async_trait]
impl server::Handler for SshHandler {
type Error = russh::Error;
async fn auth_publickey(
&mut self,
user: &str,
public_key: &russh::keys::ssh_key::PublicKey,
) -> Result<server::Auth, Self::Error> {
if !self.auth_limiter.check() {
let remote_addr_display = self
.remote_addr
.map_or("unknown".to_string(), |a| a.to_string());
let fingerprint = format!("{}", public_key.fingerprint(HashAlg::Sha256));
tracing::info!(
remote_addr = %remote_addr_display,
user = user,
key_fingerprint = %fingerprint,
result = "reject",
"auth attempt"
);
return Ok(server::Auth::Reject {
proceed_with_methods: None,
});
}
let fingerprint = format!("{}", public_key.fingerprint(HashAlg::Sha256));
let remote_addr_display = self
.remote_addr
.map_or("unknown".to_string(), |a| a.to_string());
let identity = self
.identity_provider
.resolve_from_fingerprint(&fingerprint);
match identity {
Some(id) => {
self.authenticated_identity = Some(id);
tracing::info!(
remote_addr = %remote_addr_display,
user = user,
key_fingerprint = %fingerprint,
result = "accept",
"auth attempt"
);
Ok(server::Auth::Accept)
}
None => {
self.auth_limiter.on_failure();
tracing::info!(
remote_addr = %remote_addr_display,
user = user,
key_fingerprint = %fingerprint,
result = "reject",
"auth attempt"
);
Ok(server::Auth::Reject {
proceed_with_methods: None,
})
}
}
}
async fn channel_open_direct_tcpip(
&mut self,
channel: Channel<server::Msg>,
host_to_connect: &str,
port_to_connect: u32,
originator_address: &str,
originator_port: u32,
_session: &mut server::Session,
) -> Result<bool, Self::Error> {
if host_to_connect.starts_with(ALKNET_PREFIX) {
if host_to_connect == ALKNET_CONTROL_DESTINATION && self.has_control_channel_bridge() {
let event_tx = self.bridge_event_tx.take().unwrap();
let envelope_rx = self.bridge_envelope_rx.take().unwrap();
let identity = self.authenticated_identity.clone();
tokio::spawn(async move {
let stream = channel.into_stream();
let (read_half, write_half) = tokio::io::split(stream);
run_control_channel_bridge(
read_half,
write_half,
identity,
event_tx,
envelope_rx,
)
.await;
});
let _ = (originator_address, originator_port);
return Ok(true);
}
if self.control_channel_router.has_handler() {
if let Some(handler) = self.control_channel_router.take_handler() {
let stream: Box<dyn DuplexStream> = Box::new(channel.into_stream());
tokio::spawn(async move {
handler.handle_channel(stream).await;
});
}
let _ = (originator_address, originator_port);
return Ok(true);
}
return Ok(false);
}
let identity = self
.authenticated_identity
.clone()
.unwrap_or_else(|| Identity {
id: String::new(),
scopes: vec![],
resources: std::collections::HashMap::new(),
});
let policy = self.dynamic.load();
let allowed = policy.forwarding.check(
host_to_connect,
port_to_connect as u16,
&identity,
self.transport.clone(),
);
if !allowed {
tracing::info!(
remote_addr = ?self.remote_addr,
target = %format!("{host_to_connect}:{port_to_connect}"),
identity = %identity.id,
transport = %self.transport,
"forwarding denied by policy"
);
return Ok(false);
}
let target_host = host_to_connect.to_string();
let target_port = port_to_connect;
let proxy_config =
self.outbound_proxy
.clone()
.unwrap_or(crate::server::handler::ProxyConfig {
mode: crate::server::handler::ProxyMode::Direct,
});
tokio::spawn(async move {
let target = match format!("{target_host}:{target_port}")
.parse::<std::net::SocketAddr>()
{
Ok(addr) => addr,
Err(_) => {
match tokio::net::lookup_host((&target_host[..], target_port as u16)).await {
Ok(mut addrs) => match addrs.next() {
Some(addr) => addr,
None => return,
},
Err(_) => return,
}
}
};
crate::server::channel_proxy::proxy_channel(
channel.into_stream(),
target,
&proxy_config,
)
.await;
});
let _ = (originator_address, originator_port);
Ok(true)
}
async fn channel_open_session(
&mut self,
_channel: Channel<server::Msg>,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
"rejected session channel (shell/exec not supported)"
);
let _ = session;
Ok(false)
}
async fn channel_open_x11(
&mut self,
_channel: Channel<server::Msg>,
_originator_address: &str,
_originator_port: u32,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
"rejected x11 channel"
);
let _ = session;
Ok(false)
}
async fn channel_open_forwarded_tcpip(
&mut self,
_channel: Channel<server::Msg>,
host_to_connect: &str,
port_to_connect: u32,
_originator_address: &str,
_originator_port: u32,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
target = %format!("{host_to_connect}:{port_to_connect}"),
"rejected forwarded-tcpip channel (remote port forwarding not supported)"
);
let _ = session;
Ok(false)
}
async fn exec_request(
&mut self,
channel: ChannelId,
data: &[u8],
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
data_len = data.len(),
"rejected exec request on channel (shell/exec not supported)"
);
let _ = session.channel_failure(channel);
Ok(())
}
async fn shell_request(
&mut self,
channel: ChannelId,
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
"rejected shell request on channel"
);
let _ = session.channel_failure(channel);
Ok(())
}
async fn subsystem_request(
&mut self,
channel: ChannelId,
name: &str,
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
subsystem = name,
"rejected subsystem request on channel"
);
let _ = session.channel_failure(channel);
Ok(())
}
async fn pty_request(
&mut self,
channel: ChannelId,
term: &str,
col_width: u32,
row_height: u32,
pix_width: u32,
pix_height: u32,
modes: &[(russh::Pty, u32)],
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
term = term,
"rejected pty request on channel"
);
let _ = (col_width, row_height, pix_width, pix_height, modes);
let _ = session.channel_failure(channel);
Ok(())
}
async fn env_request(
&mut self,
channel: ChannelId,
variable_name: &str,
variable_value: &str,
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
variable = variable_name,
"rejected env request on channel"
);
let _ = variable_value;
let _ = session.channel_failure(channel);
Ok(())
}
async fn x11_request(
&mut self,
channel: ChannelId,
single_connection: bool,
x11_auth_protocol: &str,
x11_auth_cookie: &str,
x11_screen_number: u32,
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
"rejected x11 request on channel"
);
let _ = (
single_connection,
x11_auth_protocol,
x11_auth_cookie,
x11_screen_number,
);
let _ = session.channel_failure(channel);
Ok(())
}
async fn agent_request(
&mut self,
channel: ChannelId,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
channel = %channel,
"rejected agent forwarding request on channel"
);
let _ = session;
Ok(false)
}
async fn tcpip_forward(
&mut self,
address: &str,
port: &mut u32,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
address = address,
port = *port,
"rejected tcpip-forward request (remote port forwarding not supported)"
);
let _ = session;
Ok(false)
}
async fn cancel_tcpip_forward(
&mut self,
address: &str,
port: u32,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
let _ = (address, port, session);
Ok(false)
}
async fn streamlocal_forward(
&mut self,
socket_path: &str,
session: &mut server::Session,
) -> Result<bool, Self::Error> {
tracing::warn!(
remote_addr = ?self.remote_addr,
socket_path = socket_path,
"rejected streamlocal-forward request"
);
let _ = session;
Ok(false)
}
async fn signal(
&mut self,
channel: ChannelId,
signal: russh::Sig,
session: &mut server::Session,
) -> Result<(), Self::Error> {
tracing::debug!(
remote_addr = ?self.remote_addr,
channel = %channel,
signal = ?signal,
"received signal on channel (ignored)"
);
let _ = session;
Ok(())
}
}
pub struct SshInterface {
config: Arc<Config>,
dynamic: Arc<ArcSwap<DynamicConfig>>,
connection_limiter: Arc<ConnectionRateLimiter>,
outbound_proxy: Option<crate::server::handler::ProxyConfig>,
max_auth_attempts: usize,
}
impl SshInterface {
pub fn new(config: Arc<Config>, dynamic: Arc<ArcSwap<DynamicConfig>>) -> Self {
Self {
config,
dynamic,
connection_limiter: Arc::new(ConnectionRateLimiter::new(0)),
outbound_proxy: None,
max_auth_attempts: 10,
}
}
pub fn with_connection_limiter(mut self, limiter: Arc<ConnectionRateLimiter>) -> Self {
self.connection_limiter = limiter;
self
}
pub fn with_outbound_proxy(
mut self,
proxy: Option<crate::server::handler::ProxyConfig>,
) -> Self {
self.outbound_proxy = proxy;
self
}
pub fn with_max_auth_attempts(mut self, max: usize) -> Self {
self.max_auth_attempts = max;
self
}
pub fn config(&self) -> &Arc<Config> {
&self.config
}
pub fn dynamic(&self) -> &Arc<ArcSwap<DynamicConfig>> {
&self.dynamic
}
async fn accept_inner(
&self,
stream: Box<dyn TransportStream>,
ssh_config: &crate::interface::SshInterfaceConfig,
remote_addr: Option<SocketAddr>,
transport: TransportKind,
) -> Result<SshSession> {
let identity_provider = Arc::clone(&ssh_config.auth);
let _forwarding = Arc::clone(&ssh_config.forwarding);
let (event_tx, event_rx) = mpsc::channel::<InterfaceEvent>(256);
let (envelope_tx, envelope_rx) = mpsc::channel::<EventEnvelope>(256);
let handler = SshHandler::new(
Arc::clone(&self.dynamic),
identity_provider,
self.outbound_proxy.clone(),
remote_addr,
transport,
Arc::clone(&self.connection_limiter),
self.max_auth_attempts,
)
.with_bridge_channels(event_tx, envelope_rx);
let running = server::run_stream(Arc::clone(&self.config), stream, handler).await?;
let handle = running.handle();
let join = tokio::spawn(async {
let _ = running.await;
});
Ok(SshSession {
handle,
_join: join,
event_rx,
envelope_tx,
})
}
}
#[async_trait]
impl StreamInterface for SshInterface {
type Session = SshSession;
async fn accept(
&self,
stream: Box<dyn TransportStream>,
config: &StreamInterfaceConfig,
) -> Result<Self::Session> {
let ssh_config = match config {
StreamInterfaceConfig::Ssh(c) => c,
StreamInterfaceConfig::RawFraming(_) => {
return Err(anyhow::anyhow!("SshInterface received RawFramingConfig"));
}
};
self.accept_inner(stream, ssh_config, None, TransportKind::Tcp)
.await
}
}
pub struct SshSession {
handle: server::Handle,
_join: tokio::task::JoinHandle<()>,
event_rx: mpsc::Receiver<InterfaceEvent>,
envelope_tx: mpsc::Sender<EventEnvelope>,
}
impl SshSession {
pub fn handle(&self) -> &server::Handle {
&self.handle
}
}
#[async_trait]
impl InterfaceSession for SshSession {
async fn recv(&mut self) -> Option<InterfaceEvent> {
self.event_rx.recv().await
}
async fn send(&mut self, envelope: EventEnvelope) -> Result<()> {
self.envelope_tx
.send(envelope)
.await
.map_err(|_| anyhow::anyhow!("control channel bridge closed"))
}
}
async fn run_control_channel_bridge<R, W>(
read_half: R,
write_half: W,
identity: Option<Identity>,
event_tx: mpsc::Sender<InterfaceEvent>,
mut envelope_rx: mpsc::Receiver<EventEnvelope>,
) where
R: tokio::io::AsyncRead + Unpin,
W: tokio::io::AsyncWrite + Unpin,
{
let mut reader = FrameFramedReader::new(read_half);
let mut writer = FrameFramedWriter::new(write_half);
loop {
tokio::select! {
frame = reader.read_frame() => {
match frame {
Ok(Some(envelope)) => {
let event = match &identity {
Some(id) => InterfaceEvent::with_identity(envelope, id.clone()),
None => InterfaceEvent::new(envelope),
};
if event_tx.send(event).await.is_err() {
return;
}
}
Ok(None) => return,
Err(_) => return,
}
}
envelope = envelope_rx.recv() => {
match envelope {
Some(envelope) => {
if writer.write_frame(&envelope).await.is_err() {
return;
}
}
None => return,
}
}
}
}
}
pub struct ControlChannelBridge {
identity: Option<Identity>,
}
impl ControlChannelBridge {
pub fn new(identity: Option<Identity>) -> Self {
Self { identity }
}
}
#[async_trait]
impl ControlChannelHandler for ControlChannelBridge {
async fn handle_channel(&self, stream: Box<dyn DuplexStream>) {
let (event_tx, _event_rx) = mpsc::channel::<InterfaceEvent>(256);
let (_envelope_tx, envelope_rx) = mpsc::channel::<EventEnvelope>(256);
let identity = self.identity.clone();
let (read_half, write_half) = tokio::io::split(stream);
tokio::spawn(run_control_channel_bridge(
read_half,
write_half,
identity,
event_tx,
envelope_rx,
));
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::call::frame::{FrameFramedReader, FrameFramedWriter};
use tokio::io::duplex;
#[test]
fn ssh_interface_constructs_with_config() {
let config = Arc::new(Config {
keys: vec![russh::keys::PrivateKey::random(
&mut rand_core::OsRng,
russh::keys::Algorithm::Ed25519,
)
.unwrap()],
..Default::default()
});
let dynamic = Arc::new(ArcSwap::new(Arc::new(DynamicConfig::default())));
let iface = SshInterface::new(config, dynamic);
assert!(iface.config().keys.len() >= 1);
}
#[test]
fn ssh_interface_builder_pattern() {
let config = Arc::new(Config {
keys: vec![russh::keys::PrivateKey::random(
&mut rand_core::OsRng,
russh::keys::Algorithm::Ed25519,
)
.unwrap()],
..Default::default()
});
let dynamic = Arc::new(ArcSwap::new(Arc::new(DynamicConfig::default())));
let limiter = Arc::new(ConnectionRateLimiter::new(5));
let iface = SshInterface::new(config, dynamic)
.with_connection_limiter(limiter)
.with_max_auth_attempts(3);
assert!(iface.config().keys.len() >= 1);
}
#[test]
fn ssh_handler_auth_delegates_to_identity_provider() {
use std::collections::HashMap;
struct MockProvider {
identities: HashMap<String, Identity>,
}
impl IdentityProvider for MockProvider {
fn resolve_from_fingerprint(&self, fp: &str) -> Option<Identity> {
self.identities.get(fp).cloned()
}
fn resolve_from_token(&self, _t: &crate::auth::AuthToken) -> Option<Identity> {
None
}
}
let mut ids = HashMap::new();
ids.insert(
"SHA256:testkey".to_string(),
Identity {
id: "SHA256:testkey".to_string(),
scopes: vec!["admin".to_string()],
resources: HashMap::new(),
},
);
let provider: Arc<dyn IdentityProvider> = Arc::new(MockProvider { identities: ids });
let dynamic = Arc::new(ArcSwap::new(Arc::new(DynamicConfig::default())));
let limiter = Arc::new(ConnectionRateLimiter::new(0));
let handler = SshHandler::new(
dynamic,
provider,
None,
None,
TransportKind::Tcp,
limiter,
10,
);
assert!(handler.authenticated_identity.is_none());
}
#[test]
fn ssh_handler_connection_rate_limiting() {
let dynamic = Arc::new(ArcSwap::new(Arc::new(DynamicConfig::default())));
let provider: Arc<dyn IdentityProvider> = Arc::new(
crate::auth::identity::ConfigIdentityProvider::new(Arc::clone(&dynamic)),
);
let limiter = Arc::new(ConnectionRateLimiter::new(1));
let addr: SocketAddr = "10.0.0.1:22".parse().unwrap();
let h1 = SshHandler::new(
Arc::clone(&dynamic),
Arc::clone(&provider),
None,
Some(addr),
TransportKind::Tcp,
Arc::clone(&limiter),
10,
);
assert!(h1.connection_allowed);
let h2 = SshHandler::new(
dynamic,
provider,
None,
Some(addr),
TransportKind::Tcp,
limiter,
10,
);
assert!(!h2.connection_allowed);
}
#[tokio::test]
async fn ssh_interface_rejects_raw_framing_config() {
let config = Arc::new(Config {
keys: vec![russh::keys::PrivateKey::random(
&mut rand_core::OsRng,
russh::keys::Algorithm::Ed25519,
)
.unwrap()],
..Default::default()
});
let dynamic = Arc::new(ArcSwap::new(Arc::new(DynamicConfig::default())));
let iface = SshInterface::new(config, dynamic);
let (_client, server) = tokio::io::duplex(1024);
let stream: Box<dyn TransportStream> = Box::new(server);
let raw_config = StreamInterfaceConfig::RawFraming(crate::interface::RawFramingConfig {
auth: Arc::new(crate::auth::ConfigIdentityProvider::new(Arc::new(
ArcSwap::new(Arc::new(DynamicConfig::default())),
))),
});
let result = iface.accept(stream, &raw_config).await;
assert!(result.is_err());
}
#[tokio::test]
async fn ssh_session_round_trip_event_envelope() {
let (client, server) = duplex(4096);
let (event_tx, mut event_rx) = mpsc::channel::<InterfaceEvent>(256);
let (envelope_tx, envelope_rx) = mpsc::channel::<EventEnvelope>(256);
let identity = Identity {
id: "SHA256:test".to_string(),
scopes: vec![],
resources: std::collections::HashMap::new(),
};
let identity_clone = identity.clone();
let (server_read, server_write) = tokio::io::split(server);
tokio::spawn(run_control_channel_bridge(
server_read,
server_write,
Some(identity_clone),
event_tx,
envelope_rx,
));
let (client_read, client_write) = tokio::io::split(client);
let mut client_reader = FrameFramedReader::new(client_read);
let mut client_writer = FrameFramedWriter::new(client_write);
let envelope = EventEnvelope::call_requested("req-1", serde_json::json!({"op": "test"}));
client_writer.write_frame(&envelope).await.unwrap();
let received_event =
tokio::time::timeout(std::time::Duration::from_secs(2), event_rx.recv())
.await
.unwrap()
.unwrap();
assert_eq!(received_event.envelope, envelope);
assert_eq!(received_event.identity.as_ref().unwrap().id, "SHA256:test");
let response = EventEnvelope::call_responded("req-1", serde_json::json!({"result": 42}));
envelope_tx.send(response.clone()).await.unwrap();
let read_back = tokio::time::timeout(
std::time::Duration::from_secs(2),
client_reader.read_frame(),
)
.await
.unwrap()
.unwrap()
.unwrap();
assert_eq!(read_back, response);
}
#[tokio::test]
async fn ssh_session_recv_without_identity() {
let (client, server) = duplex(4096);
let (event_tx, mut event_rx) = mpsc::channel::<InterfaceEvent>(256);
let (_envelope_tx, envelope_rx) = mpsc::channel::<EventEnvelope>(256);
let (server_read, server_write) = tokio::io::split(server);
tokio::spawn(run_control_channel_bridge(
server_read,
server_write,
None,
event_tx,
envelope_rx,
));
let (client_read, client_write) = tokio::io::split(client);
let mut client_writer = FrameFramedWriter::new(client_write);
let _client_reader = FrameFramedReader::new(client_read);
let envelope = EventEnvelope::call_requested("req-2", serde_json::json!({"op": "no-id"}));
client_writer.write_frame(&envelope).await.unwrap();
let received_event =
tokio::time::timeout(std::time::Duration::from_secs(2), event_rx.recv())
.await
.unwrap()
.unwrap();
assert_eq!(received_event.envelope, envelope);
assert!(received_event.identity.is_none());
}
#[tokio::test]
async fn control_channel_router_with_handler_routes_data() {
let called = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let called_clone = called.clone();
struct TrackingHandler {
called: std::sync::Arc<std::sync::atomic::AtomicBool>,
}
#[async_trait]
impl ControlChannelHandler for TrackingHandler {
async fn handle_channel(&self, _stream: Box<dyn DuplexStream>) {
self.called.store(true, std::sync::atomic::Ordering::SeqCst);
}
}
let router = ControlChannelRouter::with_handler(Box::new(TrackingHandler {
called: called_clone,
}));
assert!(router.has_handler());
let (_client, server) = duplex(64);
let stream: Box<dyn DuplexStream> = Box::new(server);
let result = router.route(stream).await;
assert!(result.is_ok());
assert!(called.load(std::sync::atomic::Ordering::SeqCst));
}
}