use std::collections::{BTreeMap, VecDeque};
use std::hash::{BuildHasher, Hasher};
use std::io;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use renet::{ChannelConfig, ConnectionConfig, RenetClient, SendType};
use serde::{Serialize, de::DeserializeOwned};
pub const MAX_MESSAGE_BYTES: usize = 1024 * 1024;
const CHANNEL_MEMORY: usize = 4 * MAX_MESSAGE_BYTES;
const HEADER: usize = 21;
const MAGIC: &[u8; 4] = b"CRM1";
const HELLO: u8 = 1;
const WELCOME: u8 = 2;
const DATA: u8 = 3;
const BYE: u8 = 4;
const FULL: u8 = 5;
const HEARTBEAT: u8 = 6;
const INTERVAL: Duration = Duration::from_millis(100);
const EVENT_LIMIT: usize = 4096;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Lane {
Reliable,
ReliableUnordered,
Unreliable,
}
impl Lane {
fn channel(self) -> u8 {
match self {
Self::Reliable => 0,
Self::Unreliable => 1,
Self::ReliableUnordered => 2,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PeerId(u64);
impl PeerId {
pub const HOST: Self = Self(0);
pub const fn to_bits(self) -> u64 {
self.0
}
pub const fn from_bits(bits: u64) -> Self {
Self(bits)
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum MessageEvent {
Connected(PeerId),
Disconnected(PeerId, String),
Message {
peer: PeerId,
lane: Lane,
bytes: Vec<u8>,
},
}
impl MessageEvent {
pub fn decode<T: DeserializeOwned>(&self) -> Option<T> {
match self {
Self::Message { bytes, .. } => serde_json::from_slice(bytes).ok(),
_ => None,
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct MessageStats {
pub packets_sent: u64,
pub bytes_sent: u64,
pub packets_received: u64,
pub bytes_received: u64,
}
struct Link {
addr: SocketAddr,
client_nonce: u64,
server_nonce: u64,
connected: bool,
last_seen: Duration,
last_sent: Duration,
renet: RenetClient,
next_unreliable: u64,
received_unreliable: Option<u64>,
}
impl Link {
fn new(addr: SocketAddr, client_nonce: u64, server_nonce: u64, now: Duration) -> Self {
let channels = vec![
ChannelConfig {
channel_id: 0,
max_memory_usage_bytes: CHANNEL_MEMORY,
send_type: SendType::ReliableOrdered { resend_time: INTERVAL },
},
ChannelConfig {
channel_id: 2,
max_memory_usage_bytes: CHANNEL_MEMORY,
send_type: SendType::ReliableUnordered { resend_time: INTERVAL },
},
ChannelConfig {
channel_id: 1,
max_memory_usage_bytes: CHANNEL_MEMORY,
send_type: SendType::Unreliable,
},
];
let mut renet = RenetClient::new(ConnectionConfig {
available_bytes_per_tick: (2 * MAX_MESSAGE_BYTES) as u64,
server_channels_config: channels.clone(),
client_channels_config: channels,
});
renet.set_connected();
Self {
addr,
client_nonce,
server_nonce,
connected: server_nonce != 0,
last_seen: now,
last_sent: now,
renet,
next_unreliable: 0,
received_unreliable: None,
}
}
}
type Inbox = VecDeque<(SocketAddr, Vec<u8>)>;
#[cfg(feature = "steam")]
pub(super) trait PacketTransport: Send {
fn send(&mut self, to: SocketAddr, bytes: &[u8]) -> io::Result<bool>;
fn recv(&mut self, buf: &mut [u8]) -> io::Result<Option<(SocketAddr, usize)>>;
#[cfg(feature = "steam")]
fn lobby(&self) -> Option<&super::SteamLobby> {
None
}
}
enum Wire {
#[cfg(feature = "steam")]
Custom(Box<dyn PacketTransport>),
Udp(UdpSocket),
Local {
me: usize,
inboxes: Arc<Mutex<Vec<Inbox>>>,
drop_one_in: u32,
sent: u64,
},
}
impl Wire {
fn addr(&self) -> io::Result<SocketAddr> {
match self {
#[cfg(feature = "steam")]
Self::Custom(_) => Err(io::Error::other("no UDP address")),
Self::Udp(s) => s.local_addr(),
Self::Local { me, .. } => Ok(local_addr(*me)),
}
}
fn send(&mut self, addr: SocketAddr, bytes: &[u8]) -> io::Result<bool> {
match self {
#[cfg(feature = "steam")]
Self::Custom(t) => t.send(addr, bytes),
Self::Udp(s) => match s.send_to(bytes, addr) {
Ok(_) => Ok(true),
Err(e) if e.kind() == io::ErrorKind::WouldBlock => Ok(false),
Err(e) => Err(e),
},
Self::Local {
me,
inboxes,
drop_one_in,
sent,
} => {
*sent += 1;
if *drop_one_in != 0 && sent.is_multiple_of(u64::from(*drop_one_in)) {
return Ok(true);
}
let mut queues = inboxes.lock().expect("loopback inbox");
if let Some(queue) = queues.get_mut(usize::from(addr.port() - 1))
&& queue.len() < 4096
{
queue.push_back((local_addr(*me), bytes.to_vec()));
}
Ok(true)
}
}
}
fn recv(&mut self, buf: &mut [u8]) -> io::Result<Option<(SocketAddr, usize)>> {
match self {
#[cfg(feature = "steam")]
Self::Custom(t) => t.recv(buf),
Self::Udp(s) => match s.recv_from(buf) {
Ok((n, addr)) => Ok(Some((addr, n))),
Err(e) if e.kind() == io::ErrorKind::WouldBlock => Ok(None),
Err(e) => Err(e),
},
Self::Local { me, inboxes, .. } => {
Ok(inboxes.lock().expect("loopback inbox")[*me]
.pop_back()
.map(|(addr, bytes)| {
buf[..bytes.len()].copy_from_slice(&bytes);
(addr, bytes.len())
}))
}
}
}
}
fn local_addr(index: usize) -> SocketAddr {
SocketAddr::from(([127, 0, 0, 1], (index + 1) as u16))
}
fn nonce() -> u64 {
let mut h = std::collections::hash_map::RandomState::new().build_hasher();
h.write_u8(1);
h.finish().max(1)
}
fn invalid(message: &str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidInput, message)
}
pub struct NetMessages {
wire: Wire,
hosting: bool,
max_clients: usize,
links: BTreeMap<PeerId, Link>,
next_peer: u64,
now: Duration,
last_poll: crate::Instant,
timeout: Duration,
events: VecDeque<MessageEvent>,
event_bytes: usize,
stats: MessageStats,
scratch: Vec<u8>,
}
impl NetMessages {
#[cfg(feature = "steam")]
pub(super) fn transported(
transport: Box<dyn PacketTransport>,
host: Option<SocketAddr>,
max_clients: usize,
) -> Self {
let mut net = Self::new(Wire::Custom(transport), host.is_none(), max_clients);
if let Some(addr) = host {
net.links.insert(PeerId::HOST, Link::new(addr, nonce(), 0, net.now));
}
net
}
#[cfg(feature = "steam")]
pub(super) fn transport_peer_addr(&self, peer: PeerId) -> Option<SocketAddr> {
self.links.get(&peer).filter(|l| l.connected).map(|l| l.addr)
}
#[cfg(feature = "steam")]
pub(super) fn transport_lobby(&self) -> Option<&super::SteamLobby> {
match &self.wire {
Wire::Custom(t) => t.lobby(),
_ => None,
}
}
fn new(wire: Wire, hosting: bool, max_clients: usize) -> Self {
Self {
wire,
hosting,
max_clients,
links: BTreeMap::new(),
next_peer: 1,
now: Duration::ZERO,
last_poll: crate::Instant::now(),
timeout: Duration::from_secs(10),
events: VecDeque::new(),
event_bytes: 0,
stats: MessageStats::default(),
scratch: Vec::with_capacity(1500),
}
}
pub fn host(addr: impl ToSocketAddrs, max_clients: usize) -> io::Result<Self> {
if !(1..=1024).contains(&max_clients) {
return Err(invalid("max_clients must be 1..=1024"));
}
let socket = UdpSocket::bind(addr)?;
socket.set_nonblocking(true)?;
Ok(Self::new(Wire::Udp(socket), true, max_clients))
}
pub fn join(addr: impl ToSocketAddrs) -> io::Result<Self> {
let addr = addr
.to_socket_addrs()?
.next()
.ok_or_else(|| invalid("no host address"))?;
let bind = if addr.is_ipv4() { "0.0.0.0:0" } else { "[::]:0" };
let socket = UdpSocket::bind(bind)?;
socket.set_nonblocking(true)?;
let mut net = Self::new(Wire::Udp(socket), false, 1);
net.links.insert(PeerId::HOST, Link::new(addr, nonce(), 0, net.now));
Ok(net)
}
pub fn loopback(players: usize) -> Vec<Self> {
Self::loopback_lossy(players, 0)
}
pub fn loopback_lossy(players: usize, drop_one_in: u32) -> Vec<Self> {
assert!((2..=1025).contains(&players), "players must be 2..=1025");
let inboxes = Arc::new(Mutex::new((0..players).map(|_| VecDeque::new()).collect()));
(0..players)
.map(|me| {
let mut net = Self::new(
Wire::Local {
me,
inboxes: inboxes.clone(),
drop_one_in,
sent: me as u64,
},
me == 0,
players - 1,
);
if me != 0 {
net.links
.insert(PeerId::HOST, Link::new(local_addr(0), nonce(), 0, net.now));
}
net
})
.collect()
}
pub fn timeout(mut self, timeout: Duration) -> Self {
assert!(!timeout.is_zero(), "timeout must be positive");
self.timeout = timeout;
self
}
pub fn local_addr(&self) -> Option<SocketAddr> {
match self.wire {
Wire::Udp(_) => self.wire.addr().ok(),
_ => None,
}
}
pub fn send(&mut self, peer: PeerId, lane: Lane, bytes: &[u8]) -> io::Result<()> {
if bytes.len() > MAX_MESSAGE_BYTES {
return Err(invalid("message exceeds 1 MiB"));
}
let link = self
.links
.get_mut(&peer)
.filter(|l| l.connected)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "peer is not connected"))?;
let extra = if lane == Lane::Unreliable { 8 } else { 0 };
if !link.renet.can_send_message(lane.channel(), bytes.len() + extra) {
return Err(io::Error::new(io::ErrorKind::WouldBlock, "message lane is full"));
}
let mut payload = Vec::with_capacity(bytes.len() + extra);
if lane == Lane::Unreliable {
payload.extend_from_slice(&link.next_unreliable.to_le_bytes());
link.next_unreliable = link
.next_unreliable
.checked_add(1)
.ok_or_else(|| invalid("sequence exhausted"))?;
}
payload.extend_from_slice(bytes);
link.renet.send_message(lane.channel(), payload);
Ok(())
}
pub fn send_serde<T: Serialize>(&mut self, peer: PeerId, lane: Lane, value: &T) -> io::Result<()> {
let bytes = serde_json::to_vec(value).map_err(io::Error::other)?;
self.send(peer, lane, &bytes)
}
pub fn send_to_all(&mut self, lane: Lane, bytes: &[u8]) -> io::Result<()> {
if bytes.len() > MAX_MESSAGE_BYTES {
return Err(invalid("message exceeds 1 MiB"));
}
let peers: Vec<_> = self
.links
.iter()
.filter(|(_, l)| l.connected)
.map(|(&p, _)| p)
.collect();
for peer in peers {
self.send(peer, lane, bytes)?;
}
Ok(())
}
pub fn poll(&mut self) -> io::Result<()> {
let delta = self.last_poll.elapsed();
self.poll_elapsed(delta)
}
pub fn poll_elapsed(&mut self, elapsed: Duration) -> io::Result<()> {
self.last_poll = crate::Instant::now();
self.now = self.now.saturating_add(elapsed);
if self.events.len() >= EVENT_LIMIT || self.event_bytes >= CHANNEL_MEMORY {
return Err(io::Error::new(
io::ErrorKind::WouldBlock,
"drain message events before polling",
));
}
for link in self.links.values_mut() {
link.renet.update(elapsed);
}
let mut buf = [0u8; 1500];
for _ in 0..256 {
let Some((addr, n)) = self.wire.recv(&mut buf)? else {
break;
};
self.stats.packets_received += 1;
self.stats.bytes_received += n as u64;
self.receive(addr, &buf[..n])?;
}
let mut links = std::mem::take(&mut self.links);
let result = (|| {
for (&peer, link) in &mut links {
if self.now.saturating_sub(link.last_seen) >= self.timeout {
self.events
.push_back(MessageEvent::Disconnected(peer, "peer timed out".into()));
link.renet.disconnect();
continue;
}
if let Some(reason) = link.renet.disconnect_reason() {
self.events
.push_back(MessageEvent::Disconnected(peer, reason.to_string()));
continue;
}
if link.connected {
for lane in [Lane::Reliable, Lane::ReliableUnordered, Lane::Unreliable] {
while self.events.len() < EVENT_LIMIT && self.event_bytes < CHANNEL_MEMORY {
let Some(message) = link.renet.receive_message(lane.channel()) else {
break;
};
let bytes = if lane == Lane::Unreliable {
let Some(seq) = message.get(..8) else { continue };
let seq = u64::from_le_bytes(seq.try_into().expect("sequence"));
if link.received_unreliable.is_some_and(|last| seq <= last) {
continue;
}
link.received_unreliable = Some(seq);
&message[8..]
} else {
&message[..]
};
if bytes.len() > MAX_MESSAGE_BYTES {
link.renet.disconnect();
self.events
.push_back(MessageEvent::Disconnected(peer, "message exceeds 1 MiB".into()));
break;
}
self.event_bytes += bytes.len();
self.events.push_back(MessageEvent::Message {
peer,
lane,
bytes: bytes.to_vec(),
});
}
}
for packet in link.renet.get_packets_to_send() {
self.transmit(link, DATA, &packet)?;
}
}
if self.now.saturating_sub(link.last_sent) >= INTERVAL {
self.transmit(link, if link.connected { HEARTBEAT } else { HELLO }, &[])?;
link.last_sent = self.now;
}
}
Ok(())
})();
links.retain(|_, link| !link.renet.is_disconnected());
self.links = links;
result
}
fn transmit(&mut self, link: &Link, kind: u8, data: &[u8]) -> io::Result<()> {
self.scratch.clear();
self.scratch.extend_from_slice(MAGIC);
self.scratch.push(kind);
self.scratch.extend_from_slice(&link.client_nonce.to_le_bytes());
self.scratch.extend_from_slice(&link.server_nonce.to_le_bytes());
self.scratch.extend_from_slice(data);
if self.wire.send(link.addr, &self.scratch)? {
self.stats.packets_sent += 1;
self.stats.bytes_sent += self.scratch.len() as u64;
}
Ok(())
}
fn receive(&mut self, addr: SocketAddr, data: &[u8]) -> io::Result<()> {
if data.len() < HEADER || &data[..4] != MAGIC {
return Ok(());
}
let kind = data[4];
let cn = u64::from_le_bytes(data[5..13].try_into().expect("nonce"));
let sn = u64::from_le_bytes(data[13..21].try_into().expect("nonce"));
if self.hosting && kind == HELLO && cn != 0 && sn == 0 {
if let Some(link) = self.links.values().find(|l| l.addr == addr) {
if link.client_nonce == cn {
let peer = *self.links.iter().find(|(_, l)| l.addr == addr).expect("link").0;
let link = self.links.remove(&peer).expect("link");
let result = self.transmit(&link, WELCOME, &[]);
self.links.insert(peer, link);
result?;
}
return Ok(());
}
let link = Link::new(addr, cn, nonce(), self.now);
if self.links.len() >= self.max_clients {
let rejected = Link::new(addr, cn, 0, self.now);
return self.transmit(&rejected, FULL, &[]);
}
let peer = PeerId(self.next_peer);
self.next_peer = self
.next_peer
.checked_add(1)
.ok_or_else(|| io::Error::other("peer IDs exhausted"))?;
self.transmit(&link, WELCOME, &[])?;
self.links.insert(peer, link);
self.events.push_back(MessageEvent::Connected(peer));
return Ok(());
}
let Some((&peer, _)) = self.links.iter().find(|(_, l)| l.addr == addr && l.client_nonce == cn) else {
return Ok(());
};
let link = self.links.get_mut(&peer).expect("link");
if !self.hosting && !link.connected && kind == FULL && sn == 0 {
self.links.remove(&peer);
self.events
.push_back(MessageEvent::Disconnected(peer, "host is full".into()));
return Ok(());
}
if !self.hosting && !link.connected && kind == WELCOME && sn != 0 {
link.server_nonce = sn;
link.connected = true;
self.events.push_back(MessageEvent::Connected(peer));
}
if !link.connected || sn != link.server_nonce {
return Ok(());
}
if !matches!(kind, WELCOME | DATA | BYE | HEARTBEAT) {
return Ok(());
}
link.last_seen = self.now;
match kind {
DATA => link.renet.process_packet(&data[HEADER..]),
BYE => {
self.links.remove(&peer);
self.events
.push_back(MessageEvent::Disconnected(peer, "peer disconnected".into()));
}
_ => {}
}
Ok(())
}
pub fn next_event(&mut self) -> Option<MessageEvent> {
let event = self.events.pop_front()?;
if let MessageEvent::Message { bytes, .. } = &event {
self.event_bytes = self.event_bytes.saturating_sub(bytes.len());
}
Some(event)
}
pub fn events(&mut self) -> impl Iterator<Item = MessageEvent> + '_ {
self.event_bytes = 0;
self.events.drain(..)
}
pub fn disconnect(&mut self, peer: PeerId) -> io::Result<()> {
let link = self
.links
.remove(&peer)
.ok_or_else(|| io::Error::new(io::ErrorKind::NotConnected, "unknown peer"))?;
self.events
.push_back(MessageEvent::Disconnected(peer, "disconnected locally".into()));
for _ in 0..3 {
self.transmit(&link, BYE, &[])?;
}
Ok(())
}
pub fn stats(&self) -> MessageStats {
self.stats
}
}