#![cfg_attr(not(feature = "std"), no_std)]
extern crate alloc;
mod error;
mod registry;
pub mod codec;
pub mod iter;
pub mod proto;
use alloc::vec::Vec;
use core::fmt;
use core::hash::{Hash, Hasher};
use core::ops::Deref;
use core::str::FromStr;
use once_cell::race::OnceBox;
use std::net::{SocketAddrV4, SocketAddrV6};
use tinyvec::{Array, ArrayVec, TinyVec};
use crate::proto::{DnsAddr, Ip4, Ip6, Tcp};
pub use error::Error;
use ockam_core::env::FromString;
pub use registry::{Registry, RegistryBuilder};
fn default_registry() -> &'static Registry {
static INSTANCE: OnceBox<Registry> = OnceBox::new();
INSTANCE.get_or_init(Box::<Registry>::default)
}
pub trait Protocol<'a>: Sized {
const CODE: Code;
const PREFIX: &'static str;
fn read_str(input: Checked<&'a str>) -> Result<Self, Error>;
fn write_str(&self, f: &mut fmt::Formatter) -> Result<(), Error>;
fn read_bytes(input: Checked<&'a [u8]>) -> Result<Self, Error>;
fn write_bytes(&self, buf: &mut dyn Buffer);
}
pub trait Codec: Send + Sync {
fn split_str<'a>(
&self,
prefix: &str,
input: &'a str,
) -> Result<(Checked<&'a str>, &'a str), Error>;
fn split_bytes<'a>(
&self,
code: Code,
input: &'a [u8],
) -> Result<(Checked<&'a [u8]>, &'a [u8]), Error>;
fn is_valid_bytes(&self, code: Code, value: Checked<&[u8]>) -> bool;
fn write_bytes(&self, val: &ProtoValue, buf: &mut dyn Buffer) -> Result<(), Error>;
fn transcode_str(
&self,
prefix: &str,
value: Checked<&str>,
buf: &mut dyn Buffer,
) -> Result<(), Error>;
fn transcode_bytes(
&self,
code: Code,
value: Checked<&[u8]>,
f: &mut fmt::Formatter,
) -> Result<(), Error>;
}
pub trait Buffer: AsRef<[u8]> {
fn extend_with(&mut self, buf: &[u8]);
}
impl Buffer for Vec<u8> {
fn extend_with(&mut self, buf: &[u8]) {
self.extend_from_slice(buf)
}
}
impl<A: tinyvec::Array<Item = u8>> Buffer for TinyVec<A> {
fn extend_with(&mut self, buf: &[u8]) {
self.extend_from_slice(buf)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Checked<T>(pub T);
impl<T> Deref for Checked<T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.0
}
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Code(u32);
impl fmt::Display for Code {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl Code {
pub const fn new(n: u32) -> Self {
Code(n)
}
}
impl From<Code> for u32 {
fn from(c: Code) -> Self {
c.0
}
}
#[derive(Debug, Clone)]
pub struct ProtoValue<'a> {
code: Code,
data: Bytes<'a>,
}
#[derive(Debug, Clone)]
enum Bytes<'a> {
Slice(Checked<&'a [u8]>),
Owned(Checked<TinyVec<[u8; 28]>>),
}
impl<'a> ProtoValue<'a> {
pub fn code(&self) -> Code {
self.code
}
pub fn data(&self) -> Checked<&[u8]> {
match &self.data {
Bytes::Slice(s) => *s,
Bytes::Owned(v) => Checked(v),
}
}
pub fn cast<P: Protocol<'a>>(&'a self) -> Option<P> {
if self.code != P::CODE {
return None;
}
P::read_bytes(self.data()).ok()
}
pub fn to_owned<'b>(&self) -> ProtoValue<'b> {
match &self.data {
Bytes::Slice(Checked(s)) => ProtoValue {
code: self.code,
data: Bytes::Owned(Checked(TinyVec::Heap(Vec::from(*s)))),
},
Bytes::Owned(Checked(v)) => ProtoValue {
code: self.code,
data: Bytes::Owned(Checked(v.clone())),
},
}
}
}
impl<'a> AsRef<[u8]> for ProtoValue<'a> {
fn as_ref(&self) -> &[u8] {
&self.data()
}
}
#[derive(Debug)]
pub struct MultiAddr {
dat: TinyVec<[u8; 28]>,
off: usize,
reg: Registry,
}
impl Default for MultiAddr {
fn default() -> Self {
MultiAddr::new(default_registry().clone())
}
}
impl PartialEq for MultiAddr {
fn eq(&self, other: &Self) -> bool {
self.as_ref().eq(other.as_ref())
}
}
impl Eq for MultiAddr {}
impl Hash for MultiAddr {
fn hash<H: Hasher>(&self, state: &mut H) {
self.as_ref().hash(state)
}
}
impl Clone for MultiAddr {
fn clone(&self) -> Self {
if self.off > 0 {
MultiAddr {
dat: match &self.dat {
TinyVec::Inline(a) => TinyVec::Inline({
let mut b = ArrayVec::default();
b.extend_from_slice(&a[self.off..]);
b
}),
TinyVec::Heap(v) => TinyVec::Heap({
let mut w = Vec::with_capacity(v.len() - self.off);
w.extend_from_slice(&v[self.off..]);
w
}),
},
off: 0,
reg: self.reg.clone(),
}
} else {
MultiAddr {
dat: self.dat.clone(),
off: self.off,
reg: self.reg.clone(),
}
}
}
}
impl FromString for MultiAddr {
fn from_string(s: &str) -> ockam_core::Result<Self> {
Self::from_str(s).map_err(|_| {
ockam_core::Error::new(
ockam_core::errcode::Origin::Core,
ockam_core::errcode::Kind::Internal,
"MultiAddr parse error",
)
})
}
}
impl MultiAddr {
pub fn new(r: Registry) -> Self {
MultiAddr {
dat: TinyVec::new(),
off: 0,
reg: r,
}
}
pub fn registry(&self) -> &Registry {
&self.reg
}
pub fn try_from_str(input: &str, r: Registry) -> Result<Self, Error> {
let iter = iter::StrIter::with_registry(input, r.clone());
let mut b = TinyVec::new();
for pair in iter {
let (prefix, value) = pair?;
let codec = r
.get_by_prefix(prefix)
.ok_or_else(|| Error::unregistered_prefix(prefix))?;
codec.transcode_str(prefix, value, &mut b)?;
}
Ok(MultiAddr {
dat: b,
off: 0,
reg: r,
})
}
pub fn try_from_bytes(input: &[u8], r: Registry) -> Result<Self, Error> {
let iter = iter::BytesIter::with_registry(input, r.clone());
let mut b = TinyVec::new();
for item in iter {
let (_, code, value) = item?;
let codec = r
.get_by_code(code)
.ok_or_else(|| Error::unregistered(code))?;
if !codec.is_valid_bytes(code, value) {
return Err(Error::invalid_proto(code));
}
}
b.extend_from_slice(input);
Ok(MultiAddr {
dat: b,
off: 0,
reg: r,
})
}
#[cfg(feature = "cbor")]
pub fn try_from_cbor(input: &[u8], r: Registry) -> Result<Self, Error> {
let bytes = minicbor::Decoder::new(input)
.bytes()
.map_err(|e| Error::message(format!("invalid cbor: {e}")))?;
Self::try_from_bytes(bytes, r)
}
pub fn is_empty(&self) -> bool {
self.as_ref().is_empty()
}
pub fn len(&self) -> usize {
self.as_ref().len()
}
pub fn push_back<'a, P: Protocol<'a>>(&mut self, p: P) -> Result<(), Error> {
if self.reg.get_by_code(P::CODE).is_none() {
return Err(Error::unregistered(P::CODE));
}
debug_assert!(self.reg.get_by_prefix(P::PREFIX).is_some());
p.write_bytes(&mut self.dat);
Ok(())
}
pub fn push_back_value(&mut self, p: &ProtoValue) -> Result<(), Error> {
if let Some(codec) = self.reg.get_by_code(p.code()) {
codec.write_bytes(p, &mut self.dat)
} else {
Err(Error::unregistered(p.code()))
}
}
pub fn push_front<'a, P: Protocol<'a>>(&mut self, p: P) -> Result<(), Error> {
if self.reg.get_by_code(P::CODE).is_none() {
return Err(Error::unregistered(P::CODE));
}
debug_assert!(self.reg.get_by_prefix(P::PREFIX).is_some());
let mut dat = TinyVec::new();
p.write_bytes(&mut dat);
dat.extend_from_slice(&self.dat); self.dat = dat;
Ok(())
}
pub fn push_front_value(&mut self, p: &ProtoValue) -> Result<(), Error> {
if let Some(codec) = self.reg.get_by_code(p.code()) {
let mut dat = TinyVec::new();
codec.write_bytes(p, &mut dat)?;
dat.extend_from_slice(&self.dat); self.dat = dat;
Ok(())
} else {
Err(Error::unregistered(p.code()))
}
}
pub fn pop_back<'b>(&mut self) -> Option<ProtoValue<'b>> {
let iter = ValidBytesIter(iter::BytesIter::with_registry(
&self.dat[self.off..],
self.reg.clone(),
));
if let Some((o, c, Checked(p))) = iter.last() {
debug_assert!(self.dat.ends_with(p));
let dlen = self.len();
let plen = p.len();
let val = split_off(&mut self.dat, self.off + dlen - plen);
self.dat.truncate(self.off + o);
Some(ProtoValue {
code: c,
data: Bytes::Owned(Checked(val)),
})
} else {
None
}
}
pub fn pop_front(&mut self) -> Option<ProtoValue> {
let mut iter = ValidBytesIter(iter::BytesIter::with_registry(
&self.dat[self.off..],
self.reg.clone(),
));
if let Some((_, c, Checked(p))) = iter.next() {
self.off += iter.0.offset();
let val = &self.dat[self.off - p.len()..self.off];
debug_assert_eq!(val, p);
Some(ProtoValue {
code: c,
data: Bytes::Slice(Checked(val)),
})
} else {
None
}
}
pub fn drop_first(&mut self) {
let mut iter = ValidBytesIter(iter::BytesIter::with_registry(
self.as_ref(),
self.reg.clone(),
));
if iter.next().is_some() {
self.off += iter.0.offset()
}
}
pub fn drop_last(&mut self) {
let iter = ValidBytesIter(iter::BytesIter::with_registry(
self.as_ref(),
self.reg.clone(),
));
if let Some((o, _, _)) = iter.last() {
self.dat.truncate(self.off + o)
}
}
pub fn first(&self) -> Option<ProtoValue> {
self.iter().next()
}
pub fn last(&self) -> Option<ProtoValue> {
self.iter().last()
}
pub fn iter(&self) -> ProtoIter {
ProtoIter(ValidBytesIter(iter::BytesIter::with_registry(
self.as_ref(),
self.reg.clone(),
)))
}
pub fn shrink_to_fit(&mut self) {
self.dat.shrink_to_fit()
}
pub fn try_extend<'a, T>(&mut self, iter: T) -> Result<(), Error>
where
T: IntoIterator<Item = ProtoValue<'a>>,
{
for p in iter.into_iter() {
self.push_back_value(&p)?
}
Ok(())
}
pub fn try_with<'a, T>(mut self, iter: T) -> Result<Self, Error>
where
T: IntoIterator<Item = ProtoValue<'a>>,
{
self.try_extend(iter)?;
Ok(self)
}
pub fn starts_with(&self, code: Code) -> bool {
self.matches(0, &[code.into()])
}
pub fn matches<'a, I>(&self, start: usize, codes: I) -> bool
where
I: IntoIterator<Item = &'a Match>,
I::IntoIter: ExactSizeIterator,
{
let codes = codes.into_iter();
let mut n = codes.len();
for (p, c) in self.iter().skip(start).zip(codes) {
n -= 1;
match c {
Match::Val(c) => {
if p.code() != *c {
return false;
}
}
Match::Any(cs) => {
if !cs.contains(&p.code()) {
return false;
}
}
}
}
n == 0
}
pub fn split(&self, at: usize) -> (MultiAddr, MultiAddr) {
let mut iter = self.iter();
let a = MultiAddr::default()
.try_with((&mut iter).take(at))
.expect("valid address");
let b = MultiAddr::default().try_with(iter).expect("valid address");
(a, b)
}
pub fn concat_mut(&mut self, other: &MultiAddr) -> Result<(), Error> {
for proto in other.iter() {
self.push_back_value(&proto)?;
}
Ok(())
}
pub fn concat(self, other: &MultiAddr) -> Result<MultiAddr, Error> {
let mut addr = self;
addr.concat_mut(other)?;
Ok(addr)
}
pub fn to_socket_addr(&self) -> Result<String, Error> {
let mut it = self.iter().peekable();
while let Some(p) = it.next() {
match p.code() {
Ip4::CODE => {
let ip4 = p.cast::<Ip4>().unwrap();
let port = it.next().unwrap().cast::<Tcp>().unwrap();
return Ok(SocketAddrV4::new(*ip4, *port).to_string());
}
Ip6::CODE => {
let ip6 = p.cast::<Ip6>().unwrap();
let port = it.next().unwrap().cast::<Tcp>().unwrap();
return Ok(SocketAddrV6::new(*ip6, *port, 0, 0).to_string());
}
DnsAddr::CODE => {
let host = p.cast::<DnsAddr>().unwrap();
if let Some(p) = it.peek() {
if p.code() == Tcp::CODE {
let port = p.cast::<Tcp>().unwrap();
return Ok(format!("{}:{}", &*host, *port));
}
}
}
other => {
return Err(Error::invalid_proto(other));
}
}
}
Err(Error::message("No socket address found"))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Match {
Val(Code),
Any(TinyVec<[Code; 4]>),
}
impl Match {
pub fn code(c: Code) -> Self {
Self::Val(c)
}
pub fn any<I: IntoIterator<Item = Code>>(cs: I) -> Self {
Self::Any(cs.into_iter().collect())
}
}
impl From<Code> for Match {
fn from(c: Code) -> Self {
Match::code(c)
}
}
impl fmt::Display for MultiAddr {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
for proto in self.iter() {
let codec = self.reg.get_by_code(proto.code()).expect("valid code");
if let Err(e) = codec.transcode_bytes(proto.code(), proto.data(), f) {
if let error::ErrorImpl::Format(e) = e.into_impl() {
return Err(e);
}
}
}
Ok(())
}
}
impl<'a> IntoIterator for &'a MultiAddr {
type Item = ProtoValue<'a>;
type IntoIter = ProtoIter<'a>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl TryFrom<&str> for MultiAddr {
type Error = Error;
fn try_from(value: &str) -> Result<Self, Self::Error> {
MultiAddr::try_from_str(value, default_registry().clone())
}
}
impl TryFrom<&[u8]> for MultiAddr {
type Error = Error;
fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
MultiAddr::try_from_bytes(value, default_registry().clone())
}
}
impl FromStr for MultiAddr {
type Err = Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
s.try_into()
}
}
impl AsRef<[u8]> for MultiAddr {
fn as_ref(&self) -> &[u8] {
&self.dat[self.off..]
}
}
#[cfg(feature = "serde")]
impl serde::Serialize for MultiAddr {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
if s.is_human_readable() {
s.serialize_str(&self.to_string())
} else {
s.serialize_bytes(self.as_ref())
}
}
}
#[cfg(feature = "serde")]
impl<'de> serde::Deserialize<'de> for MultiAddr {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
if d.is_human_readable() {
let s = <&'de str>::deserialize(d)?;
MultiAddr::try_from(s).map_err(serde::de::Error::custom)
} else {
let b = <&'de [u8]>::deserialize(d)?;
MultiAddr::try_from(b).map_err(serde::de::Error::custom)
}
}
}
#[cfg(feature = "cbor")]
impl<C> minicbor::Encode<C> for MultiAddr {
fn encode<W>(
&self,
e: &mut minicbor::Encoder<W>,
_: &mut C,
) -> Result<(), minicbor::encode::Error<W::Error>>
where
W: minicbor::encode::Write,
{
e.bytes(self.as_ref())?.ok()
}
}
#[cfg(feature = "cbor")]
impl<'b, C> minicbor::Decode<'b, C> for MultiAddr {
fn decode(d: &mut minicbor::Decoder<'b>, _: &mut C) -> Result<Self, minicbor::decode::Error> {
MultiAddr::try_from(d.bytes()?)
.map_err(|e| minicbor::decode::Error::message(format!("{e}")))
}
}
#[derive(Debug)]
pub struct ProtoIter<'a>(ValidBytesIter<'a>);
impl<'a> Iterator for ProtoIter<'a> {
type Item = ProtoValue<'a>;
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|(_, c, p)| ProtoValue {
code: c,
data: Bytes::Slice(p),
})
}
}
#[derive(Debug)]
struct ValidBytesIter<'a>(iter::BytesIter<'a>);
impl<'a> Iterator for ValidBytesIter<'a> {
type Item = (usize, Code, Checked<&'a [u8]>);
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|x| x.expect("valid protocol"))
}
}
fn split_off<A>(v: &mut TinyVec<A>, at: usize) -> TinyVec<A>
where
A: Array<Item = u8>,
{
match v {
TinyVec::Inline(a) => TinyVec::Inline(a.split_off(at)),
TinyVec::Heap(v) => {
if v.len() - at <= A::CAPACITY {
let mut a = ArrayVec::default();
a.extend_from_slice(&v[at..]);
v.truncate(at);
TinyVec::Inline(a)
} else {
TinyVec::Heap(v.split_off(at))
}
}
}
}
#[cfg(test)]
mod tests {
use tinyvec::TinyVec;
#[test]
fn split_off() {
let mut t: TinyVec<[u8; 5]> = TinyVec::new();
t.extend_from_slice(b"hello");
assert!(t.is_inline());
t.extend_from_slice(b"world");
assert!(t.is_heap());
let mut v = t.clone();
let a = v.split_off(5);
assert!(a.is_heap());
let b = super::split_off(&mut t, 5);
assert!(b.is_inline());
assert_eq!(a, b);
assert_eq!(v, t);
}
}