use std::collections::{HashMap, VecDeque};
use std::io::{self, Read, Write};
use crate::snapshot_frame::cap_count;
pub type SockId = u32;
pub const EAGAIN: i32 = 11;
pub const EBADF: i32 = 9;
pub const EINVAL: i32 = 22;
pub const EPIPE: i32 = 32;
pub const EISCONN: i32 = 106;
pub const ENOTCONN: i32 = 107;
pub const ECONNREFUSED: i32 = 111;
pub const EADDRINUSE: i32 = 98;
const RX_CAP: usize = 256 * 1024;
pub const SHUT_RD: i32 = 0;
pub const SHUT_WR: i32 = 1;
pub const SHUT_RDWR: i32 = 2;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct Endpoint {
pub v6: bool,
pub ip: [u8; 16],
pub port: u16,
}
impl Endpoint {
pub fn v4(a: [u8; 4], port: u16) -> Endpoint {
let mut ip = [0u8; 16];
ip[..4].copy_from_slice(&a);
Endpoint {
v6: false,
ip,
port,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SockState {
Unbound,
Bound,
Listen,
Established,
Closed,
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct Sock {
state: SockState,
local: Option<Endpoint>,
remote: Option<Endpoint>,
peer: Option<SockId>,
rx: VecDeque<u8>,
rd_shut: bool,
peer_wr_shut: bool,
backlog: VecDeque<SockId>,
refs: u32,
}
impl Sock {
fn new() -> Sock {
Sock {
state: SockState::Unbound,
local: None,
remote: None,
peer: None,
rx: VecDeque::new(),
rd_shut: false,
peer_wr_shut: false,
backlog: VecDeque::new(),
refs: 1,
}
}
}
#[derive(Debug, Default)]
pub struct LoopNet {
socks: HashMap<SockId, Sock>,
listeners: HashMap<Endpoint, SockId>,
next_id: SockId,
next_ephemeral: u16,
}
impl LoopNet {
pub fn new() -> LoopNet {
LoopNet {
socks: HashMap::new(),
listeners: HashMap::new(),
next_id: 1,
next_ephemeral: 49152,
}
}
fn alloc_id(&mut self) -> SockId {
let id = self.next_id;
self.next_id = self.next_id.checked_add(1).expect("SockId space exhausted");
id
}
fn alloc_ephemeral(&mut self) -> u16 {
let p = self.next_ephemeral;
self.next_ephemeral = if self.next_ephemeral >= 65535 {
49152
} else {
self.next_ephemeral + 1
};
p
}
pub fn socket(&mut self) -> SockId {
let id = self.alloc_id();
self.socks.insert(id, Sock::new());
id
}
fn get(&self, id: SockId) -> Result<&Sock, i32> {
self.socks.get(&id).ok_or(EBADF)
}
fn get_mut(&mut self, id: SockId) -> Result<&mut Sock, i32> {
self.socks.get_mut(&id).ok_or(EBADF)
}
pub fn bind(&mut self, id: SockId, mut ep: Endpoint) -> Result<(), i32> {
if ep.port == 0 {
let mut chosen = 0u16;
for _ in 0..16384 {
let p = self.alloc_ephemeral();
let cand = Endpoint { port: p, ..ep };
if !self.listeners.contains_key(&cand) {
chosen = p;
break;
}
}
if chosen == 0 {
return Err(EADDRINUSE); }
ep.port = chosen;
}
if self.listeners.contains_key(&ep) {
return Err(EADDRINUSE);
}
let s = self.get_mut(id)?;
if s.state != SockState::Unbound {
return Err(EINVAL);
}
s.local = Some(ep);
s.state = SockState::Bound;
Ok(())
}
pub fn listen(&mut self, id: SockId) -> Result<(), i32> {
let ep = {
let s = self.get_mut(id)?;
match s.state {
SockState::Bound | SockState::Listen => {}
_ => return Err(EINVAL),
}
s.state = SockState::Listen;
s.local.ok_or(EINVAL)?
};
self.listeners.insert(ep, id);
Ok(())
}
pub fn peer_of(&self, id: SockId) -> Option<SockId> {
self.socks.get(&id).and_then(|s| s.peer)
}
pub fn connect(&mut self, id: SockId, dst: Endpoint) -> Result<SockId, i32> {
{
let s = self.get(id)?;
if s.state == SockState::Established {
return Err(EISCONN);
}
if !matches!(s.state, SockState::Unbound | SockState::Bound) {
return Err(EINVAL);
}
}
let listener_id = self
.listeners
.get(&dst)
.copied()
.or_else(|| {
let wild = Endpoint {
v6: dst.v6,
ip: [0u8; 16],
port: dst.port,
};
self.listeners.get(&wild).copied()
})
.or_else(|| {
if dst.v6 {
None
} else {
let wild_v6 = Endpoint {
v6: true,
ip: [0u8; 16],
port: dst.port,
};
self.listeners.get(&wild_v6).copied()
}
})
.ok_or(ECONNREFUSED)?;
let local = match self.get(id)?.local {
Some(ep) => ep,
None => {
let port = self.alloc_ephemeral();
if dst.v6 {
let mut ip = [0u8; 16];
ip[15] = 1; Endpoint { v6: true, ip, port }
} else {
Endpoint::v4([127, 0, 0, 1], port)
}
}
};
let server_id = self.alloc_id();
let mut server = Sock::new();
server.state = SockState::Established;
server.local = Some(dst);
server.remote = Some(local);
server.peer = Some(id);
self.socks.insert(server_id, server);
{
let c = self.get_mut(id)?;
c.state = SockState::Established;
c.local = Some(local);
c.remote = Some(dst);
c.peer = Some(server_id);
}
self.get_mut(listener_id)?.backlog.push_back(server_id);
Ok(listener_id)
}
pub fn connect_first(&mut self, id: SockId) -> Result<SockId, i32> {
let dst = self.listeners.keys().copied().next().ok_or(ECONNREFUSED)?;
self.connect(id, dst)
}
pub fn accept(&mut self, listener_id: SockId) -> Result<SockId, i32> {
let s = self.get_mut(listener_id)?;
if s.state != SockState::Listen {
return Err(EINVAL);
}
s.backlog.pop_front().ok_or(EAGAIN)
}
pub fn send(&mut self, id: SockId, buf: &[u8]) -> Result<usize, i32> {
let peer = {
let s = self.get(id)?;
if s.state != SockState::Established {
return Err(ENOTCONN);
}
s.peer
};
let peer = peer.ok_or(EPIPE)?;
let p = match self.socks.get_mut(&peer) {
Some(p) if !p.rd_shut && p.state == SockState::Established => p,
_ => return Err(EPIPE),
};
let space = RX_CAP.saturating_sub(p.rx.len());
if space == 0 {
return Err(EAGAIN);
}
let n = buf.len().min(space);
p.rx.extend(&buf[..n]);
Ok(n)
}
pub fn recv(&mut self, id: SockId, max: usize) -> Result<Vec<u8>, i32> {
let s = self.get_mut(id)?;
if s.state != SockState::Established && s.state != SockState::Closed {
return Err(ENOTCONN);
}
if s.rx.is_empty() {
if s.peer_wr_shut || s.peer.is_none() {
return Ok(Vec::new()); }
return Err(EAGAIN);
}
let n = max.min(s.rx.len());
Ok(s.rx.drain(..n).collect())
}
pub fn shutdown(&mut self, id: SockId, how: i32) -> Result<(), i32> {
let peer = {
let s = self.get_mut(id)?;
if how == SHUT_RD || how == SHUT_RDWR {
s.rd_shut = true;
s.rx.clear();
}
s.peer
};
if (how == SHUT_WR || how == SHUT_RDWR) && peer.is_some() {
if let Some(p) = self.socks.get_mut(&peer.unwrap()) {
p.peer_wr_shut = true;
}
}
Ok(())
}
pub fn dup(&mut self, id: SockId) {
if let Some(s) = self.socks.get_mut(&id) {
s.refs += 1;
}
}
pub fn close(&mut self, id: SockId) -> Result<bool, i32> {
{
let s = self.socks.get_mut(&id).ok_or(EBADF)?;
if s.refs > 1 {
s.refs -= 1;
return Ok(false); }
}
let s = self.socks.get(&id).ok_or(EBADF)?;
let peer = s.peer;
let local = s.local;
let is_listener = s.state == SockState::Listen;
if let Some(pid) = peer {
if let Some(p) = self.socks.get_mut(&pid) {
p.peer_wr_shut = true; p.peer = None; }
}
if is_listener {
if let Some(ep) = local {
self.listeners.remove(&ep);
}
}
self.socks.remove(&id);
Ok(true)
}
pub fn readable(&self, id: SockId) -> bool {
match self.socks.get(&id) {
Some(s) if s.state == SockState::Listen => !s.backlog.is_empty(),
Some(s) => {
!s.rx.is_empty()
|| s.peer_wr_shut
|| (s.state == SockState::Established && s.peer.is_none())
}
None => false,
}
}
pub fn writable(&self, id: SockId) -> bool {
let s = match self.socks.get(&id) {
Some(s) if s.state == SockState::Established => s,
_ => return false,
};
match s.peer.and_then(|p| self.socks.get(&p)) {
Some(p) => !p.rd_shut && p.rx.len() < RX_CAP,
None => false,
}
}
pub fn local_addr(&self, id: SockId) -> Option<Endpoint> {
self.socks.get(&id).and_then(|s| s.local)
}
pub fn peer_addr(&self, id: SockId) -> Option<Endpoint> {
self.socks.get(&id).and_then(|s| s.remote)
}
pub fn capture(&self) -> LoopSnapshot {
let mut socks: Vec<SockRecord> = self
.socks
.iter()
.map(|(&id, s)| SockRecord {
id,
state: s.state,
local: s.local,
remote: s.remote,
peer: s.peer,
rx: s.rx.iter().copied().collect(),
rd_shut: s.rd_shut,
peer_wr_shut: s.peer_wr_shut,
backlog: s.backlog.iter().copied().collect(),
})
.collect();
socks.sort_by_key(|r| r.id); LoopSnapshot {
next_id: self.next_id,
next_ephemeral: self.next_ephemeral,
socks,
}
}
pub fn restore(snap: &LoopSnapshot) -> LoopNet {
let mut net = LoopNet {
socks: HashMap::new(),
listeners: HashMap::new(),
next_id: snap.next_id,
next_ephemeral: snap.next_ephemeral,
};
for r in &snap.socks {
let s = Sock {
state: r.state,
local: r.local,
remote: r.remote,
peer: r.peer,
rx: r.rx.iter().copied().collect(),
rd_shut: r.rd_shut,
peer_wr_shut: r.peer_wr_shut,
backlog: r.backlog.iter().copied().collect(),
refs: 1, };
if s.state == SockState::Listen {
if let Some(ep) = s.local {
net.listeners.insert(ep, r.id);
}
}
net.socks.insert(r.id, s);
}
net
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SockRecord {
pub id: SockId,
pub state: SockState,
pub local: Option<Endpoint>,
pub remote: Option<Endpoint>,
pub peer: Option<SockId>,
pub rx: Vec<u8>,
pub rd_shut: bool,
pub peer_wr_shut: bool,
pub backlog: Vec<SockId>,
}
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct LoopSnapshot {
pub next_id: SockId,
pub next_ephemeral: u16,
pub socks: Vec<SockRecord>,
}
const LOOP_MAGIC: [u8; 8] = *b"SMLOOP01";
const LOOP_VERSION: u32 = 1;
const LOOP_BLOB_MAX: usize = 64 << 20;
#[inline]
fn rb<const N: usize, R: Read>(r: &mut R) -> io::Result<[u8; N]> {
let mut x = [0u8; N];
r.read_exact(&mut x)?;
Ok(x)
}
#[inline]
fn r_u16<R: Read>(r: &mut R) -> io::Result<u16> {
Ok(u16::from_le_bytes(rb::<2, _>(r)?))
}
#[inline]
fn r_u32<R: Read>(r: &mut R) -> io::Result<u32> {
Ok(u32::from_le_bytes(rb::<4, _>(r)?))
}
#[inline]
fn ierr(m: &str) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, m.to_string())
}
fn w_ep<W: Write>(w: &mut W, ep: &Option<Endpoint>) -> io::Result<()> {
match ep {
Some(e) => {
w.write_all(&[1u8, e.v6 as u8])?;
w.write_all(&e.ip)?;
w.write_all(&e.port.to_le_bytes())
}
None => w.write_all(&[0u8]),
}
}
fn r_ep<R: Read>(r: &mut R) -> io::Result<Option<Endpoint>> {
if rb::<1, _>(r)?[0] == 0 {
return Ok(None);
}
let v6 = rb::<1, _>(r)?[0] != 0;
let ip = rb::<16, _>(r)?;
let port = r_u16(r)?;
Ok(Some(Endpoint { v6, ip, port }))
}
impl SockState {
fn tag(&self) -> u8 {
match self {
SockState::Unbound => 0,
SockState::Bound => 1,
SockState::Listen => 2,
SockState::Established => 3,
SockState::Closed => 4,
}
}
fn from_tag(t: u8) -> io::Result<SockState> {
Ok(match t {
0 => SockState::Unbound,
1 => SockState::Bound,
2 => SockState::Listen,
3 => SockState::Established,
4 => SockState::Closed,
_ => return Err(ierr("bad SockState tag")),
})
}
}
impl SockRecord {
fn write_to<W: Write>(&self, w: &mut W) -> io::Result<()> {
w.write_all(&self.id.to_le_bytes())?;
w.write_all(&[self.state.tag()])?;
w_ep(w, &self.local)?;
w_ep(w, &self.remote)?;
match self.peer {
Some(p) => {
w.write_all(&[1u8])?;
w.write_all(&p.to_le_bytes())?;
}
None => w.write_all(&[0u8])?,
}
w.write_all(&(self.rx.len() as u32).to_le_bytes())?;
w.write_all(&self.rx)?;
w.write_all(&[self.rd_shut as u8, self.peer_wr_shut as u8])?;
w.write_all(&(self.backlog.len() as u32).to_le_bytes())?;
for &b in &self.backlog {
w.write_all(&b.to_le_bytes())?;
}
Ok(())
}
fn read_from<R: Read>(r: &mut R) -> io::Result<SockRecord> {
let id = r_u32(r)?;
let state = SockState::from_tag(rb::<1, _>(r)?[0])?;
let local = r_ep(r)?;
let remote = r_ep(r)?;
let peer = if rb::<1, _>(r)?[0] != 0 {
Some(r_u32(r)?)
} else {
None
};
let rxn = r_u32(r)? as usize;
if rxn > LOOP_BLOB_MAX {
return Err(ierr("loop rx blob too large"));
}
let mut rx = vec![0u8; rxn];
r.read_exact(&mut rx)?;
let [rd_shut, peer_wr_shut] = rb::<2, _>(r)?;
let n = r_u32(r)? as usize;
let cap = cap_count(n, 4, LOOP_BLOB_MAX);
let mut backlog = Vec::with_capacity(cap);
for _ in 0..n {
backlog.push(r_u32(r)?);
}
Ok(SockRecord {
id,
state,
local,
remote,
peer,
rx,
rd_shut: rd_shut != 0,
peer_wr_shut: peer_wr_shut != 0,
backlog,
})
}
}
impl LoopSnapshot {
pub fn write_to<W: Write>(&self, w: &mut W) -> io::Result<()> {
w.write_all(&LOOP_MAGIC)?;
w.write_all(&LOOP_VERSION.to_le_bytes())?;
w.write_all(&self.next_id.to_le_bytes())?;
w.write_all(&self.next_ephemeral.to_le_bytes())?;
w.write_all(&(self.socks.len() as u32).to_le_bytes())?;
for s in &self.socks {
s.write_to(w)?;
}
Ok(())
}
pub fn read_from<R: Read>(r: &mut R) -> io::Result<LoopSnapshot> {
if rb::<8, _>(r)? != LOOP_MAGIC {
return Err(ierr("loop snapshot magic mismatch"));
}
if r_u32(r)? != LOOP_VERSION {
return Err(ierr("loop snapshot version skew (hard refusal)"));
}
let next_id = r_u32(r)?;
let next_ephemeral = r_u16(r)?;
let n = r_u32(r)? as usize;
let cap = cap_count(n, 18, LOOP_BLOB_MAX);
let mut socks = Vec::with_capacity(cap);
for _ in 0..n {
socks.push(SockRecord::read_from(r)?);
}
Ok(LoopSnapshot {
next_id,
next_ephemeral,
socks,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn ep(p: u16) -> Endpoint {
Endpoint::v4([127, 0, 0, 1], p)
}
fn connected() -> (LoopNet, SockId, SockId) {
let mut net = LoopNet::new();
let lis = net.socket();
net.bind(lis, ep(5432)).unwrap();
net.listen(lis).unwrap();
let cli = net.socket();
net.connect(cli, ep(5432)).unwrap();
let srv = net.accept(lis).unwrap();
(net, cli, srv)
}
#[test]
fn connect_accept_bidirectional_stream() {
let (mut net, cli, srv) = connected();
assert_eq!(net.send(cli, b"ping").unwrap(), 4);
assert_eq!(net.recv(srv, 64).unwrap(), b"ping");
assert_eq!(net.send(srv, b"pong!").unwrap(), 5);
assert_eq!(net.recv(cli, 64).unwrap(), b"pong!");
assert_eq!(net.peer_addr(cli).unwrap(), ep(5432));
assert_eq!(net.local_addr(srv).unwrap(), ep(5432));
assert_eq!(net.peer_addr(srv).unwrap(), net.local_addr(cli).unwrap());
}
#[test]
fn recv_empty_is_eagain_then_eof_after_peer_shutdown() {
let (mut net, cli, srv) = connected();
assert_eq!(net.recv(srv, 64), Err(EAGAIN)); net.send(cli, b"last").unwrap();
net.shutdown(cli, SHUT_WR).unwrap();
assert_eq!(net.recv(srv, 64).unwrap(), b"last"); assert_eq!(net.recv(srv, 64).unwrap(), b""); }
#[test]
fn connect_with_no_listener_is_refused() {
let mut net = LoopNet::new();
let cli = net.socket();
assert_eq!(net.connect(cli, ep(9999)), Err(ECONNREFUSED));
}
#[test]
fn connect_first_routes_to_registered_listener() {
let mut net = LoopNet::new();
let lis = net.socket();
net.bind(lis, ep(8080)).unwrap();
net.listen(lis).unwrap();
let cli = net.socket();
assert_eq!(net.connect_first(cli).unwrap(), lis);
let srv = net.accept(lis).unwrap();
net.send(cli, b"first").unwrap();
assert_eq!(net.recv(srv, 16).unwrap(), b"first");
}
#[test]
fn ipv4_loopback_reaches_ipv6_wildcard_listener() {
let mut net = LoopNet::new();
let lis = net.socket();
net.bind(
lis,
Endpoint {
v6: true,
ip: [0u8; 16],
port: 18080,
},
)
.unwrap();
net.listen(lis).unwrap();
let cli = net.socket();
assert_eq!(net.connect(cli, ep(18080)).unwrap(), lis);
let srv = net.accept(lis).unwrap();
net.send(cli, b"GET / HTTP/1.0\r\n\r\n").unwrap();
assert_eq!(net.recv(srv, 64).unwrap(), b"GET / HTTP/1.0\r\n\r\n");
}
#[test]
fn send_after_peer_close_is_epipe_and_peer_reads_eof() {
let (mut net, cli, srv) = connected();
net.send(cli, b"buffered").unwrap();
net.close(srv).unwrap();
assert_eq!(net.send(cli, b"x"), Err(EPIPE));
assert_eq!(net.recv(cli, 64).unwrap(), b""); }
#[test]
fn flow_control_backpressure_and_writable() {
let (mut net, cli, srv) = connected();
let big = vec![0u8; RX_CAP];
assert_eq!(net.send(cli, &big).unwrap(), RX_CAP);
assert!(!net.writable(cli), "peer rx full ⇒ not writable");
assert_eq!(net.send(cli, b"more"), Err(EAGAIN)); let _ = net.recv(srv, RX_CAP).unwrap();
assert!(net.writable(cli), "drained ⇒ writable again");
}
#[test]
fn readability_predicate() {
let mut net = LoopNet::new();
let lis = net.socket();
net.bind(lis, ep(80)).unwrap();
net.listen(lis).unwrap();
assert!(!net.readable(lis), "no pending conn");
let cli = net.socket();
net.connect(cli, ep(80)).unwrap();
assert!(net.readable(lis), "listener readable ⇒ accept won't block");
let srv = net.accept(lis).unwrap();
assert!(!net.readable(srv), "fresh conn, no data");
net.send(cli, b"q").unwrap();
assert!(net.readable(srv), "data buffered ⇒ readable");
}
#[test]
fn capture_restore_round_trip_preserves_live_connection() {
let (mut net, cli, srv) = connected();
net.send(cli, b"client->server in flight").unwrap();
net.send(srv, b"server->client in flight").unwrap();
net.shutdown(srv, SHUT_WR).unwrap();
let snap = net.capture();
let mut buf_round = LoopNet::restore(&snap).capture();
buf_round.socks.sort_by_key(|r| r.id);
assert_eq!(snap, buf_round, "snapshot is restore-stable");
drop(net);
let mut r = LoopNet::restore(&snap);
assert_eq!(r.recv(cli, 64).unwrap(), b"server->client in flight");
assert_eq!(
r.recv(cli, 64).unwrap(),
b"",
"peer_wr_shut survived restore ⇒ EOF"
);
assert_eq!(r.recv(srv, 64).unwrap(), b"client->server in flight");
assert_eq!(r.send(cli, b"after restore").unwrap(), 13);
assert_eq!(r.recv(srv, 64).unwrap(), b"after restore");
assert_eq!(r.peer_addr(cli).unwrap(), ep(5432));
}
#[test]
fn listener_and_backlog_survive_restore() {
let mut net = LoopNet::new();
let lis = net.socket();
net.bind(lis, ep(8080)).unwrap();
net.listen(lis).unwrap();
let cli = net.socket();
net.connect(cli, ep(8080)).unwrap(); let snap = net.capture();
let mut r = LoopNet::restore(&snap);
let srv = r.accept(lis).expect("backlog survived");
r.send(cli, b"hi").unwrap();
assert_eq!(r.recv(srv, 64).unwrap(), b"hi");
let cli2 = r.socket();
assert!(
r.connect(cli2, ep(8080)).is_ok(),
"listener re-indexed after restore"
);
}
#[test]
fn loop_snapshot_bytes_round_trip() {
let (mut net, cli, srv) = connected();
net.send(cli, b"in flight c->s").unwrap();
net.send(srv, b"in flight s->c").unwrap();
net.shutdown(srv, SHUT_WR).unwrap();
let lis = net.socket();
net.bind(lis, ep(9000)).unwrap();
net.listen(lis).unwrap();
let c2 = net.socket();
net.connect(c2, ep(9000)).unwrap();
let snap = net.capture();
let mut buf = Vec::new();
snap.write_to(&mut buf).unwrap();
let back = LoopSnapshot::read_from(&mut std::io::Cursor::new(&buf[..])).unwrap();
assert_eq!(snap, back, "LoopSnapshot survives the byte round-trip");
let mut r = LoopNet::restore(&back);
assert_eq!(r.recv(cli, 64).unwrap(), b"in flight s->c");
assert_eq!(
r.recv(cli, 64).unwrap(),
b"",
"half-close survived the byte round-trip"
);
assert_eq!(r.recv(srv, 64).unwrap(), b"in flight c->s");
let srv2 = r.accept(lis).expect("backlog survived the byte round-trip");
r.send(c2, b"hi2").unwrap();
assert_eq!(r.recv(srv2, 64).unwrap(), b"hi2");
}
#[test]
fn loop_snapshot_rejects_bad_magic_and_version() {
let mut buf = Vec::new();
LoopNet::new().capture().write_to(&mut buf).unwrap();
let mut bad_ver = buf.clone();
bad_ver[8] ^= 0xff; assert!(LoopSnapshot::read_from(&mut std::io::Cursor::new(&bad_ver[..])).is_err());
let mut bad_magic = buf.clone();
bad_magic[0] ^= 0xff;
assert!(LoopSnapshot::read_from(&mut std::io::Cursor::new(&bad_magic[..])).is_err());
}
}