use std::collections::HashMap;
use std::fmt;
use bb_ir::types::{Storage, TypeNode, TYPE_MULTIADDRESS};
use crate::ids::{ComponentRef, NodeSiteId, PeerId};
pub type Multiaddress = Address;
#[derive(Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]
pub enum Protocol {
P2p(PeerId),
Site(NodeSiteId),
Component(ComponentRef),
Op(String),
}
const CODE_P2P: u64 = 421;
const CODE_SITE: u64 = 0xE2;
const CODE_COMPONENT: u64 = 0xE3;
const CODE_OP: u64 = 0xE4;
impl Protocol {
pub const fn code(&self) -> u64 {
match self {
Protocol::P2p(_) => CODE_P2P,
Protocol::Site(_) => CODE_SITE,
Protocol::Component(_) => CODE_COMPONENT,
Protocol::Op(_) => CODE_OP,
}
}
fn write_to(&self, out: &mut Vec<u8>) {
let mut code_buf = unsigned_varint::encode::u64_buffer();
out.extend_from_slice(unsigned_varint::encode::u64(self.code(), &mut code_buf));
match self {
Protocol::P2p(p) => {
let mh_bytes = p.to_bytes();
let mut len_buf = unsigned_varint::encode::usize_buffer();
out.extend_from_slice(unsigned_varint::encode::usize(mh_bytes.len(), &mut len_buf));
out.extend_from_slice(&mh_bytes);
}
Protocol::Site(s) => out.extend_from_slice(&s.as_u64().to_be_bytes()),
Protocol::Component(c) => out.extend_from_slice(&c.as_u32().to_be_bytes()),
Protocol::Op(name) => {
let bytes = name.as_bytes();
let mut len_buf = unsigned_varint::encode::usize_buffer();
out.extend_from_slice(unsigned_varint::encode::usize(bytes.len(), &mut len_buf));
out.extend_from_slice(bytes);
}
}
}
fn read_from(buf: &[u8]) -> Result<(Self, usize), AddressError> {
let (code, rest) =
unsigned_varint::decode::u64(buf).map_err(|_| AddressError::Truncated)?;
let code_len = buf.len() - rest.len();
let (prot, payload_len) = match code {
CODE_P2P => {
let (mh_len, after_len) =
unsigned_varint::decode::usize(rest).map_err(|_| AddressError::Truncated)?;
let mh_len_bytes = rest.len() - after_len.len();
let mh_bytes = after_len.get(0..mh_len).ok_or(AddressError::Truncated)?;
let peer =
PeerId::from_bytes(mh_bytes).map_err(|_| AddressError::InvalidValue {
protocol: "p2p".into(),
value: format!("malformed multihash ({mh_len} bytes)"),
})?;
(Protocol::P2p(peer), mh_len_bytes + mh_len)
}
CODE_SITE => {
let bytes: [u8; 8] = rest
.get(0..8)
.ok_or(AddressError::Truncated)?
.try_into()
.expect("8-byte slice");
(
Protocol::Site(NodeSiteId::from(u64::from_be_bytes(bytes))),
8,
)
}
CODE_COMPONENT => {
let bytes: [u8; 4] = rest
.get(0..4)
.ok_or(AddressError::Truncated)?
.try_into()
.expect("4-byte slice");
(
Protocol::Component(ComponentRef::from(u32::from_be_bytes(bytes))),
4,
)
}
CODE_OP => {
let (str_len, after_len) =
unsigned_varint::decode::usize(rest).map_err(|_| AddressError::Truncated)?;
let len_bytes = rest.len() - after_len.len();
let str_bytes = after_len.get(0..str_len).ok_or(AddressError::Truncated)?;
let name = std::str::from_utf8(str_bytes)
.map_err(|_| AddressError::InvalidValue {
protocol: "op".into(),
value: format!("non-utf8 {str_len} bytes"),
})?
.to_string();
(Protocol::Op(name), len_bytes + str_len)
}
other => {
return Err(AddressError::UnknownCode((other & 0xFF) as u8));
}
};
Ok((prot, code_len + payload_len))
}
}
#[derive(Clone, Debug, Default, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]
pub struct Address {
segments: Vec<Protocol>,
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum AddressError {
Truncated,
UnknownCode(u8),
MalformedString(String),
InvalidValue {
protocol: String,
value: String,
},
}
impl fmt::Display for AddressError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
AddressError::Truncated => write!(f, "Address: buffer truncated"),
AddressError::UnknownCode(c) => write!(f, "Address: unknown protocol code 0x{c:x}"),
AddressError::MalformedString(s) => write!(f, "Address: malformed string `{s}`"),
AddressError::InvalidValue { protocol, value } => {
write!(
f,
"Address: invalid value `{value}` for protocol `{protocol}`"
)
}
}
}
}
impl std::error::Error for AddressError {}
impl Address {
pub fn empty() -> Self {
Self::default()
}
fn with(mut self, segment: Protocol) -> Self {
self.segments.push(segment);
self
}
pub fn p2p(self, peer: PeerId) -> Self {
self.with(Protocol::P2p(peer))
}
pub fn site(self, site: NodeSiteId) -> Self {
self.with(Protocol::Site(site))
}
pub fn component(self, c: ComponentRef) -> Self {
self.with(Protocol::Component(c))
}
pub fn op(self, name: impl Into<String>) -> Self {
self.with(Protocol::Op(name.into()))
}
pub fn segments(&self) -> &[Protocol] {
&self.segments
}
pub fn peer_id(&self) -> Option<PeerId> {
self.segments.iter().find_map(|p| match p {
Protocol::P2p(id) => Some(*id),
_ => None,
})
}
pub fn site_id(&self) -> Option<NodeSiteId> {
self.segments.iter().find_map(|p| match p {
Protocol::Site(id) => Some(*id),
_ => None,
})
}
pub fn component_ref(&self) -> Option<ComponentRef> {
self.segments.iter().find_map(|p| match p {
Protocol::Component(c) => Some(*c),
_ => None,
})
}
pub fn op_name(&self) -> Option<&str> {
self.segments.iter().find_map(|p| match p {
Protocol::Op(name) => Some(name.as_str()),
_ => None,
})
}
pub fn to_bytes(&self) -> Vec<u8> {
let mut buf = Vec::new();
for seg in &self.segments {
seg.write_to(&mut buf);
}
buf
}
pub fn from_bytes(mut bytes: &[u8]) -> Result<Self, AddressError> {
let mut segments: Vec<Protocol> = Vec::new();
while !bytes.is_empty() {
let (seg, consumed) = Protocol::read_from(bytes)?;
segments.push(seg);
bytes = &bytes[consumed..];
}
Ok(Address { segments })
}
pub fn parse_str(s: &str) -> Result<Self, AddressError> {
let trimmed = s.trim();
if trimmed.is_empty() {
return Ok(Self::empty());
}
if !trimmed.starts_with('/') {
return Err(AddressError::MalformedString(s.to_string()));
}
let mut parts = trimmed[1..].split('/').peekable();
let mut segments: Vec<Protocol> = Vec::new();
while parts.peek().is_some() {
let protocol = parts
.next()
.ok_or_else(|| AddressError::MalformedString(s.to_string()))?;
let value = parts
.next()
.ok_or_else(|| AddressError::MalformedString(s.to_string()))?;
let seg = match protocol {
"p2p" => {
let mh_bytes = bs58::decode(value)
.into_vec()
.map_err(|_| invalid(protocol, value))?;
Protocol::P2p(
PeerId::from_bytes(&mh_bytes).map_err(|_| invalid(protocol, value))?,
)
}
"site" => Protocol::Site(NodeSiteId::from(
value.parse::<u64>().map_err(|_| invalid(protocol, value))?,
)),
"component" => Protocol::Component(ComponentRef::from(
value.parse::<u32>().map_err(|_| invalid(protocol, value))?,
)),
"op" => Protocol::Op(value.to_string()),
other => return Err(invalid(other, value)),
};
segments.push(seg);
}
Ok(Address { segments })
}
}
fn invalid(protocol: &str, value: &str) -> AddressError {
AddressError::InvalidValue {
protocol: protocol.into(),
value: value.into(),
}
}
impl fmt::Display for Address {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for seg in &self.segments {
match seg {
Protocol::P2p(id) => write!(f, "/p2p/{id}")?,
Protocol::Site(id) => write!(f, "/site/{}", id.as_u64())?,
Protocol::Component(c) => write!(f, "/component/{}", c.as_u32())?,
Protocol::Op(name) => write!(f, "/op/{}", name)?,
}
}
Ok(())
}
}
impl std::str::FromStr for Address {
type Err = AddressError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Self::parse_str(s)
}
}
impl Storage for Address {
const TYPE: &'static TypeNode = &TYPE_MULTIADDRESS;
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AddressBookError {
UnknownPeer(PeerId),
EmptyAddressList,
Full {
cap: usize,
},
AllocationFailed {
requested: usize,
},
}
impl fmt::Display for AddressBookError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
AddressBookError::UnknownPeer(p) => {
write!(f, "AddressBook: peer {p} not registered")
}
AddressBookError::EmptyAddressList => {
write!(f, "AddressBook: add_peer requires a non-empty address list")
}
AddressBookError::Full { cap } => {
write!(f, "AddressBook: at cap {cap}, new peer rejected")
}
AddressBookError::AllocationFailed { requested } => {
write!(
f,
"AddressBook: dedup reservation for {requested} addresses failed"
)
}
}
}
}
impl std::error::Error for AddressBookError {}
struct AddressEntry {
addresses: Vec<Address>,
ref_count: u64,
}
pub const DEFAULT_ADDRESS_BOOK_CAP: usize = 16_384;
pub struct AddressBook {
entries: HashMap<PeerId, AddressEntry>,
cap: usize,
}
impl Default for AddressBook {
fn default() -> Self {
Self {
entries: HashMap::new(),
cap: DEFAULT_ADDRESS_BOOK_CAP,
}
}
}
impl AddressBook {
pub fn new() -> Self {
Self::default()
}
pub fn set_cap(&mut self, cap: usize) {
self.cap = cap.max(1);
}
pub fn add_peer(
&mut self,
peer: PeerId,
addresses: Vec<Address>,
) -> Result<(), AddressBookError> {
if addresses.is_empty() {
return Err(AddressBookError::EmptyAddressList);
}
match self.entries.get_mut(&peer) {
Some(entry) => {
entry.ref_count = entry.ref_count.saturating_add(1);
for addr in addresses {
if !entry.addresses.contains(&addr) {
entry.addresses.push(addr);
}
}
}
None => {
if self.entries.len() >= self.cap {
return Err(AddressBookError::Full { cap: self.cap });
}
let mut dedup: Vec<Address> = Vec::new();
crate::fallible::try_reserve_exact(&mut dedup, addresses.len()).map_err(|_| {
AddressBookError::AllocationFailed {
requested: addresses.len(),
}
})?;
for addr in addresses {
if !dedup.contains(&addr) {
dedup.push(addr);
}
}
self.entries.insert(
peer,
AddressEntry {
addresses: dedup,
ref_count: 1,
},
);
}
}
Ok(())
}
pub fn drop_peer(&mut self, peer: PeerId) -> Result<(), AddressBookError> {
let Some(entry) = self.entries.get_mut(&peer) else {
return Err(AddressBookError::UnknownPeer(peer));
};
entry.ref_count = entry.ref_count.saturating_sub(1);
if entry.ref_count == 0 {
self.entries.remove(&peer);
}
Ok(())
}
pub fn register_address(
&mut self,
peer: PeerId,
address: Address,
) -> Result<(), AddressBookError> {
let Some(entry) = self.entries.get_mut(&peer) else {
return Err(AddressBookError::UnknownPeer(peer));
};
if !entry.addresses.contains(&address) {
entry.addresses.push(address);
}
Ok(())
}
pub fn forget_address(
&mut self,
peer: PeerId,
address: &Address,
) -> Result<(), AddressBookError> {
let Some(entry) = self.entries.get_mut(&peer) else {
return Err(AddressBookError::UnknownPeer(peer));
};
entry.addresses.retain(|a| a != address);
Ok(())
}
pub fn lookup(&self, peer: PeerId) -> Option<&[Address]> {
let entry = self.entries.get(&peer)?;
if entry.addresses.is_empty() {
None
} else {
Some(entry.addresses.as_slice())
}
}
pub fn lookup_first(&self, peer: PeerId) -> Option<&Address> {
self.lookup(peer).and_then(|addrs| addrs.first())
}
pub fn ref_count(&self, peer: PeerId) -> u64 {
self.entries.get(&peer).map(|e| e.ref_count).unwrap_or(0)
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (PeerId, &[Address], u64)> {
self.entries
.iter()
.map(|(p, e)| (*p, e.addresses.as_slice(), e.ref_count))
}
pub fn restore_entry(&mut self, peer: PeerId, addresses: Vec<Address>, ref_count: u64) {
self.entries.insert(
peer,
AddressEntry {
addresses,
ref_count,
},
);
}
}