294 lines
8.6 KiB
Rust
294 lines
8.6 KiB
Rust
use bytes::Buf;
|
|
use bytes::BytesMut;
|
|
|
|
use crate::actor::*;
|
|
use crate::error::Error;
|
|
use crate::schemas::internal_protocol::*;
|
|
use crate::schemas::tunnel_relay;
|
|
|
|
use futures::Sink;
|
|
use futures::SinkExt;
|
|
use futures::Stream;
|
|
use futures::StreamExt;
|
|
|
|
use preserves::value::BinarySource;
|
|
use preserves::value::BytesBinarySource;
|
|
use preserves::value::DomainDecode;
|
|
use preserves::value::DomainEncode;
|
|
use preserves::value::IOValue;
|
|
use preserves::value::Map;
|
|
use preserves::value::NoEmbeddedDomainCodec;
|
|
use preserves::value::PackedReader;
|
|
use preserves::value::PackedWriter;
|
|
use preserves::value::Reader;
|
|
use preserves::value::Writer;
|
|
use preserves_schema::support::lazy_static;
|
|
|
|
use std::convert::TryFrom;
|
|
use std::io;
|
|
use std::pin::Pin;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic;
|
|
|
|
use tokio::io::AsyncRead;
|
|
use tokio::io::AsyncReadExt;
|
|
use tokio::io::AsyncWrite;
|
|
use tokio::io::AsyncWriteExt;
|
|
|
|
use tracing;
|
|
|
|
struct WireSymbol {
|
|
oid: Oid,
|
|
obj: Arc<Ref>,
|
|
ref_count: atomic::AtomicUsize,
|
|
}
|
|
|
|
struct Membrane {
|
|
oid_map: Map<Oid, Arc<WireSymbol>>,
|
|
ref_map: Map<Arc<Ref>, Arc<WireSymbol>>,
|
|
}
|
|
|
|
struct Membranes {
|
|
exported: Membrane,
|
|
imported: Membrane,
|
|
}
|
|
|
|
pub enum Input {
|
|
Packets(Pin<Box<dyn Stream<Item = Result<Vec<u8>, Error>> + Send>>),
|
|
Bytes(Pin<Box<dyn AsyncRead + Send>>),
|
|
}
|
|
|
|
pub enum Output {
|
|
Packets(Pin<Box<dyn Sink<Vec<u8>, Error = Error> + Send>>),
|
|
Bytes(Pin<Box<dyn AsyncWrite + Send>>),
|
|
}
|
|
|
|
// There are other kinds of relay. This one has exactly two participants connected to each other.
|
|
pub struct TunnelRelay
|
|
{
|
|
input_buffer: BytesMut,
|
|
o: Output,
|
|
inbound_assertions: Map</* remote */ Handle, (/* local */ Handle, Vec<Arc<WireSymbol>>)>,
|
|
outbound_assertions: Map<Handle, Vec<Arc<WireSymbol>>>,
|
|
membranes: Membranes,
|
|
pending_outbound: Vec<TurnEvent>,
|
|
}
|
|
|
|
lazy_static! {
|
|
static ref INERT_REF: Arc<Ref> = {
|
|
struct InertEntity;
|
|
impl crate::actor::Entity for InertEntity {}
|
|
let mut ac = crate::actor::Actor::new();
|
|
ac.create(InertEntity)
|
|
};
|
|
}
|
|
|
|
//---------------------------------------------------------------------------
|
|
|
|
impl WireSymbol {
|
|
fn acquire(&self) {
|
|
self.ref_count.fetch_add(1, atomic::Ordering::Relaxed);
|
|
}
|
|
|
|
fn release(&self) -> bool {
|
|
self.ref_count.fetch_sub(1, atomic::Ordering::Relaxed) == 1
|
|
}
|
|
}
|
|
|
|
impl Membrane {
|
|
fn new() -> Self {
|
|
Membrane {
|
|
oid_map: Map::new(),
|
|
ref_map: Map::new(),
|
|
}
|
|
}
|
|
|
|
fn insert(&mut self, oid: Oid, obj: Arc<Ref>) -> Arc<WireSymbol> {
|
|
let ws = Arc::new(WireSymbol {
|
|
oid: oid.clone(),
|
|
obj: Arc::clone(&obj),
|
|
ref_count: atomic::AtomicUsize::new(0),
|
|
});
|
|
self.oid_map.insert(oid, Arc::clone(&ws));
|
|
self.ref_map.insert(obj, Arc::clone(&ws));
|
|
ws
|
|
}
|
|
|
|
fn release(&mut self, ws: &Arc<WireSymbol>) {
|
|
if ws.release() {
|
|
self.oid_map.remove(&ws.oid);
|
|
self.ref_map.remove(&ws.obj);
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TunnelRelay {
|
|
pub fn run(ac: &mut Actor, i: Input, o: Output) -> ActorResult {
|
|
let tr = ac.create(TunnelRelay {
|
|
input_buffer: BytesMut::with_capacity(1024),
|
|
o,
|
|
inbound_assertions: Map::new(),
|
|
outbound_assertions: Map::new(),
|
|
membranes: Membranes {
|
|
exported: Membrane::new(),
|
|
imported: Membrane::new(),
|
|
},
|
|
pending_outbound: Vec::new(),
|
|
});
|
|
ac.add_exit_hook(&tr.target);
|
|
ac.linked_task(tracing::info_span!("reader"), input_loop(i, tr));
|
|
Ok(())
|
|
}
|
|
|
|
fn handle_packet(&mut self, p: Packet) -> ActorResult {
|
|
match p {
|
|
Packet::Error(b) => {
|
|
tracing::info!(message = debug(b.message.clone()),
|
|
detail = debug(b.detail.clone()),
|
|
"received Error from peer");
|
|
Err(*b)
|
|
},
|
|
Packet::Turn(b) => {
|
|
let Turn(events) = *b;
|
|
for TurnEvent { oid, event } in events {
|
|
tracing::info!(oid = debug(oid), event = debug(event))
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn send_event(&mut self, oid: Oid, event: Event) -> bool {
|
|
let need_flush = self.pending_outbound.is_empty();
|
|
self.pending_outbound.push(TurnEvent { oid, event });
|
|
need_flush
|
|
}
|
|
|
|
pub fn decode_packet(&mut self, bs: &[u8]) -> Result<Packet, Error> {
|
|
let mut src = BytesBinarySource::new(bs);
|
|
Ok(Packet::try_from(&src.packed::<_, _Any, _>(&mut self.membranes).demand_next(false)?)?)
|
|
}
|
|
|
|
fn encode_packet(&mut self, p: Packet) -> Result<Vec<u8>, Error> {
|
|
Ok(PackedWriter::encode::<_, _Any, _>(&mut self.membranes, &_Any::from(&p))?)
|
|
}
|
|
|
|
pub async fn send_packet(&mut self, p: Packet) -> ActorResult {
|
|
let bs = self.encode_packet(p)?;
|
|
match &mut self.o {
|
|
Output::Packets(sink) => Ok(sink.send(bs).await?),
|
|
Output::Bytes(w) => Ok(w.write_all(&bs).await?),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl DomainDecode<_Ptr> for Membranes {
|
|
fn decode_embedded<'de, 'src, S: BinarySource<'de>>(
|
|
&mut self,
|
|
src: &'src mut S,
|
|
_read_annotations: bool,
|
|
) -> io::Result<_Ptr> {
|
|
let v: IOValue = PackedReader::new(src, NoEmbeddedDomainCodec).demand_next(false)?;
|
|
Ok(Arc::new(_Dom::try_from(&v)?))
|
|
}
|
|
}
|
|
|
|
impl DomainEncode<_Ptr> for Membranes {
|
|
fn encode_embedded<W: Writer>(
|
|
&mut self,
|
|
w: &mut W,
|
|
d: &_Ptr,
|
|
) -> io::Result<()> {
|
|
w.write(&mut NoEmbeddedDomainCodec, &IOValue::from(d.as_ref()))
|
|
}
|
|
}
|
|
|
|
pub async fn input_loop(
|
|
i: Input,
|
|
relay: Arc<Ref>,
|
|
) -> ActorResult {
|
|
fn s<M: Into<_Any>>(relay: &Arc<Ref>, m: M) {
|
|
relay.external_event(Event::Message(Box::new(Message { body: Assertion(m.into()) })))
|
|
}
|
|
|
|
match i {
|
|
Input::Packets(mut src) => {
|
|
loop {
|
|
match src.next().await {
|
|
None => {
|
|
s(&relay, &tunnel_relay::Input::Eof);
|
|
return Ok(());
|
|
}
|
|
Some(bs) => s(&relay, &tunnel_relay::Input::Packet { bs: bs? }),
|
|
}
|
|
}
|
|
}
|
|
Input::Bytes(mut r) => {
|
|
let mut buf = BytesMut::with_capacity(1024);
|
|
loop {
|
|
buf.reserve(1024);
|
|
let n = r.read_buf(&mut buf).await?;
|
|
match n {
|
|
0 => {
|
|
s(&relay, &tunnel_relay::Input::Eof);
|
|
return Ok(());
|
|
}
|
|
_ => {
|
|
while buf.has_remaining() {
|
|
let bs = buf.chunk();
|
|
let n = bs.len();
|
|
s(&relay, &tunnel_relay::Input::Segment { bs: bs.to_vec() });
|
|
buf.advance(n);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Entity for TunnelRelay {
|
|
fn message(&mut self, t: &mut Activation, m: _Any) -> ActorResult {
|
|
if let Ok(m) = tunnel_relay::Input::try_from(&m) {
|
|
match m {
|
|
tunnel_relay::Input::Eof => {
|
|
tracing::info!("eof");
|
|
t.actor.shutdown();
|
|
}
|
|
tunnel_relay::Input::Packet { bs } => {
|
|
let p = self.decode_packet(&bs)?;
|
|
self.handle_packet(p)?
|
|
}
|
|
tunnel_relay::Input::Segment { bs } => {
|
|
self.input_buffer.extend_from_slice(&bs);
|
|
loop {
|
|
let (e, count) = {
|
|
let mut src = BytesBinarySource::new(&self.input_buffer);
|
|
let mut r = src.packed::<_, _Any, _>(&mut self.membranes);
|
|
let e = r.next(false)?;
|
|
(e, r.source.index)
|
|
};
|
|
match e {
|
|
None => break,
|
|
Some(item) => {
|
|
self.input_buffer.advance(count);
|
|
self.handle_packet(Packet::try_from(&item)?)?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn exit_hook(&mut self, _t: &mut Activation, exit_status: &ActorResult) -> BoxFuture<ActorResult> {
|
|
if let Err(e) = exit_status {
|
|
let e = e.clone();
|
|
Box::pin(self.send_packet(Packet::Error(Box::new(e))))
|
|
} else {
|
|
Box::pin(ready(Ok(())))
|
|
}
|
|
}
|
|
}
|