<commit_before>#[macro_use]
extern crate glium;
extern crate rand;
#[macro_use]
extern crate slog;
extern crate slog_async;
extern crate slog_term;
#[macro_use(define_roles)]
extern crate smithay;
extern crate wayland_server;

mod helpers;

use glium::Surface;
use helpers::{init_shell, GliumDrawer, MyWindowMap};
use slog::{Drain, Logger};
use smithay::backend::graphics::egl::EGLGraphicsBackend;
use smithay::backend::input::{self, Event, InputBackend, InputHandler, KeyboardKeyEvent, PointerButtonEvent,
                              PointerMotionAbsoluteEvent};
use smithay::backend::winit;
use smithay::compositor::{SubsurfaceRole, TraversalAction};
use smithay::compositor::roles::Role;
use smithay::seat::{KbdHandle, PointerHandle, Seat};
use smithay::shm::init_shm_global;
use std::cell::RefCell;
use std::rc::Rc;
use wayland_server::protocol::wl_pointer;

struct WinitInputHandler {
    log: Logger,
    pointer: PointerHandle,
    keyboard: KbdHandle,
    window_map: Rc<RefCell<MyWindowMap>>,
    pointer_location: (f64, f64),
    serial: u32,
}

impl WinitInputHandler {
    fn next_serial(&mut self) -> u32 {
        self.serial += 1;
        self.serial
    }
}

impl InputHandler<winit::WinitInputBackend> for WinitInputHandler {
    fn on_seat_created(&mut self, _: &input::Seat) {
        /* never happens with winit */
    }
    fn on_seat_destroyed(&mut self, _: &input::Seat) {
        /* never happens with winit */
    }
    fn on_seat_changed(&mut self, _: &input::Seat) {
        /* never happens with winit */
    }
    fn on_keyboard_key(&mut self, _: &input::Seat, evt: winit::WinitKeyboardInputEvent) {
        let keycode = evt.key_code() - 8;
        let state = evt.state();
        debug!(self.log, "key"; "keycode" => keycode, "state" => format!("{:?}", state));
        let serial = self.next_serial();
        self.keyboard.input(keycode, state, serial, |_, _| true);
    }
    fn on_pointer_move(&mut self, _: &input::Seat, _: input::UnusedEvent) {
        /* never happens with winit */
    }
    fn on_pointer_move_absolute(&mut self, _: &input::Seat, evt: winit::WinitMouseMovedEvent) {
        // on winit, mouse events are already in pixel coordinates
        let (x, y) = evt.position();
        self.pointer_location = (x, y);
        let serial = self.next_serial();
        let under = self.window_map.borrow().get_surface_under((x, y));
        self.pointer.motion(
            under.as_ref().map(|&(ref s, (x, y))| (s, x, y)),
            serial,
            evt.time(),
        );
    }
    fn on_pointer_button(&mut self, _: &input::Seat, evt: winit::WinitMouseInputEvent) {
        let serial = self.next_serial();
        let button = match evt.button() {
            input::MouseButton::Left => 0x110,
            input::MouseButton::Right => 0x111,
            input::MouseButton::Middle => 0x112,
            input::MouseButton::Other(b) => b as u32,
        };
        let state = match evt.state() {
            input::MouseButtonState::Pressed => {
                // change the keyboard focus
                let under = self.window_map
                    .borrow_mut()
                    .get_surface_and_bring_to_top(self.pointer_location);
                self.keyboard
                    .set_focus(under.as_ref().map(|&(ref s, _)| s), serial);
                wl_pointer::ButtonState::Pressed
            }
            input::MouseButtonState::Released => wl_pointer::ButtonState::Released,
        };
        self.pointer.button(button, state, serial, evt.time());
    }
    fn on_pointer_axis(&mut self, _: &input::Seat, _: winit::WinitMouseWheelEvent) {
        /* not done in this example */
    }
    fn on_touch_down(&mut self, _: &input::Seat, _: winit::WinitTouchStartedEvent) {
        /* not done in this example */
    }
    fn on_touch_motion(&mut self, _: &input::Seat, _: winit::WinitTouchMovedEvent) {
        /* not done in this example */
    }
    fn on_touch_up(&mut self, _: &input::Seat, _: winit::WinitTouchEndedEvent) {
        /* not done in this example */
    }
    fn on_touch_cancel(&mut self, _: &input::Seat, _: winit::WinitTouchCancelledEvent) {
        /* not done in this example */
    }
    fn on_touch_frame(&mut self, _: &input::Seat, _: input::UnusedEvent) {
        /* never happens with winit */
    }
    fn on_input_config_changed(&mut self, _: &mut ()) {
        /* never happens with winit */
    }
}

fn main() {
    // A logger facility, here we use the terminal for this example
    let log = Logger::root(
        slog_async::Async::default(slog_term::term_full().fuse()).fuse(),
        o!(),
    );

    // Initialize a simple backend for testing
    let (renderer, mut input) = winit::init(log.clone()).unwrap();

    let (mut display, mut event_loop) = wayland_server::create_display();

    /*
     * Initialize the globals
     */

    init_shm_global(&mut event_loop, vec![], log.clone());

    let (compositor_token, _shell_state_token, window_map) = init_shell(&mut event_loop, log.clone());

    let (seat_token, _) = Seat::new(&mut event_loop, "winit".into(), log.clone());

    let pointer = event_loop.state().get_mut(&seat_token).add_pointer();
    let keyboard = event_loop
        .state()
        .get_mut(&seat_token)
        .add_keyboard("", "fr", "oss", None, 1000, 500)
        .expect("Failed to initialize the keyboard");

    /*
     * Initialize glium
     */
    let drawer = GliumDrawer::from(renderer);

    input.set_handler(WinitInputHandler {
        log: log.clone(),
        pointer,
        keyboard,
        window_map: window_map.clone(),
        pointer_location: (0.0, 0.0),
        serial: 0,
    });

    /*
     * Add a listening socket:
     */
    let name = display.add_socket_auto().unwrap().into_string().unwrap();
    println!("Listening on socket: {}", name);

    loop {
        input.dispatch_new_events().unwrap();

        let mut frame = drawer.draw();
        frame.clear(None, Some((0.8, 0.8, 0.9, 1.0)), false, None, None);
        // redraw the frame, in a simple but inneficient way
        {
            let screen_dimensions = drawer.get_framebuffer_dimensions();
            window_map
                .borrow()
                .with_windows_from_bottom_to_top(|toplevel_surface, initial_place| {
                    if let Some(wl_surface) = toplevel_surface.get_surface() {
                        // this surface is a root of a subsurface tree that needs to be drawn
                        compositor_token
                            .with_surface_tree_upward(
                                wl_surface,
                                initial_place,
                                |_surface, attributes, role, &(mut x, mut y)| {
                                    if let Some((ref contents, (w, h))) = attributes.user_data.buffer {
                                        // there is actually something to draw !
                                        if let Ok(subdata) = Role::<SubsurfaceRole>::data(role) {
                                            x += subdata.x;
                                            y += subdata.y;
                                        }
                                        drawer.render(
                                            &mut frame,
                                            contents,
                                            (w, h),
                                            (x, y),
                                            screen_dimensions,
                                        );
                                        TraversalAction::DoChildren((x, y))
                                    } else {
                                        // we are not display, so our children are neither
                                        TraversalAction::SkipChildren
                                    }
                                },
                            )
                            .unwrap();
                    }
                });
        }
        frame.finish().unwrap();

        event_loop.dispatch(Some(16)).unwrap();
        display.flush_clients();

        window_map.borrow_mut().refresh();
    }
}
<commit_msg>fix examples<commit_after>#[macro_use]
extern crate glium;
extern crate rand;
#[macro_use]
extern crate slog;
extern crate slog_async;
extern crate slog_term;
#[macro_use(define_roles)]
extern crate smithay;
extern crate wayland_server;

mod helpers;

use glium::Surface;
use helpers::{init_shell, GliumDrawer, MyWindowMap};
use slog::{Drain, Logger};
use smithay::backend::graphics::egl::EGLGraphicsBackend;
use smithay::backend::input::{self, Event, InputBackend, InputHandler, KeyboardKeyEvent, PointerButtonEvent,
                              PointerMotionAbsoluteEvent};
use smithay::backend::winit;
use smithay::compositor::{SubsurfaceRole, TraversalAction};
use smithay::compositor::roles::Role;
use smithay::seat::{KeyboardHandle, PointerHandle, Seat};
use smithay::shm::init_shm_global;
use std::cell::RefCell;
use std::rc::Rc;
use wayland_server::protocol::wl_pointer;

struct WinitInputHandler {
    log: Logger,
    pointer: PointerHandle,
    keyboard: KeyboardHandle,
    window_map: Rc<RefCell<MyWindowMap>>,
    pointer_location: (f64, f64),
    serial: u32,
}

impl WinitInputHandler {
    fn next_serial(&mut self) -> u32 {
        self.serial += 1;
        self.serial
    }
}

impl InputHandler<winit::WinitInputBackend> for WinitInputHandler {
    fn on_seat_created(&mut self, _: &input::Seat) {
        /* never happens with winit */
    }
    fn on_seat_destroyed(&mut self, _: &input::Seat) {
        /* never happens with winit */
    }
    fn on_seat_changed(&mut self, _: &input::Seat) {
        /* never happens with winit */
    }
    fn on_keyboard_key(&mut self, _: &input::Seat, evt: winit::WinitKeyboardInputEvent) {
        let keycode = evt.key_code() - 8;
        let state = evt.state();
        debug!(self.log, "key"; "keycode" => keycode, "state" => format!("{:?}", state));
        let serial = self.next_serial();
        self.keyboard.input(keycode, state, serial, |_, _| true);
    }
    fn on_pointer_move(&mut self, _: &input::Seat, _: input::UnusedEvent) {
        /* never happens with winit */
    }
    fn on_pointer_move_absolute(&mut self, _: &input::Seat, evt: winit::WinitMouseMovedEvent) {
        // on winit, mouse events are already in pixel coordinates
        let (x, y) = evt.position();
        self.pointer_location = (x, y);
        let serial = self.next_serial();
        let under = self.window_map.borrow().get_surface_under((x, y));
        self.pointer.motion(
            under.as_ref().map(|&(ref s, (x, y))| (s, x, y)),
            serial,
            evt.time(),
        );
    }
    fn on_pointer_button(&mut self, _: &input::Seat, evt: winit::WinitMouseInputEvent) {
        let serial = self.next_serial();
        let button = match evt.button() {
            input::MouseButton::Left => 0x110,
            input::MouseButton::Right => 0x111,
            input::MouseButton::Middle => 0x112,
            input::MouseButton::Other(b) => b as u32,
        };
        let state = match evt.state() {
            input::MouseButtonState::Pressed => {
                // change the keyboard focus
                let under = self.window_map
                    .borrow_mut()
                    .get_surface_and_bring_to_top(self.pointer_location);
                self.keyboard
                    .set_focus(under.as_ref().map(|&(ref s, _)| s), serial);
                wl_pointer::ButtonState::Pressed
            }
            input::MouseButtonState::Released => wl_pointer::ButtonState::Released,
        };
        self.pointer.button(button, state, serial, evt.time());
    }
    fn on_pointer_axis(&mut self, _: &input::Seat, _: winit::WinitMouseWheelEvent) {
        /* not done in this example */
    }
    fn on_touch_down(&mut self, _: &input::Seat, _: winit::WinitTouchStartedEvent) {
        /* not done in this example */
    }
    fn on_touch_motion(&mut self, _: &input::Seat, _: winit::WinitTouchMovedEvent) {
        /* not done in this example */
    }
    fn on_touch_up(&mut self, _: &input::Seat, _: winit::WinitTouchEndedEvent) {
        /* not done in this example */
    }
    fn on_touch_cancel(&mut self, _: &input::Seat, _: winit::WinitTouchCancelledEvent) {
        /* not done in this example */
    }
    fn on_touch_frame(&mut self, _: &input::Seat, _: input::UnusedEvent) {
        /* never happens with winit */
    }
    fn on_input_config_changed(&mut self, _: &mut ()) {
        /* never happens with winit */
    }
}

fn main() {
    // A logger facility, here we use the terminal for this example
    let log = Logger::root(
        slog_async::Async::default(slog_term::term_full().fuse()).fuse(),
        o!(),
    );

    // Initialize a simple backend for testing
    let (renderer, mut input) = winit::init(log.clone()).unwrap();

    let (mut display, mut event_loop) = wayland_server::create_display();

    /*
     * Initialize the globals
     */

    init_shm_global(&mut event_loop, vec![], log.clone());

    let (compositor_token, _shell_state_token, window_map) = init_shell(&mut event_loop, log.clone());

    let (seat_token, _) = Seat::new(&mut event_loop, "winit".into(), log.clone());

    let pointer = event_loop.state().get_mut(&seat_token).add_pointer();
    let keyboard = event_loop
        .state()
        .get_mut(&seat_token)
        .add_keyboard("", "fr", "oss", None, 1000, 500)
        .expect("Failed to initialize the keyboard");

    /*
     * Initialize glium
     */
    let drawer = GliumDrawer::from(renderer);

    input.set_handler(WinitInputHandler {
        log: log.clone(),
        pointer,
        keyboard,
        window_map: window_map.clone(),
        pointer_location: (0.0, 0.0),
        serial: 0,
    });

    /*
     * Add a listening socket:
     */
    let name = display.add_socket_auto().unwrap().into_string().unwrap();
    println!("Listening on socket: {}", name);

    loop {
        input.dispatch_new_events().unwrap();

        let mut frame = drawer.draw();
        frame.clear(None, Some((0.8, 0.8, 0.9, 1.0)), false, None, None);
        // redraw the frame, in a simple but inneficient way
        {
            let screen_dimensions = drawer.get_framebuffer_dimensions();
            window_map
                .borrow()
                .with_windows_from_bottom_to_top(|toplevel_surface, initial_place| {
                    if let Some(wl_surface) = toplevel_surface.get_surface() {
                        // this surface is a root of a subsurface tree that needs to be drawn
                        compositor_token
                            .with_surface_tree_upward(
                                wl_surface,
                                initial_place,
                                |_surface, attributes, role, &(mut x, mut y)| {
                                    if let Some((ref contents, (w, h))) = attributes.user_data.buffer {
                                        // there is actually something to draw !
                                        if let Ok(subdata) = Role::<SubsurfaceRole>::data(role) {
                                            x += subdata.x;
                                            y += subdata.y;
                                        }
                                        drawer.render(
                                            &mut frame,
                                            contents,
                                            (w, h),
                                            (x, y),
                                            screen_dimensions,
                                        );
                                        TraversalAction::DoChildren((x, y))
                                    } else {
                                        // we are not display, so our children are neither
                                        TraversalAction::SkipChildren
                                    }
                                },
                            )
                            .unwrap();
                    }
                });
        }
        frame.finish().unwrap();

        event_loop.dispatch(Some(16)).unwrap();
        display.flush_clients();

        window_map.borrow_mut().refresh();
    }
}
<|endoftext|>


<commit_before>use crate::{
    channel_closer::ChannelCloser,
    consumer_canceler::ConsumerCanceler,
    consumer_status::{ConsumerState, ConsumerStatus},
    executor::Executor,
    internal_rpc::InternalRPCHandle,
    message::{Delivery, DeliveryResult},
    types::ShortString,
    BasicProperties, Channel, Error, Result,
};
use crossbeam_channel::{Receiver, Sender};
use futures_core::stream::Stream;
use log::trace;
use parking_lot::Mutex;
use std::{
    fmt,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
};

pub trait ConsumerDelegate: Send + Sync {
    fn on_new_delivery(&self, delivery: DeliveryResult)
        -> Pin<Box<dyn Future<Output = ()> + Send>>;
    fn drop_prefetched_messages(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(async move {})
    }
}

impl<
        F: Future<Output = ()> + Send + 'static,
        DeliveryHandler: Fn(DeliveryResult) -> F + Send + Sync + 'static,
    > ConsumerDelegate for DeliveryHandler
{
    fn on_new_delivery(
        &self,
        delivery: DeliveryResult,
    ) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(self(delivery))
    }
}

/// Continuously consumes message from a Queue.
///
/// A consumer represents a stream of messages created from
/// the [basic.consume](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.consume) AMQP command.
/// It continuously receives messages from the queue, as opposed to the
/// [basic.get](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.get) command, which
/// retrieves only a single message.
///
/// A consumer is obtained by calling [`Channel::basic_consume`] with the queue name.
///
/// New messages from this consumer can be accessed by obtaining the iterator from the consumer.
/// This iterator returns new messages and the associated channel in the form of a
/// [`DeliveryResult`] for as long as the consumer is subscribed to the queue.
///
/// It is also possible to set a delegate to be spawned via [`set_delegate`].
///
/// ## Message acknowledgment
///
/// There are two ways for acknowledging a message:
///
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `true` while obtaining the consumer from
///   [`Channel::basic_consume`], the server implicitely acknowledges each message after it has been
///   sent.
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `false`, a message has to be explicitely
///   acknowledged or rejected with [`Channel::basic_ack`],
///   [`Channel::basic_reject`] or [`Channel::basic_nack`]. See the documentation at [`Delivery`]
///   for further information.
///
/// Also see the RabbitMQ documentation about
/// [Acknowledgement Modes](https://www.rabbitmq.com/consumers.html#acknowledgement-modes).
///
/// ## Consumer Prefetch
///
/// To limit the maximum number of unacknowledged messages arriving, you can call [`Channel::basic_qos`]
/// before creating the consumer.
///
/// Also see the RabbitMQ documentation about
/// [Consumer Prefetch](https://www.rabbitmq.com/consumer-prefetch.html).
///
/// ## Cancel subscription
///
/// To stop receiving messages, call [`Channel::basic_cancel`] with the consumer tag of this
/// consumer.
///
///
/// ## Example
/// ```rust,no_run
/// use lapin::{options::*, types::FieldTable, Connection, ConnectionProperties, Result};
/// use futures_executor::LocalPool;
/// use futures_util::stream::StreamExt;
/// use std::future::Future;
///
/// let addr = std::env::var("AMQP_ADDR").unwrap_or_else(|_| "amqp://127.0.0.1:5672/%2f".into());
///
/// let res: Result<()> = LocalPool::new().run_until(async {
///     let conn = Connection::connect(
///         &addr,
///         ConnectionProperties::default().with_default_executor(8),
///     )
///     .await?;
///     let channel = conn.create_channel().await?;
///     let mut consumer = channel
///         .basic_consume(
///             "hello",
///             "my_consumer",
///             BasicConsumeOptions::default(),
///             FieldTable::default(),
///         )
///         .await?;
///
///     while let Some(delivery) = consumer.next().await {
///         let (channel, delivery) = delivery.expect("error in consumer");
///         channel
///             .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
///             .await?;
///     }
///     Ok(())
/// });
/// ```
///
/// [`Channel::basic_consume`]: ./struct.Channel.html#method.basic_consume
/// [`Channel::basic_qos`]: ./struct.Channel.html#method.basic_qos
/// [`Channel::basic_ack`]: ./struct.Channel.html#method.basic_ack
/// [`Channel::basic_reject`]: ./struct.Channel.html#method.basic_reject
/// [`Channel::basic_nack`]: ./struct.Channel.html#method.basic_nack
/// [`Channel::basic_cancel`]: ./struct.Channel.html#method.basic_cancel
/// [`DeliveryResult`]: ./message/type.DeliveryResult.html
/// [`BasicConsumeOptions::no_ack`]: ./options/struct.BasicConsumeOptions.html#structfield.no_ack
/// [`set_delegate`]: #method.set_delegate
#[derive(Clone)]
pub struct Consumer {
    inner: Arc<Mutex<ConsumerInner>>,
    status: ConsumerStatus,
    channel_closer: Option<Arc<ChannelCloser>>,
    consumer_canceler: Option<Arc<ConsumerCanceler>>,
}

impl Consumer {
    pub(crate) fn new(
        consumer_tag: ShortString,
        executor: Arc<dyn Executor>,
        channel_closer: Option<Arc<ChannelCloser>>,
    ) -> Self {
        let status = ConsumerStatus::default();
        Self {
            inner: Arc::new(Mutex::new(ConsumerInner::new(
                status.clone(),
                consumer_tag,
                executor,
            ))),
            status,
            channel_closer,
            consumer_canceler: None,
        }
    }

    pub(crate) fn external(&self, channel_id: u16, internal_rpc_handle: InternalRPCHandle) -> Self {
        Self {
            inner: self.inner.clone(),
            status: self.status.clone(),
            channel_closer: None,
            consumer_canceler: Some(Arc::new(ConsumerCanceler::new(
                channel_id,
                self.tag().to_string(),
                self.status.clone(),
                internal_rpc_handle,
            ))),
        }
    }

    /// Gets the consumer tag.
    ///
    /// If no consumer tag was specified when obtaining the consumer from the channel,
    /// this contains the server generated consumer tag.
    pub fn tag(&self) -> ShortString {
        self.inner.lock().tag.clone()
    }

    /// Gets the current state of the Consumer.
    pub fn state(&self) -> ConsumerState {
        self.status.state()
    }

    /// Automatically spawns the delegate on the executor for each message.
    ///
    /// Enables parallel handling of the messages.
    pub fn set_delegate<D: ConsumerDelegate + 'static>(&self, delegate: D) -> Result<()> {
        let mut inner = self.inner.lock();
        let mut status = self.status.lock();
        while let Some(delivery) = inner.next_delivery() {
            inner.executor.spawn(delegate.on_new_delivery(delivery))?;
        }
        inner.delegate = Some(Arc::new(Box::new(delegate)));
        status.set_delegate();
        Ok(())
    }

    pub(crate) fn start_new_delivery(&mut self, delivery: Delivery) {
        self.inner.lock().current_message = Some(delivery)
    }

    pub(crate) fn set_delivery_properties(&mut self, properties: BasicProperties) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.properties = properties;
        }
    }

    pub(crate) fn receive_delivery_content(&mut self, payload: Vec<u8>) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.receive_content(payload);
        }
    }

    pub(crate) fn new_delivery_complete(&mut self, channel: Channel) -> Result<()> {
        let mut inner = self.inner.lock();
        if let Some(delivery) = inner.current_message.take() {
            inner.new_delivery(channel, delivery)?;
        }
        Ok(())
    }

    pub(crate) fn drop_prefetched_messages(&self) -> Result<()> {
        self.inner.lock().drop_prefetched_messages()
    }

    pub(crate) fn cancel(&self) -> Result<()> {
        self.inner.lock().cancel()
    }

    pub(crate) fn set_error(&self, error: Error) -> Result<()> {
        self.inner.lock().set_error(error)
    }
}

struct ConsumerInner {
    status: ConsumerStatus,
    current_message: Option<Delivery>,
    deliveries_in: Sender<DeliveryResult>,
    deliveries_out: Receiver<DeliveryResult>,
    task: Option<Waker>,
    tag: ShortString,
    delegate: Option<Arc<Box<dyn ConsumerDelegate>>>,
    executor: Arc<dyn Executor>,
}

pub struct ConsumerIterator {
    receiver: Receiver<DeliveryResult>,
}

impl Iterator for ConsumerIterator {
    type Item = Result<(Channel, Delivery)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok().and_then(Result::transpose)
    }
}

impl IntoIterator for Consumer {
    type Item = Result<(Channel, Delivery)>;
    type IntoIter = ConsumerIterator;

    fn into_iter(self) -> Self::IntoIter {
        ConsumerIterator {
            receiver: self.inner.lock().deliveries_out.clone(),
        }
    }
}

impl fmt::Debug for Consumer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Consumer");
        if let Some(inner) = self.inner.try_lock() {
            debug
                .field("tag", &inner.tag)
                .field("executor", &inner.executor)
                .field("task", &inner.task);
        }
        if let Some(status) = self.status.try_lock() {
            debug.field("state", &status.state());
        }
        debug.finish()
    }
}

impl ConsumerInner {
    fn new(status: ConsumerStatus, consumer_tag: ShortString, executor: Arc<dyn Executor>) -> Self {
        let (sender, receiver) = crossbeam_channel::unbounded();
        Self {
            status,
            current_message: None,
            deliveries_in: sender,
            deliveries_out: receiver,
            task: None,
            tag: consumer_tag,
            delegate: None,
            executor,
        }
    }

    fn next_delivery(&mut self) -> Option<DeliveryResult> {
        self.deliveries_out.try_recv().ok()
    }

    fn new_delivery(&mut self, channel: Channel, delivery: Delivery) -> Result<()> {
        trace!("new_delivery; consumer_tag={}", self.tag);
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor
                .spawn(delegate.on_new_delivery(Ok(Some((channel, delivery)))))?;
        } else {
            self.deliveries_in
                .send(Ok(Some((channel, delivery))))
                .expect("failed to send delivery to consumer");
        }
        if let Some(task) = self.task.as_ref() {
            task.wake_by_ref();
        }
        Ok(())
    }

    fn drop_prefetched_messages(&mut self) -> Result<()> {
        trace!("drop_prefetched_messages; consumer_tag={}", self.tag);
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.drop_prefetched_messages())?;
        }
        while self.next_delivery().is_some() {}
        Ok(())
    }

    fn cancel(&mut self) -> Result<()> {
        trace!("cancel; consumer_tag={}", self.tag);
        let mut status = self.status.lock();
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Ok(None)))?;
        } else {
            self.deliveries_in
                .send(Ok(None))
                .expect("failed to send cancel to consumer");
        }
        if let Some(task) = self.task.take() {
            task.wake();
        }
        status.cancel();
        Ok(())
    }

    fn set_error(&mut self, error: Error) -> Result<()> {
        trace!("set_error; consumer_tag={}", self.tag);
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Err(error)))?;
        } else {
            self.deliveries_in
                .send(Err(error))
                .expect("failed to send error to consumer");
        }
        self.cancel()
    }
}

impl Stream for Consumer {
    type Item = Result<(Channel, Delivery)>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        trace!("consumer poll_next");
        let mut inner = self.inner.lock();
        trace!(
            "consumer poll; acquired inner lock, consumer_tag={}",
            inner.tag
        );
        inner.task = Some(cx.waker().clone());
        if let Some(delivery) = inner.next_delivery() {
            match delivery {
                Ok(Some((channel, delivery))) => {
                    trace!(
                        "delivery; channel={}, consumer_tag={}, delivery_tag={:?}",
                        channel.id(),
                        inner.tag,
                        delivery.delivery_tag
                    );
                    Poll::Ready(Some(Ok((channel, delivery))))
                }
                Ok(None) => {
                    trace!("consumer canceled; consumer_tag={}", inner.tag);
                    Poll::Ready(None)
                }
                Err(error) => Poll::Ready(Some(Err(error))),
            }
        } else {
            trace!("delivery; status=NotReady, consumer_tag={}", inner.tag);
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod futures_tests {
    use super::*;
    use crate::executor::DefaultExecutor;

    use std::task::{Context, Poll};

    use futures_test::task::new_count_waker;
    use futures_util::stream::StreamExt;

    #[test]
    fn stream_on_cancel() {
        let (waker, awoken_count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            Arc::new(DefaultExecutor::default()),
            None,
        );

        assert_eq!(awoken_count.get(), 0);
        assert_eq!(consumer.poll_next_unpin(&mut cx), Poll::Pending);

        consumer.cancel().unwrap();

        assert_eq!(awoken_count.get(), 1);
        assert_eq!(consumer.poll_next_unpin(&mut cx), Poll::Ready(None));
    }

    #[test]
    fn stream_on_error() {
        let (waker, awoken_count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            Arc::new(DefaultExecutor::default()),
            None,
        );

        assert_eq!(awoken_count.get(), 0);
        assert_eq!(consumer.poll_next_unpin(&mut cx), Poll::Pending);

        consumer.set_error(Error::ChannelsLimitReached).unwrap();

        assert_eq!(awoken_count.get(), 1);
        assert_eq!(
            consumer.poll_next_unpin(&mut cx),
            Poll::Ready(Some(Err(Error::ChannelsLimitReached)))
        );
    }
}
<commit_msg>consumer: dont cancel when converting to iterator<commit_after>use crate::{
    channel_closer::ChannelCloser,
    consumer_canceler::ConsumerCanceler,
    consumer_status::{ConsumerState, ConsumerStatus},
    executor::Executor,
    internal_rpc::InternalRPCHandle,
    message::{Delivery, DeliveryResult},
    types::ShortString,
    BasicProperties, Channel, Error, Result,
};
use crossbeam_channel::{Receiver, Sender};
use futures_core::stream::Stream;
use log::trace;
use parking_lot::Mutex;
use std::{
    fmt,
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Waker},
};

pub trait ConsumerDelegate: Send + Sync {
    fn on_new_delivery(&self, delivery: DeliveryResult)
        -> Pin<Box<dyn Future<Output = ()> + Send>>;
    fn drop_prefetched_messages(&self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(async move {})
    }
}

impl<
        F: Future<Output = ()> + Send + 'static,
        DeliveryHandler: Fn(DeliveryResult) -> F + Send + Sync + 'static,
    > ConsumerDelegate for DeliveryHandler
{
    fn on_new_delivery(
        &self,
        delivery: DeliveryResult,
    ) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(self(delivery))
    }
}

/// Continuously consumes message from a Queue.
///
/// A consumer represents a stream of messages created from
/// the [basic.consume](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.consume) AMQP command.
/// It continuously receives messages from the queue, as opposed to the
/// [basic.get](https://www.rabbitmq.com/amqp-0-9-1-quickref.html#basic.get) command, which
/// retrieves only a single message.
///
/// A consumer is obtained by calling [`Channel::basic_consume`] with the queue name.
///
/// New messages from this consumer can be accessed by obtaining the iterator from the consumer.
/// This iterator returns new messages and the associated channel in the form of a
/// [`DeliveryResult`] for as long as the consumer is subscribed to the queue.
///
/// It is also possible to set a delegate to be spawned via [`set_delegate`].
///
/// ## Message acknowledgment
///
/// There are two ways for acknowledging a message:
///
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `true` while obtaining the consumer from
///   [`Channel::basic_consume`], the server implicitely acknowledges each message after it has been
///   sent.
/// * If the flag [`BasicConsumeOptions::no_ack`] is set to `false`, a message has to be explicitely
///   acknowledged or rejected with [`Channel::basic_ack`],
///   [`Channel::basic_reject`] or [`Channel::basic_nack`]. See the documentation at [`Delivery`]
///   for further information.
///
/// Also see the RabbitMQ documentation about
/// [Acknowledgement Modes](https://www.rabbitmq.com/consumers.html#acknowledgement-modes).
///
/// ## Consumer Prefetch
///
/// To limit the maximum number of unacknowledged messages arriving, you can call [`Channel::basic_qos`]
/// before creating the consumer.
///
/// Also see the RabbitMQ documentation about
/// [Consumer Prefetch](https://www.rabbitmq.com/consumer-prefetch.html).
///
/// ## Cancel subscription
///
/// To stop receiving messages, call [`Channel::basic_cancel`] with the consumer tag of this
/// consumer.
///
///
/// ## Example
/// ```rust,no_run
/// use lapin::{options::*, types::FieldTable, Connection, ConnectionProperties, Result};
/// use futures_executor::LocalPool;
/// use futures_util::stream::StreamExt;
/// use std::future::Future;
///
/// let addr = std::env::var("AMQP_ADDR").unwrap_or_else(|_| "amqp://127.0.0.1:5672/%2f".into());
///
/// let res: Result<()> = LocalPool::new().run_until(async {
///     let conn = Connection::connect(
///         &addr,
///         ConnectionProperties::default().with_default_executor(8),
///     )
///     .await?;
///     let channel = conn.create_channel().await?;
///     let mut consumer = channel
///         .basic_consume(
///             "hello",
///             "my_consumer",
///             BasicConsumeOptions::default(),
///             FieldTable::default(),
///         )
///         .await?;
///
///     while let Some(delivery) = consumer.next().await {
///         let (channel, delivery) = delivery.expect("error in consumer");
///         channel
///             .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
///             .await?;
///     }
///     Ok(())
/// });
/// ```
///
/// [`Channel::basic_consume`]: ./struct.Channel.html#method.basic_consume
/// [`Channel::basic_qos`]: ./struct.Channel.html#method.basic_qos
/// [`Channel::basic_ack`]: ./struct.Channel.html#method.basic_ack
/// [`Channel::basic_reject`]: ./struct.Channel.html#method.basic_reject
/// [`Channel::basic_nack`]: ./struct.Channel.html#method.basic_nack
/// [`Channel::basic_cancel`]: ./struct.Channel.html#method.basic_cancel
/// [`DeliveryResult`]: ./message/type.DeliveryResult.html
/// [`BasicConsumeOptions::no_ack`]: ./options/struct.BasicConsumeOptions.html#structfield.no_ack
/// [`set_delegate`]: #method.set_delegate
#[derive(Clone)]
pub struct Consumer {
    inner: Arc<Mutex<ConsumerInner>>,
    status: ConsumerStatus,
    channel_closer: Option<Arc<ChannelCloser>>,
    consumer_canceler: Option<Arc<ConsumerCanceler>>,
}

impl Consumer {
    pub(crate) fn new(
        consumer_tag: ShortString,
        executor: Arc<dyn Executor>,
        channel_closer: Option<Arc<ChannelCloser>>,
    ) -> Self {
        let status = ConsumerStatus::default();
        Self {
            inner: Arc::new(Mutex::new(ConsumerInner::new(
                status.clone(),
                consumer_tag,
                executor,
            ))),
            status,
            channel_closer,
            consumer_canceler: None,
        }
    }

    pub(crate) fn external(&self, channel_id: u16, internal_rpc_handle: InternalRPCHandle) -> Self {
        Self {
            inner: self.inner.clone(),
            status: self.status.clone(),
            channel_closer: None,
            consumer_canceler: Some(Arc::new(ConsumerCanceler::new(
                channel_id,
                self.tag().to_string(),
                self.status.clone(),
                internal_rpc_handle,
            ))),
        }
    }

    /// Gets the consumer tag.
    ///
    /// If no consumer tag was specified when obtaining the consumer from the channel,
    /// this contains the server generated consumer tag.
    pub fn tag(&self) -> ShortString {
        self.inner.lock().tag.clone()
    }

    /// Gets the current state of the Consumer.
    pub fn state(&self) -> ConsumerState {
        self.status.state()
    }

    /// Automatically spawns the delegate on the executor for each message.
    ///
    /// Enables parallel handling of the messages.
    pub fn set_delegate<D: ConsumerDelegate + 'static>(&self, delegate: D) -> Result<()> {
        let mut inner = self.inner.lock();
        let mut status = self.status.lock();
        while let Some(delivery) = inner.next_delivery() {
            inner.executor.spawn(delegate.on_new_delivery(delivery))?;
        }
        inner.delegate = Some(Arc::new(Box::new(delegate)));
        status.set_delegate();
        Ok(())
    }

    pub(crate) fn start_new_delivery(&mut self, delivery: Delivery) {
        self.inner.lock().current_message = Some(delivery)
    }

    pub(crate) fn set_delivery_properties(&mut self, properties: BasicProperties) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.properties = properties;
        }
    }

    pub(crate) fn receive_delivery_content(&mut self, payload: Vec<u8>) {
        if let Some(delivery) = self.inner.lock().current_message.as_mut() {
            delivery.receive_content(payload);
        }
    }

    pub(crate) fn new_delivery_complete(&mut self, channel: Channel) -> Result<()> {
        let mut inner = self.inner.lock();
        if let Some(delivery) = inner.current_message.take() {
            inner.new_delivery(channel, delivery)?;
        }
        Ok(())
    }

    pub(crate) fn drop_prefetched_messages(&self) -> Result<()> {
        self.inner.lock().drop_prefetched_messages()
    }

    pub(crate) fn cancel(&self) -> Result<()> {
        self.inner.lock().cancel()
    }

    pub(crate) fn set_error(&self, error: Error) -> Result<()> {
        self.inner.lock().set_error(error)
    }
}

struct ConsumerInner {
    status: ConsumerStatus,
    current_message: Option<Delivery>,
    deliveries_in: Sender<DeliveryResult>,
    deliveries_out: Receiver<DeliveryResult>,
    task: Option<Waker>,
    tag: ShortString,
    delegate: Option<Arc<Box<dyn ConsumerDelegate>>>,
    executor: Arc<dyn Executor>,
}

pub struct ConsumerIterator {
    receiver: Receiver<DeliveryResult>,
    _consumer_canceler: Option<Arc<ConsumerCanceler>>,
}

impl Iterator for ConsumerIterator {
    type Item = Result<(Channel, Delivery)>;

    fn next(&mut self) -> Option<Self::Item> {
        self.receiver.recv().ok().and_then(Result::transpose)
    }
}

impl IntoIterator for Consumer {
    type Item = Result<(Channel, Delivery)>;
    type IntoIter = ConsumerIterator;

    fn into_iter(self) -> Self::IntoIter {
        ConsumerIterator {
            receiver: self.inner.lock().deliveries_out.clone(),
            _consumer_canceler: self.consumer_canceler.clone(),
        }
    }
}

impl fmt::Debug for Consumer {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut debug = f.debug_struct("Consumer");
        if let Some(inner) = self.inner.try_lock() {
            debug
                .field("tag", &inner.tag)
                .field("executor", &inner.executor)
                .field("task", &inner.task);
        }
        if let Some(status) = self.status.try_lock() {
            debug.field("state", &status.state());
        }
        debug.finish()
    }
}

impl ConsumerInner {
    fn new(status: ConsumerStatus, consumer_tag: ShortString, executor: Arc<dyn Executor>) -> Self {
        let (sender, receiver) = crossbeam_channel::unbounded();
        Self {
            status,
            current_message: None,
            deliveries_in: sender,
            deliveries_out: receiver,
            task: None,
            tag: consumer_tag,
            delegate: None,
            executor,
        }
    }

    fn next_delivery(&mut self) -> Option<DeliveryResult> {
        self.deliveries_out.try_recv().ok()
    }

    fn new_delivery(&mut self, channel: Channel, delivery: Delivery) -> Result<()> {
        trace!("new_delivery; consumer_tag={}", self.tag);
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor
                .spawn(delegate.on_new_delivery(Ok(Some((channel, delivery)))))?;
        } else {
            self.deliveries_in
                .send(Ok(Some((channel, delivery))))
                .expect("failed to send delivery to consumer");
        }
        if let Some(task) = self.task.as_ref() {
            task.wake_by_ref();
        }
        Ok(())
    }

    fn drop_prefetched_messages(&mut self) -> Result<()> {
        trace!("drop_prefetched_messages; consumer_tag={}", self.tag);
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.drop_prefetched_messages())?;
        }
        while self.next_delivery().is_some() {}
        Ok(())
    }

    fn cancel(&mut self) -> Result<()> {
        trace!("cancel; consumer_tag={}", self.tag);
        let mut status = self.status.lock();
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Ok(None)))?;
        } else {
            self.deliveries_in
                .send(Ok(None))
                .expect("failed to send cancel to consumer");
        }
        if let Some(task) = self.task.take() {
            task.wake();
        }
        status.cancel();
        Ok(())
    }

    fn set_error(&mut self, error: Error) -> Result<()> {
        trace!("set_error; consumer_tag={}", self.tag);
        if let Some(delegate) = self.delegate.as_ref() {
            let delegate = delegate.clone();
            self.executor.spawn(delegate.on_new_delivery(Err(error)))?;
        } else {
            self.deliveries_in
                .send(Err(error))
                .expect("failed to send error to consumer");
        }
        self.cancel()
    }
}

impl Stream for Consumer {
    type Item = Result<(Channel, Delivery)>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        trace!("consumer poll_next");
        let mut inner = self.inner.lock();
        trace!(
            "consumer poll; acquired inner lock, consumer_tag={}",
            inner.tag
        );
        inner.task = Some(cx.waker().clone());
        if let Some(delivery) = inner.next_delivery() {
            match delivery {
                Ok(Some((channel, delivery))) => {
                    trace!(
                        "delivery; channel={}, consumer_tag={}, delivery_tag={:?}",
                        channel.id(),
                        inner.tag,
                        delivery.delivery_tag
                    );
                    Poll::Ready(Some(Ok((channel, delivery))))
                }
                Ok(None) => {
                    trace!("consumer canceled; consumer_tag={}", inner.tag);
                    Poll::Ready(None)
                }
                Err(error) => Poll::Ready(Some(Err(error))),
            }
        } else {
            trace!("delivery; status=NotReady, consumer_tag={}", inner.tag);
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod futures_tests {
    use super::*;
    use crate::executor::DefaultExecutor;

    use std::task::{Context, Poll};

    use futures_test::task::new_count_waker;
    use futures_util::stream::StreamExt;

    #[test]
    fn stream_on_cancel() {
        let (waker, awoken_count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            Arc::new(DefaultExecutor::default()),
            None,
        );

        assert_eq!(awoken_count.get(), 0);
        assert_eq!(consumer.poll_next_unpin(&mut cx), Poll::Pending);

        consumer.cancel().unwrap();

        assert_eq!(awoken_count.get(), 1);
        assert_eq!(consumer.poll_next_unpin(&mut cx), Poll::Ready(None));
    }

    #[test]
    fn stream_on_error() {
        let (waker, awoken_count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        let mut consumer = Consumer::new(
            ShortString::from("test-consumer"),
            Arc::new(DefaultExecutor::default()),
            None,
        );

        assert_eq!(awoken_count.get(), 0);
        assert_eq!(consumer.poll_next_unpin(&mut cx), Poll::Pending);

        consumer.set_error(Error::ChannelsLimitReached).unwrap();

        assert_eq!(awoken_count.get(), 1);
        assert_eq!(
            consumer.poll_next_unpin(&mut cx),
            Poll::Ready(Some(Err(Error::ChannelsLimitReached)))
        );
    }
}
<|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::{PollToken, RawDescriptor};
use serde::{Deserialize, Serialize};
use std::{
    fs::File,
    mem::{self, ManuallyDrop},
};

/// Wraps a RawDescriptor and safely closes it when self falls out of scope.
#[derive(Serialize, Deserialize, Debug, Eq)]
#[serde(transparent)]
pub struct SafeDescriptor {
    #[serde(with = "super::with_raw_descriptor")]
    pub(crate) descriptor: RawDescriptor,
}

/// Trait for forfeiting ownership of the current raw descriptor, and returning the raw descriptor
pub trait IntoRawDescriptor {
    fn into_raw_descriptor(self) -> RawDescriptor;
}

/// Trait for returning the underlying raw descriptor, without giving up ownership of the
/// descriptor.
pub trait AsRawDescriptor {
    fn as_raw_descriptor(&self) -> RawDescriptor;
}

/// A trait similar to `AsRawDescriptor` but supports an arbitrary number of descriptors.
pub trait AsRawDescriptors {
    fn as_raw_descriptors(&self) -> Vec<RawDescriptor>;
}

pub trait FromRawDescriptor {
    /// # Safety
    /// Safe only if the caller ensures nothing has access to the descriptor after passing it to
    /// `from_raw_descriptor`
    unsafe fn from_raw_descriptor(descriptor: RawDescriptor) -> Self;
}

impl AsRawDescriptor for SafeDescriptor {
    fn as_raw_descriptor(&self) -> RawDescriptor {
        self.descriptor
    }
}

impl<T> AsRawDescriptors for T
where
    T: AsRawDescriptor,
{
    fn as_raw_descriptors(&self) -> Vec<RawDescriptor> {
        vec![self.as_raw_descriptor()]
    }
}

impl IntoRawDescriptor for SafeDescriptor {
    fn into_raw_descriptor(self) -> RawDescriptor {
        let descriptor = self.descriptor;
        mem::forget(self);
        descriptor
    }
}

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

impl TryFrom<&dyn AsRawDescriptor> for SafeDescriptor {
    type Error = std::io::Error;

    /// Clones the underlying descriptor (handle), internally creating a new descriptor.
    ///
    /// WARNING: Windows does NOT support cloning/duplicating all types of handles. DO NOT use this
    /// function on IO completion ports, sockets, or pseudo-handles (except those from
    /// GetCurrentProcess or GetCurrentThread). See
    /// <https://docs.microsoft.com/en-us/windows/win32/api/handleapi/nf-handleapi-duplicatehandle>
    /// for further details.
    ///
    /// TODO(b/191800567): this API has sharp edges on Windows. We should evaluate making some
    /// adjustments to smooth those edges.
    fn try_from(rd: &dyn AsRawDescriptor) -> std::result::Result<Self, Self::Error> {
        // Safe because the underlying raw descriptor is guaranteed valid by rd's existence.
        //
        // Note that we are cloning the underlying raw descriptor since we have no guarantee of
        // its existence after this function returns.
        let rd_as_safe_desc = ManuallyDrop::new(unsafe {
            SafeDescriptor::from_raw_descriptor(rd.as_raw_descriptor())
        });

        // We have to clone rd because we have no guarantee ownership was transferred (rd is
        // borrowed).
        rd_as_safe_desc
            .try_clone()
            .map_err(|e| Self::Error::from_raw_os_error(e.errno()))
    }
}

impl From<File> for SafeDescriptor {
    fn from(f: File) -> SafeDescriptor {
        // Safe because we own the File at this point.
        unsafe { SafeDescriptor::from_raw_descriptor(f.into_raw_descriptor()) }
    }
}

/// For use cases where a simple wrapper around a RawDescriptor is needed.
/// This is a simply a wrapper and does not manage the lifetime of the descriptor.
/// Most usages should prefer SafeDescriptor or using a RawDescriptor directly
#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Descriptor(pub RawDescriptor);
impl AsRawDescriptor for Descriptor {
    fn as_raw_descriptor(&self) -> RawDescriptor {
        self.0
    }
}

/// Implement token for implementations that wish to use this struct as such
impl PollToken for Descriptor {
    fn as_raw_token(&self) -> u64 {
        self.0 as u64
    }

    fn from_raw_token(data: u64) -> Self {
        Descriptor(data as RawDescriptor)
    }
}
<commit_msg>base: improve documentation of Descriptor and AsRawDescriptor<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::{PollToken, RawDescriptor};
use serde::{Deserialize, Serialize};
use std::{
    fs::File,
    mem::{self, ManuallyDrop},
};

/// Wraps a RawDescriptor and safely closes it when self falls out of scope.
#[derive(Serialize, Deserialize, Debug, Eq)]
#[serde(transparent)]
pub struct SafeDescriptor {
    #[serde(with = "super::with_raw_descriptor")]
    pub(crate) descriptor: RawDescriptor,
}

/// Trait for forfeiting ownership of the current raw descriptor, and returning the raw descriptor
pub trait IntoRawDescriptor {
    fn into_raw_descriptor(self) -> RawDescriptor;
}

/// Trait for returning the underlying raw descriptor, without giving up ownership of the
/// descriptor.
pub trait AsRawDescriptor {
    /// Returns the underlying raw descriptor.
    ///
    /// Since the descriptor is still owned by the provider, callers should not assume that it will
    /// remain open for longer than the immediate call of this method. In particular, it is a
    /// dangerous practice to store the result of this method for future use: instead, it should be
    /// used to e.g. obtain a raw descriptor that is immediately passed to a system call.
    ///
    /// If you need to use the descriptor for a longer time (and particularly if you cannot reliably
    /// track the lifetime of the providing object), you should probably consider using
    /// [`SafeDescriptor`] (possibly along with [`trait@IntoRawDescriptor`]) to get full ownership
    /// over a descriptor pointing to the same resource.
    fn as_raw_descriptor(&self) -> RawDescriptor;
}

/// A trait similar to `AsRawDescriptor` but supports an arbitrary number of descriptors.
pub trait AsRawDescriptors {
    /// Returns the underlying raw descriptors.
    ///
    /// Please refer to the documentation of [`AsRawDescriptor::as_raw_descriptor`] for limitations
    /// and recommended use.
    fn as_raw_descriptors(&self) -> Vec<RawDescriptor>;
}

pub trait FromRawDescriptor {
    /// # Safety
    /// Safe only if the caller ensures nothing has access to the descriptor after passing it to
    /// `from_raw_descriptor`
    unsafe fn from_raw_descriptor(descriptor: RawDescriptor) -> Self;
}

impl AsRawDescriptor for SafeDescriptor {
    fn as_raw_descriptor(&self) -> RawDescriptor {
        self.descriptor
    }
}

impl<T> AsRawDescriptors for T
where
    T: AsRawDescriptor,
{
    fn as_raw_descriptors(&self) -> Vec<RawDescriptor> {
        vec![self.as_raw_descriptor()]
    }
}

impl IntoRawDescriptor for SafeDescriptor {
    fn into_raw_descriptor(self) -> RawDescriptor {
        let descriptor = self.descriptor;
        mem::forget(self);
        descriptor
    }
}

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

impl TryFrom<&dyn AsRawDescriptor> for SafeDescriptor {
    type Error = std::io::Error;

    /// Clones the underlying descriptor (handle), internally creating a new descriptor.
    ///
    /// WARNING: Windows does NOT support cloning/duplicating all types of handles. DO NOT use this
    /// function on IO completion ports, sockets, or pseudo-handles (except those from
    /// GetCurrentProcess or GetCurrentThread). See
    /// <https://docs.microsoft.com/en-us/windows/win32/api/handleapi/nf-handleapi-duplicatehandle>
    /// for further details.
    ///
    /// TODO(b/191800567): this API has sharp edges on Windows. We should evaluate making some
    /// adjustments to smooth those edges.
    fn try_from(rd: &dyn AsRawDescriptor) -> std::result::Result<Self, Self::Error> {
        // Safe because the underlying raw descriptor is guaranteed valid by rd's existence.
        //
        // Note that we are cloning the underlying raw descriptor since we have no guarantee of
        // its existence after this function returns.
        let rd_as_safe_desc = ManuallyDrop::new(unsafe {
            SafeDescriptor::from_raw_descriptor(rd.as_raw_descriptor())
        });

        // We have to clone rd because we have no guarantee ownership was transferred (rd is
        // borrowed).
        rd_as_safe_desc
            .try_clone()
            .map_err(|e| Self::Error::from_raw_os_error(e.errno()))
    }
}

impl From<File> for SafeDescriptor {
    fn from(f: File) -> SafeDescriptor {
        // Safe because we own the File at this point.
        unsafe { SafeDescriptor::from_raw_descriptor(f.into_raw_descriptor()) }
    }
}

/// For use cases where a simple wrapper around a [`RawDescriptor`] is needed, in order to e.g.
/// implement [`trait@AsRawDescriptor`].
///
/// This is a simply a wrapper and does not manage the lifetime of the descriptor. As such it is the
/// responsibility of the user to ensure that the wrapped descriptor will not be closed for as long
/// as the `Descriptor` is alive.
///
/// Most use-cases should prefer [`SafeDescriptor`] or implementing and using
/// [`trait@AsRawDescriptor`] on the type providing the descriptor. Using this wrapper usually means
/// something can be improved in your code.
///
/// Valid uses of this struct include:
/// * You only have a valid [`RawDescriptor`] and need to pass something that implements
///   [`trait@AsRawDescriptor`] to a function,
/// * You need to serialize a [`RawDescriptor`],
/// * You need [`trait@Send`] or [`trait@Sync`] for your descriptor and properly handle the case
///   where your descriptor gets closed.
///
/// Note that with the exception of the last use-case (which requires proper error checking against
/// the descriptor being closed), the `Descriptor` instance would be very short-lived.
#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct Descriptor(pub RawDescriptor);
impl AsRawDescriptor for Descriptor {
    fn as_raw_descriptor(&self) -> RawDescriptor {
        self.0
    }
}

/// Implement token for implementations that wish to use this struct as such
impl PollToken for Descriptor {
    fn as_raw_token(&self) -> u64 {
        self.0 as u64
    }

    fn from_raw_token(data: u64) -> Self {
        Descriptor(data as RawDescriptor)
    }
}
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::fs::File;
use std::string::String;
use std::path::Path;
use std::vec::Vec;
use std::io::BufReader;
use json;
use game::constants::{CHARACTER_JSON_PATH, ZOMBIE_JSON_PATH};
use tiled::Map;
use tiled;
use character::CharacterData;

pub fn load_map_file(filename: &str) -> Map {
  let file = match File::open(&Path::new(&filename)) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let reader = BufReader::new(file);
  match tiled::parse(reader) {
    Ok(m) => m,
    Err(e) => panic!("Map parse error {:?}", e)
  }
}

pub fn get_map_tile(map: &Map, layer_index: usize, x: usize, y: usize) -> u32 {
  let layer = match map.layers.get(layer_index) {
    None => panic!("Layer_index value out of index {:?}", map.layers),
    Some(ref l) => *l
  };
  let y_index = match layer.tiles.iter().rev().nth(y) {
    None => panic!("X value out of index {:?}", map.layers[0]),
    Some(ref x) => *x
  };
  let val = match y_index.get(x) {
    None => panic!("Y value out of index {:?}", layer.tiles[y]),
    Some(ref val) => **val
  };
  val
}

fn read_sprite_file(filename: &str) -> String {
  let path = Path::new(&filename);
  let mut file = match File::open(&path) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let mut buf = String::new();
  match file.read_to_string(&mut buf) {
    Ok(_) => buf,
    Err(_) => panic!("Couldn't read file {}", filename),
  }
}

pub fn load_character() -> Vec<CharacterData> {
  let mut sprites = Vec::with_capacity(512);
  let mut move_sprite_names = Vec::with_capacity(256);
  //  let mut fire_sprite_names = Vec::with_capacity(256);
  let character_json = read_sprite_file(CHARACTER_JSON_PATH);
  let character = match json::parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character JSON parse error {:?}", e),
  };

  for x in 0..15 {
    for y in 0..14 {
      move_sprite_names.push(format!("run_{}_{}", x, y));
    }
  }
  //  for x in 0..15 {
  //    for y in 0..3 {
  //      fire_sprite_names.push(format!("fire_{}_{}", x, y));
  //    }
  //  }

  for &ref move_sprite in &move_sprite_names {
    sprites.push(CharacterData::new([character["frames"][move_sprite]["frame"]["x"].as_f32().unwrap(),
                                    character["frames"][move_sprite]["frame"]["y"].as_f32().unwrap(),
                                    0.0,
                                    0.0]));
  }

  //  for &ref fire_sprite in &fire_sprite_names {
  //    sprites.push(CharacterData::new([character["frames"][fire_sprite]["frame"]["x"].as_f32().unwrap(),
  //                                    character["frames"][fire_sprite]["frame"]["y"].as_f32().unwrap(),
  //                                    0.0,
  //                                    0.0]));
  //  }
  sprites
}

#[allow(dead_code)]
pub fn load_zombie() -> Vec<(f64, f64)> {
  let mut sprites = Vec::with_capacity(256);
  let mut idle_sprite_names = Vec::with_capacity(64);
  let mut walk_sprite_names = Vec::with_capacity(64);
  let character_json = read_sprite_file(ZOMBIE_JSON_PATH);
  let zombie = match json::parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character JSON parse error {:?}", e),
  };

  for x in 0..7 {
    for y in 0..3 {
      idle_sprite_names.push(format!("idle_{}_{}", x, y));
    }
  }

  for x in 0..7 {
    for y in 0..3 {
      walk_sprite_names.push(format!("walk_{}_{}", x, y));
    }
  }

  for &ref idle_sprite in &idle_sprite_names {
    let x = zombie["frames"][idle_sprite]["frame"]["x"].as_f64().unwrap();
    let y = zombie["frames"][idle_sprite]["frame"]["y"].as_f64().unwrap();
    sprites.push((x, y));
  }

  for &ref walk_sprite in &walk_sprite_names {
    let x = zombie["frames"][walk_sprite]["frame"]["x"].as_f64().unwrap();
    let y = zombie["frames"][walk_sprite]["frame"]["y"].as_f64().unwrap();
    sprites.push((x, y));
  }
  sprites
}
<commit_msg>Take width & height values from character.json<commit_after>use std::io::prelude::*;
use std::fs::File;
use std::string::String;
use std::path::Path;
use std::vec::Vec;
use std::io::BufReader;
use json;
use game::constants::{CHARACTER_JSON_PATH, ZOMBIE_JSON_PATH};
use tiled::Map;
use tiled;
use character::CharacterData;

pub fn load_map_file(filename: &str) -> Map {
  let file = match File::open(&Path::new(&filename)) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let reader = BufReader::new(file);
  match tiled::parse(reader) {
    Ok(m) => m,
    Err(e) => panic!("Map parse error {:?}", e)
  }
}

pub fn get_map_tile(map: &Map, layer_index: usize, x: usize, y: usize) -> u32 {
  let layer = match map.layers.get(layer_index) {
    None => panic!("Layer_index value out of index {:?}", map.layers),
    Some(ref l) => *l
  };
  let y_index = match layer.tiles.iter().rev().nth(y) {
    None => panic!("X value out of index {:?}", map.layers[0]),
    Some(ref x) => *x
  };
  let val = match y_index.get(x) {
    None => panic!("Y value out of index {:?}", layer.tiles[y]),
    Some(ref val) => **val
  };
  val
}

fn read_sprite_file(filename: &str) -> String {
  let path = Path::new(&filename);
  let mut file = match File::open(&path) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let mut buf = String::new();
  match file.read_to_string(&mut buf) {
    Ok(_) => buf,
    Err(_) => panic!("Couldn't read file {}", filename),
  }
}

pub fn load_character() -> Vec<CharacterData> {
  let mut sprites = Vec::with_capacity(512);
  let mut move_sprite_names = Vec::with_capacity(256);
//  let mut fire_sprite_names = Vec::with_capacity(256);
  let character_json = read_sprite_file(CHARACTER_JSON_PATH);
  let character = match json::parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character JSON parse error {:?}", e),
  };

  for x in 0..15 {
    for y in 0..14 {
      move_sprite_names.push(format!("run_{}_{}", x, y));
    }
  }
//  for x in 0..15 {
//    for y in 0..3 {
//      fire_sprite_names.push(format!("fire_{}_{}", x, y));
//    }
//  }

  for &ref move_sprite in &move_sprite_names {
    sprites.push(CharacterData::new([
      character["frames"][move_sprite]["frame"]["x"].as_f32().unwrap(),
      character["frames"][move_sprite]["frame"]["y"].as_f32().unwrap(),
      character["frames"][move_sprite]["frame"]["w"].as_f32().unwrap(),
      character["frames"][move_sprite]["frame"]["h"].as_f32().unwrap()
    ]));
  }

//  for &ref fire_sprite in &fire_sprite_names {
//    sprites.push(CharacterData::new([
//      character["frames"][fire_sprite]["frame"]["x"].as_f32().unwrap(),
//      character["frames"][fire_sprite]["frame"]["y"].as_f32().unwrap(),
//      character["frames"][fire_sprite]["frame"]["w"].as_f32().unwrap(),
//      character["frames"][fire_sprite]["frame"]["h"].as_f32().unwrap()
//    ]));
//  }
  sprites
}

#[allow(dead_code)]
pub fn load_zombie() -> Vec<(f64, f64)> {
  let mut sprites = Vec::with_capacity(256);
  let mut idle_sprite_names = Vec::with_capacity(64);
  let mut walk_sprite_names = Vec::with_capacity(64);
  let character_json = read_sprite_file(ZOMBIE_JSON_PATH);
  let zombie = match json::parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character JSON parse error {:?}", e),
  };

  for x in 0..7 {
    for y in 0..3 {
      idle_sprite_names.push(format!("idle_{}_{}", x, y));
    }
  }

  for x in 0..7 {
    for y in 0..3 {
      walk_sprite_names.push(format!("walk_{}_{}", x, y));
    }
  }

  for &ref idle_sprite in &idle_sprite_names {
    let x = zombie["frames"][idle_sprite]["frame"]["x"].as_f64().unwrap();
    let y = zombie["frames"][idle_sprite]["frame"]["y"].as_f64().unwrap();
    sprites.push((x, y));
  }

  for &ref walk_sprite in &walk_sprite_names {
    let x = zombie["frames"][walk_sprite]["frame"]["x"].as_f64().unwrap();
    let y = zombie["frames"][walk_sprite]["frame"]["y"].as_f64().unwrap();
    sprites.push((x, y));
  }
  sprites
}
<|endoftext|>


<commit_before>use std::net::{TcpStream,Ipv4Addr};
use metainfo::Metainfo;
use tracker_response::Peer;
use std::io::{Write, Read};

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            PeerConnection::connect(peer, &info.info_hash);
        }
    }
}

struct PeerConnection {
    stream: TcpStream,
}

impl PeerConnection {
    fn connect(peer: &Peer, info_hash: &[u8]) -> Option<PeerConnection> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        match TcpStream::connect((peer.ip, peer.port)) {
            Ok(stream) => {
                let mut conn = PeerConnection { stream: stream };
                conn.handshake(info_hash);
                Some(conn)
            },
            Err(_) => {
                println!("Failed to connect to {}:{}", peer.ip, peer.port);
                None
            }
        }
    }

    fn handshake(&mut self, info_hash: &[u8]) {
        self.send_handshake(info_hash);
        self.receive_handshake();
    }

    fn send_handshake(&mut self, info_hash: &[u8]) {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        self.stream.write_all(&message).unwrap();
    }

    fn receive_handshake(&mut self) {
        let pstrlen = self.read_n(1);
        let pstr = self.read_n(pstrlen[0] as usize);
        let reserved = self.read_n(8);
        let info_hash = self.read_n(20);
        let peer_id = self.read_n(20);
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Vec<u8> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n)  => assert_eq!(bytes_to_read as usize, n),
            Err(_) => panic!("Didn't read enough"),
        }
        buf
    }
}
<commit_msg>Added error handling to downloads.<commit_after>use std::net::{TcpStream,Ipv4Addr};
use metainfo::Metainfo;
use tracker_response::Peer;
use std::{convert, io};
use std::io::{Write, Read};

pub fn download(info: &Metainfo, peers: &[Peer]) {
    for peer in peers {
        if peer.ip != Ipv4Addr::new(207, 251, 103, 46) {
            match PeerConnection::connect(peer, &info.info_hash) {
                Ok(_) => println!("Peer done"),
                Err(e) => println!("Error: {:?}", e)
            }
        }
    }
}

struct PeerConnection {
    stream: TcpStream,
}

impl PeerConnection {
    fn connect(peer: &Peer, info_hash: &[u8]) -> Result<PeerConnection, Error> {
        println!("Connecting to {}:{}", peer.ip, peer.port);
        let stream = try!(TcpStream::connect((peer.ip, peer.port)));
        let mut conn = PeerConnection { stream: stream };
        try!(conn.handshake(info_hash));
        Ok(conn)
    }

    fn handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        try!(self.send_handshake(info_hash));
        try!(self.receive_handshake());
        Ok(())
    }

    fn send_handshake(&mut self, info_hash: &[u8]) -> Result<(), Error> {
        let mut message = vec![];
        message.push(19);
        message.extend("BitTorrent protocol".as_bytes().iter().cloned());
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.push(0);
        message.extend(info_hash.iter().cloned());
        message.extend("-TZ-0000-00000000001".as_bytes().iter().cloned());
        try!(self.stream.write_all(&message));
        Ok(())
    }

    fn receive_handshake(&mut self) -> Result<(), Error> {
        let pstrlen = try!(self.read_n(1));
        let pstr = try!(self.read_n(pstrlen[0] as usize));
        let reserved = try!(self.read_n(8));
        let info_hash = try!(self.read_n(20));
        let peer_id = try!(self.read_n(20));
        Ok(())
    }

    fn read_n(&mut self, bytes_to_read: usize) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let bytes_read = (&mut self.stream).take(bytes_to_read as u64).read_to_end(&mut buf);
        match bytes_read {
            Ok(n) if n == bytes_to_read  => Ok(buf),
            Ok(_)  => Err(Error::NotEnoughData),
            Err(e) => try!(Err(e))
        }
    }
}

#[derive(Debug)]
pub enum Error {
    IoError(io::Error),
    NotEnoughData,
}

impl convert::From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        Error::IoError(err)
    }
}
<|endoftext|>


<commit_before>// Defined note types for GNU systems.

#[cfg(feature = "log")]
use log::debug;
#[cfg(feature = "alloc")]
use scroll::{Pread, Pwrite, IOread, IOwrite, SizeWith};

/// ABI information.
///
/// The descriptor consists of words:
///  * word 0: OS descriptor
///  * word 1: major version of the ABI
///  * word 2: minor version of the ABI
///  * word 3: subminor version of the ABI
pub const NT_GNU_ABI_TAG: u32 = 1;

/// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
// `NT_GNU_ABI_TAG` note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

/// Synthetic `hwcap` information.
///
/// The descriptor begins with two words:
///  * word 0: number of entries
///  * word 1: bitmask of enabled entries
///
/// Then follow variable-length entries, one byte followed by a '\0'-terminated
/// `hwcap` name string.  The byte gives the bit number to test if enabled,
/// `(1U << bit) & bitmask`.
pub const NT_GNU_HWCAP: u32 = 2;

/// Build ID bits as generated by ld --build-id.
///
/// The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

/// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_alloc! {
    use crate::error;
    use crate::container;
    use scroll::ctx;
    use alloc::vec::Vec;

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, usize), Self::Error> {
            let offset = &mut 0;
            let mut alignment = alignment;
            if alignment < 4 {
                alignment = 4;
            }
            let header: NoteHeader = {
                match alignment {
                    4|8 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<commit_msg>elf.note: added new constants for structs: prstatus, prpsinfo, siginfo_t, mapped files<commit_after>// Defined note types for GNU systems.

#[cfg(feature = "log")]
use log::debug;
#[cfg(feature = "alloc")]
use scroll::{Pread, Pwrite, IOread, IOwrite, SizeWith};

/// ABI information.
///
/// The descriptor consists of words:
///  * word 0: OS descriptor
///  * word 1: major version of the ABI
///  * word 2: minor version of the ABI
///  * word 3: subminor version of the ABI
pub const NT_GNU_ABI_TAG: u32 = 1;

/// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
// `NT_GNU_ABI_TAG` note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

/// Synthetic `hwcap` information.
///
/// The descriptor begins with two words:
///  * word 0: number of entries
///  * word 1: bitmask of enabled entries
///
/// Then follow variable-length entries, one byte followed by a '\0'-terminated
/// `hwcap` name string.  The byte gives the bit number to test if enabled,
/// `(1U << bit) & bitmask`.
pub const NT_GNU_HWCAP: u32 = 2;

/// Build ID bits as generated by ld --build-id.
///
/// The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

/// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

///Contains copy of prstatus struct.
pub const NT_PRSTATUS: u32 = 1;

///Contains copy of prpsinfo struct.
pub const NT_PRPSINFO: u32 = 3;

///Fields of siginfo_t.
pub const NT_SIGINFO: u32 = 0x53494749;

///Description of mapped files.
pub const NT_FILE: u32 = 0x46494c45;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_alloc! {
    use crate::error;
    use crate::container;
    use scroll::ctx;
    use alloc::vec::Vec;

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, usize), Self::Error> {
            let offset = &mut 0;
            let mut alignment = alignment;
            if alignment < 4 {
                alignment = 4;
            }
            let header: NoteHeader = {
                match alignment {
                    4|8 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<|endoftext|>