Compare commits
1 Commits
feat/napi/
...
feat/serve
| Author | SHA1 | Date | |
|---|---|---|---|
| f963898a05 |
@@ -60,6 +60,22 @@ pub enum ConfigError {
|
|||||||
IncompatibleOptions,
|
IncompatibleOptions,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, thiserror::Error)]
|
||||||
|
pub enum ForwardError {
|
||||||
|
#[error("invalid forward specification: {spec}")]
|
||||||
|
InvalidSpec { spec: String },
|
||||||
|
#[error("bind failed")]
|
||||||
|
BindFailed {
|
||||||
|
#[source]
|
||||||
|
source: io::Error,
|
||||||
|
},
|
||||||
|
#[error("channel open failed")]
|
||||||
|
ChannelOpenFailed {
|
||||||
|
#[source]
|
||||||
|
source: Box<dyn std::error::Error + Send + Sync>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|||||||
186
crates/wraith-core/src/server/control_channel.rs
Normal file
186
crates/wraith-core/src/server/control_channel.rs
Normal file
@@ -0,0 +1,186 @@
|
|||||||
|
use std::io;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use tokio::io::{AsyncRead, AsyncWrite};
|
||||||
|
|
||||||
|
pub const WRAITH_CONTROL_DESTINATION: &str = "wraith-control";
|
||||||
|
pub const WRAITH_PREFIX: &str = "wraith-";
|
||||||
|
|
||||||
|
pub fn is_reserved_destination(host: &str) -> bool {
|
||||||
|
host.starts_with(WRAITH_PREFIX)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub trait DuplexStream: AsyncRead + AsyncWrite + Unpin + Send {}
|
||||||
|
|
||||||
|
impl<T: AsyncRead + AsyncWrite + Unpin + Send> DuplexStream for T {}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
pub trait ControlChannelHandler: Send + Sync {
|
||||||
|
async fn handle_channel(&self, stream: Box<dyn DuplexStream>);
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ControlChannelRouter {
|
||||||
|
handler: Option<Box<dyn ControlChannelHandler>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ControlChannelRouter {
|
||||||
|
pub fn new(handler: Option<Box<dyn ControlChannelHandler>>) -> Self {
|
||||||
|
Self { handler }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn without_handler() -> Self {
|
||||||
|
Self { handler: None }
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn with_handler(handler: Box<dyn ControlChannelHandler>) -> Self {
|
||||||
|
Self {
|
||||||
|
handler: Some(handler),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn has_handler(&self) -> bool {
|
||||||
|
self.handler.is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn route(&self, stream: Box<dyn DuplexStream>) -> io::Result<()> {
|
||||||
|
match &self.handler {
|
||||||
|
Some(handler) => {
|
||||||
|
handler.handle_channel(stream).await;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
None => Err(io::Error::new(
|
||||||
|
io::ErrorKind::ConnectionRefused,
|
||||||
|
"no control channel handler configured",
|
||||||
|
)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use tokio::io::duplex;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn wraith_control_destination_constant() {
|
||||||
|
assert_eq!(WRAITH_CONTROL_DESTINATION, "wraith-control");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn wraith_prefix_constant() {
|
||||||
|
assert_eq!(WRAITH_PREFIX, "wraith-");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn reserved_destination_detected() {
|
||||||
|
assert!(is_reserved_destination("wraith-control"));
|
||||||
|
assert!(is_reserved_destination("wraith-status"));
|
||||||
|
assert!(is_reserved_destination("wraith-events"));
|
||||||
|
assert!(is_reserved_destination("wraith-"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn non_reserved_destination_passes_through() {
|
||||||
|
assert!(!is_reserved_destination("example.com"));
|
||||||
|
assert!(!is_reserved_destination("localhost"));
|
||||||
|
assert!(!is_reserved_destination("192.168.1.1"));
|
||||||
|
assert!(!is_reserved_destination("wraith.example.com"));
|
||||||
|
assert!(!is_reserved_destination(""));
|
||||||
|
assert!(!is_reserved_destination("wrait-control"));
|
||||||
|
assert!(!is_reserved_destination("WRAITH-control"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn prefix_matching_case_sensitive() {
|
||||||
|
assert!(!is_reserved_destination("Wraith-control"));
|
||||||
|
assert!(!is_reserved_destination("WRAITH-control"));
|
||||||
|
assert!(is_reserved_destination("wraith-Control"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn router_without_handler_has_no_handler() {
|
||||||
|
let router = ControlChannelRouter::without_handler();
|
||||||
|
assert!(!router.has_handler());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn router_with_handler_has_handler() {
|
||||||
|
struct DummyHandler;
|
||||||
|
#[async_trait]
|
||||||
|
impl ControlChannelHandler for DummyHandler {
|
||||||
|
async fn handle_channel(&self, _stream: Box<dyn DuplexStream>) {}
|
||||||
|
}
|
||||||
|
let router = ControlChannelRouter::with_handler(Box::new(DummyHandler));
|
||||||
|
assert!(router.has_handler());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn route_without_handler_returns_error() {
|
||||||
|
let router = ControlChannelRouter::without_handler();
|
||||||
|
let (_client, server) = duplex(64);
|
||||||
|
let stream: Box<dyn DuplexStream> = Box::new(server);
|
||||||
|
let result = router.route(stream).await;
|
||||||
|
assert!(result.is_err());
|
||||||
|
let err = result.unwrap_err();
|
||||||
|
assert_eq!(err.kind(), io::ErrorKind::ConnectionRefused);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn route_with_handler_succeeds() {
|
||||||
|
use std::sync::atomic::{AtomicBool, Ordering};
|
||||||
|
use std::sync::Arc;
|
||||||
|
|
||||||
|
struct TrackedHandler {
|
||||||
|
called: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
#[async_trait]
|
||||||
|
impl ControlChannelHandler for TrackedHandler {
|
||||||
|
async fn handle_channel(&self, _stream: Box<dyn DuplexStream>) {
|
||||||
|
self.called.store(true, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let called = Arc::new(AtomicBool::new(false));
|
||||||
|
let handler = TrackedHandler {
|
||||||
|
called: called.clone(),
|
||||||
|
};
|
||||||
|
let router = ControlChannelRouter::with_handler(Box::new(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(Ordering::SeqCst));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn route_with_handler_can_read_write() {
|
||||||
|
struct EchoHandler;
|
||||||
|
#[async_trait]
|
||||||
|
impl ControlChannelHandler for EchoHandler {
|
||||||
|
async fn handle_channel(&self, mut stream: Box<dyn DuplexStream>) {
|
||||||
|
let mut buf = [0u8; 64];
|
||||||
|
let n = stream.read(&mut buf).await.unwrap();
|
||||||
|
stream.write_all(&buf[..n]).await.unwrap();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let router = ControlChannelRouter::with_handler(Box::new(EchoHandler));
|
||||||
|
let (client, server) = duplex(64);
|
||||||
|
let stream: Box<dyn DuplexStream> = Box::new(server);
|
||||||
|
tokio::spawn(async move {
|
||||||
|
router.route(stream).await.unwrap();
|
||||||
|
});
|
||||||
|
|
||||||
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||||
|
let mut client = client;
|
||||||
|
client.write_all(b"hello").await.unwrap();
|
||||||
|
let mut buf = [0u8; 5];
|
||||||
|
client.read_exact(&mut buf).await.unwrap();
|
||||||
|
assert_eq!(&buf, b"hello");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn control_channel_destination_matches_prefix() {
|
||||||
|
assert!(is_reserved_destination(WRAITH_CONTROL_DESTINATION));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -7,8 +7,9 @@ use russh::server::{Auth, Handler, Msg, Session};
|
|||||||
use russh::Channel;
|
use russh::Channel;
|
||||||
|
|
||||||
use crate::auth::ServerAuthConfig;
|
use crate::auth::ServerAuthConfig;
|
||||||
|
use crate::server::control_channel::{
|
||||||
const WRAITH_PREFIX: &str = "wraith-";
|
ControlChannelHandler, ControlChannelRouter, WRAITH_PREFIX,
|
||||||
|
};
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub enum ProxyMode {
|
pub enum ProxyMode {
|
||||||
@@ -26,6 +27,7 @@ pub struct ServerHandler {
|
|||||||
auth_config: Arc<ServerAuthConfig>,
|
auth_config: Arc<ServerAuthConfig>,
|
||||||
outbound_proxy: Option<ProxyConfig>,
|
outbound_proxy: Option<ProxyConfig>,
|
||||||
remote_addr: Option<SocketAddr>,
|
remote_addr: Option<SocketAddr>,
|
||||||
|
control_channel_router: ControlChannelRouter,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ServerHandler {
|
impl ServerHandler {
|
||||||
@@ -38,8 +40,21 @@ impl ServerHandler {
|
|||||||
auth_config,
|
auth_config,
|
||||||
outbound_proxy,
|
outbound_proxy,
|
||||||
remote_addr,
|
remote_addr,
|
||||||
|
control_channel_router: ControlChannelRouter::without_handler(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn with_control_channel_handler(
|
||||||
|
mut self,
|
||||||
|
handler: Box<dyn ControlChannelHandler>,
|
||||||
|
) -> Self {
|
||||||
|
self.control_channel_router = ControlChannelRouter::with_handler(handler);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn control_channel_router(&self) -> &ControlChannelRouter {
|
||||||
|
&self.control_channel_router
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
@@ -98,6 +113,16 @@ impl Handler for ServerHandler {
|
|||||||
port = port_to_connect,
|
port = port_to_connect,
|
||||||
"routing to internal control channel handler"
|
"routing to internal control channel handler"
|
||||||
);
|
);
|
||||||
|
|
||||||
|
if !self.control_channel_router.has_handler() {
|
||||||
|
tracing::warn!(
|
||||||
|
host = host_to_connect,
|
||||||
|
"no control channel handler configured, rejecting channel open"
|
||||||
|
);
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
let _ = channel;
|
||||||
return Ok(true);
|
return Ok(true);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -251,12 +276,20 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn reserved_wraith_destination_routing() {
|
fn reserved_wraith_destination_routing() {
|
||||||
assert!("wraith-control".starts_with(WRAITH_PREFIX));
|
use crate::server::control_channel::is_reserved_destination;
|
||||||
assert!("wraith-status".starts_with(WRAITH_PREFIX));
|
assert!(is_reserved_destination("wraith-control"));
|
||||||
assert!("wraith-events".starts_with(WRAITH_PREFIX));
|
assert!(is_reserved_destination("wraith-status"));
|
||||||
assert!(!"example.com".starts_with(WRAITH_PREFIX));
|
assert!(is_reserved_destination("wraith-events"));
|
||||||
assert!(!"localhost".starts_with(WRAITH_PREFIX));
|
assert!(!is_reserved_destination("example.com"));
|
||||||
assert!(!"wraith.example.com".starts_with(WRAITH_PREFIX));
|
assert!(!is_reserved_destination("localhost"));
|
||||||
|
assert!(!is_reserved_destination("wraith.example.com"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn server_handler_without_control_handler_rejects_wraith_destinations() {
|
||||||
|
let auth_config = make_empty_auth_config();
|
||||||
|
let handler = ServerHandler::new(auth_config, None, None);
|
||||||
|
assert!(!handler.control_channel_router().has_handler());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -1,3 +1,8 @@
|
|||||||
|
pub mod control_channel;
|
||||||
pub mod handler;
|
pub mod handler;
|
||||||
|
|
||||||
|
pub use control_channel::{
|
||||||
|
ControlChannelHandler, ControlChannelRouter, DuplexStream, WRAITH_CONTROL_DESTINATION,
|
||||||
|
WRAITH_PREFIX, is_reserved_destination,
|
||||||
|
};
|
||||||
pub use handler::{ProxyConfig, ProxyMode, ServerHandler};
|
pub use handler::{ProxyConfig, ProxyMode, ServerHandler};
|
||||||
Reference in New Issue
Block a user