<commit_before>use crate::{
    buffer::Buffer,
    channels::Channels,
    connection_status::ConnectionState,
    frames::Frames,
    heartbeat::Heartbeat,
    internal_rpc::InternalRPCHandle,
    protocol::{self, AMQPError, AMQPHardError},
    socket_state::{SocketEvent, SocketState},
    thread::ThreadHandle,
    types::FrameSize,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use std::{
    collections::VecDeque,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
    thread::Builder as ThreadBuilder,
    time::Duration,
};
use reactor_trait::AsyncIOHandle;
use tracing::{error, trace};

const FRAMES_STORAGE: usize = 32;

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

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPCHandle,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    connection_io_loop_handle: ThreadHandle,
    stream: Pin<Box<dyn AsyncIOHandle + Send>>,
    status: Status,
    frame_size: FrameSize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(FrameSize, Option<PromiseResolver<()>>)>,
}

impl IoLoop {
    pub(crate) async fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: Pin<Box<dyn AsyncIOHandle + Send>>,
        heartbeat: Heartbeat,
    ) -> Result<Self> {
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE,
            configuration.frame_max(),
        );

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

    fn readable_waker(&self) -> Waker {
        let socket_state_handle = self.socket_state.handle();
        waker_fn::waker_fn(move || socket_state_handle.send(SocketEvent::Readable))
    }

    fn writable_waker(&self) -> Waker {
        let socket_state_handle = self.socket_state.handle();
        waker_fn::waker_fn(move || socket_state_handle.send(SocketEvent::Writable))
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max();
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer
                .grow(FRAMES_STORAGE * self.frame_size as usize);
            self.send_buffer
                .grow(FRAMES_STORAGE * self.frame_size as usize);
            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.heartbeat.start();
            }
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        match self.status {
            Status::Initial => self.finish_setup(),
            Status::Connected => Ok(true),
            Status::Stop => 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();
        let handle = self.connection_io_loop_handle.clone();
        handle.register(
            ThreadBuilder::new()
                .name("lapin-io-loop".to_owned())
                .spawn(move || {
                    let readable_waker = self.readable_waker();
                    let mut readable_context = Context::from_waker(&readable_waker);
                    let writable_waker = self.writable_waker();
                    let mut writable_context = Context::from_waker(&writable_waker);
                    let noop_waker = waker_fn::waker_fn(|| {});
                    let mut noop_context = Context::from_waker(&noop_waker);
                    while self.should_continue() {
                        if let Err(err) = self.run(
                            &mut readable_context,
                            &mut writable_context,
                            &mut noop_context,
                        ) {
                            self.critical_error(err)?;
                        }
                    }
                    self.internal_rpc.stop();
                    self.heartbeat.cancel();
                    Ok(())
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_socket_events(&mut self) {
        self.socket_state.poll_events();
    }

    fn check_connection_state(&mut self) {
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
    }

    fn run(
        &mut self,
        readable_context: &mut Context<'_>,
        writable_context: &mut Context<'_>,
        noop_context: &mut Context<'_>,
    ) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events();
        if !self.ensure_setup()? {
            return Ok(());
        }
        self.check_connection_state();
        trace!(
            can_read=%self.socket_state.readable(),
            can_write=%self.socket_state.writable(),
            has_data=%self.has_data(),
            "io_loop do_run",
        );
        if !self.can_read() && !self.can_write() && self.should_continue() {
            self.socket_state.wait();
        }
        self.poll_socket_events();
        self.attempt_flush(noop_context)?;
        self.write(writable_context)?;
        self.check_connection_state();
        if self.should_continue() {
            self.read(readable_context)?;
        }
        self.handle_frames()?;
        self.check_connection_state();
        trace!(
            can_read=%self.socket_state.readable(),
            can_write=%self.socket_state.writable(),
            has_data=%self.has_data(),
            status=?self.status,
            "io_loop do_run done",
        );
        Ok(())
    }

    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()));
            }
        }
        Err(error)
    }

    fn attempt_flush(&mut self, noop_context: &mut Context<'_>) -> Result<()> {
        let res = self.flush(noop_context);
        self.handle_io_result(res)
    }

    fn handle_io_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self.socket_state.handle_io_result(result) {
            error!(error=?e, "error doing IO");
            self.critical_error(e)?;
        }
        Ok(())
    }

    fn flush(&mut self, writable_context: &mut Context<'_>) -> Result<()> {
        let res = self.stream.as_mut().poll_flush(writable_context)?;
        self.socket_state.handle_write_poll(res);
        Ok(())
    }

    fn write(&mut self, writable_context: &mut Context<'_>) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream(writable_context);
            self.handle_io_result(res)?;
        }
        Ok(())
    }

    fn read(&mut self, readable_context: &mut Context<'_>) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream(readable_context);
            self.handle_io_result(res)?;
        }
        Ok(())
    }

    fn write_to_stream(&mut self, writable_context: &mut Context<'_>) -> Result<()> {
        self.flush(writable_context)?;
        self.serialize()?;

        let res = self
            .send_buffer
            .poll_write_to(writable_context, Pin::new(&mut self.stream))?;

        if let Some(sz) = self.socket_state.handle_write_poll(res) {
            if sz > 0 {
                self.heartbeat.update_last_write();

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

                let mut written = sz as FrameSize;
                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(writable_context)?;
            } else {
                error!("Socket was writable but we wrote 0, marking as wouldblock");
                self.socket_state.handle_write_poll::<()>(Poll::Pending);
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self, readable_context: &mut Context<'_>) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let res = self
                    .receive_buffer
                    .poll_read_from(readable_context, Pin::new(&mut self.stream))?;

                if let Some(sz) = self.socket_state.handle_read_poll(res) {
                    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.socket_state.handle_read_poll(Poll::Pending);
                    }
                }
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!(%next_msg, "will write to buffer");
            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 as FrameSize, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!(error=?e, "error generating frame");
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        Ok(())
    }

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

    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!(bytes = consumed, "received large frame");
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!(error=?e, "parse error");
                    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::InternalRPCHandle,
    protocol::{self, AMQPError, AMQPHardError},
    socket_state::{SocketEvent, SocketState},
    thread::ThreadHandle,
    types::FrameSize,
    Configuration, ConnectionStatus, Error, PromiseResolver, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError};
use reactor_trait::AsyncIOHandle;
use std::{
    collections::VecDeque,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
    thread::Builder as ThreadBuilder,
    time::Duration,
};
use tracing::{error, trace};

const FRAMES_STORAGE: usize = 32;

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

pub struct IoLoop {
    connection_status: ConnectionStatus,
    configuration: Configuration,
    channels: Channels,
    internal_rpc: InternalRPCHandle,
    frames: Frames,
    heartbeat: Heartbeat,
    socket_state: SocketState,
    connection_io_loop_handle: ThreadHandle,
    stream: Pin<Box<dyn AsyncIOHandle + Send>>,
    status: Status,
    frame_size: FrameSize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    serialized_frames: VecDeque<(FrameSize, Option<PromiseResolver<()>>)>,
}

impl IoLoop {
    pub(crate) async fn new(
        connection_status: ConnectionStatus,
        configuration: Configuration,
        channels: Channels,
        internal_rpc: InternalRPCHandle,
        frames: Frames,
        socket_state: SocketState,
        connection_io_loop_handle: ThreadHandle,
        stream: Pin<Box<dyn AsyncIOHandle + Send>>,
        heartbeat: Heartbeat,
    ) -> Result<Self> {
        let frame_size = std::cmp::max(
            protocol::constants::FRAME_MIN_SIZE,
            configuration.frame_max(),
        );

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

    fn readable_waker(&self) -> Waker {
        let socket_state_handle = self.socket_state.handle();
        waker_fn::waker_fn(move || socket_state_handle.send(SocketEvent::Readable))
    }

    fn writable_waker(&self) -> Waker {
        let socket_state_handle = self.socket_state.handle();
        waker_fn::waker_fn(move || socket_state_handle.send(SocketEvent::Writable))
    }

    fn finish_setup(&mut self) -> Result<bool> {
        if self.connection_status.connected() {
            let frame_max = self.configuration.frame_max();
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer
                .grow(FRAMES_STORAGE * self.frame_size as usize);
            self.send_buffer
                .grow(FRAMES_STORAGE * self.frame_size as usize);
            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.heartbeat.start();
            }
            self.status = Status::Connected;
        }
        Ok(true)
    }

    fn ensure_setup(&mut self) -> Result<bool> {
        match self.status {
            Status::Initial => self.finish_setup(),
            Status::Connected => Ok(true),
            Status::Stop => 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();
        let handle = self.connection_io_loop_handle.clone();
        handle.register(
            ThreadBuilder::new()
                .name("lapin-io-loop".to_owned())
                .spawn(move || {
                    let readable_waker = self.readable_waker();
                    let mut readable_context = Context::from_waker(&readable_waker);
                    let writable_waker = self.writable_waker();
                    let mut writable_context = Context::from_waker(&writable_waker);
                    let noop_waker = waker_fn::waker_fn(|| {});
                    let mut noop_context = Context::from_waker(&noop_waker);
                    while self.should_continue() {
                        if let Err(err) = self.run(
                            &mut readable_context,
                            &mut writable_context,
                            &mut noop_context,
                        ) {
                            self.critical_error(err)?;
                        }
                    }
                    self.internal_rpc.stop();
                    self.heartbeat.cancel();
                    Ok(())
                })?,
        );
        waker.wake();
        Ok(())
    }

    fn poll_socket_events(&mut self) {
        self.socket_state.poll_events();
    }

    fn check_connection_state(&mut self) {
        if self.connection_status.closed() {
            self.status = Status::Stop;
        }
    }

    fn run(
        &mut self,
        readable_context: &mut Context<'_>,
        writable_context: &mut Context<'_>,
        noop_context: &mut Context<'_>,
    ) -> Result<()> {
        trace!("io_loop run");
        self.poll_socket_events();
        if !self.ensure_setup()? {
            return Ok(());
        }
        self.check_connection_state();
        trace!(
            can_read=%self.socket_state.readable(),
            can_write=%self.socket_state.writable(),
            has_data=%self.has_data(),
            "io_loop do_run",
        );
        if !self.can_read() && !self.can_write() && self.should_continue() {
            self.socket_state.wait();
        }
        self.poll_socket_events();
        self.attempt_flush(noop_context)?;
        self.write(writable_context)?;
        self.check_connection_state();
        if self.should_continue() {
            self.read(readable_context)?;
        }
        self.handle_frames()?;
        self.check_connection_state();
        trace!(
            can_read=%self.socket_state.readable(),
            can_write=%self.socket_state.writable(),
            has_data=%self.has_data(),
            status=?self.status,
            "io_loop do_run done",
        );
        Ok(())
    }

    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()));
            }
        }
        Err(error)
    }

    fn attempt_flush(&mut self, noop_context: &mut Context<'_>) -> Result<()> {
        let res = self.flush(noop_context);
        self.handle_io_result(res)
    }

    fn handle_io_result(&mut self, result: Result<()>) -> Result<()> {
        if let Err(e) = self.socket_state.handle_io_result(result) {
            error!(error=?e, "error doing IO");
            self.critical_error(e)?;
        }
        Ok(())
    }

    fn flush(&mut self, writable_context: &mut Context<'_>) -> Result<()> {
        let res = self.stream.as_mut().poll_flush(writable_context)?;
        self.socket_state.handle_write_poll(res);
        Ok(())
    }

    fn write(&mut self, writable_context: &mut Context<'_>) -> Result<()> {
        while self.can_write() {
            let res = self.write_to_stream(writable_context);
            self.handle_io_result(res)?;
        }
        Ok(())
    }

    fn read(&mut self, readable_context: &mut Context<'_>) -> Result<()> {
        while self.can_read() {
            let res = self.read_from_stream(readable_context);
            self.handle_io_result(res)?;
        }
        Ok(())
    }

    fn write_to_stream(&mut self, writable_context: &mut Context<'_>) -> Result<()> {
        self.flush(writable_context)?;
        self.serialize()?;

        let res = self
            .send_buffer
            .poll_write_to(writable_context, Pin::new(&mut self.stream))?;

        if let Some(sz) = self.socket_state.handle_write_poll(res) {
            if sz > 0 {
                self.heartbeat.update_last_write();

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

                let mut written = sz as FrameSize;
                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(writable_context)?;
            } else {
                error!("Socket was writable but we wrote 0, marking as wouldblock");
                self.socket_state.handle_write_poll::<()>(Poll::Pending);
            }
        }
        Ok(())
    }

    fn read_from_stream(&mut self, readable_context: &mut Context<'_>) -> Result<()> {
        match self.connection_status.state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                let res = self
                    .receive_buffer
                    .poll_read_from(readable_context, Pin::new(&mut self.stream))?;

                if let Some(sz) = self.socket_state.handle_read_poll(res) {
                    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.socket_state.handle_read_poll(Poll::Pending);
                    }
                }
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        while let Some((next_msg, resolver)) = self.frames.pop(self.channels.flow()) {
            trace!(%next_msg, "will write to buffer");
            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 as FrameSize, resolver)),
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.frames.retry((next_msg, resolver));
                            break;
                        }
                        e => {
                            error!(error=?e, "error generating frame");
                            self.critical_error(Error::SerialisationError(Arc::new(e)))?;
                        }
                    }
                }
            }
        }
        Ok(())
    }

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

    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!(bytes = consumed, "received large frame");
                    let error = AMQPError::new(
                        AMQPHardError::FRAMEERROR.into(),
                        format!("frame too large: {} bytes", consumed).into(),
                    );
                    self.internal_rpc.close_connection(
                        error.get_id(),
                        error.get_message().to_string(),
                        0,
                        0,
                    );
                    self.critical_error(Error::ProtocolError(error))?;
                }
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if !e.is_incomplete() {
                    error!(error=?e, "parse error");
                    self.critical_error(Error::ParsingError(e))?;
                }
                Ok(None)
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use crate::generate_scoped_event;
use serde::{Deserialize, Serialize};
use std::{
    ffi::CString,
    mem,
    mem::MaybeUninit,
    ops::Deref,
    os::windows::io::{AsRawHandle, RawHandle},
    ptr,
    ptr::null,
    time::Duration,
};
use win_util::{SecurityAttributes, SelfRelativeSecurityDescriptor};
use winapi::{
    shared::{
        minwindef::{DWORD, FALSE, TRUE},
        winerror::WAIT_TIMEOUT,
    },
    um::{
        handleapi::DuplicateHandle,
        processthreadsapi::GetCurrentProcess,
        synchapi::{CreateEventA, OpenEventA, ResetEvent, SetEvent, WaitForSingleObject},
        winbase::WAIT_FAILED,
        winnt::{DUPLICATE_SAME_ACCESS, EVENT_MODIFY_STATE, HANDLE},
    },
};

use super::{errno_result, Error, RawDescriptor, Result};
use crate::{
    descriptor::{AsRawDescriptor, FromRawDescriptor, IntoRawDescriptor, SafeDescriptor},
    Event as CrateEvent,
};

/// A safe wrapper around Windows synchapi methods used to mimic Linux eventfd (man 2 eventfd).
/// Since the eventfd isn't using "EFD_SEMAPHORE", we don't need to keep count so we can just use
/// events.
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Event {
    event_handle: SafeDescriptor,
}

pub trait EventExt {
    fn reset(&self) -> Result<()>;
}

impl EventExt for CrateEvent {
    fn reset(&self) -> Result<()> {
        self.0.reset()
    }
}

/// Wrapper around the return value of doing a read on an EventFd which distinguishes between
/// getting a valid count of the number of times the eventfd has been written to and timing out
/// waiting for the count to be non-zero.
#[derive(Debug, PartialEq, Eq)]
pub enum EventReadResult {
    Count(u64),
    Timeout,
}

impl Event {
    pub fn new_with_manual_reset(manual_reset: bool) -> Result<Event> {
        let handle = unsafe {
            CreateEventA(
                SecurityAttributes::new_with_security_descriptor(
                    SelfRelativeSecurityDescriptor::get_singleton(),
                    /* inherit= */ false,
                )
                .as_mut(),
                if manual_reset { TRUE } else { FALSE },
                FALSE, // initial state = unsignalled
                null(),
            )
        };
        if handle.is_null() {
            return errno_result();
        }
        Ok(Event {
            event_handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
        })
    }

    pub fn create_event_with_name(name: &str) -> Result<Event> {
        let event_str = CString::new(String::from(name)).unwrap();
        let handle = unsafe {
            CreateEventA(
                SecurityAttributes::new_with_security_descriptor(
                    SelfRelativeSecurityDescriptor::get_singleton(),
                    /* inherit= */ false,
                )
                .as_mut(),
                FALSE, // manual_reset = false
                FALSE, // initial state = unsignalled
                event_str.as_ptr(),
            )
        };
        if handle.is_null() {
            return errno_result();
        }
        Ok(Event {
            event_handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
        })
    }

    pub fn new() -> Result<Event> {
        // Require manual reset
        Event::new_with_manual_reset(true)
    }

    pub fn open(name: &str) -> Result<Event> {
        let event_str = CString::new(String::from(name)).unwrap();
        let handle = unsafe { OpenEventA(EVENT_MODIFY_STATE, FALSE, event_str.as_ptr()) };
        if handle.is_null() {
            return errno_result();
        }
        Ok(Event {
            event_handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
        })
    }

    pub fn new_auto_reset() -> Result<Event> {
        Event::new_with_manual_reset(false)
    }

    pub fn write(&self, _v: u64) -> Result<()> {
        let event_result = unsafe { SetEvent(self.event_handle.as_raw_descriptor()) };
        if event_result == 0 {
            return errno_result();
        }
        Ok(())
    }

    pub fn read(&self) -> Result<u64> {
        let read_result = self.read_timeout(Duration::new(std::i64::MAX as u64, 0));
        match read_result {
            Ok(EventReadResult::Count(c)) => Ok(c),
            Ok(EventReadResult::Timeout) => Err(Error::new(WAIT_TIMEOUT)),
            Err(e) => Err(e),
        }
    }

    pub fn reset(&self) -> Result<()> {
        let res = unsafe { ResetEvent(self.event_handle.as_raw_descriptor()) };
        if res == 0 {
            errno_result()
        } else {
            Ok(())
        }
    }

    /// Blocks for a maximum of `timeout` duration until the the event is signaled. If
    /// a timeout does not occur then the count is returned as a EventReadResult::Count(1),
    /// and the event resets. If a timeout does occur then this function will return
    /// EventReadResult::Timeout.
    pub fn read_timeout(&self, timeout: Duration) -> Result<EventReadResult> {
        let wait_result = unsafe {
            WaitForSingleObject(
                self.event_handle.as_raw_descriptor(),
                timeout.as_millis() as DWORD,
            )
        };

        // We are using an infinite timeout so we can ignore WAIT_ABANDONED
        match wait_result {
            WAIT_FAILED => errno_result(),
            WAIT_TIMEOUT => Ok(EventReadResult::Timeout),
            _ => {
                // Safe because self manages the handle and we know it was valid as it
                // was just successfully waited upon. It is safe to reset a non manual reset event as well.
                match unsafe { ResetEvent(self.event_handle.as_raw_descriptor()) } {
                    0 => errno_result(),
                    _ => Ok(EventReadResult::Count(1)),
                }
            }
        }
    }

    pub fn try_clone(&self) -> Result<Event> {
        let mut event_clone: HANDLE = MaybeUninit::uninit().as_mut_ptr();
        let duplicate_result = unsafe {
            DuplicateHandle(
                GetCurrentProcess(),
                self.event_handle.as_raw_descriptor(),
                GetCurrentProcess(),
                &mut event_clone,
                0,
                0,
                DUPLICATE_SAME_ACCESS,
            )
        };
        if duplicate_result == 0 {
            return errno_result();
        }
        Ok(unsafe { Event::from_raw_descriptor(event_clone) })
    }
}

impl AsRawDescriptor for Event {
    fn as_raw_descriptor(&self) -> RawDescriptor {
        self.event_handle.as_raw_descriptor()
    }
}

impl FromRawDescriptor for Event {
    unsafe fn from_raw_descriptor(descriptor: RawDescriptor) -> Self {
        Event {
            event_handle: SafeDescriptor::from_raw_descriptor(descriptor),
        }
    }
}

impl AsRawHandle for Event {
    fn as_raw_handle(&self) -> RawHandle {
        self.as_raw_descriptor()
    }
}

impl IntoRawDescriptor for Event {
    fn into_raw_descriptor(self) -> RawDescriptor {
        self.event_handle.into_raw_descriptor()
    }
}

// Event is safe for send & Sync despite containing a raw handle to its
// file mapping object. As long as the instance to Event stays alive, this
// pointer will be a valid handle.
unsafe impl Send for Event {}
unsafe impl Sync for Event {}

generate_scoped_event!(Event);

#[cfg(test)]
mod tests {
    use super::*;
    use winapi::{
        shared::winerror::WAIT_TIMEOUT,
        um::winbase::{INFINITE, WAIT_OBJECT_0},
    };

    #[test]
    fn new() {
        Event::new().unwrap();
    }

    #[test]
    fn read_write() {
        let evt = Event::new().unwrap();
        evt.write(55).unwrap();
        assert_eq!(evt.read(), Ok(1));
    }

    #[test]
    fn read_write_auto_reset() {
        let evt = Event::new_auto_reset().unwrap();
        evt.write(55).unwrap();

        // Wait for the notification.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), INFINITE) };
        assert_eq!(result, WAIT_OBJECT_0);

        // The notification should have reset since we already received it.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), 0) };
        assert_eq!(result, WAIT_TIMEOUT);
    }

    #[test]
    fn read_write_notifies_until_read() {
        let evt = Event::new().unwrap();
        evt.write(55).unwrap();

        // Wait for the notification.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), INFINITE) };
        assert_eq!(result, WAIT_OBJECT_0);

        // The notification should still be active because read wasn't called.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), 0) };
        assert_eq!(result, WAIT_OBJECT_0);

        // Read and ensure the notification has cleared.
        evt.read().expect("Failed to read event.");
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), 0) };
        assert_eq!(result, WAIT_TIMEOUT);
    }

    #[test]
    fn clone() {
        let evt = Event::new().unwrap();
        let evt_clone = evt.try_clone().unwrap();
        evt.write(923).unwrap();
        assert_eq!(evt_clone.read(), Ok(1));
    }

    #[test]
    fn timeout() {
        let evt = Event::new().expect("failed to create event");
        assert_eq!(
            evt.read_timeout(Duration::from_millis(1))
                .expect("failed to read from event with timeout"),
            EventReadResult::Timeout
        );
    }

    #[test]
    fn scoped_event() {
        let scoped_evt = ScopedEvent::new().unwrap();
        let evt_clone: Event = scoped_evt.try_clone().unwrap();
        drop(scoped_evt);
        assert_eq!(evt_clone.read(), Ok(1));
    }

    #[test]
    fn eventfd_from_scoped_event() {
        let scoped_evt = ScopedEvent::new().unwrap();
        let evt: Event = scoped_evt.into();
        evt.write(1).unwrap();
    }
}
<commit_msg>base: add few more windows specific Event functions<commit_after>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use crate::generate_scoped_event;
use serde::{Deserialize, Serialize};
use std::{
    ffi::CString,
    mem,
    mem::MaybeUninit,
    ops::Deref,
    os::windows::io::{AsRawHandle, RawHandle},
    ptr,
    ptr::null,
    time::Duration,
};
use win_util::{SecurityAttributes, SelfRelativeSecurityDescriptor};
use winapi::{
    shared::{
        minwindef::{DWORD, FALSE, TRUE},
        winerror::WAIT_TIMEOUT,
    },
    um::{
        handleapi::DuplicateHandle,
        processthreadsapi::GetCurrentProcess,
        synchapi::{CreateEventA, OpenEventA, ResetEvent, SetEvent, WaitForSingleObject},
        winbase::WAIT_FAILED,
        winnt::{DUPLICATE_SAME_ACCESS, EVENT_MODIFY_STATE, HANDLE},
    },
};

use super::{errno_result, Error, RawDescriptor, Result};
use crate::{
    descriptor::{AsRawDescriptor, FromRawDescriptor, IntoRawDescriptor, SafeDescriptor},
    Event as CrateEvent,
};

/// A safe wrapper around Windows synchapi methods used to mimic Linux eventfd (man 2 eventfd).
/// Since the eventfd isn't using "EFD_SEMAPHORE", we don't need to keep count so we can just use
/// events.
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Event {
    event_handle: SafeDescriptor,
}

pub trait EventExt {
    fn reset(&self) -> Result<()>;
    fn new_with_manual_reset(manual_reset: bool) -> Result<CrateEvent>;
    fn new_auto_reset() -> Result<CrateEvent>;
    fn open(name: &str) -> Result<CrateEvent>;
    fn create_event_with_name(name: &str) -> Result<CrateEvent>;
}

impl EventExt for CrateEvent {
    fn reset(&self) -> Result<()> {
        self.0.reset()
    }

    fn new_with_manual_reset(manual_reset: bool) -> Result<CrateEvent> {
        Event::new_with_manual_reset(manual_reset).map(CrateEvent)
    }

    fn new_auto_reset() -> Result<CrateEvent> {
        CrateEvent::new_with_manual_reset(false)
    }

    fn open(name: &str) -> Result<CrateEvent> {
        Event::open(name).map(CrateEvent)
    }

    fn create_event_with_name(name: &str) -> Result<CrateEvent> {
        Event::create_event_with_name(name).map(CrateEvent)
    }
}

/// Wrapper around the return value of doing a read on an EventFd which distinguishes between
/// getting a valid count of the number of times the eventfd has been written to and timing out
/// waiting for the count to be non-zero.
#[derive(Debug, PartialEq, Eq)]
pub enum EventReadResult {
    Count(u64),
    Timeout,
}

impl Event {
    pub fn new_with_manual_reset(manual_reset: bool) -> Result<Event> {
        let handle = unsafe {
            CreateEventA(
                SecurityAttributes::new_with_security_descriptor(
                    SelfRelativeSecurityDescriptor::get_singleton(),
                    /* inherit= */ false,
                )
                .as_mut(),
                if manual_reset { TRUE } else { FALSE },
                FALSE, // initial state = unsignalled
                null(),
            )
        };
        if handle.is_null() {
            return errno_result();
        }
        Ok(Event {
            event_handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
        })
    }

    pub fn create_event_with_name(name: &str) -> Result<Event> {
        let event_str = CString::new(String::from(name)).unwrap();
        let handle = unsafe {
            CreateEventA(
                SecurityAttributes::new_with_security_descriptor(
                    SelfRelativeSecurityDescriptor::get_singleton(),
                    /* inherit= */ false,
                )
                .as_mut(),
                FALSE, // manual_reset = false
                FALSE, // initial state = unsignalled
                event_str.as_ptr(),
            )
        };
        if handle.is_null() {
            return errno_result();
        }
        Ok(Event {
            event_handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
        })
    }

    pub fn new() -> Result<Event> {
        // Require manual reset
        Event::new_with_manual_reset(true)
    }

    pub fn open(name: &str) -> Result<Event> {
        let event_str = CString::new(String::from(name)).unwrap();
        let handle = unsafe { OpenEventA(EVENT_MODIFY_STATE, FALSE, event_str.as_ptr()) };
        if handle.is_null() {
            return errno_result();
        }
        Ok(Event {
            event_handle: unsafe { SafeDescriptor::from_raw_descriptor(handle) },
        })
    }

    pub fn new_auto_reset() -> Result<Event> {
        Event::new_with_manual_reset(false)
    }

    pub fn write(&self, _v: u64) -> Result<()> {
        let event_result = unsafe { SetEvent(self.event_handle.as_raw_descriptor()) };
        if event_result == 0 {
            return errno_result();
        }
        Ok(())
    }

    pub fn read(&self) -> Result<u64> {
        let read_result = self.read_timeout(Duration::new(std::i64::MAX as u64, 0));
        match read_result {
            Ok(EventReadResult::Count(c)) => Ok(c),
            Ok(EventReadResult::Timeout) => Err(Error::new(WAIT_TIMEOUT)),
            Err(e) => Err(e),
        }
    }

    pub fn reset(&self) -> Result<()> {
        let res = unsafe { ResetEvent(self.event_handle.as_raw_descriptor()) };
        if res == 0 {
            errno_result()
        } else {
            Ok(())
        }
    }

    /// Blocks for a maximum of `timeout` duration until the the event is signaled. If
    /// a timeout does not occur then the count is returned as a EventReadResult::Count(1),
    /// and the event resets. If a timeout does occur then this function will return
    /// EventReadResult::Timeout.
    pub fn read_timeout(&self, timeout: Duration) -> Result<EventReadResult> {
        let wait_result = unsafe {
            WaitForSingleObject(
                self.event_handle.as_raw_descriptor(),
                timeout.as_millis() as DWORD,
            )
        };

        // We are using an infinite timeout so we can ignore WAIT_ABANDONED
        match wait_result {
            WAIT_FAILED => errno_result(),
            WAIT_TIMEOUT => Ok(EventReadResult::Timeout),
            _ => {
                // Safe because self manages the handle and we know it was valid as it
                // was just successfully waited upon. It is safe to reset a non manual reset event as well.
                match unsafe { ResetEvent(self.event_handle.as_raw_descriptor()) } {
                    0 => errno_result(),
                    _ => Ok(EventReadResult::Count(1)),
                }
            }
        }
    }

    pub fn try_clone(&self) -> Result<Event> {
        let mut event_clone: HANDLE = MaybeUninit::uninit().as_mut_ptr();
        let duplicate_result = unsafe {
            DuplicateHandle(
                GetCurrentProcess(),
                self.event_handle.as_raw_descriptor(),
                GetCurrentProcess(),
                &mut event_clone,
                0,
                0,
                DUPLICATE_SAME_ACCESS,
            )
        };
        if duplicate_result == 0 {
            return errno_result();
        }
        Ok(unsafe { Event::from_raw_descriptor(event_clone) })
    }
}

impl AsRawDescriptor for Event {
    fn as_raw_descriptor(&self) -> RawDescriptor {
        self.event_handle.as_raw_descriptor()
    }
}

impl FromRawDescriptor for Event {
    unsafe fn from_raw_descriptor(descriptor: RawDescriptor) -> Self {
        Event {
            event_handle: SafeDescriptor::from_raw_descriptor(descriptor),
        }
    }
}

impl AsRawHandle for Event {
    fn as_raw_handle(&self) -> RawHandle {
        self.as_raw_descriptor()
    }
}

impl IntoRawDescriptor for Event {
    fn into_raw_descriptor(self) -> RawDescriptor {
        self.event_handle.into_raw_descriptor()
    }
}

// Event is safe for send & Sync despite containing a raw handle to its
// file mapping object. As long as the instance to Event stays alive, this
// pointer will be a valid handle.
unsafe impl Send for Event {}
unsafe impl Sync for Event {}

generate_scoped_event!(Event);

#[cfg(test)]
mod tests {
    use super::*;
    use winapi::{
        shared::winerror::WAIT_TIMEOUT,
        um::winbase::{INFINITE, WAIT_OBJECT_0},
    };

    #[test]
    fn new() {
        Event::new().unwrap();
    }

    #[test]
    fn read_write() {
        let evt = Event::new().unwrap();
        evt.write(55).unwrap();
        assert_eq!(evt.read(), Ok(1));
    }

    #[test]
    fn read_write_auto_reset() {
        let evt = Event::new_auto_reset().unwrap();
        evt.write(55).unwrap();

        // Wait for the notification.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), INFINITE) };
        assert_eq!(result, WAIT_OBJECT_0);

        // The notification should have reset since we already received it.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), 0) };
        assert_eq!(result, WAIT_TIMEOUT);
    }

    #[test]
    fn read_write_notifies_until_read() {
        let evt = Event::new().unwrap();
        evt.write(55).unwrap();

        // Wait for the notification.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), INFINITE) };
        assert_eq!(result, WAIT_OBJECT_0);

        // The notification should still be active because read wasn't called.
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), 0) };
        assert_eq!(result, WAIT_OBJECT_0);

        // Read and ensure the notification has cleared.
        evt.read().expect("Failed to read event.");
        let result = unsafe { WaitForSingleObject(evt.as_raw_descriptor(), 0) };
        assert_eq!(result, WAIT_TIMEOUT);
    }

    #[test]
    fn clone() {
        let evt = Event::new().unwrap();
        let evt_clone = evt.try_clone().unwrap();
        evt.write(923).unwrap();
        assert_eq!(evt_clone.read(), Ok(1));
    }

    #[test]
    fn timeout() {
        let evt = Event::new().expect("failed to create event");
        assert_eq!(
            evt.read_timeout(Duration::from_millis(1))
                .expect("failed to read from event with timeout"),
            EventReadResult::Timeout
        );
    }

    #[test]
    fn scoped_event() {
        let scoped_evt = ScopedEvent::new().unwrap();
        let evt_clone: Event = scoped_evt.try_clone().unwrap();
        drop(scoped_evt);
        assert_eq!(evt_clone.read(), Ok(1));
    }

    #[test]
    fn eventfd_from_scoped_event() {
        let scoped_evt = ScopedEvent::new().unwrap();
        let evt: Event = scoped_evt.into();
        evt.write(1).unwrap();
    }
}
<|endoftext|>


<commit_before>extern crate bls;
extern crate boolean_bitfield;
extern crate ethereum_types;
extern crate ssz;

pub mod active_state;
pub mod attestation_record;
pub mod beacon_block;
pub mod chain_config;
pub mod crosslink_record;
pub mod crystallized_state;
pub mod shard_and_committee;
pub mod special_record;
pub mod validator_record;
pub mod validator_registration;

use self::boolean_bitfield::BooleanBitfield;
use self::ethereum_types::{
    H256,
    H160,
    U256
};
use std::collections::HashMap;

pub use active_state::ActiveState;
pub use attestation_record::AttestationRecord;
pub use beacon_block::BeaconBlock;
pub use chain_config::ChainConfig;
pub use crosslink_record::CrosslinkRecord;
pub use crystallized_state::CrystallizedState;
pub use shard_and_committee::ShardAndCommittee;
pub use special_record::{SpecialRecord, SpecialRecordKind};
pub use validator_record::{ValidatorRecord, ValidatorStatus};
pub use validator_registration::ValidatorRegistration;

pub type Hash256 = H256;
pub type Address = H160;
pub type EthBalance = U256;
pub type Bitfield = boolean_bitfield::BooleanBitfield;
pub type BitfieldError = boolean_bitfield::Error;

/// Maps a (slot, shard_id) to attestation_indices.
pub type AttesterMap = HashMap<(u64, u16), Vec<usize>>;

/// Maps a slot to a block proposer.
pub type ProposerMap = HashMap<u64, usize>;
<commit_msg>Remove warning about unused import<commit_after>extern crate bls;
extern crate boolean_bitfield;
extern crate ethereum_types;
extern crate ssz;

pub mod active_state;
pub mod attestation_record;
pub mod beacon_block;
pub mod chain_config;
pub mod crosslink_record;
pub mod crystallized_state;
pub mod shard_and_committee;
pub mod special_record;
pub mod validator_record;
pub mod validator_registration;

use self::ethereum_types::{
    H256,
    H160,
    U256
};
use std::collections::HashMap;

pub use active_state::ActiveState;
pub use attestation_record::AttestationRecord;
pub use beacon_block::BeaconBlock;
pub use chain_config::ChainConfig;
pub use crosslink_record::CrosslinkRecord;
pub use crystallized_state::CrystallizedState;
pub use shard_and_committee::ShardAndCommittee;
pub use special_record::{SpecialRecord, SpecialRecordKind};
pub use validator_record::{ValidatorRecord, ValidatorStatus};
pub use validator_registration::ValidatorRegistration;

pub type Hash256 = H256;
pub type Address = H160;
pub type EthBalance = U256;
pub type Bitfield = boolean_bitfield::BooleanBitfield;
pub type BitfieldError = boolean_bitfield::Error;

/// Maps a (slot, shard_id) to attestation_indices.
pub type AttesterMap = HashMap<(u64, u16), Vec<usize>>;

/// Maps a slot to a block proposer.
pub type ProposerMap = HashMap<u64, usize>;
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    cmp,
    fs::OpenOptions,
    panic,
    path::{Path, PathBuf},
};

use either::Either;
use serde_json::{from_reader, Value};
use uuid::Uuid;

use devicemapper::{
    devnode_to_devno, Bytes, Device, DmDevice, DmName, LinearDev, LinearDevTargetParams,
    LinearTargetParams, Sectors, TargetLine, IEC,
};

use crate::engine::strat_engine::{
    device::{blkdev_size, wipe_sectors},
    dm::get_dm,
    tests::{logger::init_logger, util::clean_up},
};

pub struct RealTestDev {
    dev: Either<PathBuf, LinearDev>,
}

impl RealTestDev {
    /// Construct a new test device for the given path.
    /// Wipe initial MiB to clear metadata.
    pub fn new(dev: Either<PathBuf, LinearDev>) -> RealTestDev {
        let test_dev = RealTestDev { dev };
        wipe_sectors(test_dev.as_path(), Sectors(0), Bytes(IEC::Mi).sectors()).unwrap();
        test_dev
    }

    /// Get the device node of the device.
    fn as_path(&self) -> PathBuf {
        self.dev.as_ref().either(|p| p.clone(), |l| l.devnode())
    }

    /// Teardown a real test dev
    fn teardown(self) {
        wipe_sectors(&self.as_path(), Sectors(0), Bytes(IEC::Mi).sectors()).unwrap();
        if let Some(mut ld) = self.dev.right() {
            ld.teardown(get_dm()).unwrap();
        }
    }
}

/// Shorthand for specifying devices to use for tests
/// For the minimum size, if None is given, the default 1 GiB is chosen.
pub enum DeviceLimits {
    /// Use exactly the number of devices specified
    /// The second argument is the minimum size for all devices.
    /// The third argument is the maximum size for all devices.
    Exactly(usize, Option<Sectors>, Option<Sectors>),
    /// Use at least the number of devices specified, but if there are more
    /// devices available, also use the maximum number of devices.
    /// The second argument is the minimum size for all devices.
    /// The third argument is the maximum size for all devices.
    AtLeast(usize, Option<Sectors>, Option<Sectors>),
    /// Use exactly the number of devices specified in the first argument,
    /// and the minimum of the number of devices specified and the number
    /// of devices available in the second argument.
    /// The third argument is the minimum size for all devices.
    /// The fourth argument is the maximum size for all devices.
    #[allow(dead_code)]
    Range(usize, usize, Option<Sectors>, Option<Sectors>),
}

/// Get a list of lists of devices to use for tests.
/// May return an empty list if the request is not satisfiable.
#[allow(clippy::type_complexity)]
fn get_device_runs<'a>(
    limits: &DeviceLimits,
    dev_sizes: &[(&'a Path, Sectors)],
) -> Vec<Vec<(&'a Path, Option<(Sectors, Sectors)>)>> {
    let (lower, maybe_upper, min_size, max_size) = match limits {
        DeviceLimits::Exactly(num, min_size, max_size) => (*num, Some(*num), *min_size, *max_size),
        DeviceLimits::AtLeast(num, min_size, max_size) => (*num, None, *min_size, *max_size),
        DeviceLimits::Range(lower, upper, min_size, max_size) => {
            assert!(lower < upper);
            (*lower, Some(*upper), *min_size, *max_size)
        }
    };

    let min_size = min_size.unwrap_or_else(|| Bytes(IEC::Gi).sectors());

    assert!(max_size.is_none() || Some(min_size) <= max_size);

    // Retain only devices that are larger than min_size
    let (matches, _): (Vec<(&Path, Sectors)>, Vec<(&Path, Sectors)>) =
        dev_sizes.iter().partition(|&&(_, s)| s >= min_size);

    // If there is not a sufficient number of devices large enough to match
    // the lower bound, return an empty vec. TODO: It would be possible to try
    // harder, by merging several devices into a single linear device. If this
    // turns out to be necessary, this is the correct place to do it.
    if lower > matches.len() {
        return vec![];
    }

    // Retain only devices that are less than max_size
    let (mut matches, mut too_large): (Vec<(&Path, Sectors)>, Vec<(&Path, Sectors)>) =
        if max_size.is_none() {
            (matches, vec![])
        } else {
            let max_size = max_size.expect("!max_size.is_none()");
            matches.iter().partition(|&&(_, s)| s <= max_size)
        };

    let avail = matches.len();
    let needed = maybe_upper.unwrap_or(lower);
    let must_generate = if avail > needed { 0 } else { needed - avail };
    let avail_specs = {
        let mut avail_specs = vec![];
        while avail_specs.len() < must_generate && (!too_large.is_empty() || !matches.is_empty()) {
            let (path, size) = if too_large.is_empty() {
                matches.pop().expect("!matches.is_empty()")
            } else {
                too_large.pop().expect("!too_large.is_empty()")
            };
            let mut new_pairs = vec![];
            for i in 0..size / min_size {
                new_pairs.push((path, Some((i * min_size, min_size))))
            }
            avail_specs.extend(new_pairs);
        }
        avail_specs.extend(matches.iter().map(|&(p, _)| (p, None)));
        avail_specs
    };

    let avail = avail_specs.len();

    // If there is not a sufficient number of devices small enough to match
    // the lower bound, return an empty vec.
    if lower > avail {
        return vec![];
    }

    let mut device_lists = vec![];

    device_lists.push(avail_specs.iter().take(lower).cloned().collect::<Vec<_>>());

    if lower != avail {
        match maybe_upper {
            None => device_lists.push(avail_specs.iter().take(avail).cloned().collect::<Vec<_>>()),
            Some(upper) => {
                if lower < upper {
                    device_lists.push(
                        avail_specs
                            .iter()
                            .take(cmp::min(upper, avail))
                            .cloned()
                            .collect::<Vec<_>>(),
                    )
                }
            }
        }
    }

    device_lists
}

/// Make a new LinearDev according to specs
fn make_linear_test_dev(devnode: &Path, start: Sectors, length: Sectors) -> LinearDev {
    let params = LinearTargetParams::new(
        Device::from(devnode_to_devno(devnode).unwrap().unwrap()),
        start,
    );
    let table = vec![TargetLine::new(
        Sectors(0),
        length,
        LinearDevTargetParams::Linear(params),
    )];
    LinearDev::setup(
        get_dm(),
        DmName::new(&format!("stratis_test_{}", Uuid::new_v4())).expect("valid format"),
        None,
        table.clone(),
    )
    .unwrap()
}

/// Run test on real devices, using given constraints. Constraints may result
/// in multiple invocations of the test, with differing numbers of block
/// devices.
pub fn test_with_spec<F>(limits: &DeviceLimits, test: F)
where
    F: Fn(&[&Path]) -> () + panic::RefUnwindSafe,
{
    let file = OpenOptions::new()
        .read(true)
        .open("tests/test_config.json")
        .unwrap();
    let config: Value = from_reader(&file).unwrap();
    let devpaths: Vec<_> = config
        .get("ok_to_destroy_dev_array_key")
        .unwrap()
        .as_array()
        .unwrap()
        .iter()
        .map(|x| Path::new(x.as_str().unwrap()))
        .collect();

    let dev_sizes: Vec<(&Path, Sectors)> = devpaths
        .iter()
        .map(|p| {
            (
                *p,
                blkdev_size(&OpenOptions::new().read(true).open(p).unwrap())
                    .unwrap()
                    .sectors(),
            )
        })
        .collect();

    let runs = get_device_runs(limits, &dev_sizes);

    assert!(!runs.is_empty());

    init_logger();

    for run_paths in runs {
        clean_up().unwrap();

        let mut devices: Vec<_> = run_paths
            .iter()
            .map(|&(path, spec)| match spec {
                Some((start, length)) => {
                    RealTestDev::new(Either::Right(make_linear_test_dev(path, start, length)))
                }
                None => RealTestDev::new(Either::Left(path.to_path_buf())),
            })
            .collect();

        let paths: Vec<PathBuf> = devices.iter().map(|x| x.as_path()).collect();
        let paths: Vec<&Path> = paths.iter().map(|x| x.as_path()).collect();
        let result = panic::catch_unwind(|| test(&paths));
        let tear_down = clean_up();

        result.unwrap();
        tear_down.unwrap();

        for dev in devices.drain(..) {
            dev.teardown();
        }
    }
}
<commit_msg>Fix for unnecessary clone as checked by clippy<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    cmp,
    fs::OpenOptions,
    panic,
    path::{Path, PathBuf},
};

use either::Either;
use serde_json::{from_reader, Value};
use uuid::Uuid;

use devicemapper::{
    devnode_to_devno, Bytes, Device, DmDevice, DmName, LinearDev, LinearDevTargetParams,
    LinearTargetParams, Sectors, TargetLine, IEC,
};

use crate::engine::strat_engine::{
    device::{blkdev_size, wipe_sectors},
    dm::get_dm,
    tests::{logger::init_logger, util::clean_up},
};

pub struct RealTestDev {
    dev: Either<PathBuf, LinearDev>,
}

impl RealTestDev {
    /// Construct a new test device for the given path.
    /// Wipe initial MiB to clear metadata.
    pub fn new(dev: Either<PathBuf, LinearDev>) -> RealTestDev {
        let test_dev = RealTestDev { dev };
        wipe_sectors(test_dev.as_path(), Sectors(0), Bytes(IEC::Mi).sectors()).unwrap();
        test_dev
    }

    /// Get the device node of the device.
    fn as_path(&self) -> PathBuf {
        self.dev.as_ref().either(|p| p.clone(), |l| l.devnode())
    }

    /// Teardown a real test dev
    fn teardown(self) {
        wipe_sectors(&self.as_path(), Sectors(0), Bytes(IEC::Mi).sectors()).unwrap();
        if let Some(mut ld) = self.dev.right() {
            ld.teardown(get_dm()).unwrap();
        }
    }
}

/// Shorthand for specifying devices to use for tests
/// For the minimum size, if None is given, the default 1 GiB is chosen.
pub enum DeviceLimits {
    /// Use exactly the number of devices specified
    /// The second argument is the minimum size for all devices.
    /// The third argument is the maximum size for all devices.
    Exactly(usize, Option<Sectors>, Option<Sectors>),
    /// Use at least the number of devices specified, but if there are more
    /// devices available, also use the maximum number of devices.
    /// The second argument is the minimum size for all devices.
    /// The third argument is the maximum size for all devices.
    AtLeast(usize, Option<Sectors>, Option<Sectors>),
    /// Use exactly the number of devices specified in the first argument,
    /// and the minimum of the number of devices specified and the number
    /// of devices available in the second argument.
    /// The third argument is the minimum size for all devices.
    /// The fourth argument is the maximum size for all devices.
    #[allow(dead_code)]
    Range(usize, usize, Option<Sectors>, Option<Sectors>),
}

/// Get a list of lists of devices to use for tests.
/// May return an empty list if the request is not satisfiable.
#[allow(clippy::type_complexity)]
fn get_device_runs<'a>(
    limits: &DeviceLimits,
    dev_sizes: &[(&'a Path, Sectors)],
) -> Vec<Vec<(&'a Path, Option<(Sectors, Sectors)>)>> {
    let (lower, maybe_upper, min_size, max_size) = match limits {
        DeviceLimits::Exactly(num, min_size, max_size) => (*num, Some(*num), *min_size, *max_size),
        DeviceLimits::AtLeast(num, min_size, max_size) => (*num, None, *min_size, *max_size),
        DeviceLimits::Range(lower, upper, min_size, max_size) => {
            assert!(lower < upper);
            (*lower, Some(*upper), *min_size, *max_size)
        }
    };

    let min_size = min_size.unwrap_or_else(|| Bytes(IEC::Gi).sectors());

    assert!(max_size.is_none() || Some(min_size) <= max_size);

    // Retain only devices that are larger than min_size
    let (matches, _): (Vec<(&Path, Sectors)>, Vec<(&Path, Sectors)>) =
        dev_sizes.iter().partition(|&&(_, s)| s >= min_size);

    // If there is not a sufficient number of devices large enough to match
    // the lower bound, return an empty vec. TODO: It would be possible to try
    // harder, by merging several devices into a single linear device. If this
    // turns out to be necessary, this is the correct place to do it.
    if lower > matches.len() {
        return vec![];
    }

    // Retain only devices that are less than max_size
    let (mut matches, mut too_large): (Vec<(&Path, Sectors)>, Vec<(&Path, Sectors)>) =
        if max_size.is_none() {
            (matches, vec![])
        } else {
            let max_size = max_size.expect("!max_size.is_none()");
            matches.iter().partition(|&&(_, s)| s <= max_size)
        };

    let avail = matches.len();
    let needed = maybe_upper.unwrap_or(lower);
    let must_generate = if avail > needed { 0 } else { needed - avail };
    let avail_specs = {
        let mut avail_specs = vec![];
        while avail_specs.len() < must_generate && (!too_large.is_empty() || !matches.is_empty()) {
            let (path, size) = if too_large.is_empty() {
                matches.pop().expect("!matches.is_empty()")
            } else {
                too_large.pop().expect("!too_large.is_empty()")
            };
            let mut new_pairs = vec![];
            for i in 0..size / min_size {
                new_pairs.push((path, Some((i * min_size, min_size))))
            }
            avail_specs.extend(new_pairs);
        }
        avail_specs.extend(matches.iter().map(|&(p, _)| (p, None)));
        avail_specs
    };

    let avail = avail_specs.len();

    // If there is not a sufficient number of devices small enough to match
    // the lower bound, return an empty vec.
    if lower > avail {
        return vec![];
    }

    let mut device_lists = vec![];

    device_lists.push(avail_specs.iter().take(lower).cloned().collect::<Vec<_>>());

    if lower != avail {
        match maybe_upper {
            None => device_lists.push(avail_specs.iter().take(avail).cloned().collect::<Vec<_>>()),
            Some(upper) => {
                if lower < upper {
                    device_lists.push(
                        avail_specs
                            .iter()
                            .take(cmp::min(upper, avail))
                            .cloned()
                            .collect::<Vec<_>>(),
                    )
                }
            }
        }
    }

    device_lists
}

/// Make a new LinearDev according to specs
fn make_linear_test_dev(devnode: &Path, start: Sectors, length: Sectors) -> LinearDev {
    let params = LinearTargetParams::new(
        Device::from(devnode_to_devno(devnode).unwrap().unwrap()),
        start,
    );
    let table = vec![TargetLine::new(
        Sectors(0),
        length,
        LinearDevTargetParams::Linear(params),
    )];
    LinearDev::setup(
        get_dm(),
        DmName::new(&format!("stratis_test_{}", Uuid::new_v4())).expect("valid format"),
        None,
        table,
    )
    .unwrap()
}

/// Run test on real devices, using given constraints. Constraints may result
/// in multiple invocations of the test, with differing numbers of block
/// devices.
pub fn test_with_spec<F>(limits: &DeviceLimits, test: F)
where
    F: Fn(&[&Path]) -> () + panic::RefUnwindSafe,
{
    let file = OpenOptions::new()
        .read(true)
        .open("tests/test_config.json")
        .unwrap();
    let config: Value = from_reader(&file).unwrap();
    let devpaths: Vec<_> = config
        .get("ok_to_destroy_dev_array_key")
        .unwrap()
        .as_array()
        .unwrap()
        .iter()
        .map(|x| Path::new(x.as_str().unwrap()))
        .collect();

    let dev_sizes: Vec<(&Path, Sectors)> = devpaths
        .iter()
        .map(|p| {
            (
                *p,
                blkdev_size(&OpenOptions::new().read(true).open(p).unwrap())
                    .unwrap()
                    .sectors(),
            )
        })
        .collect();

    let runs = get_device_runs(limits, &dev_sizes);

    assert!(!runs.is_empty());

    init_logger();

    for run_paths in runs {
        clean_up().unwrap();

        let mut devices: Vec<_> = run_paths
            .iter()
            .map(|&(path, spec)| match spec {
                Some((start, length)) => {
                    RealTestDev::new(Either::Right(make_linear_test_dev(path, start, length)))
                }
                None => RealTestDev::new(Either::Left(path.to_path_buf())),
            })
            .collect();

        let paths: Vec<PathBuf> = devices.iter().map(|x| x.as_path()).collect();
        let paths: Vec<&Path> = paths.iter().map(|x| x.as_path()).collect();
        let result = panic::catch_unwind(|| test(&paths));
        let tear_down = clean_up();

        result.unwrap();
        tear_down.unwrap();

        for dev in devices.drain(..) {
            dev.teardown();
        }
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
    thread::sleep,
    time::Duration,
};

use nix::mount::{umount2, MntFlags};

use devicemapper::{DevId, DmFlags, DmName, DmNameBuf, DmOptions, DM};

use crate::{
    engine::strat_engine::{
        cmd::udev_settle,
        device::blkdev_size,
        dm::{get_dm, get_dm_init},
    },
    stratis::StratisResult,
};

mod cleanup_errors {
    use std::fmt;

    use devicemapper::DmError;

    #[derive(Debug)]
    pub enum Error {
        Ioe(std::io::Error),
        Mnt(libmount::mountinfo::ParseError),
        Nix(nix::Error),
        Msg(String),
        Chained(String, Box<Error>),
        Dm(DmError),
    }

    pub type Result<T> = std::result::Result<T, Error>;

    impl From<nix::Error> for Error {
        fn from(err: nix::Error) -> Error {
            Error::Nix(err)
        }
    }

    impl From<libmount::mountinfo::ParseError> for Error {
        fn from(err: libmount::mountinfo::ParseError) -> Error {
            Error::Mnt(err)
        }
    }

    impl From<std::io::Error> for Error {
        fn from(err: std::io::Error) -> Error {
            Error::Ioe(err)
        }
    }

    impl From<String> for Error {
        fn from(err: String) -> Error {
            Error::Msg(err)
        }
    }

    impl From<DmError> for Error {
        fn from(err: DmError) -> Error {
            Error::Dm(err)
        }
    }

    impl fmt::Display for Error {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            match self {
                Error::Ioe(err) => write!(f, "IO error: {}", err),
                Error::Mnt(err) => write!(f, "mountinfo::ParseError: {}", err),
                Error::Nix(err) => write!(f, "Nix error: {}", err),
                Error::Msg(err) => write!(f, "{}", err),
                Error::Chained(msg, err) => write!(f, "{}: {}", msg, err),
                Error::Dm(err) => write!(f, "DM error: {}", err),
            }
        }
    }

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

use self::cleanup_errors::{Error, Result};

/// Attempt to remove all device mapper devices which match the stratis naming convention.
/// FIXME: Current implementation complicated by https://bugzilla.redhat.com/show_bug.cgi?id=1506287
pub fn dm_stratis_devices_remove() -> Result<()> {
    /// One iteration of removing devicemapper devices
    fn one_iteration() -> Result<(bool, Vec<DmNameBuf>)> {
        let mut progress_made = false;
        let mut remain = get_dm()
            .list_devices()
            .map_err(|e| {
                Error::Chained(
                    "failed while listing DM devices, giving up".into(),
                    Box::new(e.into()),
                )
            })?
            .iter()
            .map(|d| &d.0)
            .filter_map(|n| {
                if !n.to_string().starts_with("stratis-1") {
                    None
                } else {
                    match get_dm().device_remove(&DevId::Name(n), &DmOptions::new()) {
                        Ok(_) => {
                            progress_made = true;
                            None
                        }
                        Err(_) => Some(n.to_owned()),
                    }
                }
            })
            .collect::<Vec<_>>();

        // Retries if no progress has been made.
        if !remain.is_empty() && !progress_made {
            remain = remain
                .into_iter()
                .filter(|name| {
                    for _ in 0..3 {
                        match get_dm().device_remove(&DevId::Name(name), &DmOptions::new()) {
                            Ok(_) => {
                                progress_made = true;
                                return false;
                            }
                            Err(e) => {
                                debug!(
                                    "Failed to remove device {} on retry: {}",
                                    name.to_string(),
                                    e
                                );
                                sleep(Duration::from_secs(1));
                            }
                        }
                    }
                    true
                })
                .collect();
        }

        Ok((progress_made, remain))
    }

    /// Do one iteration of removals until progress stops. Return remaining
    /// dm devices.
    fn do_while_progress() -> Result<Vec<DmNameBuf>> {
        let mut result = one_iteration()?;
        while result.0 {
            result = one_iteration()?;
        }
        Ok(result.1)
    }

    || -> Result<()> {
        udev_settle().unwrap();
        get_dm_init().map_err(|_| Error::Msg("Unable to initialize DM".into()))?;
        do_while_progress().and_then(|remain| {
            if !remain.is_empty() {
                Err(format!("Some Stratis DM devices remaining: {:?}", remain).into())
            } else {
                Ok(())
            }
        })
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure removal of all test-generated DM devices".into(),
            Box::new(e),
        )
    })
}

/// Try and un-mount any filesystems that have the name stratis in the mount point, returning
/// immediately on the first one we are unable to unmount.
fn stratis_filesystems_unmount() -> Result<()> {
    || -> Result<()> {
        let mut mount_data = String::new();
        File::open("/proc/self/mountinfo")?.read_to_string(&mut mount_data)?;
        let parser = libmount::mountinfo::Parser::new(mount_data.as_bytes());

        for mount_point in parser
            .filter_map(|x| x.ok())
            .filter_map(|m| m.mount_point.into_owned().into_string().ok())
            .filter(|mp| mp.contains("stratis"))
        {
            umount2(&PathBuf::from(mount_point), MntFlags::MNT_DETACH)?;
        }

        Ok(())
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure all Stratis filesystems were unmounted".into(),
            Box::new(e),
        )
    })
}

/// When a unit test panics we can leave the system in an inconsistent state.  This function
/// tries to clean up by un-mounting any mounted file systems which contain the string
/// "stratis_testing" and then it tries to remove any device mapper tables which are also stratis
/// created.
pub fn clean_up() -> Result<()> {
    stratis_filesystems_unmount().and_then(|_| dm_stratis_devices_remove())
}

pub struct FailDevice {
    backing_device: PathBuf,
    test_device_name: String,
    dm_context: DM,
    size: u64,
}

impl FailDevice {
    pub fn new(backing_device: &Path, test_device_name: &str) -> StratisResult<Self> {
        let dm = DM::new()?;
        let dm_name = DmName::new(test_device_name)?;
        let dev_id = DevId::Name(dm_name);

        let size = {
            let file = File::open(backing_device)?;
            blkdev_size(&file)?
        };

        dm.device_create(dm_name, None, &DmOptions::new())?;
        dm.table_load(
            &dev_id,
            &[(
                0,
                *size.sectors(),
                "linear".to_string(),
                format!("{} 0", backing_device.display()),
            )],
        )?;
        dm.device_suspend(&dev_id, &DmOptions::new())?;

        Ok(FailDevice {
            backing_device: backing_device.to_owned(),
            test_device_name: test_device_name.to_owned(),
            dm_context: dm,
            size: *size.sectors(),
        })
    }

    pub fn as_path(&self) -> PathBuf {
        vec!["/dev/mapper", self.test_device_name.as_str()]
            .into_iter()
            .collect::<PathBuf>()
    }

    pub fn start_failing(&self, num_sectors_after_start: u64) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::new().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[
                (
                    0,
                    num_sectors_after_start,
                    "error".to_string(),
                    String::new(),
                ),
                (
                    num_sectors_after_start,
                    self.size - num_sectors_after_start,
                    "linear".to_string(),
                    format!("{} 0", self.backing_device.display()),
                ),
            ],
        )?;
        self.dm_context.device_suspend(&dev_id, &DmOptions::new())?;

        Ok(())
    }

    pub fn stop_failing(&self) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::new().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[(
                0,
                self.size,
                "linear".to_string(),
                format!("{} 0", self.backing_device.display()),
            )],
        )?;
        self.dm_context.device_suspend(&dev_id, &DmOptions::new())?;

        Ok(())
    }
}

impl Drop for FailDevice {
    fn drop(&mut self) {
        fn drop_fail(dev: &mut FailDevice) -> StratisResult<()> {
            let dev_id = DevId::Name(DmName::new(dev.test_device_name.as_str())?);

            let (dev_info, _) = dev.dm_context.table_status(&dev_id, &DmOptions::new())?;
            if dev_info.flags() & DmFlags::DM_SUSPEND == DmFlags::DM_SUSPEND {
                dev.dm_context.device_suspend(&dev_id, &DmOptions::new())?;
            }
            dev.dm_context.device_remove(&dev_id, &DmOptions::new())?;

            Ok(())
        }

        if let Err(e) = drop_fail(self) {
            warn!(
                "Teardown of test device /dev/mapper/{} failed: {}",
                self.test_device_name, e
            );
        }
    }
}
<commit_msg>Drop libmount ParseError<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
    thread::sleep,
    time::Duration,
};

use nix::mount::{umount2, MntFlags};

use devicemapper::{DevId, DmFlags, DmName, DmNameBuf, DmOptions, DM};

use crate::{
    engine::strat_engine::{
        cmd::udev_settle,
        device::blkdev_size,
        dm::{get_dm, get_dm_init},
    },
    stratis::StratisResult,
};

mod cleanup_errors {
    use std::fmt;

    use devicemapper::DmError;

    #[derive(Debug)]
    pub enum Error {
        Ioe(std::io::Error),
        Nix(nix::Error),
        Msg(String),
        Chained(String, Box<Error>),
        Dm(DmError),
    }

    pub type Result<T> = std::result::Result<T, Error>;

    impl From<nix::Error> for Error {
        fn from(err: nix::Error) -> Error {
            Error::Nix(err)
        }
    }

    impl From<std::io::Error> for Error {
        fn from(err: std::io::Error) -> Error {
            Error::Ioe(err)
        }
    }

    impl From<String> for Error {
        fn from(err: String) -> Error {
            Error::Msg(err)
        }
    }

    impl From<DmError> for Error {
        fn from(err: DmError) -> Error {
            Error::Dm(err)
        }
    }

    impl fmt::Display for Error {
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
            match self {
                Error::Ioe(err) => write!(f, "IO error: {}", err),
                Error::Nix(err) => write!(f, "Nix error: {}", err),
                Error::Msg(err) => write!(f, "{}", err),
                Error::Chained(msg, err) => write!(f, "{}: {}", msg, err),
                Error::Dm(err) => write!(f, "DM error: {}", err),
            }
        }
    }

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

use self::cleanup_errors::{Error, Result};

/// Attempt to remove all device mapper devices which match the stratis naming convention.
/// FIXME: Current implementation complicated by https://bugzilla.redhat.com/show_bug.cgi?id=1506287
pub fn dm_stratis_devices_remove() -> Result<()> {
    /// One iteration of removing devicemapper devices
    fn one_iteration() -> Result<(bool, Vec<DmNameBuf>)> {
        let mut progress_made = false;
        let mut remain = get_dm()
            .list_devices()
            .map_err(|e| {
                Error::Chained(
                    "failed while listing DM devices, giving up".into(),
                    Box::new(e.into()),
                )
            })?
            .iter()
            .map(|d| &d.0)
            .filter_map(|n| {
                if !n.to_string().starts_with("stratis-1") {
                    None
                } else {
                    match get_dm().device_remove(&DevId::Name(n), &DmOptions::new()) {
                        Ok(_) => {
                            progress_made = true;
                            None
                        }
                        Err(_) => Some(n.to_owned()),
                    }
                }
            })
            .collect::<Vec<_>>();

        // Retries if no progress has been made.
        if !remain.is_empty() && !progress_made {
            remain = remain
                .into_iter()
                .filter(|name| {
                    for _ in 0..3 {
                        match get_dm().device_remove(&DevId::Name(name), &DmOptions::new()) {
                            Ok(_) => {
                                progress_made = true;
                                return false;
                            }
                            Err(e) => {
                                debug!(
                                    "Failed to remove device {} on retry: {}",
                                    name.to_string(),
                                    e
                                );
                                sleep(Duration::from_secs(1));
                            }
                        }
                    }
                    true
                })
                .collect();
        }

        Ok((progress_made, remain))
    }

    /// Do one iteration of removals until progress stops. Return remaining
    /// dm devices.
    fn do_while_progress() -> Result<Vec<DmNameBuf>> {
        let mut result = one_iteration()?;
        while result.0 {
            result = one_iteration()?;
        }
        Ok(result.1)
    }

    || -> Result<()> {
        udev_settle().unwrap();
        get_dm_init().map_err(|_| Error::Msg("Unable to initialize DM".into()))?;
        do_while_progress().and_then(|remain| {
            if !remain.is_empty() {
                Err(format!("Some Stratis DM devices remaining: {:?}", remain).into())
            } else {
                Ok(())
            }
        })
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure removal of all test-generated DM devices".into(),
            Box::new(e),
        )
    })
}

/// Try and un-mount any filesystems that have the name stratis in the mount point, returning
/// immediately on the first one we are unable to unmount.
fn stratis_filesystems_unmount() -> Result<()> {
    || -> Result<()> {
        let mut mount_data = String::new();
        File::open("/proc/self/mountinfo")?.read_to_string(&mut mount_data)?;
        let parser = libmount::mountinfo::Parser::new(mount_data.as_bytes());

        for mount_point in parser
            .filter_map(|x| x.ok())
            .filter_map(|m| m.mount_point.into_owned().into_string().ok())
            .filter(|mp| mp.contains("stratis"))
        {
            umount2(&PathBuf::from(mount_point), MntFlags::MNT_DETACH)?;
        }

        Ok(())
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure all Stratis filesystems were unmounted".into(),
            Box::new(e),
        )
    })
}

/// When a unit test panics we can leave the system in an inconsistent state.  This function
/// tries to clean up by un-mounting any mounted file systems which contain the string
/// "stratis_testing" and then it tries to remove any device mapper tables which are also stratis
/// created.
pub fn clean_up() -> Result<()> {
    stratis_filesystems_unmount().and_then(|_| dm_stratis_devices_remove())
}

pub struct FailDevice {
    backing_device: PathBuf,
    test_device_name: String,
    dm_context: DM,
    size: u64,
}

impl FailDevice {
    pub fn new(backing_device: &Path, test_device_name: &str) -> StratisResult<Self> {
        let dm = DM::new()?;
        let dm_name = DmName::new(test_device_name)?;
        let dev_id = DevId::Name(dm_name);

        let size = {
            let file = File::open(backing_device)?;
            blkdev_size(&file)?
        };

        dm.device_create(dm_name, None, &DmOptions::new())?;
        dm.table_load(
            &dev_id,
            &[(
                0,
                *size.sectors(),
                "linear".to_string(),
                format!("{} 0", backing_device.display()),
            )],
        )?;
        dm.device_suspend(&dev_id, &DmOptions::new())?;

        Ok(FailDevice {
            backing_device: backing_device.to_owned(),
            test_device_name: test_device_name.to_owned(),
            dm_context: dm,
            size: *size.sectors(),
        })
    }

    pub fn as_path(&self) -> PathBuf {
        vec!["/dev/mapper", self.test_device_name.as_str()]
            .into_iter()
            .collect::<PathBuf>()
    }

    pub fn start_failing(&self, num_sectors_after_start: u64) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::new().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[
                (
                    0,
                    num_sectors_after_start,
                    "error".to_string(),
                    String::new(),
                ),
                (
                    num_sectors_after_start,
                    self.size - num_sectors_after_start,
                    "linear".to_string(),
                    format!("{} 0", self.backing_device.display()),
                ),
            ],
        )?;
        self.dm_context.device_suspend(&dev_id, &DmOptions::new())?;

        Ok(())
    }

    pub fn stop_failing(&self) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::new().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[(
                0,
                self.size,
                "linear".to_string(),
                format!("{} 0", self.backing_device.display()),
            )],
        )?;
        self.dm_context.device_suspend(&dev_id, &DmOptions::new())?;

        Ok(())
    }
}

impl Drop for FailDevice {
    fn drop(&mut self) {
        fn drop_fail(dev: &mut FailDevice) -> StratisResult<()> {
            let dev_id = DevId::Name(DmName::new(dev.test_device_name.as_str())?);

            let (dev_info, _) = dev.dm_context.table_status(&dev_id, &DmOptions::new())?;
            if dev_info.flags() & DmFlags::DM_SUSPEND == DmFlags::DM_SUSPEND {
                dev.dm_context.device_suspend(&dev_id, &DmOptions::new())?;
            }
            dev.dm_context.device_remove(&dev_id, &DmOptions::new())?;

            Ok(())
        }

        if let Err(e) = drop_fail(self) {
            warn!(
                "Teardown of test device /dev/mapper/{} failed: {}",
                self.test_device_name, e
            );
        }
    }
}
<|endoftext|>


<commit_before>type buf_reader = obj {
  fn read(vec[u8] buf) -> uint;
};

type buf_writer = obj {
  fn write(vec[u8] buf) -> uint;
};

fn mk_buf_reader(str s) -> buf_reader {

  obj fd_reader(int fd) {
    fn read(vec[u8] v) -> uint {
      auto len = _vec.len[u8](v);
      auto buf = _vec.buf[u8](v);
      auto count = os.libc.read(fd, buf, len);
      if (count < 0) {
        log "error filling buffer";
        log sys.rustrt.last_os_error();
        fail;
      } else {
        ret uint(count);
      }
    }
    drop {
      os.libc.close(fd);
    }
  }

  auto fd = os.libc.open(_str.buf(s), 0);
  if (fd < 0) {
    log "error opening file";
    log sys.rustrt.last_os_error();
    fail;
  }
  ret fd_reader(fd);
}
<commit_msg>Modify _io's fd_buf_reader to do something slightly more useful: produce multiple buffers on demand.<commit_after>type buf_reader = unsafe obj {
  fn read() -> vec[u8];
};

fn default_bufsz() -> uint {
  ret uint(4096);
}

fn new_buf() -> vec[u8] {
  let vec[u8] v = vec();
  let uint i = default_bufsz();
  while (i > uint(0)) {
    i -= uint(1);
    v += vec(u8(0));
  }
  // FIXME (issue #93): should be:
  // ret _vec.alloc[u8](default_bufsz());
}

fn new_buf_reader(str s) -> buf_reader {

  unsafe obj fd_buf_reader(int fd, mutable vec[u8] buf) {

    fn read() -> vec[u8] {

      // Ensure our buf is singly-referenced.
      if (_vec.rustrt.refcount[u8](buf) != uint(1)) {
        buf = new_buf();
      }

      auto len = _vec.len[u8](buf);
      auto vbuf = _vec.buf[u8](buf);
      auto count = os.libc.read(fd, vbuf, len);

      if (count < 0) {
        log "error filling buffer";
        log sys.rustrt.last_os_error();
        fail;
      } else {
        ret buf;
      }
    }

    drop {
      os.libc.close(fd);
    }
  }

  auto fd = os.libc.open(_str.buf(s), 0);
  if (fd < 0) {
    log "error opening file";
    log sys.rustrt.last_os_error();
    fail;
  }
  ret fd_buf_reader(fd, new_buf());
}
<|endoftext|>