use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use bytes::Bytes;
use rings_core::dht::Did;
use serde::de::DeserializeOwned;
use serde::Deserialize;
use serde::Serialize;
use crate::extension::ext::Ctx;
#[cfg(any(rings_native, rings_browser))]
use crate::extension::ext::EffectScope;
#[cfg(any(rings_native, rings_browser))]
use crate::extension::ext::Interpret;
use crate::extension::ext::MaybeSend;
use crate::extension::ext::Protocol;
use crate::extension::ext::Reject;
use crate::extension::ext::Scope;
use crate::extension::ext::Transition;
use crate::extension::ext::Wire;
use crate::extension::transport::EffectEnqueue;
use crate::extension::transport::Frame;
use crate::extension::transport::Initiator;
use crate::extension::transport::SessionId;
use crate::extension::transport::SessionKey;
use crate::extension::transport::TransportKind;
use crate::peer_quota::PeerQuota;
#[cfg(any(rings_native, rings_browser))]
mod control_outbox;
#[cfg(any(rings_native, rings_browser))]
use self::control_outbox::ControlOutbox;
#[cfg(all(test, rings_native))]
pub(crate) use self::control_outbox::ControlSendTestHook;
pub const TCP: &str = "tcp";
pub const UDP: &str = "udp";
pub(crate) const MAX_RELAY_SESSIONS: usize = 1_024;
pub(crate) const MAX_RELAY_SESSIONS_PER_PEER: usize = 64;
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub enum RelayCommand<T> {
RegisterService {
name: String,
target: T,
},
Accepted {
token: u64,
peer: Did,
service: String,
},
Untrack {
peer: Did,
session: SessionId,
initiator: Initiator,
},
Abort {
peer: Did,
session: SessionId,
initiator: Initiator,
},
}
pub enum RelayEvent<T> {
Command(RelayCommand<T>),
Frame {
from: Did,
frame: Frame,
},
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RelayEffect<T> {
Connect {
key: SessionKey,
target: T,
kind: TransportKind,
},
Write {
key: SessionKey,
bytes: Bytes,
},
Shutdown {
key: SessionKey,
},
Close {
key: SessionKey,
},
SendClose {
to: Did,
session: SessionId,
from_opener: bool,
},
OpenAccepted {
token: u64,
key: SessionKey,
service: String,
},
RejectAccepted {
token: u64,
},
}
#[derive(Clone)]
pub struct RelayState<T> {
services: Arc<HashMap<String, T>>,
sessions: Arc<HashSet<SessionKey>>,
session_quota: Arc<PeerQuota>,
peer_shutdown: Arc<HashSet<SessionKey>>,
next_session: u64,
}
impl<T> Default for RelayState<T> {
fn default() -> Self {
Self {
services: Arc::new(HashMap::new()),
sessions: Arc::new(HashSet::new()),
session_quota: Arc::new(PeerQuota::new(
MAX_RELAY_SESSIONS,
MAX_RELAY_SESSIONS_PER_PEER,
)),
peer_shutdown: Arc::new(HashSet::new()),
next_session: 0,
}
}
}
impl<T> RelayState<T> {
fn can_admit_session(&self, key: &SessionKey) -> bool {
self.session_quota.can_reserve(key.peer).is_ok()
}
fn insert_session(&mut self, key: SessionKey) -> bool {
if self.sessions.contains(&key) {
return false;
}
if Arc::make_mut(&mut self.session_quota)
.reserve(key.peer)
.is_err()
{
return false;
}
if Arc::make_mut(&mut self.sessions).insert(key.clone()) {
Arc::make_mut(&mut self.peer_shutdown).remove(&key);
true
} else {
let rolled_back = Arc::make_mut(&mut self.session_quota).release(key.peer);
debug_assert!(rolled_back);
false
}
}
fn remove_session(&mut self, key: &SessionKey) -> bool {
if !self.sessions.contains(key) {
return false;
}
if !Arc::make_mut(&mut self.session_quota).release(key.peer) {
debug_assert!(false, "session quota missing admitted peer {}", key.peer);
return false;
}
Arc::make_mut(&mut self.peer_shutdown).remove(key);
let removed = Arc::make_mut(&mut self.sessions).remove(key);
debug_assert!(removed);
removed
}
fn peer_can_send(&self, key: &SessionKey) -> bool {
self.sessions.contains(key) && !self.peer_shutdown.contains(key)
}
fn shutdown_peer(&mut self, key: &SessionKey, kind: TransportKind) -> bool {
kind == TransportKind::Tcp
&& self.sessions.contains(key)
&& !self.peer_shutdown.contains(key)
&& Arc::make_mut(&mut self.peer_shutdown).insert(key.clone())
}
}
#[derive(Clone)]
pub struct Relay<T> {
namespace: String,
kind: TransportKind,
config: HashMap<String, T>,
}
impl<T> Relay<T> {
pub fn tcp(config: HashMap<String, T>) -> Self {
Self {
namespace: TCP.to_string(),
kind: TransportKind::Tcp,
config,
}
}
pub fn udp(config: HashMap<String, T>) -> Self {
Self {
namespace: UDP.to_string(),
kind: TransportKind::Udp,
config,
}
}
}
impl<T> Protocol for Relay<T>
where T: Clone + DeserializeOwned + Serialize + MaybeSend + 'static
{
type State = RelayState<T>;
type Event = RelayEvent<T>;
type Effect = RelayEffect<T>;
fn namespace(&self) -> &str {
self.namespace.as_str()
}
fn init(&self) -> RelayState<T> {
RelayState {
services: Arc::new(self.config.clone()),
sessions: Arc::new(HashSet::new()),
session_quota: Arc::new(PeerQuota::new(
MAX_RELAY_SESSIONS,
MAX_RELAY_SESSIONS_PER_PEER,
)),
peer_shutdown: Arc::new(HashSet::new()),
next_session: 0,
}
}
fn decode(&self, wire: Wire<'_>) -> Result<RelayEvent<T>, Reject> {
if wire.from == wire.me {
let command = rings_codec::deserialize::<RelayCommand<T>>(wire.payload)
.map_err(|e| Reject(format!("bad relay command: {e}")))?;
Ok(RelayEvent::Command(command))
} else {
let frame = rings_codec::deserialize::<Frame>(wire.payload)
.map_err(|e| Reject(format!("bad relay frame: {e}")))?;
Ok(RelayEvent::Frame {
from: wire.from,
frame,
})
}
}
fn step(
&self,
ctx: Ctx<'_, RelayState<T>>,
event: RelayEvent<T>,
) -> Transition<RelayState<T>, RelayEffect<T>> {
match event {
RelayEvent::Command(command) => {
step_command(self.namespace.as_str(), ctx.state, command)
}
RelayEvent::Frame { from, frame } => {
step_frame(self.kind, self.namespace.as_str(), ctx.state, from, frame)
}
}
}
}
fn step_command<T: Clone>(
namespace: &str,
state: &RelayState<T>,
command: RelayCommand<T>,
) -> Transition<RelayState<T>, RelayEffect<T>> {
let mut next = state.clone();
match command {
RelayCommand::RegisterService { name, target } => {
Arc::make_mut(&mut next.services).insert(name, target);
Transition::pure(next)
}
RelayCommand::Accepted {
token,
peer,
service,
} => {
let session = SessionId(next.next_session);
let key = SessionKey::new(peer, namespace, session, Initiator::Local);
if !next.can_admit_session(&key) {
return Transition::with(next, vec![RelayEffect::RejectAccepted { token }]);
}
let Some(next_session) = next.next_session.checked_add(1) else {
return Transition::with(next, vec![RelayEffect::RejectAccepted { token }]);
};
next.next_session = next_session;
if !next.insert_session(key.clone()) {
return Transition::with(next, vec![RelayEffect::RejectAccepted { token }]);
}
Transition::with(next, vec![RelayEffect::OpenAccepted {
token,
key,
service,
}])
}
RelayCommand::Untrack {
peer,
session,
initiator,
} => {
next.remove_session(&SessionKey::new(peer, namespace, session, initiator));
Transition::pure(next)
}
RelayCommand::Abort {
peer,
session,
initiator,
} => {
let key = SessionKey::new(peer, namespace, session, initiator);
if next.remove_session(&key) {
Transition::with(next, vec![RelayEffect::SendClose {
to: peer,
session,
from_opener: matches!(initiator, Initiator::Local),
}])
} else {
Transition::pure(next)
}
}
}
}
fn step_frame<T: Clone>(
kind: TransportKind,
namespace: &str,
state: &RelayState<T>,
from: Did,
frame: Frame,
) -> Transition<RelayState<T>, RelayEffect<T>> {
match frame {
Frame::Open { session, service } => {
let key = SessionKey::new(from, namespace, session, Initiator::Remote);
if state.sessions.contains(&key) {
return Transition::pure(state.clone());
}
match state.services.get(service.as_str()) {
Some(target) => {
let mut next = state.clone();
if !next.can_admit_session(&key) {
return rejected_open(next, key);
}
let target = target.clone();
if !next.insert_session(key.clone()) {
return rejected_open(next, key);
}
Transition::with(next, vec![RelayEffect::Connect { key, target, kind }])
}
None => rejected_open(state.clone(), key),
}
}
Frame::Data {
session,
from_opener,
bytes,
} => {
let key = SessionKey::new(from, namespace, session, opener_to_initiator(from_opener));
if state.peer_can_send(&key) {
Transition::with(state.clone(), vec![RelayEffect::Write { key, bytes }])
} else {
Transition::pure(state.clone())
}
}
Frame::Shutdown {
session,
from_opener,
} => {
let key = SessionKey::new(from, namespace, session, opener_to_initiator(from_opener));
let mut next = state.clone();
if next.shutdown_peer(&key, kind) {
Transition::with(next, vec![RelayEffect::Shutdown { key }])
} else {
Transition::pure(next)
}
}
Frame::Close {
session,
from_opener,
} => {
let key = SessionKey::new(from, namespace, session, opener_to_initiator(from_opener));
if state.sessions.contains(&key) {
let mut next = state.clone();
next.remove_session(&key);
Transition::with(next, vec![RelayEffect::Close { key }])
} else {
Transition::pure(state.clone())
}
}
}
}
fn rejected_open<T>(
state: RelayState<T>,
key: SessionKey,
) -> Transition<RelayState<T>, RelayEffect<T>> {
Transition::with(state, vec![RelayEffect::SendClose {
to: key.peer,
session: key.session,
from_opener: false,
}])
}
fn opener_to_initiator(from_opener: bool) -> Initiator {
if from_opener {
Initiator::Remote
} else {
Initiator::Local
}
}
pub(crate) fn close_frame(session: SessionId, from_opener: bool) -> crate::error::Result<Bytes> {
let frame = Frame::Close {
session,
from_opener,
};
rings_codec::serialize(&frame)
.map(Bytes::from)
.map_err(|_| crate::error::Error::EncodeError)
}
#[cfg(rings_native)]
pub(crate) struct NativeRelay {
engine: Arc<crate::extension::transport::engine::TransportSessions>,
control_outbox: ControlOutbox,
}
#[cfg(rings_native)]
impl NativeRelay {
pub(crate) fn new(engine: Arc<crate::extension::transport::engine::TransportSessions>) -> Self {
Self {
engine,
control_outbox: ControlOutbox::default(),
}
}
#[cfg(all(test, rings_native))]
pub(crate) fn new_with_control_send_test_hook(
engine: Arc<crate::extension::transport::engine::TransportSessions>,
hook: Arc<ControlSendTestHook>,
) -> Self {
Self {
engine,
control_outbox: ControlOutbox::with_test_hook(hook),
}
}
}
#[cfg(rings_native)]
#[async_trait::async_trait]
impl Interpret for NativeRelay {
type Effect = RelayEffect<std::net::SocketAddr>;
async fn run(
&self,
scope: &EffectScope,
effect: RelayEffect<std::net::SocketAddr>,
) -> crate::error::Result<Vec<Bytes>> {
match effect {
RelayEffect::Connect { key, target, kind } => {
let admission =
self.engine
.clone()
.connect(scope.lifecycle(), key.clone(), target, kind);
return enqueue_feedback::<std::net::SocketAddr>(key, admission);
}
RelayEffect::Write { key, bytes } => {
let admission = self.engine.write(&key, bytes);
return enqueue_feedback::<std::net::SocketAddr>(key, admission);
}
RelayEffect::Shutdown { key } => {
let admission = self.engine.shutdown(&key);
return enqueue_feedback::<std::net::SocketAddr>(key, admission);
}
RelayEffect::Close { key } => {
self.engine.close_for_effect(&key);
}
RelayEffect::SendClose {
to,
session,
from_opener,
} => {
self.control_outbox.enqueue(
scope.lifecycle(),
to,
close_frame(session, from_opener)?,
)?;
}
RelayEffect::OpenAccepted {
token,
key,
service,
} => {
let feedback =
self.engine
.clone()
.bind_accepted(scope.lifecycle(), token, key, service);
return feedback
.map(untrack_feedback::<std::net::SocketAddr>)
.transpose()
.map(|feedback| feedback.into_iter().collect());
}
RelayEffect::RejectAccepted { token } => {
self.engine.evict_pending_for_effect(token);
}
}
Ok(Vec::new())
}
}
fn untrack_feedback<T: Serialize>(key: SessionKey) -> crate::error::Result<Bytes> {
let command = RelayCommand::<T>::Untrack {
peer: key.peer,
session: key.session,
initiator: key.initiator,
};
rings_codec::serialize(&command)
.map(Bytes::from)
.map_err(|_| crate::error::Error::EncodeError)
}
fn abort_feedback<T: Serialize>(key: SessionKey) -> crate::error::Result<Bytes> {
let command = RelayCommand::<T>::Abort {
peer: key.peer,
session: key.session,
initiator: key.initiator,
};
rings_codec::serialize(&command)
.map(Bytes::from)
.map_err(|_| crate::error::Error::EncodeError)
}
fn enqueue_feedback<T: Serialize>(
key: SessionKey,
admission: EffectEnqueue,
) -> crate::error::Result<Vec<Bytes>> {
match admission {
EffectEnqueue::Enqueued => Ok(Vec::new()),
EffectEnqueue::Missing => untrack_feedback::<T>(key).map(|feedback| vec![feedback]),
EffectEnqueue::Failed => abort_feedback::<T>(key).map(|feedback| vec![feedback]),
}
}
#[cfg(rings_browser)]
pub(crate) struct WtRelay {
engine: Arc<crate::extension::transport::wt::WtSessions>,
control_outbox: ControlOutbox,
}
#[cfg(rings_browser)]
impl WtRelay {
pub(crate) fn new(engine: Arc<crate::extension::transport::wt::WtSessions>) -> Self {
Self {
engine,
control_outbox: ControlOutbox::default(),
}
}
}
#[cfg(rings_browser)]
#[async_trait::async_trait(?Send)]
impl Interpret for WtRelay {
type Effect = RelayEffect<String>;
async fn run(
&self,
scope: &EffectScope,
effect: RelayEffect<String>,
) -> crate::error::Result<Vec<Bytes>> {
match effect {
RelayEffect::Connect { key, target, kind } => {
let admission =
self.engine
.clone()
.connect(scope.lifecycle(), key.clone(), target, kind);
return enqueue_feedback::<String>(key, admission);
}
RelayEffect::Write { key, bytes } => {
let admission = self.engine.write(scope.lifecycle(), key.clone(), bytes);
return enqueue_feedback::<String>(key, admission);
}
RelayEffect::Shutdown { key } => {
let admission = self.engine.shutdown(scope.lifecycle(), key.clone());
return enqueue_feedback::<String>(key, admission);
}
RelayEffect::Close { key } => {
self.engine.close_for_effect(&key);
}
RelayEffect::SendClose {
to,
session,
from_opener,
} => {
self.control_outbox.enqueue(
scope.lifecycle(),
to,
close_frame(session, from_opener)?,
)?;
}
RelayEffect::OpenAccepted { .. } => {
tracing::warn!("browser relay received OpenAccepted; it has no local listener");
}
RelayEffect::RejectAccepted { .. } => {
tracing::warn!("browser relay received RejectAccepted; it has no local listener");
}
}
Ok(Vec::new())
}
}
#[cfg(rings_native)]
#[derive(Clone)]
pub struct RelayHandle {
engine: Arc<crate::extension::transport::engine::TransportSessions>,
tcp: Scope,
udp: Scope,
}
#[cfg(rings_native)]
impl RelayHandle {
pub fn install(extensions: &crate::extension::ext::Extensions) -> crate::error::Result<Self> {
let engine = Arc::new(crate::extension::transport::engine::TransportSessions::new());
extensions.register_many(vec![
(Relay::tcp(HashMap::new()), NativeRelay::new(engine.clone())),
(Relay::udp(HashMap::new()), NativeRelay::new(engine.clone())),
])?;
let core = extensions.core();
Ok(Self {
engine,
tcp: Scope::new(core.clone(), TCP.to_string()),
udp: Scope::new(core, UDP.to_string()),
})
}
pub async fn open_tcp_tunnel(
&self,
local_addr: std::net::SocketAddr,
peer: Did,
service: String,
) -> crate::error::Result<()> {
self.open_tunnel(&self.tcp, local_addr, peer, service, TransportKind::Tcp)
.await
}
pub async fn relay_tcp_stream(
&self,
stream: tokio::net::TcpStream,
peer: Did,
service: String,
) -> crate::error::Result<()> {
self.engine
.clone()
.relay_tcp_stream(self.tcp.clone(), stream, peer, service)
.await;
Ok(())
}
pub async fn open_udp_tunnel(
&self,
local_addr: std::net::SocketAddr,
peer: Did,
service: String,
) -> crate::error::Result<()> {
self.open_tunnel(&self.udp, local_addr, peer, service, TransportKind::Udp)
.await
}
async fn open_tunnel(
&self,
scope: &Scope,
local_addr: std::net::SocketAddr,
peer: Did,
service: String,
kind: TransportKind,
) -> crate::error::Result<()> {
self.engine
.clone()
.listen(scope.clone(), local_addr, peer, service, kind)
.await;
Ok(())
}
pub async fn register_tcp_service(
&self,
name: String,
addr: std::net::SocketAddr,
) -> crate::error::Result<()> {
register_service(&self.tcp, name, addr).await
}
pub async fn register_udp_service(
&self,
name: String,
addr: std::net::SocketAddr,
) -> crate::error::Result<()> {
register_service(&self.udp, name, addr).await
}
}
#[cfg(any(rings_native, rings_browser))]
async fn register_service<T>(scope: &Scope, name: String, target: T) -> crate::error::Result<()>
where T: Serialize {
let command = RelayCommand::RegisterService { name, target };
let payload = rings_codec::serialize(&command).map_err(|_| crate::error::Error::EncodeError)?;
scope.inject(Bytes::from(payload)).await
}
#[cfg(rings_browser)]
#[derive(Clone)]
pub struct RelayHandle {
tcp: Scope,
udp: Scope,
}
#[cfg(rings_browser)]
impl RelayHandle {
pub fn install(extensions: &crate::extension::ext::Extensions) -> crate::error::Result<Self> {
let engine = Arc::new(crate::extension::transport::wt::WtSessions::new());
extensions.register_many(vec![
(Relay::tcp(HashMap::new()), WtRelay::new(engine.clone())),
(Relay::udp(HashMap::new()), WtRelay::new(engine)),
])?;
let core = extensions.core();
Ok(Self {
tcp: Scope::new(core.clone(), TCP.to_string()),
udp: Scope::new(core, UDP.to_string()),
})
}
pub async fn register_wt_service(&self, name: String, url: String) -> crate::error::Result<()> {
register_service(&self.tcp, name, url).await
}
pub async fn register_wt_udp_service(
&self,
name: String,
url: String,
) -> crate::error::Result<()> {
register_service(&self.udp, name, url).await
}
}
#[cfg(test)]
mod test_relay;