lavina/src/protos/irc/mod.rs

//! Client-to-Server IRC protocol.

use std::io::Write;

use nom::{
    branch::alt,
    bytes::complete::{tag, take, take_while},
    IResult,
};

type ByteVec = Vec<u8>;

/// Server-to-client message.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ServerMessage {
    /// Optional tags section, prefixed with `@`
    tags: Vec<Tag>,
    /// Optional server name, prefixed with `:`.
    sender: Option<String>,
    body: ServerMessageBody,
}
impl ServerMessage {
    pub fn write(&self, writer: &mut impl Write) -> std::io::Result<()> {
        self.body.write(writer)?;
        writer.write(b"\n")?;
        Ok(())
    }
}

pub fn server_message(input: &[u8]) -> IResult<&[u8], ServerMessage> {
    let (input, command) = server_message_body(input)?;
    let (input, _) = tag(b"\n")(input)?;

    let message = ServerMessage {
        tags: vec![],
        sender: None,
        body: command,
    };
    Ok((input, message))
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ServerMessageBody {
    Notice {
        first_target: ByteVec,
        rest_targets: Vec<ByteVec>,
        text: ByteVec,
    },
    Ping {
        token: ByteVec,
    },
    Pong {
        from: ByteVec,
        token: ByteVec,
    },
}
impl ServerMessageBody {
    pub fn write(&self, writer: &mut impl Write) -> std::io::Result<()> {
        match self {
            ServerMessageBody::Notice {
                first_target,
                rest_targets,
                text,
            } => {
                writer.write(b"NOTICE ")?;
                writer.write(&first_target)?;
                writer.write(b" :")?;
                writer.write(&text)?;
            }
            ServerMessageBody::Ping { token } => {
                writer.write(b"PING ")?;
                writer.write(&token)?;
            }
            ServerMessageBody::Pong { from, token } => {
                writer.write(b"PONG ")?;
                writer.write(&from)?;
                writer.write(b" :")?;
                writer.write(&token)?;
            }
        }
        Ok(())
    }
}

fn server_message_body(input: &[u8]) -> IResult<&[u8], ServerMessageBody> {
    alt((
        sserver_message_body_notice,
        server_message_body_ping,
        server_message_body_pong,
    ))(input)
}

fn sserver_message_body_notice(input: &[u8]) -> IResult<&[u8], ServerMessageBody> {
    let (input, _) = tag("NOTICE ")(input)?;
    let (input, first_target) = receiver(input)?;
    let (input, _) = tag(" :")(input)?;
    let (input, text) = token(input)?;

    let first_target = first_target.to_owned();
    let text = text.to_owned();
    Ok((
        input,
        ServerMessageBody::Notice {
            first_target,
            rest_targets: vec![],
            text,
        },
    ))
}

fn server_message_body_ping(input: &[u8]) -> IResult<&[u8], ServerMessageBody> {
    let (input, _) = tag("PING ")(input)?;
    let (input, token) = token(input)?;

    Ok((
        input,
        ServerMessageBody::Ping {
            token: token.to_owned(),
        },
    ))
}

fn server_message_body_pong(input: &[u8]) -> IResult<&[u8], ServerMessageBody> {
    let (input, _) = tag("PONG ")(input)?;
    let (input, from) = receiver(input)?;
    let (input, _) = tag(" :")(input)?;
    let (input, token) = token(input)?;

    Ok((
        input,
        ServerMessageBody::Pong {
            from: from.to_owned(),
            token: token.to_owned(),
        },
    ))
}

/// Single message tag value.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Tag {
    key: ByteVec,
    value: Option<u8>,
}

/// Client-to-server command.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ClientMessage {
    /// CAP. Capability-related commands.
    Capability {
        subcommand: CapabilitySubcommand,
    },
    /// PING
    Ping {
        token: ByteVec,
    },
    Pong {
        from: ByteVec,
        token: ByteVec,
    },
}

fn receiver(input: &[u8]) -> IResult<&[u8], &[u8]> {
    take_while(|i| i != b'\n' && i != b'\r' && i != b' ')(input)
}

fn token(input: &[u8]) -> IResult<&[u8], &[u8]> {
    take_while(|i| i != b'\n' && i != b'\r')(input)
}

pub fn parse_client_message(input: &[u8]) -> IResult<&[u8], ClientMessage> {
    alt((command_capability, command_ping))(input)
}

fn command_capability(input: &[u8]) -> IResult<&[u8], ClientMessage> {
    let (input, _) = tag("CAP ")(input)?;
    let (input, subcommand) = capability_subcommand(input)?;

    Ok((input, ClientMessage::Capability { subcommand }))
}

fn command_ping(input: &[u8]) -> IResult<&[u8], ClientMessage> {
    let (input, _) = tag("PING ")(input)?;
    let (input, token) = token(input)?;

    Ok((
        input,
        ClientMessage::Ping {
            token: token.to_owned(),
        },
    ))
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CapabilitySubcommand {
    /// CAP LS {code}
    List { code: [u8; 3] },
    /// CAP END
    End,
}

fn capability_subcommand(input: &[u8]) -> IResult<&[u8], CapabilitySubcommand> {
    alt((capability_subcommand_ls, capability_subcommand_end))(input)
}

fn capability_subcommand_ls(input: &[u8]) -> IResult<&[u8], CapabilitySubcommand> {
    let (input, _) = tag("LS ")(input)?;
    let (input, code) = take(3usize)(input)?;

    Ok((
        input,
        CapabilitySubcommand::List {
            code: code.try_into().unwrap(),
        },
    ))
}

fn capability_subcommand_end(input: &[u8]) -> IResult<&[u8], CapabilitySubcommand> {
    let (input, _) = tag("END")(input)?;
    Ok((input, CapabilitySubcommand::End))
}

#[cfg(test)]
mod test {
    use assert_matches::*;

    use super::*;

    #[test]
    fn test_server_message_notice() {
        let input = b"NOTICE * :*** Looking up your hostname...\n";
        let expected = ServerMessage {
            tags: vec![],
            sender: None,
            body: ServerMessageBody::Notice {
                first_target: b"*".to_vec(),
                rest_targets: vec![],
                text: b"*** Looking up your hostname...".to_vec(),
            },
        };

        let result = server_message(input);
        assert_matches!(result, Ok((_, result)) => assert_eq!(expected, result));

        let mut bytes = vec![];
        expected.write(&mut bytes).unwrap();
        assert_eq!(bytes, input);
    }

    #[test]
    fn test_server_message_pong() {
        let input = b"PONG server.example :LAG004911\n";
        let expected = ServerMessage {
            tags: vec![],
            sender: None,
            body: ServerMessageBody::Pong {
                from: b"server.example".to_vec(),
                token: b"LAG004911".to_vec(),
            },
        };

        let result = server_message(input);
        assert_matches!(result, Ok((_, result)) => assert_eq!(expected, result));

        let mut bytes = vec![];
        expected.write(&mut bytes).unwrap();
        assert_eq!(bytes, input);
    }

    #[test]
    fn test_client_message_cap_ls() {
        let input = b"CAP LS 302";
        let expected = ClientMessage::Capability {
            subcommand: CapabilitySubcommand::List { code: *b"302" },
        };

        let result = parse_client_message(input);
        assert_matches!(result, Ok((_, result)) => assert_eq!(expected, result));
    }

    #[test]
    fn test_client_message_cap_end() {
        let input = b"CAP END";
        let expected = ClientMessage::Capability {
            subcommand: CapabilitySubcommand::End,
        };

        let result = parse_client_message(input);
        assert_matches!(result, Ok((_, result)) => assert_eq!(expected, result));
    }

    #[test]
    fn test_client_message_ping() {
        let input = b"PING 1337";
        let expected = ClientMessage::Ping {
            token: b"1337".to_vec(),
        };

        let result = parse_client_message(input);
        assert_matches!(result, Ok((_, result)) => assert_eq!(expected, result));
    }
}