romp 0.5.2

use std::sync::{Arc, RwLock};

use log::*;
use log::Level::Debug;

use tokio::io::ReadHalf;
use tokio::net::TcpStream;
use tokio::prelude::*;
use tokio::prelude::task::Task;

use crate::body_parser::BodyParser;
use crate::body_parser::fixed_length::FixedLengthBodyParser;
use crate::body_parser::text::TextBodyParser;
use crate::errors::ClientError;
use crate::init::CONFIG;
use crate::message::stomp_message::Ownership;
use crate::parser::{ParserState, StompParser};
use crate::session::stomp_session::{StompSession, FLAG_WEB_SOCKETS};
use crate::message::stomp_message::StompCommand;
use crate::workflow::router;
use crate::web_socket::ws_demunge::{ws_demunge, WsState};
use crate::body_parser::ParserState::BodyFlup;

/// responsible for reading from the Tcp connection's ReadHalf

enum ReaderMode {
    ReadingCommand,
    ReadingTextBody,
    ReadingBinaryBody,
}

pub struct Reader {
    session: Arc<RwLock<StompSession>>,
    session_id: usize,
    read_half: ReadHalf<TcpStream>,
    buf: Box<[u8]>,
    // how far into buf we have analyzed
    pos: usize,
    // how much data has been pushed to buf
    end: usize,
    read_done: bool,
    web_socket: bool,
    ws: WsState,
    // q of frame lengths so we can validate frame length == STOMP msg length
    frame_lengths: Box<Vec<usize>>,
    mode: ReaderMode,
    parser: StompParser,
    // both have trait BodyParser but this does not work with stack references only for boxed vars (unknown size)
    fixed_length_body_parser: FixedLengthBodyParser,
    text_body_parser: TextBodyParser,

}

impl Reader {

    pub fn new(session: Arc<RwLock<StompSession>>, session_id: usize, read_half: ReadHalf<TcpStream>) -> Reader {
        Reader {
            session,
            session_id,
            read_half,
            buf: vec![0 as u8; CONFIG.request_client_buffer].into_boxed_slice(),
            pos: 0,
            end: 0,
            read_done: false,
            web_socket: false,
            ws: WsState::new(),
            frame_lengths: Box::new(vec![]),
            mode: ReaderMode::ReadingCommand,
            parser: StompParser::new(session_id),
            fixed_length_body_parser: FixedLengthBodyParser::new(),
            text_body_parser: TextBodyParser::new(CONFIG.max_message_size),
        }
    }

    /// called when Stomp command and headers have been read, unless its a SEND, \0 has also been read.
    fn handle_command_received(&mut self) -> Result<(), ParserState> {
        let message =  &self.parser.message;
        if let Err(ps) = self.assert_hdr_lengths() {
            return Err(ps);
        }
        match &message.command {
            StompCommand::Send | StompCommand::Message => {
                debug!("headers read, reading body");
                self.setup_body_parser();
            },
            // we read a trailing \0
            _ => {
                if ! self.assert_ws_frame_len() {
                    return Err(ParserState::InvalidCommand);
                }
                self.route_message();
            }
        }
        self.realign_buffer();
        if self.end > 0 {
            return self.process_buffer_portion();
        } else {
            return Ok(());
        }
    }

    /// Called when a message body has been completely read to memory
    fn handle_message_received(&mut self) -> Result<(), ParserState> {
        if ! self.assert_ws_frame_len() {
            return Err(ParserState::InvalidCommand);
        }
        debug!("body read");
        self.realign_buffer();

        self.route_message();

        self.mode = ReaderMode::ReadingCommand;

        if self.end > 0 {
            return self.process_buffer_portion();
        } else {
            return Ok(());
        }
    }

    fn assert_hdr_lengths(&self) -> Result<(), ParserState> {
        let message =  &self.parser.message;
        // TODO be good to exclude required headers from this count
        // 8 is bear minimum for Http Upgrade, could separate counts for STOMP and Http
        if message.count_headers() > CONFIG.max_headers {
            debug!("hdr flup count");
            return Err(ParserState::HdrFlup);
        }
        for hdr in message.headers() {
            if hdr.name.len() + 1 + hdr.value.len() > CONFIG.max_header_len {
                debug!("hdr flup len");
                return Err(ParserState::HdrFlup);
            }
        }
        Ok(())
    }

    /// check the amount we have read parsing the data is the ws.frame_len
    /// this allows failing fast, rather than trying to parse message data as websockets binary frame information.
    /// This could be optional so we allow multiple STOMP messages in a single ws frame, error detection is harder if we do
    fn assert_ws_frame_len(&mut self) -> bool {
        if self.web_socket {

            let actual;
            if self.fixed_length_body_parser.expected_len() > 0 {
                actual = self.parser.bytes_read() + self.fixed_length_body_parser.expected_len() + 1;
            } else if self.text_body_parser.bytes_read() > 0 {
                actual = self.parser.bytes_read() + self.text_body_parser.bytes_read();
            } else {
                actual = self.parser.bytes_read();
            }

            // we have not read a full frame yet, but we have read a full STOMP message, frame len is too long.
            if self.frame_lengths.len() == 0 {
                warn!("ws.frame_len too long {} > {}", self.ws.frame_len(), actual);
                if log_enabled!(Debug) {
                    debug!("header={}, text-body={}, bin-body={}", self.parser.bytes_read(), self.text_body_parser.bytes_read(), self.fixed_length_body_parser.expected_len());
                    debug!("msg_hdrs={}", self.parser.message.hdrs_as_string());
                }
                // cant recover from this
                return false;
            }

            let expected = self.frame_lengths.remove(0);

            if expected != actual {
                warn!("Invalid ws.frame_len {} != {}", expected, actual);
                if log_enabled!(Debug) {
                    debug!("header={}, text-body={}, bin-body={}", self.parser.bytes_read(), self.text_body_parser.bytes_read(), self.fixed_length_body_parser.expected_len());
                    debug!("msg_hdrs={}", self.parser.message.hdrs_as_string());
                }
                // is ws reader is happy and parser is happy we continue, must be a bug in length calculation
                return true;
            }
        }
        true
    }

    /// send fully uploaded message to router for workflow
    /// resets parser as a side effect
    fn route_message(&mut self) {

        let mut message =  self.parser.take_message(Ownership::Session);
        //debug!("StompSession::route_message() {:?}", message);

        let session = self.session.clone();
        router::route(message, session, self.web_socket);
    }

    /// setup the correct body parser depending on hdrs
    fn setup_body_parser(&mut self) {
        let message =  &self.parser.message;
        match message.get_header("content-length") {
            Some(content_length) => {
                match content_length.parse::<usize>() {
                    Ok(content_length) => {
                        if content_length > CONFIG.max_message_size {
                            // TODO could drain here, so its not fatal
                            self.session.write().unwrap().send_client_error_fatal(ClientError::BodyFlup);
                        }
                        self.mode = ReaderMode::ReadingBinaryBody;
                        self.text_body_parser.reinit(CONFIG.max_message_size);
                        self.fixed_length_body_parser.reinit(content_length);
                    },
                    Err(_e) => {
                        // protocol error here, send syntax error and bomb the tcp connection
                        self.session.write().unwrap().send_client_error_fatal(ClientError::Syntax);

                    }
                }
            },
            // no content-length means body must be \0 terminated
            None => {
                self.mode = ReaderMode::ReadingTextBody;
                self.text_body_parser.reinit(CONFIG.max_message_size);
                self.fixed_length_body_parser.reinit(0);
            }
        }

    }

    /// After reading a Command or getting to the end of a body there can be data left in buf
    /// that is as yet unparsed. This method shifts the remaining data to the front of the buffer.
    /// ensures each message has space for all the headers.
    fn realign_buffer(&mut self) {
        if self.end < self.pos {
            println!("WTF?? {} {}", self.pos, self.end);
        }
        let len = self.end - self.pos;
        self.buf.copy_within(self.pos..self.end, 0);
        self.pos = 0;
        self.end = len;
    }

    /// Called either when we have been pushed extra data from the network, or when
    /// we processed only part of the data and there is some left between self.pos and self.end
    fn process_buffer_portion(&mut self) -> Result<(), ParserState> {
        if self.pos == self.end {
            // read something but not STOMP data
            return Ok(())
        }

        let chunk = &self.buf[self.pos..self.end];

        match self.mode {

            ReaderMode::ReadingTextBody => {
                match self.text_body_parser.push(chunk, &mut self.parser.message) {
                    Ok(read) => {
                        if self.text_body_parser.is_done() {
                            self.pos = read;
                            return self.handle_message_received();
                        } else {
                            self.pos = 0;
                            self.end = 0;
                            return Ok(());
                        }
                    },
                    Err(ps) => {
                        debug!("body parser error {:?}", ps);
                        if ps == BodyFlup {
                            return Err(ParserState::BodyFlup);
                        }
                        return Err(ParserState::InvalidMessage);
                    }
                }
            },

            ReaderMode::ReadingBinaryBody => {
                match self.fixed_length_body_parser.push(chunk, &mut self.parser.message) {
                    Ok(read) => {
                        if self.fixed_length_body_parser.is_done() {
                            self.pos = read;
                            return self.handle_message_received();
                        } else {
                            self.pos = 0;
                            self.end = 0;
                            return Ok(());
                        }
                    },
                    Err(ps) => {
                        debug!("body parser error {:?}", ps);
                        if ps == BodyFlup {
                            return Err(ParserState::BodyFlup);
                        }
                        return Err(ParserState::InvalidMessage);
                    }
                }
            },

            ReaderMode::ReadingCommand => {

                match self.parser.push(&self.buf, self.pos, chunk) {
                    Ok(ParserState::Again) => {
                        self.pos += chunk.len();
                        if self.pos != self.end {
                            debug!("pos not what it ought ta be");
                        }
                        return Ok(());
                    },
                    Ok(ParserState::Done(read)) => {
                        self.pos = read;
                        return self.handle_command_received();
                    },
                    Ok(ParserState::Message(read)) => {
                        // loop
                        self.pos = read;
                        return self.handle_command_received();
                    },
                    Ok(ps) => {
                        panic!("unreachable {:?}", ps)
                    },
                    Err(ps) => {
                        return Err(ps);
                    }

                }
            }
        }

    }

    /// Called prior to returning Err() which is our last breath
    fn terminated(&self) {
        self.session.write().unwrap().read_terminated();
    }
}

impl Future for Reader {

    type Item = ();
    type Error = ();

    fn poll(&mut self) -> Result<Async<Self::Item>, Self::Error> {
        debug!("read polled"); // called on connect, data, and on zero bytes

        if self.session.read().unwrap().shutdown_pending() {
            debug!("reader closed");
            self.terminated();
            return Err(());
        }

        loop {

            if !self.read_done {

                // &mut because poll_read reads into the buffer,
                // this only writes to pos
                let remain = self.buf.len() - self.pos;

                match self.read_half.poll_read(&mut self.buf[self.pos..remain]) {
                    Ok(Async::Ready(n)) => {
                        if n == 0 {
                            // TcpConnection closed
                            self.read_done = true;
                            debug!("tcp connection closed");
                            self.session.write().unwrap().shutdown();
                        } else {

                            if self.pos != self.end {
                                debug!("pos not what it ought ta be");
                                self.pos = self.end;
                            }

                            self.end += n;

                            // WebSocket un-munging
                            if self.web_socket || self.session.read().unwrap().get_flag(FLAG_WEB_SOCKETS) {
                                if ! self.web_socket {
                                    self.web_socket = true;
                                }
                                match ws_demunge(&mut self.ws, &mut self.buf, self.pos, self.end, &mut self.frame_lengths) {
                                    Ok((new_pos, new_end)) => {
                                        self.pos = new_pos;
                                        self.end = new_end;
                                        // potentially here we did not process any STOMP data
                                    },
                                    Err(_) => {
                                        self.session.write().unwrap().read_error(ParserState::InvalidMessage);
                                        self.terminated();
                                        return Err(());
                                    }
                                }
                            }


                            match self.process_buffer_portion() {
                                Err(ps) => {
                                    debug!("reader closed {:?}", ps);
                                    self.session.write().unwrap().read_error(ps);
                                    self.terminated();
                                    return Err(());
                                },
                                Ok(()) => {
                                    self.session.read().unwrap().read_something();
                                    // loop
                                }
                            }
                        }
                    }
                    Ok(Async::NotReady) => {
                        return Ok(Async::NotReady);
                    }
                    Err(_e) => {
                        debug!("reader closed");
                        self.terminated();
                        return Err(());
                    }
                };

            }

            if self.read_done {
                debug!("reader closed");
                self.terminated();
                return Err(());
            }
        }

    }

}

// elaborate way to be able to notify read thread when we want to tell it to die

pub struct ReadKiller {
    pub die: bool,
    task: Option<Task>,
}

impl ReadKiller {

    pub fn new() -> ReadKiller {
        ReadKiller {
            die: false,
            task: None,
        }
    }

    pub fn kill(&mut self) {
        self.die = true;
        match &self.task {
            Some(task) => {
                debug!("notifying read to stop it");
                // seems this might not do anything if we get here from a read
                task.notify();
                self.task = None;
            },
            None => {}
        }
    }
}

impl Future for ReadKiller {
    type Item = ();
    type Error = ();
    fn poll(&mut self) -> Result<Async<Self::Item>, Self::Error> {
        debug!("read killer polled");
        self.task = Some(futures::task::current());
        match self.die {
            true => {
                debug!("read killer closed");
                self.task = None;
                Err(())
            },
            _ => Ok(Async::NotReady),
        }
    }
}