rust-libp2p/misc/multiaddr/src/protocol.rs

use bs58;
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use multihash::Multihash;
use std::{
    borrow::Cow,
    convert::From,
    fmt,
    io::{Cursor, Write},
    net::{Ipv4Addr, Ipv6Addr},
    str::{self, FromStr}
};
use unsigned_varint::{encode, decode};
use {Result, Error};

const DCCP: u32 = 33;
const DNS4: u32 = 54;
const DNS6: u32 = 55;
const HTTP: u32 = 480;
const HTTPS: u32 = 443;
const IP4: u32 = 4;
const IP6: u32 = 41;
const P2P_WEBRTC_DIRECT: u32 = 276;
const P2P_WEBRTC_STAR: u32 = 275;
const P2P_WEBSOCKET_STAR: u32 = 479;
const MEMORY: u32 = 777;
const ONION: u32 = 444;
const P2P: u32 = 421;
const P2P_CIRCUIT: u32 = 290;
const QUIC: u32 = 460;
const SCTP: u32 = 132;
const TCP: u32 = 6;
const UDP: u32 = 17;
const UDT: u32 = 301;
const UNIX: u32 = 400;
const UTP: u32 = 302;
const WS: u32 = 477;
const WSS: u32 = 478;

/// `Protocol` describes all possible multiaddress protocols.
#[derive(PartialEq, Eq, Clone, Debug)]
pub enum Protocol<'a> {
    Dccp(u16),
    Dns4(Cow<'a, str>),
    Dns6(Cow<'a, str>),
    Http,
    Https,
    Ip4(Ipv4Addr),
    Ip6(Ipv6Addr),
    P2pWebRtcDirect,
    P2pWebRtcStar,
    P2pWebSocketStar,
    Memory,
    Onion(Cow<'a, [u8]>),
    P2p(Multihash),
    P2pCircuit,
    Quic,
    Sctp(u16),
    Tcp(u16),
    Udp(u16),
    Udt,
    Unix(Cow<'a, str>),
    Utp,
    Ws,
    Wss
}

impl<'a> Protocol<'a> {
    /// Parse a protocol value from the given iterator of string slices.
    ///
    /// The parsing only consumes the minimum amount of string slices necessary to
    /// produce a well-formed protocol. The same iterator can thus be used to parse
    /// a sequence of protocols in succession. It is up to client code to check
    /// that iteration has finished whenever appropriate.
    pub fn from_str_parts<I>(mut iter: I) -> Result<Self>
    where
        I: Iterator<Item=&'a str>
    {
        match iter.next().ok_or(Error::InvalidProtocolString)? {
            "ip4" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Ip4(Ipv4Addr::from_str(s)?))
            }
            "tcp" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Tcp(s.parse()?))
            }
            "udp" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Udp(s.parse()?))
            }
            "dccp" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Dccp(s.parse()?))
            }
            "ip6" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Ip6(Ipv6Addr::from_str(s)?))
            }
            "dns4" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Dns4(Cow::Borrowed(s)))
            }
            "dns6" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Dns6(Cow::Borrowed(s)))
            }
            "sctp" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Sctp(s.parse()?))
            }
            "udt" => Ok(Protocol::Udt),
            "utp" => Ok(Protocol::Utp),
            "unix" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                Ok(Protocol::Unix(Cow::Borrowed(s)))
            }
            "p2p" => {
                let s = iter.next().ok_or(Error::InvalidProtocolString)?;
                let decoded = bs58::decode(s).into_vec()?;
                Ok(Protocol::P2p(Multihash::from_bytes(decoded)?))
            }
            "http" => Ok(Protocol::Http),
            "https" => Ok(Protocol::Https),
            "onion" => unimplemented!(), // FIXME
            "quic" => Ok(Protocol::Quic),
            "ws" => Ok(Protocol::Ws),
            "wss" => Ok(Protocol::Wss),
            "p2p-websocket-star" => Ok(Protocol::P2pWebSocketStar),
            "p2p-webrtc-star" => Ok(Protocol::P2pWebRtcStar),
            "p2p-webrtc-direct" => Ok(Protocol::P2pWebRtcDirect),
            "p2p-circuit" => Ok(Protocol::P2pCircuit),
            "memory" => Ok(Protocol::Memory),
            _ => Err(Error::UnknownProtocolString)
        }
    }

    /// Parse a single `Protocol` value from its byte slice representation,
    /// returning the protocol as well as the remaining byte slice.
    pub fn from_bytes(input: &'a [u8]) -> Result<(Self, &'a [u8])> {
        fn split_at(n: usize, input: &[u8]) -> Result<(&[u8], &[u8])> {
            if input.len() < n {
                return Err(Error::DataLessThanLen)
            }
            Ok(input.split_at(n))
        }
        let (id, input) = decode::u32(input)?;
        match id {
            DCCP => {
                let (data, rest) = split_at(2, input)?;
                let mut rdr = Cursor::new(data);
                let num = rdr.read_u16::<BigEndian>()?;
                Ok((Protocol::Dccp(num), rest))
            }
            DNS4 => {
                let (n, input) = decode::usize(input)?;
                let (data, rest) = split_at(n, input)?;
                Ok((Protocol::Dns4(Cow::Borrowed(str::from_utf8(data)?)), rest))
            }
            DNS6 => {
                let (n, input) = decode::usize(input)?;
                let (data, rest) = split_at(n, input)?;
                Ok((Protocol::Dns6(Cow::Borrowed(str::from_utf8(data)?)), rest))
            }
            HTTP => Ok((Protocol::Http, input)),
            HTTPS => Ok((Protocol::Https, input)),
            IP4 => {
                let (data, rest) = split_at(4, input)?;
                Ok((Protocol::Ip4(Ipv4Addr::new(data[0], data[1], data[2], data[3])), rest))
            }
            IP6 => {
                let (data, rest) = split_at(16, input)?;
                let mut rdr = Cursor::new(data);
                let mut seg = [0_u16; 8];

                for x in seg.iter_mut() {
                    *x = rdr.read_u16::<BigEndian>()?;
                }

                let addr = Ipv6Addr::new(seg[0],
                                         seg[1],
                                         seg[2],
                                         seg[3],
                                         seg[4],
                                         seg[5],
                                         seg[6],
                                         seg[7]);

                Ok((Protocol::Ip6(addr), rest))
            }
            P2P_WEBRTC_DIRECT => Ok((Protocol::P2pWebRtcDirect, input)),
            P2P_WEBRTC_STAR => Ok((Protocol::P2pWebRtcStar, input)),
            P2P_WEBSOCKET_STAR => Ok((Protocol::P2pWebSocketStar, input)),
            MEMORY => Ok((Protocol::Memory, input)),
            ONION => unimplemented!(), // FIXME
            P2P => {
                let (n, input) = decode::usize(input)?;
                let (data, rest) = split_at(n, input)?;
                Ok((Protocol::P2p(Multihash::from_bytes(data.to_owned())?), rest))
            }
            P2P_CIRCUIT => Ok((Protocol::P2pCircuit, input)),
            QUIC => Ok((Protocol::Quic, input)),
            SCTP => {
                let (data, rest) = split_at(2, input)?;
                let mut rdr = Cursor::new(data);
                let num = rdr.read_u16::<BigEndian>()?;
                Ok((Protocol::Sctp(num), rest))
            }
            TCP => {
                let (data, rest) = split_at(2, input)?;
                let mut rdr = Cursor::new(data);
                let num = rdr.read_u16::<BigEndian>()?;
                Ok((Protocol::Tcp(num), rest))
            }
            UDP => {
                let (data, rest) = split_at(2, input)?;
                let mut rdr = Cursor::new(data);
                let num = rdr.read_u16::<BigEndian>()?;
                Ok((Protocol::Udp(num), rest))
            }
            UDT => Ok((Protocol::Udt, input)),
            UNIX => {
                let (n, input) = decode::usize(input)?;
                let (data, rest) = split_at(n, input)?;
                Ok((Protocol::Unix(Cow::Borrowed(str::from_utf8(data)?)), rest))
            }
            UTP => Ok((Protocol::Utp, input)),
            WS => Ok((Protocol::Ws, input)),
            WSS => Ok((Protocol::Wss, input)),
            _ => Err(Error::UnknownProtocolId(id))
        }
    }

    /// Encode this protocol by writing its binary representation into
    /// the given `Write` impl.
    pub fn write_bytes<W: Write>(&self, w: &mut W) -> Result<()> {
        let mut buf = encode::u32_buffer();
        match self {
            Protocol::Ip4(addr) => {
                w.write_all(encode::u32(IP4, &mut buf))?;
                w.write_all(&addr.octets())?
            }
            Protocol::Ip6(addr) => {
                w.write_all(encode::u32(IP6, &mut buf))?;
                for &segment in &addr.segments() {
                    w.write_u16::<BigEndian>(segment)?
                }
            }
            Protocol::Tcp(port) => {
                w.write_all(encode::u32(TCP, &mut buf))?;
                w.write_u16::<BigEndian>(*port)?
            }
            Protocol::Udp(port) => {
                w.write_all(encode::u32(UDP, &mut buf))?;
                w.write_u16::<BigEndian>(*port)?
            }
            Protocol::Dccp(port) => {
                w.write_all(encode::u32(DCCP, &mut buf))?;
                w.write_u16::<BigEndian>(*port)?
            }
            Protocol::Sctp(port) => {
                w.write_all(encode::u32(SCTP, &mut buf))?;
                w.write_u16::<BigEndian>(*port)?
            }
            Protocol::Dns4(s) => {
                w.write_all(encode::u32(DNS4, &mut buf))?;
                let bytes = s.as_bytes();
                w.write_all(encode::usize(bytes.len(), &mut encode::usize_buffer()))?;
                w.write_all(&bytes)?
            }
            Protocol::Dns6(s) => {
                w.write_all(encode::u32(DNS6, &mut buf))?;
                let bytes = s.as_bytes();
                w.write_all(encode::usize(bytes.len(), &mut encode::usize_buffer()))?;
                w.write_all(&bytes)?
            }
            Protocol::Unix(s) => {
                w.write_all(encode::u32(UNIX, &mut buf))?;
                let bytes = s.as_bytes();
                w.write_all(encode::usize(bytes.len(), &mut encode::usize_buffer()))?;
                w.write_all(&bytes)?
            }
            Protocol::P2p(multihash) => {
                w.write_all(encode::u32(P2P, &mut buf))?;
                let bytes = multihash.as_bytes();
                w.write_all(encode::usize(bytes.len(), &mut encode::usize_buffer()))?;
                w.write_all(&bytes)?
            }
            Protocol::Onion(_) => unimplemented!(), // FIXME
            Protocol::Quic => w.write_all(encode::u32(QUIC, &mut buf))?,
            Protocol::Utp => w.write_all(encode::u32(UTP, &mut buf))?,
            Protocol::Udt => w.write_all(encode::u32(UDT, &mut buf))?,
            Protocol::Http => w.write_all(encode::u32(HTTP, &mut buf))?,
            Protocol::Https => w.write_all(encode::u32(HTTPS, &mut buf))?,
            Protocol::Ws => w.write_all(encode::u32(WS, &mut buf))?,
            Protocol::Wss => w.write_all(encode::u32(WSS, &mut buf))?,
            Protocol::P2pWebSocketStar => w.write_all(encode::u32(P2P_WEBSOCKET_STAR, &mut buf))?,
            Protocol::P2pWebRtcStar => w.write_all(encode::u32(P2P_WEBRTC_STAR, &mut buf))?,
            Protocol::P2pWebRtcDirect => w.write_all(encode::u32(P2P_WEBRTC_DIRECT, &mut buf))?,
            Protocol::P2pCircuit => w.write_all(encode::u32(P2P_CIRCUIT, &mut buf))?,
            Protocol::Memory => w.write_all(encode::u32(MEMORY, &mut buf))?
        }
        Ok(())
    }

    /// Turn this `Protocol` into one that owns its data, thus being valid for any lifetime.
    pub fn acquire<'b>(self) -> Protocol<'b> {
        use Protocol::*;
        match self {
            Dccp(a) => Dccp(a),
            Dns4(cow) => Dns4(Cow::Owned(cow.into_owned())),
            Dns6(cow) => Dns6(Cow::Owned(cow.into_owned())),
            Http => Http,
            Https => Https,
            Ip4(a) => Ip4(a),
            Ip6(a) => Ip6(a),
            P2pWebRtcDirect => P2pWebRtcDirect,
            P2pWebRtcStar => P2pWebRtcStar,
            P2pWebSocketStar => P2pWebSocketStar,
            Memory => Memory,
            Onion(cow) => Onion(Cow::Owned(cow.into_owned())),
            P2p(a) => P2p(a),
            P2pCircuit => P2pCircuit,
            Quic => Quic,
            Sctp(a) => Sctp(a),
            Tcp(a) => Tcp(a),
            Udp(a) => Udp(a),
            Udt => Udt,
            Unix(cow) => Unix(Cow::Owned(cow.into_owned())),
            Utp => Utp,
            Ws => Ws,
            Wss => Wss
        }
    }
}

impl<'a> fmt::Display for Protocol<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use Protocol::*;
        match self {
            Dccp(port) => write!(f, "/dccp/{}", port),
            Dns4(s) => write!(f, "/dns4/{}", s),
            Dns6(s) => write!(f, "/dns6/{}", s),
            Http => f.write_str("/http"),
            Https => f.write_str("/https"),
            Ip4(addr) => write!(f, "/ip4/{}", addr),
            Ip6(addr) => write!(f, "/ip6/{}", addr),
            P2pWebRtcDirect => f.write_str("/p2p-webrtc-direct"),
            P2pWebRtcStar => f.write_str("/p2p-webrtc-star"),
            P2pWebSocketStar => f.write_str("/p2p-websocket-star"),
            Memory => f.write_str("/memory"),
            Onion(_) => unimplemented!(), // FIXME!
            P2p(c) => write!(f, "/p2p/{}", bs58::encode(c.as_bytes()).into_string()),
            P2pCircuit => f.write_str("/p2p-circuit"),
            Quic => f.write_str("/quic"),
            Sctp(port) => write!(f, "/sctp/{}", port),
            Tcp(port) => write!(f, "/tcp/{}", port),
            Udp(port) => write!(f, "/udp/{}", port),
            Udt => f.write_str("/udt"),
            Unix(s) => write!(f, "/unix/{}", s),
            Utp => f.write_str("/utp"),
            Ws => f.write_str("/ws"),
            Wss => f.write_str("/wss"),
        }
    }
}

impl<'a> From<Ipv4Addr> for Protocol<'a> {
    #[inline]
    fn from(addr: Ipv4Addr) -> Self {
        Protocol::Ip4(addr)
    }
}

impl<'a> From<Ipv6Addr> for Protocol<'a> {
    #[inline]
    fn from(addr: Ipv6Addr) -> Self {
        Protocol::Ip6(addr)
    }
}