<commit_before>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::{HandshakeResult, MidHandshakeTlsStream, TcpStream},
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{debug, error, trace};
use std::{
    collections::VecDeque,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    SocketConnected,
    SocketWritable,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    handshake: Option<MidHandshakeTlsStream>,
    stream: Option<TcpStream>,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

macro_rules! stream_mut {
    ($self:ident) => {
        $self.stream.as_mut().unwrap()
    };
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let (mut stream, mut handshake) = match stream {
            Ok(stream) => (Some(stream), None),
            Err(error) => (None, Some(error.into_mid_handshake_tls_stream()?)),
        };
        let source = stream
            .as_mut()
            .unwrap_or_else(|| handshake.as_mut().unwrap().get_mut());
        let slot = reactor.register(source, socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            handshake,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn stream(&self) -> &TcpStream {
        self.stream
            .as_ref()
            .unwrap_or_else(|| self.handshake.as_ref().unwrap().get_ref())
    }

    fn ensure_connected(&mut self) -> Result<bool> {
        match self.stream().peer_addr() {
            Ok(peer) => {
                debug!("Connecting to {}", peer);
                self.status = Status::SocketConnected;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected = err.kind() {
                    Ok(false)
                } else {
                    Err(err.into())
                }
            }
        }
    }

    fn ensure_writable(&mut self) -> Result<bool> {
        match self.stream().is_writable() {
            Ok(()) => {
                self.status = Status::SocketWritable;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected | io::ErrorKind::WouldBlock = err.kind() {
                    Ok(false)
                } else {
                    Err(err.into())
                }
            }
        }
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        loop {
            let success = match self.status {
                Status::Initial => self.ensure_connected(),
                Status::SocketConnected => self.ensure_writable(),
                Status::SocketWritable => return self.finish_setup(),
                Status::Connected => return Ok(true),
                Status::Stop => return Ok(false),
            }?;
            if !success {
                return Ok(false);
            }
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable() && self.receive_buffer.available_space() > 0
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown()?;
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if let Some(handshake) = self.handshake.take() {
            match handshake.handshake() {
                Ok(stream) => self.stream = Some(stream),
                Err(error) => self.handshake = Some(error.into_mid_handshake_tls_stream()?),
            }
            if self.handshake.is_some() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        self.reactor.shutdown()?;
        Err(error)
    }

    fn handle_read_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_read_result(result, &*self.reactor, self.slot)
        {
            error!("error reading: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        stream_mut!(self).flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        if self.socket_state.writable() {
            let res = self.flush();
            self.handle_write_result(res)?;
        }
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            self.handle_read_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(stream_mut!(self))?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let sz = self.receive_buffer.read_from(stream_mut!(self))?;

                if sz > 0 {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                } else {
                    error!("Socket was readable but we read 0, marking as wouldblock");
                    self.handle_read_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
                }
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<commit_msg>rustfmt<commit_after>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPC,
    protocol::{self, AMQPError, AMQPHardError},
    reactor::{ReactorBuilder, ReactorHandle, Slot},
    socket_state::SocketState,
    tcp::{HandshakeResult, MidHandshakeTlsStream, TcpStream},
    thread::ThreadHandle,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use log::{debug, error, trace};
use std::{
    collections::VecDeque,
    io::{self, Write},
    sync::Arc,
    thread::Builder as ThreadBuilder,
    time::Duration,
};

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    SocketConnected,
    SocketWritable,
    Connected,
    Stop,
}

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPC,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    reactor: Box<dyn ReactorHandle + Send>,
    reactor_thread_handle: ThreadHandle,
    connection_io_loop_handle: ThreadHandle,
    handshake: Option<MidHandshakeTlsStream>,
    stream: Option<TcpStream>,
    slot: Slot,
    status: Status,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(u64, Option<PromiseResolver<()>>)>,
}

macro_rules! stream_mut {
    ($self:ident) => {
        $self.stream.as_mut().unwrap()
    };
}

impl IoLoop {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPC,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: HandshakeResult,
        reactor_builder: &dyn ReactorBuilder,
    ) -> Result<Self> {
        let heartbeat = Heartbeat::new(channels.clone());
        let mut reactor = reactor_builder.build(heartbeat.clone())?;
        let reactor_handle = reactor.handle();
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE as usize,
            configuration.frame_max() as usize,
        );
        let (mut stream, mut handshake) = match stream {
            Ok(stream) => (Some(stream), None),
            Err(error) => (None, Some(error.into_mid_handshake_tls_stream()?)),
        };
        let source = stream
            .as_mut()
            .unwrap_or_else(|| handshake.as_mut().unwrap().get_mut());
        let slot = reactor.register(source, socket_state.handle())?;

        let reactor_thread_handle = reactor.start()?;

        Ok(Self {
            connection_status,
            configuration,
            channels,
            internal_rpc,
            frames,
            heartbeat,
            socket_state,
            reactor: reactor_handle,
            reactor_thread_handle,
            connection_io_loop_handle,
            stream,
            handshake,
            slot,
            status: Status::Initial,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            serialized_frames: VecDeque::default(),
        })
    }

    fn stream(&self) -> &TcpStream {
        self.stream
            .as_ref()
            .unwrap_or_else(|| self.handshake.as_ref().unwrap().get_ref())
    }

    fn ensure_connected(&mut self) -> Result<bool> {
        match self.stream().peer_addr() {
            Ok(peer) => {
                debug!("Connecting to {}", peer);
                self.status = Status::SocketConnected;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected = err.kind() {
                    Ok(false)
                } else {
                    Err(err.into())
                }
            }
        }
    }

    fn ensure_writable(&mut self) -> Result<bool> {
        match self.stream().is_writable() {
            Ok(()) => {
                self.status = Status::SocketWritable;
                Ok(true)
            }
            Err(err) => {
                if let io::ErrorKind::NotConnected | io::ErrorKind::WouldBlock = err.kind() {
                    Ok(false)
                } else {
                    Err(err.into())
                }
            }
        }
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.configuration.heartbeat();
            if heartbeat != 0 {
                let heartbeat = Duration::from_millis(u64::from(heartbeat) * 500); // * 1000 (ms) / 2 (half the negotiated timeout)
                self.heartbeat.set_timeout(heartbeat);
                self.reactor.start_heartbeat();
            }
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        loop {
            let success = match self.status {
                Status::Initial => self.ensure_connected(),
                Status::SocketConnected => self.ensure_writable(),
                Status::SocketWritable => return self.finish_setup(),
                Status::Connected => return Ok(true),
                Status::Stop => return Ok(false),
            }?;
            if !success {
                return Ok(false);
            }
        }
    }

    fn has_data(&self) -> bool {
        self.frames.has_pending()
            || self.send_buffer.available_data() > 0
            || !self.serialized_frames.is_empty()
    }

    fn can_write(&mut self) -> bool {
        self.socket_state.writable() && self.has_data() && !self.connection_status.blocked()
    }

    fn can_read(&mut self) -> bool {
        self.socket_state.readable() && self.receive_buffer.available_space() > 0
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        (self.status != Status::Connected
            || self.connection_status.connected()
            || self.connection_status.closing())
            && self.status != Status::Stop
            && !self.connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.socket_state.handle();
        self.connection_io_loop_handle.clone().register(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    while self.should_continue() {
                        if let Err(err) = self.run() {
                            self.critical_error(err)?;
                        }
                    }
                    self.heartbeat.cancel();
                    self.reactor.shutdown()?;
                    self.reactor_thread_handle.wait("reactor")
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_internal_rpc(&self) -> Result<()> {
        self.internal_rpc.poll(&self.channels)
    }

    fn poll_socket_events(&mut self) -> Result<()> {
        self.socket_state.poll_events();
        if self.socket_state.error() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionAborted).into())?;
        }
        if self.socket_state.closed() {
            self.critical_error(io::Error::from(io::ErrorKind::ConnectionReset).into())?;
        }
        self.poll_internal_rpc()
    }

    fn run(&mut self) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events()?;
        if !self.ensure_setup()? {
            return Ok(());
        }
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data()
        );
        if !self.can_read() && !self.can_write() {
            self.socket_state.wait();
        }
        self.poll_socket_events()?;
        if let Some(handshake) = self.handshake.take() {
            match handshake.handshake() {
                Ok(stream) => self.stream = Some(stream),
                Err(error) => self.handshake = Some(error.into_mid_handshake_tls_stream()?),
            }
            if self.handshake.is_some() {
                // We hit WOULDBLOCK while handshaking, wait for the next socket event
                return Ok(());
            }
        }
        self.write()?;
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
        if self.should_continue() {
            self.read()?;
        }
        self.handle_frames()?;
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.socket_state.readable(),
            self.socket_state.writable(),
            self.has_data(),
            self.status
        );
        self.poll_internal_rpc()
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let Some(resolver) = self.connection_status.connection_resolver() {
            resolver.swear(Err(error.clone()));
        }
        self.status = Status::Stop;
        self.channels.set_connection_error(error.clone())?;
        for (_, resolver) in std::mem::take(&mut self.serialized_frames) {
            if let Some(resolver) = resolver {
                resolver.swear(Err(error.clone()));
            }
        }
        self.reactor.shutdown()?;
        Err(error)
    }

    fn handle_read_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_read_result(result, &*self.reactor, self.slot)
        {
            error!("error reading: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn handle_write_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self
            .socket_state
            .handle_write_result(result, &*self.reactor, self.slot)
        {
            error!("error writing: {:?}", e);
            self.critical_error(e)?;
        }
        self.poll_internal_rpc()
    }

    fn flush(&mut self) -> Result<()> {
        stream_mut!(self).flush()?;
        self.poll_internal_rpc()
    }

    fn write(&mut self) -> Result<()> {
        if self.socket_state.writable() {
            let res = self.flush();
            self.handle_write_result(res)?;
        }
        while self.can_write() {
            let res = self.write_to_stream();
            self.handle_write_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn read(&mut self) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream();
            self.handle_read_result(res)?;
        }
        self.poll_internal_rpc()
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.flush()?;
        self.serialize()?;

        let sz = self.send_buffer.write_to(stream_mut!(self))?;

        if sz > 0 {
            self.heartbeat.update_last_write();

            trace!("wrote {} bytes", sz);
            self.send_buffer.consume(sz);

            let mut written = sz as u64;
            while written > 0 {
                if let Some((to_write, resolver)) = self.serialized_frames.pop_front() {
                    if written < to_write {
                        self.serialized_frames
                            .push_front((to_write - written, resolver));
                        trace!("{} to write to complete this frame", to_write - written);
                        written = 0;
                    } else {
                        if let Some(resolver) = resolver {
                            resolver.swear(Ok(()));
                        }
                        written -= to_write;
                    }
                } else {
                    error!(
                        "We've written {} but didn't expect to write anything",
                        written
                    );
                    break;
                }
            }

            if self.send_buffer.available_data() > 0 {
                // We didn't write all the data yet
                trace!("Still {} to write", self.send_buffer.available_data());
            }

            self.flush()?;
        } else {
            error!("Socket was writable but we wrote 0, marking as wouldblock");
            self.handle_write_result(Err(io::Error::from(io::ErrorKind::WouldBlock).into()))?;
        }
        self.poll_internal_rpc()
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let sz = self.receive_buffer.read_from(stream_mut!(self))?;

                if sz > 0 {
                    trace!("read {} bytes", sz);
                    self.receive_buffer.fill(sz);
                } else {
                    error!("Socket was readable but we read 0, marking as wouldblock");
                    self.handle_read_result(
                        Err(io::Error::from(io::ErrorKind::WouldBlock).into()),
                    )?;
                }
                self.poll_internal_rpc()
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!("will write to buffer: {}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(sz) => self.serialized_frames.push_back((sz, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        self.poll_internal_rpc()
    }

    fn handle_frames(&mut self) -> Result<()> {
        while self.can_parse() {
            if let Some(frame) = self.parse()? {
                self.channels.handle_frame(frame)?;
            } else {
                break;
            }
        }
        self.poll_internal_rpc()
    }

    fn parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.parsing_context()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.offset(i);
                let frame_max = self.configuration.frame_max() as usize;
                if frame_max > 0 && consumed > frame_max {
                    error!("received large ({} bytes) frame", consumed);
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.handle().close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.poll_internal_rpc()?;
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!("parse error: {:?}", e);
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>


<commit_before>//! IPC for way-cooler

use std::thread;
use std::env;
use std::socket_path::PathBuf;
use std::fs;

use unix_socket::UnixListener;

mod channel;
mod command;
mod event;

#[cfg(test)]
mod tests;

/// Versions are incremented.
pub const VERSION: u64 = 0u64; // Increment to 1 on release.

/// Very much not cross-platform!
/// Submit an issue when Wayland is ported to Windoze.
pub const SOCKET_FOLDER: &'static str = "/var/run/way-cooler/";
/// Socket over which synchronous communication is made with clients.
pub const COMMAND_SOCKET: &'static str = "command";
/// Socket over which events are sent to clients.
pub const EVENT_SOCKET: &'static str = "event";
/// Folder in which sockets are created
pub const PATH_VAR: &'static str = "WAY_COOLER_SOCKET_FOLDER";

/// We need random folder names to place sockets in, but they don't need
/// to be _that_ random.
pub fn unique_ish_id() -> u32 {
    use std::hash::{Hash, Hasher, SipHasher};
    use std::time::Instant;

    // If you shift a u64 hash right by this you get a "checksum",
    // a number which retains some of the entropy of the hash but
    // is small enough to fit a more comfortable file name.
    const MAGIC_SHIFT_NUMBER: u64 = 0b110000;

    // Instant doesn't implement hash, and it's supposed to be an opaque
    // struct, but it does implement debug...
    let now = Instant::now();
    let mut hasher = SipHasher::new();
    format!("{:?}", now).hash(&mut hasher);
    (hasher.finish() >> MAGIC_SHIFT_NUMBER) as u32
}

/// Initialize the IPC socket.
pub fn init() {
    let id = unique_ish_id();
    info!("Starting IPC with unique ID {}", id);

    let mut socket_path = PathBuf::from(SOCKET_FOLDER);
    socket_path.push(id.to_string());

    if let Err(ioerr) = fs::create_dir_all(socket_path.clone()) {
        // How can we handle not having a socket?
        // In the future, we could log and continue.
        // We could have a config option to not create/create-if-possible
        error!("Unable to create temp folder: {:?}", ioerr);
        return;
    }
    let command_socket = UnixListener::bind(socket_path.join(COMMAND_SOCKET))
        .expect("Unable to open command socket!");

    let event_socket = UnixListener::bind(socket_path.join(EVENT_SOCKET))
        .expect("Unable to open event socket!");

    env::set_var(PATH_VAR, socket_path.clone());

    debug!("IPC initialized, now listening for clients.");

    let _server_handle = thread::Builder::new()
        .name("Command socket listener".to_string())
        .spawn(move || { command_thread(command_socket) });

    let _event_handle = thread::Builder::new()
        .name("Event socket listener".to_string())
        .spawn(move || { event_thread(event_socket) });

    trace!("IPC initialized.");
}

fn command_thread(socket: UnixListener) {
    for stream in socket.incoming() {
        trace!("Sever: new connection: {:?}", stream);
        match stream {
            Ok(mut stream) => {
                info!("Command: connected to {:?}", stream);
                let _handle = thread::Builder::new()
                    .name("IPC server helper".to_string())
                    .spawn(move || command::thread(&mut stream));
            },
            Err(err) => {
                info!("Error receiving a stream: {}", err);
            }
        }
    }
}

fn event_thread(socket: UnixListener) {
    for stream in socket.incoming() {
        trace!("Event: new connection: {:?}", stream);
        match stream {
            Ok(mut stream) => {
                info!("Event: connected to {:?}", stream);
                let _handle = thread::Builder::new()
                    .name("IPC evemt helper".to_string())
                    .spawn(move || event::thread(&mut stream));
            },
            Err(err) => {
                info!("Error receiving a stream: {}", err);
            }
        }
    }
}
<commit_msg>Fix find+replace problem<commit_after>//! IPC for way-cooler

use std::thread;
use std::env;
use std::path::PathBuf;
use std::fs;

use unix_socket::UnixListener;

mod channel;
mod command;
mod event;

#[cfg(test)]
mod tests;

/// Versions are incremented.
pub const VERSION: u64 = 0u64; // Increment to 1 on release.

/// Very much not cross-platform!
/// Submit an issue when Wayland is ported to Windoze.
pub const SOCKET_FOLDER: &'static str = "/var/run/way-cooler/";
/// Socket over which synchronous communication is made with clients.
pub const COMMAND_SOCKET: &'static str = "command";
/// Socket over which events are sent to clients.
pub const EVENT_SOCKET: &'static str = "event";
/// Folder in which sockets are created
pub const PATH_VAR: &'static str = "WAY_COOLER_SOCKET_FOLDER";

/// We need random folder names to place sockets in, but they don't need
/// to be _that_ random.
pub fn unique_ish_id() -> u32 {
    use std::hash::{Hash, Hasher, SipHasher};
    use std::time::Instant;

    // If you shift a u64 hash right by this you get a "checksum",
    // a number which retains some of the entropy of the hash but
    // is small enough to fit a more comfortable file name.
    const MAGIC_SHIFT_NUMBER: u64 = 0b110000;

    // Instant doesn't implement hash, and it's supposed to be an opaque
    // struct, but it does implement debug...
    let now = Instant::now();
    let mut hasher = SipHasher::new();
    format!("{:?}", now).hash(&mut hasher);
    (hasher.finish() >> MAGIC_SHIFT_NUMBER) as u32
}

/// Initialize the IPC socket.
pub fn init() {
    let id = unique_ish_id();
    info!("Starting IPC with unique ID {}", id);

    let mut socket_path = PathBuf::from(SOCKET_FOLDER);
    socket_path.push(id.to_string());

    if let Err(ioerr) = fs::create_dir_all(socket_path.clone()) {
        // How can we handle not having a socket?
        // In the future, we could log and continue.
        // We could have a config option to not create/create-if-possible
        error!("Unable to create temp folder: {:?}", ioerr);
        return;
    }
    let command_socket = UnixListener::bind(socket_path.join(COMMAND_SOCKET))
        .expect("Unable to open command socket!");

    let event_socket = UnixListener::bind(socket_path.join(EVENT_SOCKET))
        .expect("Unable to open event socket!");

    env::set_var(PATH_VAR, socket_path.clone());

    debug!("IPC initialized, now listening for clients.");

    let _server_handle = thread::Builder::new()
        .name("Command socket listener".to_string())
        .spawn(move || { command_thread(command_socket) });

    let _event_handle = thread::Builder::new()
        .name("Event socket listener".to_string())
        .spawn(move || { event_thread(event_socket) });

    trace!("IPC initialized.");
}

fn command_thread(socket: UnixListener) {
    for stream in socket.incoming() {
        trace!("Sever: new connection: {:?}", stream);
        match stream {
            Ok(mut stream) => {
                info!("Command: connected to {:?}", stream);
                let _handle = thread::Builder::new()
                    .name("IPC server helper".to_string())
                    .spawn(move || command::thread(&mut stream));
            },
            Err(err) => {
                info!("Error receiving a stream: {}", err);
            }
        }
    }
}

fn event_thread(socket: UnixListener) {
    for stream in socket.incoming() {
        trace!("Event: new connection: {:?}", stream);
        match stream {
            Ok(mut stream) => {
                info!("Event: connected to {:?}", stream);
                let _handle = thread::Builder::new()
                    .name("IPC evemt helper".to_string())
                    .spawn(move || event::thread(&mut stream));
            },
            Err(err) => {
                info!("Error receiving a stream: {}", err);
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;

use super::Register;

type AbsoluteIndex = usize;

#[derive(Debug, Copy, Clone)]
pub enum Instruction {
    Inbox,
    Outbox,
    CopyFrom(Register),
    CopyTo(Register),
    BumpUp(Register),
    BumpDown(Register),
    Add(Register),
    Sub(Register),
    Jump(AbsoluteIndex),
    JumpIfZero(AbsoluteIndex),
    JumpIfNegative(AbsoluteIndex),
    NoOp,
}

// Clamped at [-999, 999]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct NumberValue(i16);

impl NumberValue {
    fn from_char(c: char) -> Result<NumberValue, StepError> {
        NumberValue::clamp(c as i16)
    }

    fn clamp(v: i16) -> Result<NumberValue, StepError> {
        if v > 999 {
            Err(StepError::Overflow)
        } else if v < -999 {
            Err(StepError::Underflow)
        } else {
            Ok(NumberValue(v))
        }
    }

    fn add(self, other: NumberValue) -> Result<NumberValue, StepError> {
        NumberValue::clamp(self.0 + other.0)
    }

    fn sub(self, other: NumberValue) -> Result<NumberValue, StepError> {
        NumberValue::clamp(self.0 - other.0)
    }

    fn is_zero(self) -> bool { self.0 == 0 }
    fn is_negative(self) -> bool { self.0 < 0 }

    fn increment(self) -> Result<NumberValue, StepError> {
        NumberValue::clamp(self.0 + 1)
    }

    fn decrement(self) -> Result<NumberValue, StepError> {
        NumberValue::clamp(self.0 - 1)
    }

    fn into_u8(self) -> u8 {
        self.0 as u8 // Should have error of some kind?
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Tile {
    Number(NumberValue),
    Letter(char),
}

impl Tile {
    pub fn num(i: i16) -> Tile {
        Tile::Number(NumberValue::clamp(i).unwrap())
    }
}

#[derive(Debug, Copy, Clone)]
pub enum StepError {
    EndOfProgram,
    IndirectThroughNil,
    IndirectThroughNegative,
    IndirectThroughLetter,
    OutputNil,
    CopyFromNil,
    CopyToNil,
    BumpNil,
    BumpLetter,
    AddWithNil,
    AddToNil,
    AddWithLetter,
    SubFromNil,
    SubWithNil,
    SubCrossTypes,
    JumpZeroNil,
    JumpNegativeNil,
    Underflow,
    Overflow,
}

pub type Input = Vec<Tile>;
pub type Output = Vec<Tile>;
pub type Registers = BTreeMap<u8, Tile>;

#[derive(Debug, Clone)]
pub struct Machine {
    program: Vec<Instruction>,
    input: Input,
    output: Output,
    pc: usize,
    accumulator: Option<Tile>,
    registers: Registers,
}

impl Machine {
    pub fn new<I>(program: I, mut input: Input, registers: Registers) -> Machine
        where I: IntoIterator<Item = Instruction>
    {
        // We want to pop off the front, so flip it around for efficiency.
        input.reverse();

        Machine {
            program: program.into_iter().collect(),
            input: input,
            output: Vec::new(),
            pc: 0,
            accumulator: None,
            registers: registers,
        }
    }

    fn deref_target(&self, r: Register) -> Result<u8, StepError> {
        match r {
            Register::Direct(r) => Ok(r),
            Register::Indirect(r) => match self.registers.get(&r) {
                None => Err(StepError::IndirectThroughNil),
                Some(&Tile::Number(v)) if v.is_negative() => Err(StepError::IndirectThroughNegative),
                Some(&Tile::Number(v)) => Ok(v.into_u8()), // Should have max # registers?
                Some(&Tile::Letter(..)) => Err(StepError::IndirectThroughLetter),
            },
        }
    }

    pub fn into_output(self) -> Output { self.output }

    pub fn step(&mut self) -> Result<(), StepError> {
        use self::Instruction::*;

        // println!("PC: {}", self.pc);
        // println!("Instr: {:?}", self.program[self.pc]);
        // println!("Acc: {:?}", self.accumulator);

        match self.program[self.pc] {
            Inbox => {
                match self.input.pop() {
                    Some(v) => self.accumulator = Some(v),
                    None => return Err(StepError::EndOfProgram),
                }
            },
            Outbox => {
                match self.accumulator {
                    Some(v) => self.output.push(v),
                    None => return Err(StepError::OutputNil),
                }
            },
            CopyFrom(r) => {
                let r = try!(self.deref_target(r));
                let v = try!(self.registers.get(&r).ok_or(StepError::CopyFromNil));
                self.accumulator = Some(*v);
            },
            CopyTo(r) => {
                match self.accumulator {
                    Some(v) => {
                        let r = try!(self.deref_target(r));
                        self.registers.insert(r, v);
                    },
                    None => return Err(StepError::CopyToNil),
                }
            },
            BumpUp(r) => {
                let r = try!(self.deref_target(r));
                let v = match self.registers.get_mut(&r) {
                    None => return Err(StepError::BumpNil),
                    Some(&mut Tile::Number(ref mut v)) => {
                        *v = try!(v.increment());
                        *v
                    },
                    Some(&mut Tile::Letter(..)) => return Err(StepError::BumpLetter)
                };
                self.accumulator = Some(Tile::Number(v))
            },
            BumpDown(r) => {
                let r = try!(self.deref_target(r));
                let v = match self.registers.get_mut(&r) {
                    None => return Err(StepError::BumpNil),
                    Some(&mut Tile::Number(ref mut v)) => {
                        *v = try!(v.decrement());
                        *v
                    },
                    Some(&mut Tile::Letter(..)) => return Err(StepError::BumpLetter)
                };
                self.accumulator = Some(Tile::Number(v))
            },
            Add(r) => {
                let r = try!(self.deref_target(r));
                let v = match (self.accumulator, self.registers.get(&r)) {
                    (None, _) => return Err(StepError::AddToNil),
                    (_, None) => return Err(StepError::AddWithNil),
                    (Some(Tile::Number(a)), Some(&Tile::Number(v))) => try!(a.add(v)),
                    _ => return Err(StepError::AddWithLetter),
                };
                self.accumulator = Some(Tile::Number(v));
            },
            Sub(r) => {
                let r = try!(self.deref_target(r));
                let v = match (self.accumulator, self.registers.get(&r)) {
                    (None, _) => return Err(StepError::SubFromNil),
                    (_, None) => return Err(StepError::SubWithNil),
                    (Some(Tile::Number(a)), Some(&Tile::Number(v))) => try!(a.sub(v)),
                    (Some(Tile::Letter(a)), Some(&Tile::Letter(v))) => {
                        let a = try!(NumberValue::from_char(a));
                        let v = try!(NumberValue::from_char(v));
                        try!(a.sub(v))
                    },
                    _ => return Err(StepError::SubCrossTypes),
                };
                self.accumulator = Some(Tile::Number(v))
            },
            Jump(i) => {
                self.pc = i;
                return Ok(());
            },
            JumpIfZero(i) => {
                match self.accumulator {
                    None => return Err(StepError::JumpZeroNil),
                    Some(v) => match v {
                        Tile::Number(v) if v.is_zero() => {
                            self.pc = i;
                            return Ok(());
                        }
                        Tile::Number(..) |
                        Tile::Letter(..) => {}, // noop
                    }
                }
            },
            JumpIfNegative(i) => {
                match self.accumulator {
                    None => return Err(StepError::JumpNegativeNil),
                    Some(v) => match v {
                        Tile::Number(v) if v.is_negative() => {
                            self.pc = i;
                            return Ok(());
                        }
                        Tile::Number(..) |
                        Tile::Letter(..) => {}, // noop
                    }
                }
            },
            NoOp => {},
        }

        self.pc += 1;

        if self.pc >= self.program.len() {
            Err(StepError::EndOfProgram)
        } else {
            Ok(())
        }
    }

    pub fn run(&mut self) -> Result<(), StepError> {
        loop {
            match self.step() {
                Ok(..) => continue,
                Err(StepError::EndOfProgram) => return Ok(()),
                Err(e) => return Err(e),
            }
        }
    }
}
<commit_msg>Remove self-namespacing<commit_after>use std::collections::BTreeMap;

use super::Register;

type AbsoluteIndex = usize;

#[derive(Debug, Copy, Clone)]
pub enum Instruction {
    Inbox,
    Outbox,
    CopyFrom(Register),
    CopyTo(Register),
    BumpUp(Register),
    BumpDown(Register),
    Add(Register),
    Sub(Register),
    Jump(AbsoluteIndex),
    JumpIfZero(AbsoluteIndex),
    JumpIfNegative(AbsoluteIndex),
    NoOp,
}

// Clamped at [-999, 999]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct NumberValue(i16);

impl NumberValue {
    fn from_char(c: char) -> Result<NumberValue, Error> {
        NumberValue::clamp(c as i16)
    }

    fn clamp(v: i16) -> Result<NumberValue, Error> {
        if v > 999 {
            Err(Error::Overflow)
        } else if v < -999 {
            Err(Error::Underflow)
        } else {
            Ok(NumberValue(v))
        }
    }

    fn add(self, other: NumberValue) -> Result<NumberValue, Error> {
        NumberValue::clamp(self.0 + other.0)
    }

    fn sub(self, other: NumberValue) -> Result<NumberValue, Error> {
        NumberValue::clamp(self.0 - other.0)
    }

    fn is_zero(self) -> bool { self.0 == 0 }
    fn is_negative(self) -> bool { self.0 < 0 }

    fn increment(self) -> Result<NumberValue, Error> {
        NumberValue::clamp(self.0 + 1)
    }

    fn decrement(self) -> Result<NumberValue, Error> {
        NumberValue::clamp(self.0 - 1)
    }

    fn into_u8(self) -> u8 {
        self.0 as u8 // Should have error of some kind?
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Tile {
    Number(NumberValue),
    Letter(char),
}

impl Tile {
    pub fn num(i: i16) -> Tile {
        Tile::Number(NumberValue::clamp(i).unwrap())
    }
}

#[derive(Debug, Copy, Clone)]
pub enum Error {
    EndOfProgram,
    IndirectThroughNil,
    IndirectThroughNegative,
    IndirectThroughLetter,
    OutputNil,
    CopyFromNil,
    CopyToNil,
    BumpNil,
    BumpLetter,
    AddWithNil,
    AddToNil,
    AddWithLetter,
    SubFromNil,
    SubWithNil,
    SubCrossTypes,
    JumpZeroNil,
    JumpNegativeNil,
    Underflow,
    Overflow,
}

pub type Input = Vec<Tile>;
pub type Output = Vec<Tile>;
pub type Registers = BTreeMap<u8, Tile>;

#[derive(Debug, Clone)]
pub struct Machine {
    program: Vec<Instruction>,
    input: Input,
    output: Output,
    pc: usize,
    accumulator: Option<Tile>,
    registers: Registers,
}

impl Machine {
    pub fn new<I>(program: I, mut input: Input, registers: Registers) -> Machine
        where I: IntoIterator<Item = Instruction>
    {
        // We want to pop off the front, so flip it around for efficiency.
        input.reverse();

        Machine {
            program: program.into_iter().collect(),
            input: input,
            output: Vec::new(),
            pc: 0,
            accumulator: None,
            registers: registers,
        }
    }

    fn deref_target(&self, r: Register) -> Result<u8, Error> {
        match r {
            Register::Direct(r) => Ok(r),
            Register::Indirect(r) => match self.registers.get(&r) {
                None => Err(Error::IndirectThroughNil),
                Some(&Tile::Number(v)) if v.is_negative() => Err(Error::IndirectThroughNegative),
                Some(&Tile::Number(v)) => Ok(v.into_u8()), // Should have max # registers?
                Some(&Tile::Letter(..)) => Err(Error::IndirectThroughLetter),
            },
        }
    }

    pub fn into_output(self) -> Output { self.output }

    pub fn step(&mut self) -> Result<(), Error> {
        use self::Instruction::*;

        // println!("PC: {}", self.pc);
        // println!("Instr: {:?}", self.program[self.pc]);
        // println!("Acc: {:?}", self.accumulator);

        match self.program[self.pc] {
            Inbox => {
                match self.input.pop() {
                    Some(v) => self.accumulator = Some(v),
                    None => return Err(Error::EndOfProgram),
                }
            },
            Outbox => {
                match self.accumulator {
                    Some(v) => self.output.push(v),
                    None => return Err(Error::OutputNil),
                }
            },
            CopyFrom(r) => {
                let r = try!(self.deref_target(r));
                let v = try!(self.registers.get(&r).ok_or(Error::CopyFromNil));
                self.accumulator = Some(*v);
            },
            CopyTo(r) => {
                match self.accumulator {
                    Some(v) => {
                        let r = try!(self.deref_target(r));
                        self.registers.insert(r, v);
                    },
                    None => return Err(Error::CopyToNil),
                }
            },
            BumpUp(r) => {
                let r = try!(self.deref_target(r));
                let v = match self.registers.get_mut(&r) {
                    None => return Err(Error::BumpNil),
                    Some(&mut Tile::Number(ref mut v)) => {
                        *v = try!(v.increment());
                        *v
                    },
                    Some(&mut Tile::Letter(..)) => return Err(Error::BumpLetter)
                };
                self.accumulator = Some(Tile::Number(v))
            },
            BumpDown(r) => {
                let r = try!(self.deref_target(r));
                let v = match self.registers.get_mut(&r) {
                    None => return Err(Error::BumpNil),
                    Some(&mut Tile::Number(ref mut v)) => {
                        *v = try!(v.decrement());
                        *v
                    },
                    Some(&mut Tile::Letter(..)) => return Err(Error::BumpLetter)
                };
                self.accumulator = Some(Tile::Number(v))
            },
            Add(r) => {
                let r = try!(self.deref_target(r));
                let v = match (self.accumulator, self.registers.get(&r)) {
                    (None, _) => return Err(Error::AddToNil),
                    (_, None) => return Err(Error::AddWithNil),
                    (Some(Tile::Number(a)), Some(&Tile::Number(v))) => try!(a.add(v)),
                    _ => return Err(Error::AddWithLetter),
                };
                self.accumulator = Some(Tile::Number(v));
            },
            Sub(r) => {
                let r = try!(self.deref_target(r));
                let v = match (self.accumulator, self.registers.get(&r)) {
                    (None, _) => return Err(Error::SubFromNil),
                    (_, None) => return Err(Error::SubWithNil),
                    (Some(Tile::Number(a)), Some(&Tile::Number(v))) => try!(a.sub(v)),
                    (Some(Tile::Letter(a)), Some(&Tile::Letter(v))) => {
                        let a = try!(NumberValue::from_char(a));
                        let v = try!(NumberValue::from_char(v));
                        try!(a.sub(v))
                    },
                    _ => return Err(Error::SubCrossTypes),
                };
                self.accumulator = Some(Tile::Number(v))
            },
            Jump(i) => {
                self.pc = i;
                return Ok(());
            },
            JumpIfZero(i) => {
                match self.accumulator {
                    None => return Err(Error::JumpZeroNil),
                    Some(v) => match v {
                        Tile::Number(v) if v.is_zero() => {
                            self.pc = i;
                            return Ok(());
                        }
                        Tile::Number(..) |
                        Tile::Letter(..) => {}, // noop
                    }
                }
            },
            JumpIfNegative(i) => {
                match self.accumulator {
                    None => return Err(Error::JumpNegativeNil),
                    Some(v) => match v {
                        Tile::Number(v) if v.is_negative() => {
                            self.pc = i;
                            return Ok(());
                        }
                        Tile::Number(..) |
                        Tile::Letter(..) => {}, // noop
                    }
                }
            },
            NoOp => {},
        }

        self.pc += 1;

        if self.pc >= self.program.len() {
            Err(Error::EndOfProgram)
        } else {
            Ok(())
        }
    }

    pub fn run(&mut self) -> Result<(), Error> {
        loop {
            match self.step() {
                Ok(..) => continue,
                Err(Error::EndOfProgram) => return Ok(()),
                Err(e) => return Err(e),
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io::{Cursor, Read, Write};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::iter::once;
use std::collections::HashMap;

use tag::Tag;
use error::Error;

///
/// A Roughtime protocol message; a map of u32 tags to arbitrary byte-strings.
///
#[derive(Debug, Clone)]
pub struct RtMessage {
    tags: Vec<Tag>,
    values: Vec<Vec<u8>>,
}

impl RtMessage {
    /// Construct a new RtMessage
    ///
    /// ## Arguments
    ///
    /// * `num_fields` - Reserve space for this many fields.
    ///
    pub fn new(num_fields: u32) -> Self {
        RtMessage {
            tags: Vec::with_capacity(num_fields as usize),
            values: Vec::with_capacity(num_fields as usize),
        }
    }

    /// Construct a new RtMessage from the on-the-wire representation in `bytes`
    ///
    /// ## Arguments
    ///
    /// * `bytes` - On-the-wire representation
    ///
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
        let bytes_len = bytes.len();

        if bytes_len < 4 {
            return Err(Error::MessageTooShort);
        } else if bytes_len % 4 != 0 {
            return Err(Error::InvalidAlignment(bytes_len as u32));
        }

        let mut msg = Cursor::new(bytes);
        let num_tags = msg.read_u32::<LittleEndian>()?;

        match num_tags {
            0 => Ok(RtMessage::new(0)),
            1 => RtMessage::single_tag_message(bytes, &mut msg),
            _ => RtMessage::multi_tag_message(num_tags, bytes, &mut msg),
        }
    }

    /// Internal function to create a single tag message
    fn single_tag_message(bytes: &[u8], msg: &mut Cursor<&[u8]>) -> Result<Self, Error> {
        if bytes.len() < 8 {
            return Err(Error::MessageTooShort);
        }

        let pos = msg.position() as usize;
        msg.set_position((pos + 4) as u64);

        let mut value = Vec::new();
        msg.read_to_end(&mut value)?;

        let tag = Tag::from_wire(&bytes[pos..pos + 4])?;
        let mut rt_msg = RtMessage::new(1);
        rt_msg.add_field(tag, &value)?;

        Ok(rt_msg)
    }

    /// Internal function to create a multiple tag message
    fn multi_tag_message(
        num_tags: u32,
        bytes: &[u8],
        msg: &mut Cursor<&[u8]>,
    ) -> Result<Self, Error> {
        let bytes_len = bytes.len();
        let mut offsets = Vec::with_capacity((num_tags - 1) as usize);

        for _ in 0..num_tags - 1 {
            let offset = msg.read_u32::<LittleEndian>()?;

            if offset % 4 != 0 {
                return Err(Error::InvalidAlignment(offset));
            } else if offset > bytes_len as u32 {
                return Err(Error::InvalidOffsetValue(offset));
            }

            offsets.push(offset as usize);
        }

        let mut buf = [0; 4];
        let mut tags = Vec::with_capacity(num_tags as usize);

        for _ in 0..num_tags {
            if msg.read_exact(&mut buf).is_err() {
                return Err(Error::MessageTooShort);
            }

            let tag = Tag::from_wire(&buf)?;

            if let Some(last_tag) = tags.last() {
                if tag <= *last_tag {
                    return Err(Error::TagNotStrictlyIncreasing(tag));
                }
            }

            tags.push(tag);
        }

        // All offsets are relative to the end of the header,
        // which is our current position
        let header_end = msg.position() as usize;

        // Compute the end of the last value,
        // as an offset from the end of the header
        let msg_end = bytes.len() - header_end;

        let mut rt_msg = RtMessage::new(num_tags);

        for (tag, (value_start, value_end)) in tags.into_iter().zip(
            once(&0)
                .chain(offsets.iter())
                .zip(offsets.iter().chain(once(&msg_end))),
        ) {
            let start_idx = header_end + value_start;
            let end_idx = header_end + value_end;

            if end_idx > bytes_len || start_idx > end_idx {
                return Err(Error::InvalidValueLength(tag, end_idx as u32));
            }

            let value = bytes[start_idx..end_idx].to_vec();
            rt_msg.add_field(tag, &value)?;
        }

        Ok(rt_msg)
    }

    /// Add a field to this `RtMessage`
    ///
    /// ## Arguments
    ///
    /// * `tag` - The [`Tag`](enum.Tag.html) to add. Tags must be added in **strictly
    ///   increasing order**, violating this will result in a
    ///   [`Error::TagNotStrictlyIncreasing`](enum.Error.html).
    ///
    /// * `value` - Value for the tag.
    ///
    pub fn add_field(&mut self, tag: Tag, value: &[u8]) -> Result<(), Error> {
        if let Some(last_tag) = self.tags.last() {
            if tag <= *last_tag {
                return Err(Error::TagNotStrictlyIncreasing(tag));
            }
        }

        self.tags.push(tag);
        self.values.push(value.to_vec());

        Ok(())
    }

    /// Returns the number of tag/value pairs in the message
    pub fn num_fields(&self) -> u32 {
        self.tags.len() as u32
    }

    /// Returns a slice of the tags in the message
    pub fn tags(&self) -> &[Tag] {
        &self.tags
    }

    /// Returns a slice of the values in the message
    pub fn values(&self) -> &[Vec<u8>] {
        &self.values
    }

    /// Converts the message into a `HashMap` mapping each tag to its value
    pub fn into_hash_map(self) -> HashMap<Tag, Vec<u8>> {
        self.tags.into_iter().zip(self.values.into_iter()).collect()
    }

    /// Encode this message into its on-the-wire representation.
    pub fn encode(&self) -> Result<Vec<u8>, Error> {
        let num_tags = self.tags.len();
        let mut out = Vec::with_capacity(self.encoded_size());

        // number of tags
        out.write_u32::<LittleEndian>(num_tags as u32)?;

        // offset(s) to values, IFF there are two or more tags
        if num_tags > 1 {
            let mut offset_sum = self.values[0].len();

            for val in &self.values[1..] {
                out.write_u32::<LittleEndian>(offset_sum as u32)?;
                offset_sum += val.len();
            }
        }

        // write tags
        for tag in &self.tags {
            out.write_all(tag.wire_value())?;
        }

        // write values
        for value in &self.values {
            out.write_all(value)?;
        }

        // check we wrote exactly what we expected
        assert_eq!(out.len(), self.encoded_size(), "unexpected length");

        Ok(out)
    }

    /// Returns the length in bytes of this message's on-the-wire representation.
    pub fn encoded_size(&self) -> usize {
        let num_tags = self.tags.len();
        let tags_size = 4 * num_tags;
        let offsets_size = if num_tags < 2 { 0 } else { 4 * (num_tags - 1) };
        let values_size: usize = self.values.iter().map(|v| v.len()).sum();

        4 + tags_size + offsets_size + values_size
    }

    /// Adds a PAD tag to the end of this message, with a length
    /// set such that the final encoded size of this message is 1KB
    ///
    /// If the encoded size of this message is already >= 1KB,
    /// this method does nothing
    pub fn pad_to_kilobyte(&mut self) {
        let size = self.encoded_size();
        if size >= 1024 {
            return;
        }

        let mut padding_needed = 1024 - size;
        if self.tags.len() == 1 {
            // If we currently only have one tag, adding a padding tag will cause
            // a 32-bit offset values to be written
            padding_needed -= 4;
        }
        padding_needed -= Tag::PAD.wire_value().len();
        let padding = vec![0; padding_needed];

        self.add_field(Tag::PAD, &padding).unwrap();

        assert_eq!(self.encoded_size(), 1024);
    }
}

#[cfg(test)]
mod test {
    use std::io::{Cursor, Read};
    use byteorder::{LittleEndian, ReadBytesExt};
    use message::*;
    use tag::Tag;

    #[test]
    fn empty_message_size() {
        let msg = RtMessage::new(0);

        assert_eq!(msg.num_fields(), 0);
        // Empty message is 4 bytes, a single num_tags value
        assert_eq!(msg.encoded_size(), 4);
    }

    #[test]
    fn single_field_message_size() {
        let mut msg = RtMessage::new(1);
        msg.add_field(Tag::NONC, "1234".as_bytes()).unwrap();

        assert_eq!(msg.num_fields(), 1);
        // Single tag message is 4 (num_tags) + 4 (NONC) + 4 (value)
        assert_eq!(msg.encoded_size(), 12);
    }

    #[test]
    fn two_field_message_size() {
        let mut msg = RtMessage::new(2);
        msg.add_field(Tag::NONC, "1234".as_bytes()).unwrap();
        msg.add_field(Tag::PAD, "abcd".as_bytes()).unwrap();

        assert_eq!(msg.num_fields(), 2);
        // Two tag message
        //   4 num_tags
        //   8 (NONC, PAD) tags
        //   4 PAD offset
        //   8 values
        assert_eq!(msg.encoded_size(), 24);
    }

    #[test]
    fn empty_message_encoding() {
        let msg = RtMessage::new(0);
        let mut encoded = Cursor::new(msg.encode().unwrap());

        assert_eq!(encoded.read_u32::<LittleEndian>().unwrap(), 0);
    }

    #[test]
    fn single_field_message_encoding() {
        let value = vec![b'a'; 64];
        let mut msg = RtMessage::new(1);

        msg.add_field(Tag::CERT, &value).unwrap();

        let mut encoded = Cursor::new(msg.encode().unwrap());

        // num tags
        assert_eq!(encoded.read_u32::<LittleEndian>().unwrap(), 1);

        // CERT tag
        let mut cert = [0u8; 4];
        encoded.read_exact(&mut cert).unwrap();
        assert_eq!(cert, Tag::CERT.wire_value());

        // CERT value
        let mut read_val = vec![0u8; 64];
        encoded.read_exact(&mut read_val).unwrap();
        assert_eq!(value, read_val);

        // Entire message was read
        assert_eq!(encoded.position(), 72);

        // Round-trip single-tag message
        RtMessage::from_bytes(&msg.encode().unwrap()).unwrap();
    }

    #[test]
    fn two_field_message_encoding() {
        let dele_value = vec![b'a'; 24];
        let maxt_value = vec![b'z'; 32];

        let mut msg = RtMessage::new(2);
        msg.add_field(Tag::DELE, &dele_value).unwrap();
        msg.add_field(Tag::MAXT, &maxt_value).unwrap();

        let mut encoded = Cursor::new(msg.encode().unwrap());
        // Wire encoding
        //   4 num_tags
        //   8 (DELE, MAXT) tags
        //   4 MAXT offset
        //  24 DELE value
        //  32 MAXT value

        // num tags
        assert_eq!(encoded.read_u32::<LittleEndian>().unwrap(), 2);

        // Offset past DELE value to start of MAXT value
        assert_eq!(
            encoded.read_u32::<LittleEndian>().unwrap(),
            dele_value.len() as u32
        );

        // DELE tag
        let mut dele = [0u8; 4];
        encoded.read_exact(&mut dele).unwrap();
        assert_eq!(dele, Tag::DELE.wire_value());

        // MAXT tag
        let mut maxt = [0u8; 4];
        encoded.read_exact(&mut maxt).unwrap();
        assert_eq!(maxt, Tag::MAXT.wire_value());

        // DELE value
        let mut read_dele_val = vec![0u8; 24];
        encoded.read_exact(&mut read_dele_val).unwrap();
        assert_eq!(dele_value, read_dele_val);

        // MAXT value
        let mut read_maxt_val = vec![0u8; 32];
        encoded.read_exact(&mut read_maxt_val).unwrap();
        assert_eq!(maxt_value, read_maxt_val);

        // Everything was read
        assert_eq!(encoded.position() as usize, msg.encoded_size());

        // Round-trip multi-tag message
        RtMessage::from_bytes(&msg.encode().unwrap()).unwrap();
    }

    #[test]
    fn from_bytes_zero_tags() {
        let bytes = [0, 0, 0, 0];
        let msg = RtMessage::from_bytes(&bytes).unwrap();

        assert_eq!(msg.num_fields(), 0);
    }

    #[test]
    #[should_panic(expected = "InvalidAlignment")]
    fn from_bytes_offset_past_end_of_message() {
        let mut msg = RtMessage::new(2);
        msg.add_field(Tag::NONC, "1111".as_bytes()).unwrap();
        msg.add_field(Tag::PAD, "aaaaaaaaa".as_bytes()).unwrap();

        let mut bytes = msg.encode().unwrap();
        // set the PAD value offset to beyond end of the message
        bytes[4] = 128;

        RtMessage::from_bytes(&bytes).unwrap();
    }

    #[test]
    #[should_panic(expected = "InvalidAlignment")]
    fn from_bytes_too_few_bytes_for_tags() {
        // Header says two tags (8 bytes) but truncate first tag at 2 bytes
        let bytes = &[0x02, 0, 0, 0, 4, 0, 0, 0, 0, 0];
        RtMessage::from_bytes(bytes).unwrap();
    }

}
<commit_msg>limit message tag count to 1024<commit_after>// Copyright 2017 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::io::{Cursor, Read, Write};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::iter::once;
use std::collections::HashMap;

use tag::Tag;
use error::Error;

///
/// A Roughtime protocol message; a map of u32 tags to arbitrary byte-strings.
///
#[derive(Debug, Clone)]
pub struct RtMessage {
    tags: Vec<Tag>,
    values: Vec<Vec<u8>>,
}

impl RtMessage {
    /// Construct a new RtMessage
    ///
    /// ## Arguments
    ///
    /// * `num_fields` - Reserve space for this many fields.
    ///
    pub fn new(num_fields: u32) -> Self {
        RtMessage {
            tags: Vec::with_capacity(num_fields as usize),
            values: Vec::with_capacity(num_fields as usize),
        }
    }

    /// Construct a new RtMessage from the on-the-wire representation in `bytes`
    ///
    /// ## Arguments
    ///
    /// * `bytes` - On-the-wire representation
    ///
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
        let bytes_len = bytes.len();

        if bytes_len < 4 {
            return Err(Error::MessageTooShort);
        } else if bytes_len % 4 != 0 {
            return Err(Error::InvalidAlignment(bytes_len as u32));
        }

        let mut msg = Cursor::new(bytes);
        let num_tags = msg.read_u32::<LittleEndian>()?;

        match num_tags {
            0 => Ok(RtMessage::new(0)),
            1 => RtMessage::single_tag_message(bytes, &mut msg),
            2 ... 1024 => RtMessage::multi_tag_message(num_tags, bytes, &mut msg),
            _ => Err(Error::InvalidNumTags(num_tags)),
        }
    }

    /// Internal function to create a single tag message
    fn single_tag_message(bytes: &[u8], msg: &mut Cursor<&[u8]>) -> Result<Self, Error> {
        if bytes.len() < 8 {
            return Err(Error::MessageTooShort);
        }

        let pos = msg.position() as usize;
        msg.set_position((pos + 4) as u64);

        let mut value = Vec::new();
        msg.read_to_end(&mut value)?;

        let tag = Tag::from_wire(&bytes[pos..pos + 4])?;
        let mut rt_msg = RtMessage::new(1);
        rt_msg.add_field(tag, &value)?;

        Ok(rt_msg)
    }

    /// Internal function to create a multiple tag message
    fn multi_tag_message(
        num_tags: u32,
        bytes: &[u8],
        msg: &mut Cursor<&[u8]>,
    ) -> Result<Self, Error> {
        let bytes_len = bytes.len();
        let mut offsets = Vec::with_capacity((num_tags - 1) as usize);

        for _ in 0..num_tags - 1 {
            let offset = msg.read_u32::<LittleEndian>()?;

            if offset % 4 != 0 {
                return Err(Error::InvalidAlignment(offset));
            } else if offset > bytes_len as u32 {
                return Err(Error::InvalidOffsetValue(offset));
            }

            offsets.push(offset as usize);
        }

        let mut buf = [0; 4];
        let mut tags = Vec::with_capacity(num_tags as usize);

        for _ in 0..num_tags {
            if msg.read_exact(&mut buf).is_err() {
                return Err(Error::MessageTooShort);
            }

            let tag = Tag::from_wire(&buf)?;

            if let Some(last_tag) = tags.last() {
                if tag <= *last_tag {
                    return Err(Error::TagNotStrictlyIncreasing(tag));
                }
            }

            tags.push(tag);
        }

        // All offsets are relative to the end of the header,
        // which is our current position
        let header_end = msg.position() as usize;

        // Compute the end of the last value,
        // as an offset from the end of the header
        let msg_end = bytes.len() - header_end;

        let mut rt_msg = RtMessage::new(num_tags);

        for (tag, (value_start, value_end)) in tags.into_iter().zip(
            once(&0)
                .chain(offsets.iter())
                .zip(offsets.iter().chain(once(&msg_end))),
        ) {
            let start_idx = header_end + value_start;
            let end_idx = header_end + value_end;

            if end_idx > bytes_len || start_idx > end_idx {
                return Err(Error::InvalidValueLength(tag, end_idx as u32));
            }

            let value = bytes[start_idx..end_idx].to_vec();
            rt_msg.add_field(tag, &value)?;
        }

        Ok(rt_msg)
    }

    /// Add a field to this `RtMessage`
    ///
    /// ## Arguments
    ///
    /// * `tag` - The [`Tag`](enum.Tag.html) to add. Tags must be added in **strictly
    ///   increasing order**, violating this will result in a
    ///   [`Error::TagNotStrictlyIncreasing`](enum.Error.html).
    ///
    /// * `value` - Value for the tag.
    ///
    pub fn add_field(&mut self, tag: Tag, value: &[u8]) -> Result<(), Error> {
        if let Some(last_tag) = self.tags.last() {
            if tag <= *last_tag {
                return Err(Error::TagNotStrictlyIncreasing(tag));
            }
        }

        self.tags.push(tag);
        self.values.push(value.to_vec());

        Ok(())
    }

    /// Returns the number of tag/value pairs in the message
    pub fn num_fields(&self) -> u32 {
        self.tags.len() as u32
    }

    /// Returns a slice of the tags in the message
    pub fn tags(&self) -> &[Tag] {
        &self.tags
    }

    /// Returns a slice of the values in the message
    pub fn values(&self) -> &[Vec<u8>] {
        &self.values
    }

    /// Converts the message into a `HashMap` mapping each tag to its value
    pub fn into_hash_map(self) -> HashMap<Tag, Vec<u8>> {
        self.tags.into_iter().zip(self.values.into_iter()).collect()
    }

    /// Encode this message into its on-the-wire representation.
    pub fn encode(&self) -> Result<Vec<u8>, Error> {
        let num_tags = self.tags.len();
        let mut out = Vec::with_capacity(self.encoded_size());

        // number of tags
        out.write_u32::<LittleEndian>(num_tags as u32)?;

        // offset(s) to values, IFF there are two or more tags
        if num_tags > 1 {
            let mut offset_sum = self.values[0].len();

            for val in &self.values[1..] {
                out.write_u32::<LittleEndian>(offset_sum as u32)?;
                offset_sum += val.len();
            }
        }

        // write tags
        for tag in &self.tags {
            out.write_all(tag.wire_value())?;
        }

        // write values
        for value in &self.values {
            out.write_all(value)?;
        }

        // check we wrote exactly what we expected
        assert_eq!(out.len(), self.encoded_size(), "unexpected length");

        Ok(out)
    }

    /// Returns the length in bytes of this message's on-the-wire representation.
    pub fn encoded_size(&self) -> usize {
        let num_tags = self.tags.len();
        let tags_size = 4 * num_tags;
        let offsets_size = if num_tags < 2 { 0 } else { 4 * (num_tags - 1) };
        let values_size: usize = self.values.iter().map(|v| v.len()).sum();

        4 + tags_size + offsets_size + values_size
    }

    /// Adds a PAD tag to the end of this message, with a length
    /// set such that the final encoded size of this message is 1KB
    ///
    /// If the encoded size of this message is already >= 1KB,
    /// this method does nothing
    pub fn pad_to_kilobyte(&mut self) {
        let size = self.encoded_size();
        if size >= 1024 {
            return;
        }

        let mut padding_needed = 1024 - size;
        if self.tags.len() == 1 {
            // If we currently only have one tag, adding a padding tag will cause
            // a 32-bit offset values to be written
            padding_needed -= 4;
        }
        padding_needed -= Tag::PAD.wire_value().len();
        let padding = vec![0; padding_needed];

        self.add_field(Tag::PAD, &padding).unwrap();

        assert_eq!(self.encoded_size(), 1024);
    }
}

#[cfg(test)]
mod test {
    use std::io::{Cursor, Read};
    use byteorder::{LittleEndian, ReadBytesExt};
    use message::*;
    use tag::Tag;

    #[test]
    fn empty_message_size() {
        let msg = RtMessage::new(0);

        assert_eq!(msg.num_fields(), 0);
        // Empty message is 4 bytes, a single num_tags value
        assert_eq!(msg.encoded_size(), 4);
    }

    #[test]
    fn single_field_message_size() {
        let mut msg = RtMessage::new(1);
        msg.add_field(Tag::NONC, "1234".as_bytes()).unwrap();

        assert_eq!(msg.num_fields(), 1);
        // Single tag message is 4 (num_tags) + 4 (NONC) + 4 (value)
        assert_eq!(msg.encoded_size(), 12);
    }

    #[test]
    fn two_field_message_size() {
        let mut msg = RtMessage::new(2);
        msg.add_field(Tag::NONC, "1234".as_bytes()).unwrap();
        msg.add_field(Tag::PAD, "abcd".as_bytes()).unwrap();

        assert_eq!(msg.num_fields(), 2);
        // Two tag message
        //   4 num_tags
        //   8 (NONC, PAD) tags
        //   4 PAD offset
        //   8 values
        assert_eq!(msg.encoded_size(), 24);
    }

    #[test]
    fn empty_message_encoding() {
        let msg = RtMessage::new(0);
        let mut encoded = Cursor::new(msg.encode().unwrap());

        assert_eq!(encoded.read_u32::<LittleEndian>().unwrap(), 0);
    }

    #[test]
    fn single_field_message_encoding() {
        let value = vec![b'a'; 64];
        let mut msg = RtMessage::new(1);

        msg.add_field(Tag::CERT, &value).unwrap();

        let mut encoded = Cursor::new(msg.encode().unwrap());

        // num tags
        assert_eq!(encoded.read_u32::<LittleEndian>().unwrap(), 1);

        // CERT tag
        let mut cert = [0u8; 4];
        encoded.read_exact(&mut cert).unwrap();
        assert_eq!(cert, Tag::CERT.wire_value());

        // CERT value
        let mut read_val = vec![0u8; 64];
        encoded.read_exact(&mut read_val).unwrap();
        assert_eq!(value, read_val);

        // Entire message was read
        assert_eq!(encoded.position(), 72);

        // Round-trip single-tag message
        RtMessage::from_bytes(&msg.encode().unwrap()).unwrap();
    }

    #[test]
    fn two_field_message_encoding() {
        let dele_value = vec![b'a'; 24];
        let maxt_value = vec![b'z'; 32];

        let mut msg = RtMessage::new(2);
        msg.add_field(Tag::DELE, &dele_value).unwrap();
        msg.add_field(Tag::MAXT, &maxt_value).unwrap();

        let mut encoded = Cursor::new(msg.encode().unwrap());
        // Wire encoding
        //   4 num_tags
        //   8 (DELE, MAXT) tags
        //   4 MAXT offset
        //  24 DELE value
        //  32 MAXT value

        // num tags
        assert_eq!(encoded.read_u32::<LittleEndian>().unwrap(), 2);

        // Offset past DELE value to start of MAXT value
        assert_eq!(
            encoded.read_u32::<LittleEndian>().unwrap(),
            dele_value.len() as u32
        );

        // DELE tag
        let mut dele = [0u8; 4];
        encoded.read_exact(&mut dele).unwrap();
        assert_eq!(dele, Tag::DELE.wire_value());

        // MAXT tag
        let mut maxt = [0u8; 4];
        encoded.read_exact(&mut maxt).unwrap();
        assert_eq!(maxt, Tag::MAXT.wire_value());

        // DELE value
        let mut read_dele_val = vec![0u8; 24];
        encoded.read_exact(&mut read_dele_val).unwrap();
        assert_eq!(dele_value, read_dele_val);

        // MAXT value
        let mut read_maxt_val = vec![0u8; 32];
        encoded.read_exact(&mut read_maxt_val).unwrap();
        assert_eq!(maxt_value, read_maxt_val);

        // Everything was read
        assert_eq!(encoded.position() as usize, msg.encoded_size());

        // Round-trip multi-tag message
        RtMessage::from_bytes(&msg.encode().unwrap()).unwrap();
    }

    #[test]
    fn from_bytes_zero_tags() {
        let bytes = [0, 0, 0, 0];
        let msg = RtMessage::from_bytes(&bytes).unwrap();

        assert_eq!(msg.num_fields(), 0);
    }

    #[test]
    #[should_panic(expected = "InvalidAlignment")]
    fn from_bytes_offset_past_end_of_message() {
        let mut msg = RtMessage::new(2);
        msg.add_field(Tag::NONC, "1111".as_bytes()).unwrap();
        msg.add_field(Tag::PAD, "aaaaaaaaa".as_bytes()).unwrap();

        let mut bytes = msg.encode().unwrap();
        // set the PAD value offset to beyond end of the message
        bytes[4] = 128;

        RtMessage::from_bytes(&bytes).unwrap();
    }

    #[test]
    #[should_panic(expected = "InvalidAlignment")]
    fn from_bytes_too_few_bytes_for_tags() {
        // Header says two tags (8 bytes) but truncate first tag at 2 bytes
        let bytes = &[0x02, 0, 0, 0, 4, 0, 0, 0, 0, 0];
        RtMessage::from_bytes(bytes).unwrap();
    }

}
<|endoftext|>


<commit_before>/*
 * Copyright 2020 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use once_cell::sync::Lazy;
use prometheus::{
    core::Collector, CounterVec, HistogramOpts, HistogramVec, IntCounterVec, Opts, Registry,
    DEFAULT_BUCKETS,
};

pub use prometheus::Result;

/// "event" is used as a label for Metrics that can apply to both Filter
/// `read` and `write` executions.
pub const DIRECTION_LABEL: &str = "event";

/// Label value for [DIRECTION_LABEL] for `read` events
pub const READ_DIRECTION_LABEL: &str = "read";
/// Label value for [DIRECTION_LABEL] for `write` events
pub const WRITE_DIRECTION_LABEL: &str = "write";

/// Returns the [prometheus::Registry] containing all the metrics
/// registered in Quilkin.
pub fn registry() -> &'static Registry {
    static REGISTRY: Lazy<Registry> = Lazy::new(Registry::default);

    &*REGISTRY
}

/// Start the histogram bucket at a quarter of a millisecond, as number below a millisecond are
/// what we are aiming for, but some granularity below a millisecond is useful for performance
/// profiling.
const BUCKET_START: f64 = 0.00025;

const BUCKET_FACTOR: f64 = 2.0;

/// At an exponential factor of 2.0 (BUCKET_FACTOR), 13 iterations gets us to just over 1 second.
/// Any processing that occurs over a second is far too long, so we end bucketing there as we don't
/// care about granularity past 1 second.
const BUCKET_COUNT: usize = 13;

pub(crate) static PROCESSING_TIME: Lazy<HistogramVec> = Lazy::new(|| {
    prometheus::register_histogram_vec_with_registry! {
        prometheus::histogram_opts! {
            "packets_processing_duration_seconds",
            "Total processing time for a packet",
            prometheus::exponential_buckets(BUCKET_START, BUCKET_FACTOR, BUCKET_COUNT).unwrap(),
        },
        &[DIRECTION_LABEL],
        registry(),
    }
    .unwrap()
});

pub(crate) static PACKETS_TOTAL: Lazy<IntCounterVec> = Lazy::new(|| {
    prometheus::register_int_counter_vec_with_registry! {
        prometheus::opts! {
            "packets_total",
            "Total number of packets",
        },
        &[DIRECTION_LABEL],
        registry(),
    }
    .unwrap()
});

pub(crate) static PACKETS_SIZE: Lazy<CounterVec> = Lazy::new(|| {
    prometheus::register_counter_vec_with_registry! {
        prometheus::opts! {
            "packets_size",
            "The total size of received packets",
        },
        &[DIRECTION_LABEL],
        registry(),
    }
    .unwrap()
});

pub(crate) static PACKETS_DROPPED: Lazy<IntCounterVec> = Lazy::new(|| {
    prometheus::register_int_counter_vec_with_registry! {
        prometheus::opts! {
            "packets_dropped",
            "Total number of dropped packets",
        },
        &[DIRECTION_LABEL, "reason"],
        registry(),
    }
    .unwrap()
});

/// Create a generic metrics options.
/// Use [filter_opts] instead if the intended target is a filter.
pub fn opts(name: &str, subsystem: &str, description: &str) -> Opts {
    Opts::new(name, description)
        .namespace("quilkin")
        .subsystem(subsystem)
}

pub fn histogram_opts(
    name: &str,
    subsystem: &str,
    description: &str,
    buckets: Option<Vec<f64>>,
) -> HistogramOpts {
    HistogramOpts {
        common_opts: opts(name, subsystem, description),
        buckets: buckets.unwrap_or_else(|| Vec::from(DEFAULT_BUCKETS as &'static [f64])),
    }
}

/// Create a generic metrics options for a filter.
pub fn filter_opts(name: &str, filter_name: &str, description: &str) -> Opts {
    opts(name, &format!("filter_{filter_name}"), description)
}

/// Registers the current metric collector with the provided registry.
/// Returns an error if a collector with the same name has already
/// been registered.
fn register_metric<T: Collector + Sized + 'static>(
    registry: &Registry,
    collector: T,
) -> Result<()> {
    registry.register(Box::new(collector))
}

pub trait CollectorExt: Collector + Clone + Sized + 'static {
    /// Registers the current metric collector with the provided registry
    /// if not already registered.
    fn register_if_not_exists(self) -> Result<Self> {
        match register_metric(crate::metrics::registry(), self.clone()) {
            Ok(_) | Err(prometheus::Error::AlreadyReg) => Ok(self),
            Err(err) => Err(err),
        }
    }
}

impl<C: Collector + Clone + 'static> CollectorExt for C {}
<commit_msg>Fix metrics namespace not always being present (#596)<commit_after>/*
 * Copyright 2020 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use once_cell::sync::Lazy;
use prometheus::{
    core::Collector, CounterVec, HistogramOpts, HistogramVec, IntCounterVec, Opts, Registry,
    DEFAULT_BUCKETS,
};

pub use prometheus::Result;

/// "event" is used as a label for Metrics that can apply to both Filter
/// `read` and `write` executions.
pub const DIRECTION_LABEL: &str = "event";

/// Label value for [DIRECTION_LABEL] for `read` events
pub const READ_DIRECTION_LABEL: &str = "read";
/// Label value for [DIRECTION_LABEL] for `write` events
pub const WRITE_DIRECTION_LABEL: &str = "write";

/// Returns the [prometheus::Registry] containing all the metrics
/// registered in Quilkin.
pub fn registry() -> &'static Registry {
    static REGISTRY: Lazy<Registry> =
        Lazy::new(|| Registry::new_custom(Some("quilkin".into()), None).unwrap());

    &*REGISTRY
}

/// Start the histogram bucket at a quarter of a millisecond, as number below a millisecond are
/// what we are aiming for, but some granularity below a millisecond is useful for performance
/// profiling.
const BUCKET_START: f64 = 0.00025;

const BUCKET_FACTOR: f64 = 2.0;

/// At an exponential factor of 2.0 (BUCKET_FACTOR), 13 iterations gets us to just over 1 second.
/// Any processing that occurs over a second is far too long, so we end bucketing there as we don't
/// care about granularity past 1 second.
const BUCKET_COUNT: usize = 13;

pub(crate) static PROCESSING_TIME: Lazy<HistogramVec> = Lazy::new(|| {
    prometheus::register_histogram_vec_with_registry! {
        prometheus::histogram_opts! {
            "packets_processing_duration_seconds",
            "Total processing time for a packet",
            prometheus::exponential_buckets(BUCKET_START, BUCKET_FACTOR, BUCKET_COUNT).unwrap(),
        },
        &[DIRECTION_LABEL],
        registry(),
    }
    .unwrap()
});

pub(crate) static PACKETS_TOTAL: Lazy<IntCounterVec> = Lazy::new(|| {
    prometheus::register_int_counter_vec_with_registry! {
        prometheus::opts! {
            "packets_total",
            "Total number of packets",
        },
        &[DIRECTION_LABEL],
        registry(),
    }
    .unwrap()
});

pub(crate) static PACKETS_SIZE: Lazy<CounterVec> = Lazy::new(|| {
    prometheus::register_counter_vec_with_registry! {
        prometheus::opts! {
            "packets_size",
            "The total size of received packets",
        },
        &[DIRECTION_LABEL],
        registry(),
    }
    .unwrap()
});

pub(crate) static PACKETS_DROPPED: Lazy<IntCounterVec> = Lazy::new(|| {
    prometheus::register_int_counter_vec_with_registry! {
        prometheus::opts! {
            "packets_dropped",
            "Total number of dropped packets",
        },
        &[DIRECTION_LABEL, "reason"],
        registry(),
    }
    .unwrap()
});

/// Create a generic metrics options.
/// Use [filter_opts] instead if the intended target is a filter.
pub fn opts(name: &str, subsystem: &str, description: &str) -> Opts {
    Opts::new(name, description).subsystem(subsystem)
}

pub fn histogram_opts(
    name: &str,
    subsystem: &str,
    description: &str,
    buckets: Option<Vec<f64>>,
) -> HistogramOpts {
    HistogramOpts {
        common_opts: opts(name, subsystem, description),
        buckets: buckets.unwrap_or_else(|| Vec::from(DEFAULT_BUCKETS as &'static [f64])),
    }
}

/// Create a generic metrics options for a filter.
pub fn filter_opts(name: &str, filter_name: &str, description: &str) -> Opts {
    opts(name, &format!("filter_{filter_name}"), description)
}

/// Registers the current metric collector with the provided registry.
/// Returns an error if a collector with the same name has already
/// been registered.
fn register_metric<T: Collector + Sized + 'static>(
    registry: &Registry,
    collector: T,
) -> Result<()> {
    registry.register(Box::new(collector))
}

pub trait CollectorExt: Collector + Clone + Sized + 'static {
    /// Registers the current metric collector with the provided registry
    /// if not already registered.
    fn register_if_not_exists(self) -> Result<Self> {
        match register_metric(crate::metrics::registry(), self.clone()) {
            Ok(_) | Err(prometheus::Error::AlreadyReg) => Ok(self),
            Err(err) => Err(err),
        }
    }
}

impl<C: Collector + Clone + 'static> CollectorExt for C {}
<|endoftext|>