333 lines
13 KiB
Rust
333 lines
13 KiB
Rust
use std::sync::Arc;
|
|
use tokio::task::spawn;
|
|
|
|
#[derive(Clone, Copy, Default)]
|
|
pub struct Config {
|
|
pub verbosity: u8,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct Routes {
|
|
pub config: Config,
|
|
binds: Binds,
|
|
wires: Wires,
|
|
state: State,
|
|
}
|
|
|
|
#[derive(Clone, Copy, PartialEq, Eq)]
|
|
pub struct Prot(ProtType, bool);
|
|
macro_rules! protocols {
|
|
($($p:ident:$t:ident),*) => {
|
|
#[derive(Clone, Copy, PartialEq, Eq)]
|
|
pub enum ProtType {
|
|
$($t,)*
|
|
}
|
|
pub enum ProtRx {
|
|
$($t($crate::$p::Rx),)*
|
|
}
|
|
pub enum ProtTx {
|
|
$($t($crate::$p::Tx),)*
|
|
}
|
|
#[derive(Default)]
|
|
struct Binds {
|
|
$($p: Bind<$crate::$p::Bind>,)*
|
|
}
|
|
#[derive(Default)]
|
|
struct Wires {
|
|
$($p: Wire<$crate::$p::Wire>,)*
|
|
}
|
|
#[derive(Default)]
|
|
struct State {
|
|
$($p: $crate::$p::State,)*
|
|
}
|
|
use $crate::{Rx, Tx};
|
|
impl Tx for ProtTx {
|
|
async fn tx(&mut self, buf: &[u8]) -> Result<(), String> {
|
|
match self {
|
|
$(
|
|
ProtTx::$t($crate::$p::Tx::Wait(tx)) => tx.tx(buf).await,
|
|
ProtTx::$t($crate::$p::Tx::Bind(tx)) => tx.tx(buf).await,
|
|
ProtTx::$t($crate::$p::Tx::Wire(tx)) => tx.tx(buf).await,
|
|
)*
|
|
}
|
|
}
|
|
}
|
|
async fn prot_fwd(r: &mut ProtRx, tx: &mut impl Tx) -> Result<bool, String> {
|
|
use $crate::{BUF_SIZE, RxT};
|
|
match r {
|
|
$(
|
|
ProtRx::$t($crate::$p::Rx::Wait(r)) => {
|
|
let mut buf = [0u8; BUF_SIZE];
|
|
let buf = r.rx(&mut buf).await?;
|
|
let buf = buf.buf();
|
|
if buf.is_empty() {
|
|
return Ok(false);
|
|
}
|
|
tx.tx(buf).await?;
|
|
}
|
|
ProtRx::$t($crate::$p::Rx::Bind(r)) => {
|
|
let mut buf = [0u8; BUF_SIZE];
|
|
let buf = r.rx(&mut buf).await?;
|
|
let buf = buf.buf();
|
|
if buf.is_empty() {
|
|
return Ok(false);
|
|
}
|
|
tx.tx(buf).await?;
|
|
}
|
|
ProtRx::$t($crate::$p::Rx::Wire(r)) => {
|
|
let mut buf = [0u8; BUF_SIZE];
|
|
let buf = r.rx(&mut buf).await?;
|
|
let buf = buf.buf();
|
|
if buf.is_empty() {
|
|
return Ok(false);
|
|
}
|
|
tx.tx(buf).await?;
|
|
}
|
|
,)*
|
|
}
|
|
Ok(true)
|
|
}
|
|
async fn prot_drop(r: &mut ProtRx) -> Result<bool, String> {
|
|
use $crate::{BUF_SIZE, RxT};
|
|
match r {
|
|
$(
|
|
ProtRx::$t($crate::$p::Rx::Wait(r)) => {
|
|
r.rx(&mut [0u8; BUF_SIZE]).await.map(|buf| !buf.buf().is_empty())
|
|
}
|
|
ProtRx::$t($crate::$p::Rx::Bind(r)) => {
|
|
r.rx(&mut [0u8; BUF_SIZE]).await.map(|buf| !buf.buf().is_empty())
|
|
}
|
|
ProtRx::$t($crate::$p::Rx::Wire(r)) => {
|
|
r.rx(&mut [0u8; BUF_SIZE]).await.map(|buf| !buf.buf().is_empty())
|
|
}
|
|
)*
|
|
}
|
|
}
|
|
impl Routes {
|
|
pub fn has(&self, id: (Prot, usize)) -> bool {
|
|
match (id.0.wire(), id.0.prot()) {
|
|
$((true, ProtType::$t) => self.wires.$p.get(id.1).is_some(),
|
|
(false, ProtType::$t) => self.binds.$p.get(id.1).is_some(),)*
|
|
}
|
|
}
|
|
pub async fn bind(
|
|
&mut self,
|
|
prot: impl AsRef<str>,
|
|
arg: impl AsRef<str>,
|
|
) -> Result<(Prot, usize), String> {
|
|
use $crate::Wire;
|
|
match prot.as_ref().to_lowercase().as_str() {
|
|
$(stringify!($p) => {
|
|
let bind = crate::$p::Bind::parse(arg.as_ref()).await?;
|
|
Ok((ProtType::$t.bind(), self.binds.$p.add(bind)))
|
|
},)*
|
|
_ => return Err(format!("unknown/unsupported protocol: `{}`", prot.as_ref())),
|
|
}
|
|
}
|
|
pub async fn wire(
|
|
&mut self,
|
|
prot: impl AsRef<str>,
|
|
arg: impl AsRef<str>,
|
|
) -> Result<(Prot, usize), String> {
|
|
use $crate::Wire;
|
|
match prot.as_ref().to_lowercase().as_str() {
|
|
$(stringify!($p) => {
|
|
let wire = crate::$p::Wire::parse(arg.as_ref()).await?;
|
|
Ok((ProtType::$t.wire(), self.wires.$p.add(wire)))
|
|
},)*
|
|
_ => return Err(format!("unknown/unsupported protocol: `{}`", prot.as_ref())),
|
|
}
|
|
}
|
|
pub fn pipe(&mut self, bind: (Prot, usize), wire: (Prot, usize)) -> Result<(), String> {
|
|
if !bind.0.bind() {
|
|
Err("can't pipe data from a wire (:), you need to use a bind (@).")?;
|
|
}
|
|
if !self.has(wire) {
|
|
Err("internal error: bad wire index")?;
|
|
}
|
|
match bind.0.prot() {
|
|
$(ProtType::$t => self.binds
|
|
.$p
|
|
.pipe(bind.1, wire)
|
|
.ok_or("internal error: bad bind index")?,)*
|
|
}
|
|
Ok(())
|
|
}
|
|
pub async fn execute(self) -> Vec<tokio::task::JoinHandle<()>> {
|
|
use $crate::{Bind, Rx, Tx, RxT, BUF_SIZE};
|
|
let mut tasks = Vec::new();
|
|
let config = self.config;
|
|
let routes = Arc::new(self);
|
|
$({
|
|
for i in 0..routes.binds.$p.0.len() {
|
|
let routes = Arc::clone(&routes);
|
|
tasks.push(spawn(async move {
|
|
let (bind, _) = &routes.binds.$p.0[i];
|
|
while let Ok((mut rx, mut tx)) = bind.wait().await {
|
|
if config.verbosity >= 2 {
|
|
eprintln!("[binds] new connection.");
|
|
}
|
|
let routes = Arc::clone(&routes);
|
|
spawn(async move {
|
|
let (_, wires) = &routes.binds.$p.0[i];
|
|
// connect all wires
|
|
let mut rxs = Vec::new();
|
|
let mut txs = Vec::new();
|
|
if config.verbosity >= 3 {
|
|
eprintln!("[wires] connecting...");
|
|
}
|
|
for (j, wire) in wires.iter().enumerate() {
|
|
match routes.wire_open(*wire).await {
|
|
Ok((r, t)) => {
|
|
rxs.push(r);
|
|
txs.push(t);
|
|
}
|
|
Err(e) => if config.verbosity >= if j == 0 { 1 } else { 3 } {
|
|
eprintln!("[wires] failed to connect: {e}");
|
|
}
|
|
}
|
|
}
|
|
if config.verbosity >= 2 {
|
|
eprintln!("[wires] connected to {}/{} wires.", txs.len(), wires.len());
|
|
}
|
|
// forward data to all wires
|
|
spawn(async move {
|
|
let mut buf = [0u8; BUF_SIZE];
|
|
let mut any_ok = 1;
|
|
while any_ok > 0 {
|
|
any_ok = 0;
|
|
let buf = rx.rx(&mut buf).await?;
|
|
let buf = buf.buf();
|
|
if buf.is_empty() {
|
|
break;
|
|
}
|
|
for t in &mut txs {
|
|
if t.tx(buf).await.is_ok() {
|
|
any_ok += 1;
|
|
}
|
|
}
|
|
if config.verbosity >= 4 {
|
|
eprintln!("[wires] forwarded {} bytes through {}/{} wires.", buf.len(), any_ok, txs.len());
|
|
}
|
|
}
|
|
Ok::<(), String>(())
|
|
});
|
|
// get the first wire
|
|
rxs.reverse();
|
|
let r = rxs.pop();
|
|
// drop data from other wires
|
|
spawn(async move {
|
|
let mut any_ok = true;
|
|
while any_ok {
|
|
any_ok = false;
|
|
for r in &mut rxs {
|
|
if let Ok(true) = prot_drop(r).await {
|
|
any_ok = true;
|
|
}
|
|
}
|
|
}
|
|
});
|
|
// forward data from first wire
|
|
if let Some(mut r) = r {
|
|
while let Ok(true) = prot_fwd(&mut r, &mut tx).await {
|
|
if config.verbosity >= 4 {
|
|
eprintln!("[wires] forwarded response.");
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}));
|
|
}
|
|
})*
|
|
tasks
|
|
}
|
|
async fn wire_open(&self, id: (Prot, usize)) -> Result<(ProtRx, ProtTx), String> {
|
|
use $crate::Wire;
|
|
Ok(match (id.0.wire(), id.0.prot()) {
|
|
$((true, ProtType::$t) => {
|
|
let (r, t) = self.wires.$p.get(id.1).unwrap().open().await?;
|
|
(ProtRx::$t($crate::$p::Rx::Wire(r)), ProtTx::$t($crate::$p::Tx::Wire(t)))
|
|
}
|
|
(false, ProtType::$t) => {
|
|
let (r, t) = self.binds.$p.get(id.1).unwrap().open().await?;
|
|
(ProtRx::$t($crate::$p::Rx::Bind(r)), ProtTx::$t($crate::$p::Tx::Bind(t)))
|
|
}
|
|
)*})
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
protocols!(tcp:Tcp, udp:Udp, ws:Ws);
|
|
|
|
struct Bind<T>(Vec<(T, Vec<(Prot, usize)>)>);
|
|
struct Wire<T>(Vec<(T,)>);
|
|
|
|
impl Routes {
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
}
|
|
|
|
impl<T> Bind<T> {
|
|
fn get(&self, id: usize) -> Option<&T> {
|
|
self.0.get(id).map(|(t, _)| t)
|
|
}
|
|
fn get_mut(&mut self, id: usize) -> Option<&mut T> {
|
|
self.0.get_mut(id).map(|(t, _)| t)
|
|
}
|
|
fn add(&mut self, bind: T) -> usize {
|
|
self.0.push((bind, Vec::new()));
|
|
self.0.len() - 1
|
|
}
|
|
fn pipe(&mut self, id: usize, wire: (Prot, usize)) -> Option<()> {
|
|
let (_, pipes) = self.0.get_mut(id)?;
|
|
pipes.push(wire);
|
|
Some(())
|
|
}
|
|
}
|
|
impl<T> Wire<T> {
|
|
fn get(&self, id: usize) -> Option<&T> {
|
|
self.0.get(id).map(|(t,)| t)
|
|
}
|
|
fn get_mut(&mut self, id: usize) -> Option<&mut T> {
|
|
self.0.get_mut(id).map(|(t,)| t)
|
|
}
|
|
fn add(&mut self, wire: T) -> usize {
|
|
self.0.push((wire,));
|
|
self.0.len() - 1
|
|
}
|
|
}
|
|
|
|
impl ProtType {
|
|
fn bind(self) -> Prot {
|
|
Prot(self, false)
|
|
}
|
|
fn wire(self) -> Prot {
|
|
Prot(self, true)
|
|
}
|
|
}
|
|
impl Prot {
|
|
fn prot(&self) -> ProtType {
|
|
self.0
|
|
}
|
|
fn bind(&self) -> bool {
|
|
!self.1
|
|
}
|
|
fn wire(&self) -> bool {
|
|
self.1
|
|
}
|
|
}
|
|
|
|
impl<T> Default for Bind<T> {
|
|
fn default() -> Self {
|
|
Self(Vec::new())
|
|
}
|
|
}
|
|
impl<T> Default for Wire<T> {
|
|
fn default() -> Self {
|
|
Self(Vec::new())
|
|
}
|
|
}
|