koibumi-core 0.0.9

use std::{
    collections::HashMap,
    convert::{TryFrom, TryInto},
    fmt,
    io::{self, Read, Write},
};

use crate::io::{ReadFrom, WriteTo};

/// The type of a Bitmessage message.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct Command([u8; 12]);

impl Command {
    /// The "error" command.
    pub const ERROR: Command = Command::new(*b"error\0\0\0\0\0\0\0");

    /// The "getdata" command.
    pub const GETDATA: Command = Command::new(*b"getdata\0\0\0\0\0");

    /// The "inv" command.
    pub const INV: Command = Command::new(*b"inv\0\0\0\0\0\0\0\0\0");

    /// The "dinv" command.
    pub const DINV: Command = Command::new(*b"dinv\0\0\0\0\0\0\0\0");

    /// The "object" command.
    pub const OBJECT: Command = Command::new(*b"object\0\0\0\0\0\0");

    /// The "addr" command.
    pub const ADDR: Command = Command::new(*b"addr\0\0\0\0\0\0\0\0");

    /// The "portcheck" command.
    pub const PORTCHECK: Command = Command::new(*b"portcheck\0\0\0");

    /// The "ping" command.
    pub const PING: Command = Command::new(*b"ping\0\0\0\0\0\0\0\0");

    /// The "pong" command.
    pub const PONG: Command = Command::new(*b"pong\0\0\0\0\0\0\0\0");

    /// The "verack" command.
    pub const VERACK: Command = Command::new(*b"verack\0\0\0\0\0\0");

    /// The "version" command.
    pub const VERSION: Command = Command::new(*b"version\0\0\0\0\0");

    /// Constructs a command from a byte array.
    pub const fn new(bytes: [u8; 12]) -> Self {
        Self(bytes)
    }
}

impl AsRef<[u8]> for Command {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

/// An error which can be returned when parsing a message type.
///
/// This error is used as the error type for
/// the [`Command::new()`] method and
/// the `TryFrom` implementation for [`Command`].
///
/// [`Command::new()`]: struct.Command.html#method.new
/// [`Command`]: struct.Command.html
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum ParseCommandError {
    /// The input was too long to construct a message type.
    /// The maximum length allowed and the actual length of the input
    /// are returned as payloads of this variant.
    TooLong {
        /// The maximum length allowed.
        max: usize,
        /// The actual length of the input.
        len: usize,
    },
}

impl fmt::Display for ParseCommandError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ParseCommandError::TooLong { max, len } => {
                write!(f, "length must be <={}, but {}", max, len)
            }
        }
    }
}

impl std::error::Error for ParseCommandError {}

impl TryFrom<&[u8]> for Command {
    type Error = ParseCommandError;

    fn try_from(bytes: &[u8]) -> Result<Self, <Self as TryFrom<&[u8]>>::Error> {
        let mut b: [u8; 12] = [0; 12];
        if bytes.len() <= 12 {
            b[0..bytes.len()].copy_from_slice(bytes);
            Ok(Self(b))
        } else {
            Err(Self::Error::TooLong {
                max: 12,
                len: bytes.len(),
            })
        }
    }
}

impl WriteTo for Command {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        w.write_all(&self.0)
    }
}

impl ReadFrom for Command {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let mut bytes: [u8; 12] = [0; 12];
        r.read_exact(&mut bytes)?;
        Ok(Self(bytes))
    }
}

#[test]
fn test_command_write_to() {
    let test: Command = b"hello".as_ref().try_into().unwrap();
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [b'h', b'e', b'l', b'l', b'o', 0, 0, 0, 0, 0, 0, 0];
    assert_eq!(bytes, expected);
}

#[test]
fn test_command_read_from() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([b'h', b'e', b'l', b'l', b'o', 0, 0, 0, 0, 0, 0, 0]);
    let test = Command::read_from(&mut bytes).unwrap();
    let expected: Command = b"hello".as_ref().try_into().unwrap();
    assert_eq!(test, expected);
}

/// Known types of Bitmessage messages.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum CommandKind {
    /// The "error" command.
    Error,

    /// The "getdata" command.
    Getdata,

    /// The "inv" command.
    Inv,

    /// The "dinv" command.
    Dinv,

    /// The "object" command.
    Object,

    /// The "addr" command.
    Addr,

    /// The "portcheck" command.
    Portcheck,

    /// The "ping" command.
    Ping,

    /// The "pong" command.
    Pong,

    /// The "verack" command.
    Verack,

    /// The "version" command.
    Version,
}

lazy_static! {
    static ref STRING_TO_COMMAND_MAP: HashMap<Command, CommandKind> = {
        let mut m = HashMap::new();
        m.insert(
            b"version".as_ref().try_into().unwrap(),
            CommandKind::Version,
        );
        m.insert(b"verack".as_ref().try_into().unwrap(), CommandKind::Verack);
        m.insert(b"addr".as_ref().try_into().unwrap(), CommandKind::Addr);
        m.insert(b"inv".as_ref().try_into().unwrap(), CommandKind::Inv);
        m.insert(
            b"getdata".as_ref().try_into().unwrap(),
            CommandKind::Getdata,
        );
        m.insert(b"object".as_ref().try_into().unwrap(), CommandKind::Object);
        m.insert(b"error".as_ref().try_into().unwrap(), CommandKind::Error);
        m.insert(b"ping".as_ref().try_into().unwrap(), CommandKind::Ping);
        m.insert(b"pong".as_ref().try_into().unwrap(), CommandKind::Pong);
        m
    };
}

/// The error type returned when a conversion from a Bitmessage command type
/// to a known command type fails.
///
/// This error is used as the error type for the `TryFrom` implementation
/// for `CommandKind`.
///
/// The source command type is returned as a payload of this error type.
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct TryFromCommandError(Command);

impl fmt::Display for TryFromCommandError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "unknown command {}",
            String::from_utf8_lossy((self.0).0.as_ref())
        )
    }
}

impl std::error::Error for TryFromCommandError {}

impl TryFrom<Command> for CommandKind {
    type Error = TryFromCommandError;

    fn try_from(s: Command) -> Result<Self, <Self as TryFrom<Command>>::Error> {
        match STRING_TO_COMMAND_MAP.get(&s) {
            Some(c) => Ok(*c),
            None => Err(TryFromCommandError(s)),
        }
    }
}

impl From<CommandKind> for Command {
    fn from(c: CommandKind) -> Self {
        match c {
            CommandKind::Error => Command::ERROR,
            CommandKind::Getdata => Command::GETDATA,
            CommandKind::Inv => Command::INV,
            CommandKind::Dinv => Command::DINV,
            CommandKind::Object => Command::OBJECT,
            CommandKind::Addr => Command::ADDR,
            CommandKind::Portcheck => Command::PORTCHECK,
            CommandKind::Ping => Command::PING,
            CommandKind::Pong => Command::PONG,
            CommandKind::Verack => Command::VERACK,
            CommandKind::Version => Command::VERSION,
        }
    }
}

impl From<&CommandKind> for Command {
    fn from(c: &CommandKind) -> Self {
        Command::from(*c)
    }
}

impl WriteTo for CommandKind {
    fn write_to(&self, w: &mut dyn Write) -> io::Result<()> {
        let s: Command = self.clone().into();
        s.write_to(w)
    }
}

impl ReadFrom for CommandKind {
    fn read_from(r: &mut dyn Read) -> io::Result<Self>
    where
        Self: Sized,
    {
        let s = Command::read_from(r)?;
        match Self::try_from(s) {
            Ok(c) => Ok(c),
            Err(str) => Err(io::Error::new(io::ErrorKind::Other, str)),
        }
    }
}

#[test]
fn test_command_kind_write_to() {
    let test = CommandKind::Version;
    let mut bytes = Vec::new();
    test.write_to(&mut bytes).unwrap();
    let expected = [b'v', b'e', b'r', b's', b'i', b'o', b'n', 0, 0, 0, 0, 0];
    assert_eq!(bytes, expected);
}

#[test]
fn test_command_kind_read_from() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([b'v', b'e', b'r', b's', b'i', b'o', b'n', 0, 0, 0, 0, 0]);
    let test = CommandKind::read_from(&mut bytes).unwrap();
    let expected = CommandKind::Version;
    assert_eq!(test, expected);
}

#[test]
#[should_panic]
fn test_command_kind_read_from_fail() {
    use std::io::Cursor;

    let mut bytes = Cursor::new([b'h', b'e', b'l', b'l', b'o', 0, 0, 0, 0, 0, 0, 0]);
    let test = CommandKind::read_from(&mut bytes).unwrap();
    let expected = CommandKind::Version;
    assert_eq!(test, expected);
}