init (tcp, udp, ws)
This commit is contained in:
Mark
332
src/routes.rs
Normal file
332
src/routes.rs
Normal file
@@ -0,0 +1,332 @@
|
||||
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())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user