romp 0.5.2

//! STOMP protocol push parser, parses command and headers, not the body.

/// readonly bytes pushed in from the network.
/// Parser treats u8 as char as US-ASCII, we accept utf-8 but all protocol attributes are US-ASCII.
/// Also reads HTTP GET to support upgrading the stream from HTTP to WebSockets

use log::*;

use crate::message::stomp_message::*;
use crate::message::stomp_message::StompCommand;


/// State of the Struct
#[derive(Debug, PartialEq)]
pub enum ParserState {
    /// finished reading the whole message
    Done(usize),
    /// read all the headers now comes the body (value is how much data was read)
    Message(usize),
    /// awaiting more input
    Again,
    /// Bombed parsing the command (first line)
    InvalidCommand,
    /// bombed parsing something else
    InvalidMessage,
    /// flupped a buffer
    BodyFlup,
    /// reached max headers
    HdrFlup,
}

/// state of the loop
#[derive(Debug, PartialEq)]
enum State {
    Start,
    Command,
    CommandLf,

    HdrStart,
    HdrName,
    HdrValue,

    AlmostDone,
    MsgRead,
}

/// A push parser, externally a buffer is being filled with the STOMP message and headers as it arrive
/// we process the data as it comes off the stream and save state between chunks of data arriving.
#[derive(Debug)]
pub struct StompParser {
    /// count of bytes read parsing the stream, can be > message len since optional whitespace is discarded.
    bytes_read: usize,
    state: State,
    pub message: StompMessage,
    cmd_start: usize,
    request_line_start: usize,
    hdr_name_start: usize,
    hdr_name_end: usize,
    hdr_value_start: usize,
    hdr_value_end: usize,
}

impl StompParser {

    pub fn new(session_id: usize) -> StompParser {
        let mut m = StompMessage::new(Ownership::Parser);
        m.session_id = Some(session_id);
        StompParser {
            bytes_read: 0,
            state: State::Start,
            message: m,
            cmd_start: 0,
            request_line_start: 0,
            hdr_name_start: 0,
            hdr_name_end: 0,
            hdr_value_start: 0,
            hdr_value_end: 0,
        }
    }

    /// take ownership of the parsed message
    pub fn take_message(&mut self, owner: Ownership) -> StompMessage {
        self.bytes_read = 0;
        self.state = State::Start;
        self.cmd_start = 0;
        self.request_line_start = 0;
        self.hdr_name_start = 0;
        self.hdr_name_end = 0;
        self.hdr_value_start = 0;
        self.hdr_value_end = 0;

        self.message.take(owner)
    }

    /// reset wiping the message
    pub fn reset(&mut self) {
        self.bytes_read = 0;
        self.state = State::Start;
        self.cmd_start = 0;
        self.request_line_start = 0;
        self.hdr_name_start = 0;
        self.hdr_name_end = 0;
        self.hdr_value_start = 0;
        self.hdr_value_end = 0;

        let session_id = self.message.id;
        self.message = StompMessage::new(Ownership::Parser);
        self.message.session_id = Some(session_id);
    }

    pub fn bytes_read(&self) -> usize {
        self.bytes_read
    }

    /// Required a complete rewrite from xtomp because rust does not have for loops or p++ and can not index arrays easily
    ///
    /// `buffer` the full buffer we are loading
    /// `pos` position in the outer buffer (not the chunk)
    /// `chunk` slice of input buffer we are reading
    pub fn push(&mut self, buffer: &[u8], mut pos: usize, chunk: &[u8]) -> Result<ParserState, ParserState> {

        // debug!("READ chunk '{}'", String::from_utf8_lossy(chunk));

        for c in chunk {
            self.bytes_read += 1;
            //println!("loop {}, {}", *c as char, String::from_utf8_lossy(buffer));

            let ch: u8 = *c;

            if ch == b'\0' && self.state != State::AlmostDone {
                debug!("early frame termination");
                return Err(ParserState::InvalidCommand);
            }

            match self.state {

                State::Start => {
                    if ch == b'\n' || ch == b'\r' {
                        // TODO if all we got was a heart-beat we should reset buffer pos here
                        // continue
                    } else if ch < b'A' || ch > b'Z' {
                        return Err(ParserState::InvalidCommand);
                    } else {
                        self.cmd_start = pos;
                        self.state = State::Command;
                        self.message.command = StompCommand::Unknown;
                    }
                }

                State::Command => {
// TODO only allow non A-Z for HTTP GET
//                    if ch == b'\n' || ch == b'\r' {
//
//                    }
//                    else if ch < b'A' || ch > b'Z' {
//                        return Err(ParserState::InvalidCommand);
//                    }

                    let c = self.cmd_start;

                    // amount of command statement read
                    let cmd_read = pos - c;
                    let com = &buffer[c..c + cmd_read];

                    //println!("Reading Command {} cmd_start={} cmd_read={}", ch, self.cmd_start, cmd_read);

                    if cmd_read == 0 {
                        // continue
                    }
                    else if cmd_read == 1 {
                        // continue
                    }
                    else if cmd_read == 2 {
                        // continue
                    }
                    else if cmd_read == 3 {  //  ACK, GET
                        if com.eq(b"ACK") {
                            self.message.command = StompCommand::Ack;
                        }
                    }
                    else if cmd_read == 4 {  //  NACK, SEND
                        if com.eq(b"NACK") {
                            self.message.command = StompCommand::Nack;
                        }
                        if com.eq(b"SEND") {
                            self.message.command = StompCommand::Send;
                        }
                        if com.eq(b"GET ") {
                            self.message.command = StompCommand::Get;
                            self.message.message_type = MessageType::Http;
                            self.request_line_start = pos - 4;
                        }
                    }
                    else if cmd_read == 5 {  //  BEGIN, ABORT, ERROR, STOMP
                        if com.eq(b"BEGIN") {
                            self.message.command = StompCommand::Begin;
                        }
                        if com.eq(b"ABORT") {
                            self.message.command = StompCommand::Abort;
                        }
                        if com.eq(b"ERROR") {
                            self.message.command = StompCommand::Error;
                        }
                        if com.eq(b"STOMP") {
                            self.message.command = StompCommand::Stomp;
                        }
                    }
                    else if cmd_read == 6 {  //  COMMIT
                        if com.eq(b"COMMIT") {
                            self.message.command = StompCommand::Commit;
                        }
                    }
                    else if cmd_read == 7 {  //  CONNECT, MESSAGE, RECEIPT
                        if com.eq(b"CONNECT") {
                            self.message.command = StompCommand::Connect;
                        }
                        if com.eq(b"MESSAGE") {
                            self.message.command = StompCommand::Message;
                        }
                        if com.eq(b"RECEIPT") {
                            self.message.command = StompCommand::Receipt;
                        }
                    }
                    else if cmd_read == 8 {  //  none
                        // continue
                    }
                    else if cmd_read == 9 {  //  SUBSCRIBE, CONNECTED
                        if com.eq(b"SUBSCRIBE") {
                            self.message.command = StompCommand::Subscribe;
                        }
                        if com.eq(b"CONNECTED") {
                            self.message.command = StompCommand::Connected;
                        }
                    }
                    else if cmd_read == 10 {  //  DISCONNECT
                        if com.eq(b"DISCONNECT") {
                            self.message.command = StompCommand::Disconnect;
                        }
                    }
                    else if cmd_read == 11 {  //  UNSUBSCRIBE
                        if com.eq(b"UNSUBSCRIBE") {
                            self.message.command = StompCommand::Unsubscribe;
                        }
                    }
                    else if cmd_read == 12 {
                        if StompCommand::Get == self.message.command {
                            // HTTP GET can be long
                        } else {
                            return Err(ParserState::InvalidCommand);
                        }
                    } // end command types

                    match self.is_command_done(ch) {
                        Ok(state) => {
                            self.state = state;
                            self.hdr_name_start = pos;
                            if StompCommand::Get == self.message.command {
                                self.request_line_done(buffer, pos);
                            }
                        },
                        Err(ParserState::Again) => {
                            // continue
                        },
                        Err(e) => {
                            warn!("invalid command '{}'", String::from_utf8_lossy(com));
                            return Err(e);
                        }
                    }
                }

                State::CommandLf => {
                    if ch == b'\r' {
                        // continue;
                    } else if ch == b'\n' {
                        self.state = State::HdrStart;
                        self.hdr_name_start = pos;
                        // continue;
                    } else {
                        return Err(ParserState::InvalidCommand);
                    }
                }

                State::HdrStart => {
                    if ch == b' ' || ch == b'\t' || ch == b'\r' {
                        // continue (ignore leading whitespace)
                    } else if ch == b'\n' {
                        // two LF marks end of headers
                        match self.message.command {
                            StompCommand::Message | StompCommand::Send => {
                                // has a body, return without reading the trailing \0
                                self.state = State::MsgRead;
                                return Ok(ParserState::Message(pos + 1));
                            }
                            StompCommand::Get => {
                                // has no trailing \0
                                self.state = State::MsgRead;
                                return Ok(ParserState::Message(pos + 1));
                            }
                            _ => {
                                // wait for \0
                                self.state = State::AlmostDone;
                            }
                        }
                    } else {
                        self.hdr_name_start = pos;
                        self.state = State::HdrName;
                    }
                }

                State::HdrName => {
                    if ch == b':' {
                        self.hdr_name_end = pos;
                        self.hdr_value_start = pos + 1;
                        self.state = State::HdrValue;
                    } else if ch == b'\r' || ch == b'\n' {
                        return Err(ParserState::InvalidMessage);
                    }
                }

                // TODO ignore whitespace between name and value and :
                State::HdrValue => {
                    if ch == b'\r' {
                        self.hdr_value_end = pos;
                        // continue
                    } else if ch == b'\n' {
                        if self.hdr_value_end == 0 {
                            self.hdr_value_end = pos;
                        };
                        self.header_done(buffer);

                        self.hdr_name_start = pos  - 1;
                        self.state = State::HdrStart;
                    }
                }

                State::AlmostDone => {
                    if ch == b'\0' {
                        self.state = State::MsgRead;
                        debug!("DONE {} {}", pos, self.bytes_read);
                        return Ok(ParserState::Done(pos + 1));
                    } else {
                        return Err(ParserState::InvalidCommand);
                    }
                }

                State::MsgRead => {
                    panic!("unreachable in parser")
                }
            }

            pos = pos + 1;
        } // end for

        Ok(ParserState::Again)
    }

    /// return Err(ParserState::Again) until a command is read, command is the first line of the message
    fn is_command_done(&self, ch: u8) -> Result<State, ParserState> {
        if ch == b'\r' || ch == b'\n' {
            // println!("Read a command {:?}", self.message.command);
            if StompCommand::Unknown == self.message.command {
                return Err(ParserState::InvalidCommand);
            }
            if ch == b'\r' {
                return Ok(State::CommandLf);
            }
            if ch == b'\n' {
                return Ok(State::HdrStart);
            }
        }
        Err(ParserState::Again)
    }

    /// parsed a header name:value
    fn header_done(&mut self, buffer: &[u8]) {

        self.message.add_header_clone(
            &String::from_utf8_lossy(&buffer[self.hdr_name_start..self.hdr_name_end]),
            &String::from_utf8_lossy(&buffer[self.hdr_value_start..self.hdr_value_end]));

        self.hdr_name_start = 0;
        self.hdr_name_end = 0;
        self.hdr_value_start = 0;
        self.hdr_value_end = 0;
    }

    fn request_line_done(&mut self, buffer: &[u8], pos: usize) {
        //println!("read request line '{}'", String::from_utf8_lossy(&buffer[self.request_line_start..pos]) );
        self.message.add_header_clone(
            &String::from("request-line"),
            &String::from_utf8_lossy(&buffer[self.request_line_start..pos]));
    }
}




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

    #[test]
    fn test_happy() {
        let mut p = StompParser::new(0);
        let buffer = b"ACK\nhdr1:value1\nhdr2:value2\n\n\0";
        let mut pos = 0;
        let mut end = 10;
        let mut chunk = &buffer[pos..end];
        match p.push(buffer, pos, chunk) {
            Ok(ParserState::Again) => {
                if StompCommand::Ack == p.message.command {
                    println!("parsed ACK");
                } else {
                    panic!("wrong command")
                }

                println!("p.message.command {:?}", p.message.command);
                println!("chunk parsed");
            }
            Err(_) => {
                println!("Parser error {:?}", p);
                panic!("parser error");
            }
            _  => panic!("unexpected state")
        };

        let read = 10;
        pos = end;
        end += read;
        chunk = &buffer[pos..end];
        match p.push(buffer, pos, chunk) {
            Ok(ParserState::Again) => println!("chunk parsed"),
            Err(_) => panic!("parser error"),
            _  => panic!("unexpected state")
        };

        pos = end;
        end = buffer.len();
        chunk = &buffer[pos..end];
        match p.push(buffer, pos, chunk) {
            Ok(ParserState::Done(_)) => println!("all parsed"),
            Err(_) => panic!("parser error"),
            _  => panic!("unexpected state")
        };
        println!("Parsed Message {:?}", p.message);
    }

    #[test]
    fn test_happy_login() {
        let mut p = StompParser::new(0);
        let buffer = b"STOMP\nlogin:xtomp\npasscode:passcode\n\n\0";
        match p.push(buffer, 0, buffer) {
            Ok(ParserState::Done(38)) => {
                if StompCommand::Stomp == p.message.command {}
                else { panic!("wrong command"); }
                match p.message.get_header("login") {
                    Some(value) => assert!( value.eq("xtomp")),
                    None => panic!("missing header")
                }
                match p.message.get_header("passcode") {
                    Some(value) => assert!( value.eq("passcode")),
                    None => panic!("missing header")
                }
            },
            Ok(ParserState::Done(_)) => {
                panic!("reporting wrong length")
            }
            Err(_) => {
                println!("Parser error {:?}", p);
                panic!("parser error")
            },
            _  => panic!("unexpected state")
        };
    }

    #[test]
    fn test_happy_body() {
        let mut p = StompParser::new(0);
        let buffer = b"SEND\ndestination:memtop-a\n\nsome text follows\0";
        match p.push(buffer, 0, buffer) {
            Ok(ParserState::Message(read)) => {
                assert_eq!(27, read);
                if StompCommand::Send == p.message.command {}
                else { panic!("wrong command"); }
                match p.message.get_header("destination") {
                    Some(value) => assert!( value.eq("memtop-a")),
                    None => panic!("missing header")
                }
                if State::MsgRead == p.state {}
                else { panic!("wrong parser state"); }
            }
            Err(_) => {
                println!("Parser error {:?}", p);
                panic!("parser error")
            },
            _  => panic!("unexpected state")
        };
    }

    #[test]
    fn test_telnet_line_endings_login() {
        let mut p = StompParser::new(0);
        let buffer = b"STOMP\r\nlogin:xtomp\r\npasscode:passcode\r\n\r\n\0";
        match p.push(buffer, 0, buffer) {
            Ok(ParserState::Done(_)) => {
                if StompCommand::Stomp == p.message.command {}
                else { panic!("wrong command"); }
                match p.message.get_header("login") {
                    Some(value) => assert!( value.eq("xtomp")),
                    None => panic!("missing header")
                }
                match p.message.get_header("passcode") {
                    Some(value) => assert!( value.eq("passcode")),
                    None => panic!("missing header")
                }
            }
            Err(_) => {
                println!("Parser error {:?}", p);
                panic!("parser error")
            },
            _  => panic!("unexpected state")
        };
    }


    #[test]
    fn test_invalid_command() {
        let buffer = b"WIBBLE\n\n\0";
        if let Err(ParserState::InvalidCommand) = StompParser::new(0).push(buffer, 0, buffer) {
            // expected
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_non_text_command() {
        let buffer = b"*uck*you\n\n\0";
        if let Err(ParserState::InvalidCommand) = StompParser::new(0).push(buffer, 0, buffer) {
            // expected
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_wrong_case_command() {
        let buffer = b"stomp\n\n\0";
        if let Err(ParserState::InvalidCommand) = StompParser::new(0).push(buffer, 0, buffer) {
            // expected
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_heart_beats() {
        let mut p = StompParser::new(0);
        // leading \n ignored
        let buffer = b"\n\n\nSTOMP\n\n\0";
        if let Ok(_) = p.push(buffer, 0, buffer) {
            assert_eq!(StompCommand::Stomp, p.message.command);
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_early_zero_termination_in_cmd() {
        let buffer = b"STOMP\0";
        if let Err(ParserState::InvalidCommand) = StompParser::new(0).push(buffer, 0, buffer) {
            // expected
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_early_zero_termination_in_hdr_val() {
        let buffer = b"STOMP\n\nhdr:value\0";
        if let Err(ParserState::InvalidCommand) = StompParser::new(0).push(buffer, 0, buffer) {
            // expected
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_early_zero_termination_in_hdr_name() {
        let buffer = b"STOMP\n\nhdr:value\n\0";
        if let Err(ParserState::InvalidCommand) = StompParser::new(0).push(buffer, 0, buffer) {
            // expected
        } else {
            panic!("unexpected state")
        }
    }

    #[test]
    fn test_http_get() {
        let mut p = StompParser::new(0);
        let buffer = b"GET /foo/baa HTTP/1.1\ndestination:memtop-a\n\n";
        match p.push(buffer, 0, buffer) {
            Ok(ParserState::Message(read)) => {
                assert_eq!(44, read);
                if StompCommand::Get == p.message.command {}
                else { panic!("wrong command"); }
                match p.message.get_header("destination") {
                    Some(value) => assert!( value.eq("memtop-a")),
                    None => panic!("missing header")
                }
                if State::MsgRead == p.state {}
                else { panic!("wrong parser state"); }
            }
            Err(_) => {
                println!("Parser error {:?}", p);
                panic!("parser error")
            },
            _  => panic!("unexpected state")
        };
    }
}