<commit_before>//! The `ztext` module contains encoding functions to encode text in z-ascii characters.
//!

use super::zbytes::Bytes;

pub static ALPHABET: [char; 78] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',

    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',

    ' ', '\n', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.',
    ',', '!', '?', '_', '#', '\'','"', '/', '\\','-', ':', '(', ')'];

/// encodes an string to z-characters
/// and returns the length of the used bytes
/// TODO: Only works with lower case letters
///
/// # Examples
///
/// ```ignore
/// let data = Bytes{bytes: Vec::new()};
/// let byteLength = data.encode("hello");
/// ```
pub fn encode(data: &mut Bytes, content: &str, unicode_table: &Vec<u16>) -> u16 {
    let zchars: Vec<u8> = string_to_zchar(content, unicode_table);

    let mut two_bytes: u16 = 0;
    let len = zchars.len();
    for i in 0..len {
        let zasci_id =zchars[i];

        two_bytes |= shift(zasci_id as u16, i as u8);

        if i % 3 == 2 {
            data.write_u16(two_bytes, pos_to_index(i));
            two_bytes = 0;
        }

        // end of string
        if i == len -1 {
            if i % 3 != 2 {
                for j in (i % 3) + 1..3 {
                    two_bytes |= shift(0x05 as u16, j as u8);
                }

                data.write_u16(two_bytes, pos_to_index(i));
            }

            // end bit is written to the first bit of the next to last byte
            let end_index: usize = data.bytes.len() - 2 as usize;
            let mut end_byte: u8 = data.bytes[end_index];

            end_byte |= 0x80;
            data.write_byte(end_byte, end_index as usize);
        }
    }

    data.bytes.len() as u16
}

/// reads the content and converts it to a zasci vector
fn string_to_zchar(content: &str, unicode_table: &Vec<u16>) -> Vec<u8> {
    //let string_bytes = content.to_string().into_bytes();
    let mut zchars: Vec<u8> = Vec::new();

    for character in content.chars() {

        let mut byte: u8 = character as u8;
        let alpha_index = pos_in_alpha(byte as u8);
        if character as u16 <= 126 && alpha_index != -1 {

            if byte == 0x0A {
                // newline
                zchars.push(0x05);
                zchars.push(7);
            } else if byte == 0x20 {
                // space
                zchars.push(0x00);
            } else {
                if alpha_index > 51 {
                    // in A2
                    zchars.push(0x05);
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else if alpha_index < 26 {
                    // in A0
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else {
                    // in A1
                    zchars.push(0x04);
                    zchars.push(alpha_index as u8 % 26 + 6);
                }
            }
        } else {
            // not in alphabet or unicode

            // to change alphabet
            zchars.push(0x05);

            // for special char (10 bit z-ascii)
            zchars.push(0x06);

            //let mut byte: u8;
            if character as u16 <= 126 {
                // not in alphabet, but still ascii
                byte = character as u8;
            } else {
                // unicode
                let unicode_index = pos_in_unicode(character as u16, unicode_table);
                byte = unicode_index as u8 + 155;
            }

            zchars.push(byte >> 5);
            zchars.push(byte & 0x1f);
        }
    }
    zchars
}

/// shifts the z-char in a 2 bytes-array to the right position
/// shift_length has 3 possibilities: 10, 5, 0
/// an z-char use 5 bytes
/// look into two bytes:
/// ____byte one____ ____byte two____
/// 1 2 3 4 5 6  7 8 1 2 3  4 5 6 7 8
///   ^          ^          1
///   1 2 3 4 5  1 2 3 4 5  1 2 3 4 5
///   1. zchar   2. zchar   3. zchar
///   10         5          0
fn shift(zchar: u16, position: u8) -> u16 {
    let shift_length = 10 - (position % 3) * 5;
    zchar << shift_length
}

/// Returns the location of the character of the specified index in the zcode character array
///
/// # Examples
///
fn pos_in_alpha(letter: u8) -> i8 {
    for i in 0..ALPHABET.len() {
        if ALPHABET[i] as u8 == letter {
            return i as i8
        }
    }

    return -1
}

/// returns the position in the unicode-table
pub fn pos_in_unicode(letter: u16, unicode_table: &Vec<u16>) -> i8 {
    for (i, character) in unicode_table.iter().enumerate() {
        if *character == letter {
            return i as i8
        }
    }

    return -1
}

/// position in the vector from the position of an character in the string
/// every 3 zchars encode() writes 2 bytes
/// for example "helloworld"
/// "h": nothing
/// "e": nothing
/// "l": write 2 bytes to position 0   (2 * 2/3))
/// "l": nothing
/// "o": nothing
/// "w": write 2 bytes to position 2   (2 * 5/3))
/// ...
///
/// # Examples
///
fn pos_to_index(position: usize) -> usize {
    2 * (position / 3)
}

#[test]
fn test_pos_in_alpha() {
    assert_eq!(pos_in_alpha('a' as u8), 0);
    assert_eq!(pos_in_alpha('b' as u8), 1);
    assert_eq!(pos_in_alpha('c' as u8), 2);
    assert_eq!(pos_in_alpha('A' as u8), 26);
    assert_eq!(pos_in_alpha('B' as u8), 27);
    assert_eq!(pos_in_alpha('C' as u8), 28);
}

#[test]
fn test_pos_to_index() {
    assert_eq!(pos_to_index(5), 2);
}

#[test]
fn test_shift() {
    assert_eq!(shift(6,2), 6);
    assert_eq!(shift(26,5), 26);
}

#[test]
fn test_string_to_zchar() {
    let mut vec: Vec<u16> = Vec::new();
    assert_eq!(string_to_zchar("i am a string, please test me, no unicode",&vec), vec![14, 0, 6, 18, 0, 6, 0, 24, 25, 23, 14, 19, 12, 5, 19, 0, 21, 17, 10, 6, 24, 10, 0, 25, 10, 24, 25, 0, 18, 10, 5, 19, 0, 19, 20, 0, 26, 19, 14, 8, 20, 9, 10]);
    vec.push('€' as u16);
    assert_eq!(string_to_zchar("nasty char: €",&vec), vec![19, 6, 24, 25, 30, 0, 8, 13, 6, 23, 5, 29, 0, 5, 6, 4, 27]);
}<commit_msg>u8 to usize to avoid overflow<commit_after>//! The `ztext` module contains encoding functions to encode text in z-ascii characters.
//!

use super::zbytes::Bytes;

pub static ALPHABET: [char; 78] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
    'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',

    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
    'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',

    ' ', '\n', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.',
    ',', '!', '?', '_', '#', '\'','"', '/', '\\','-', ':', '(', ')'];

/// encodes an string to z-characters
/// and returns the length of the used bytes
/// TODO: Only works with lower case letters
///
/// # Examples
///
/// ```ignore
/// let data = Bytes{bytes: Vec::new()};
/// let byteLength = data.encode("hello");
/// ```
pub fn encode(data: &mut Bytes, content: &str, unicode_table: &Vec<u16>) -> u16 {
    let zchars: Vec<u8> = string_to_zchar(content, unicode_table);

    let mut two_bytes: u16 = 0;
    let len = zchars.len();
    for i in 0..len {
        let zasci_id =zchars[i];

        two_bytes |= shift(zasci_id as u16, i);

        if i % 3 == 2 {
            data.write_u16(two_bytes, pos_to_index(i));
            two_bytes = 0;
        }

        // end of string
        if i == len -1 {
            if i % 3 != 2 {
                for j in (i % 3) + 1..3 {
                    two_bytes |= shift(0x05 as u16, j);
                }

                data.write_u16(two_bytes, pos_to_index(i));
            }

            // end bit is written to the first bit of the next to last byte
            let end_index: usize = data.bytes.len() - 2 as usize;
            let mut end_byte: u8 = data.bytes[end_index];

            end_byte |= 0x80;
            data.write_byte(end_byte, end_index as usize);
        }
    }

    data.bytes.len() as u16
}

/// reads the content and converts it to a zasci vector
fn string_to_zchar(content: &str, unicode_table: &Vec<u16>) -> Vec<u8> {
    //let string_bytes = content.to_string().into_bytes();
    let mut zchars: Vec<u8> = Vec::new();

    for character in content.chars() {

        let mut byte: u8 = character as u8;
        let alpha_index = pos_in_alpha(byte as u8);
        if character as u16 <= 126 && alpha_index != -1 {

            if byte == 0x0A {
                // newline
                zchars.push(0x05);
                zchars.push(7);
            } else if byte == 0x20 {
                // space
                zchars.push(0x00);
            } else {
                if alpha_index > 51 {
                    // in A2
                    zchars.push(0x05);
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else if alpha_index < 26 {
                    // in A0
                    zchars.push(alpha_index as u8 % 26 + 6);
                } else {
                    // in A1
                    zchars.push(0x04);
                    zchars.push(alpha_index as u8 % 26 + 6);
                }
            }
        } else {
            // not in alphabet or unicode

            // to change alphabet
            zchars.push(0x05);

            // for special char (10 bit z-ascii)
            zchars.push(0x06);

            //let mut byte: u8;
            if character as u16 <= 126 {
                // not in alphabet, but still ascii
                byte = character as u8;
            } else {
                // unicode
                let unicode_index = pos_in_unicode(character as u16, unicode_table);
                byte = unicode_index as u8 + 155;
            }

            zchars.push(byte >> 5);
            zchars.push(byte & 0x1f);
        }
    }
    zchars
}

/// shifts the z-char in a 2 bytes-array to the right position
/// shift_length has 3 possibilities: 10, 5, 0
/// an z-char use 5 bytes
/// look into two bytes:
/// ____byte one____ ____byte two____
/// 1 2 3 4 5 6  7 8 1 2 3  4 5 6 7 8
///   ^          ^          1
///   1 2 3 4 5  1 2 3 4 5  1 2 3 4 5
///   1. zchar   2. zchar   3. zchar
///   10         5          0
fn shift(zchar: u16, position: usize) -> u16 {
    let shift_length = 10 - (position % 3) as u16 * 5;
    zchar << shift_length
}

/// Returns the location of the character of the specified index in the zcode character array
///
/// # Examples
///
fn pos_in_alpha(letter: u8) -> i8 {
    for i in 0..ALPHABET.len() {
        if ALPHABET[i] as u8 == letter {
            return i as i8
        }
    }

    return -1
}

/// returns the position in the unicode-table
pub fn pos_in_unicode(letter: u16, unicode_table: &Vec<u16>) -> i8 {
    for (i, character) in unicode_table.iter().enumerate() {
        if *character == letter {
            return i as i8
        }
    }

    return -1
}

/// position in the vector from the position of an character in the string
/// every 3 zchars encode() writes 2 bytes
/// for example "helloworld"
/// "h": nothing
/// "e": nothing
/// "l": write 2 bytes to position 0   (2 * 2/3))
/// "l": nothing
/// "o": nothing
/// "w": write 2 bytes to position 2   (2 * 5/3))
/// ...
///
/// # Examples
///
fn pos_to_index(position: usize) -> usize {
    2 * (position / 3)
}

#[test]
fn test_pos_in_alpha() {
    assert_eq!(pos_in_alpha('a' as u8), 0);
    assert_eq!(pos_in_alpha('b' as u8), 1);
    assert_eq!(pos_in_alpha('c' as u8), 2);
    assert_eq!(pos_in_alpha('A' as u8), 26);
    assert_eq!(pos_in_alpha('B' as u8), 27);
    assert_eq!(pos_in_alpha('C' as u8), 28);
}

#[test]
fn test_pos_to_index() {
    assert_eq!(pos_to_index(5), 2);
}

#[test]
fn test_shift() {
    assert_eq!(shift(6,2), 6);
    assert_eq!(shift(26,5), 26);
}

#[test]
fn test_string_to_zchar() {
    let mut vec: Vec<u16> = Vec::new();
    assert_eq!(string_to_zchar("i am a string, please test me, no unicode",&vec), vec![14, 0, 6, 18, 0, 6, 0, 24, 25, 23, 14, 19, 12, 5, 19, 0, 21, 17, 10, 6, 24, 10, 0, 25, 10, 24, 25, 0, 18, 10, 5, 19, 0, 19, 20, 0, 26, 19, 14, 8, 20, 9, 10]);
    vec.push('€' as u16);
    assert_eq!(string_to_zchar("nasty char: €",&vec), vec![19, 6, 24, 25, 30, 0, 8, 13, 6, 23, 5, 29, 0, 5, 6, 4, 27]);
}<|endoftext|>


<commit_before>mod control_handle;
mod control_base;
mod window;
mod button;
mod check_box;
mod radio_button;
mod text_input;
mod label;
mod image_frame;

#[cfg(feature = "textbox")]
mod text_box;

#[cfg(feature = "rich-textbox")]
mod rich_text_box;

#[cfg(feature = "status-bar")]
mod status_bar;

#[cfg(feature = "tooltip")]
mod tooltip;

#[cfg(feature = "trackbar")]
mod track_bar;

#[cfg(feature = "menu")]
mod menu;

#[cfg(feature = "timer")]
mod timer;

#[cfg(feature = "notice")]
mod notice;

#[cfg(feature = "combobox")]
mod combo_box;

#[cfg(feature = "listbox")]
mod list_box;

#[cfg(feature = "datetime-picker")]
mod date_picker;

#[cfg(feature = "progress-bar")]
mod progress_bar;

#[cfg(feature = "tabs")]
mod tabs;

#[cfg(feature = "tree-view")]
mod treeview;

#[cfg(all(feature = "tree-view-iterator", feature = "tree-view") )]
mod treeview_iterator;

#[cfg(feature = "tray-notification")]
mod tray_notification;

#[cfg(feature = "message-window")]
mod message_window;

#[cfg(feature = "list-view")]
mod list_view;

#[cfg(feature = "number-select")]
mod number_select;

#[cfg(feature = "extern-canvas")]
mod extern_canvas;

#[cfg(feature = "frame")]
mod frame;

#[cfg(feature = "scroll-bar")]
mod scroll_bar;

mod handle_from_control;

pub use control_handle::ControlHandle;
pub use control_base::{ControlBase, HwndBuilder, TimerBuilder as BaseTimerBuilder, OtherBuilder};
pub use window::{Window, WindowFlags};
pub use button::{Button, ButtonFlags, ButtonBuilder};
pub use check_box::{CheckBox, CheckBoxState, CheckBoxFlags};
pub use radio_button::{RadioButton, RadioButtonState, RadioButtonFlags};
pub use text_input::{TextInput, TextInputFlags};
pub use label::{Label, LabelFlags};
pub use image_frame::{ImageFrame, ImageFrameFlags};

#[cfg(feature = "textbox")]
pub use text_box::{TextBox, TextBoxFlags};

#[cfg(feature = "rich-textbox")]
pub use rich_text_box::{RichTextBox, RichTextBoxFlags};

#[cfg(feature = "status-bar")]
pub use status_bar::StatusBar;

#[cfg(feature = "tooltip")]
pub use tooltip::{Tooltip, TooltipIcon};

#[cfg(feature = "trackbar")]
pub use track_bar::{TrackBar, TrackBarFlags};

#[cfg(feature = "menu")]
pub use menu::{Menu, MenuBuilder, MenuItem, MenuSeparator, MenuItemBuilder};

#[cfg(feature = "menu")]
pub use control_base::HmenuBuilder;

#[cfg(feature = "timer")]
pub use timer::{Timer, TimerBuilder};

#[cfg(feature = "notice")]
pub use notice::{Notice, NoticeSender, NoticeBuilder};

#[cfg(feature = "combobox")]
pub use combo_box::{ComboBox, ComboBoxFlags, ComboBoxBuilder};

#[cfg(feature = "listbox")]
pub use list_box::{ListBox, ListBoxFlags, ListBoxBuilder};

#[cfg(feature = "datetime-picker")]
pub use date_picker::{DatePicker, DatePickerValue, DatePickerFlags, DatePickerBuilder};

#[cfg(feature = "progress-bar")]
pub use progress_bar::{ProgressBar, ProgressBarState, ProgressBarFlags, ProgressBarBuilder};

#[cfg(feature = "tabs")]
pub use tabs::{TabsContainer, Tab, TabsContainerFlags, TabsContainerBuilder, TabBuilder};

#[cfg(feature = "tree-view")]
pub use treeview::{TreeView, TreeViewBuilder, TreeItem, TreeInsert, TreeItemAction, ExpandState, TreeItemState, TreeViewFlags};

#[cfg(all(feature = "tree-view-iterator", feature = "tree-view") )]
pub use treeview_iterator::TreeViewIterator;

#[cfg(feature = "tray-notification")]
pub use tray_notification::{TrayNotificationFlags, TrayNotification, TrayNotificationBuilder};

#[cfg(feature = "message-window")]
pub use message_window::{MessageWindow, MessageWindowBuilder};

#[cfg(feature = "list-view")]
pub use list_view::{ListView, ListViewStyle, ListViewBuilder, ListViewFlags, ListViewExFlags, InsertListViewItem, ListViewItem, InsertListViewColumn, ListViewColumn};

#[cfg(all(feature="list-view", feature="image-list"))]
pub use list_view::ListViewImageListType;

#[cfg(feature = "number-select")]
pub use number_select::{NumberSelect, NumberSelectBuilder, NumberSelectFlags};

#[cfg(feature = "extern-canvas")]
pub use extern_canvas::{ExternCanvas, ExternCanvasBuilder, ExternCanvasFlags};

#[cfg(feature = "frame")]
pub use frame::{Frame, FrameBuilder, FrameFlags};

#[cfg(feature = "scroll-bar")]
pub use scroll_bar::{ScrollBar, ScrollBarBuilder, ScrollBarFlags};

pub use handle_from_control::*;
<commit_msg>export missing builder in the docs<commit_after>mod control_handle;
mod control_base;
mod window;
mod button;
mod check_box;
mod radio_button;
mod text_input;
mod label;
mod image_frame;

#[cfg(feature = "textbox")]
mod text_box;

#[cfg(feature = "rich-textbox")]
mod rich_text_box;

#[cfg(feature = "status-bar")]
mod status_bar;

#[cfg(feature = "tooltip")]
mod tooltip;

#[cfg(feature = "trackbar")]
mod track_bar;

#[cfg(feature = "menu")]
mod menu;

#[cfg(feature = "timer")]
mod timer;

#[cfg(feature = "notice")]
mod notice;

#[cfg(feature = "combobox")]
mod combo_box;

#[cfg(feature = "listbox")]
mod list_box;

#[cfg(feature = "datetime-picker")]
mod date_picker;

#[cfg(feature = "progress-bar")]
mod progress_bar;

#[cfg(feature = "tabs")]
mod tabs;

#[cfg(feature = "tree-view")]
mod treeview;

#[cfg(all(feature = "tree-view-iterator", feature = "tree-view") )]
mod treeview_iterator;

#[cfg(feature = "tray-notification")]
mod tray_notification;

#[cfg(feature = "message-window")]
mod message_window;

#[cfg(feature = "list-view")]
mod list_view;

#[cfg(feature = "number-select")]
mod number_select;

#[cfg(feature = "extern-canvas")]
mod extern_canvas;

#[cfg(feature = "frame")]
mod frame;

#[cfg(feature = "scroll-bar")]
mod scroll_bar;

mod handle_from_control;

pub use control_handle::ControlHandle;
pub use control_base::{ControlBase, HwndBuilder, TimerBuilder as BaseTimerBuilder, OtherBuilder};
pub use window::{Window, WindowBuilder, WindowFlags};
pub use button::{Button, ButtonBuilder, ButtonFlags};
pub use check_box::{CheckBox, CheckBoxBuilder, CheckBoxState, CheckBoxFlags};
pub use radio_button::{RadioButton, RadioButtonBuilder, RadioButtonState, RadioButtonFlags};
pub use text_input::{TextInput, TextInputBuilder, TextInputFlags};
pub use label::{Label, LabelBuilder, LabelFlags};
pub use image_frame::{ImageFrame, ImageFrameBuilder, ImageFrameFlags};

#[cfg(feature = "textbox")]
pub use text_box::{TextBox, TextBoxBuilder, TextBoxFlags};

#[cfg(feature = "rich-textbox")]
pub use rich_text_box::{RichTextBox, RichTextBoxBuilder, RichTextBoxFlags};

#[cfg(feature = "status-bar")]
pub use status_bar::{StatusBar, StatusBarBuilder};

#[cfg(feature = "tooltip")]
pub use tooltip::{Tooltip, TooltipBuilder, TooltipIcon};

#[cfg(feature = "trackbar")]
pub use track_bar::{TrackBar, TrackBarBuilder, TrackBarFlags};

#[cfg(feature = "menu")]
pub use menu::{Menu, MenuBuilder, MenuItem, MenuSeparator, MenuItemBuilder};

#[cfg(feature = "menu")]
pub use control_base::HmenuBuilder;

#[cfg(feature = "timer")]
pub use timer::{Timer, TimerBuilder};

#[cfg(feature = "notice")]
pub use notice::{Notice, NoticeSender, NoticeBuilder};

#[cfg(feature = "combobox")]
pub use combo_box::{ComboBox, ComboBoxFlags, ComboBoxBuilder};

#[cfg(feature = "listbox")]
pub use list_box::{ListBox, ListBoxFlags, ListBoxBuilder};

#[cfg(feature = "datetime-picker")]
pub use date_picker::{DatePicker, DatePickerValue, DatePickerFlags, DatePickerBuilder};

#[cfg(feature = "progress-bar")]
pub use progress_bar::{ProgressBar, ProgressBarState, ProgressBarFlags, ProgressBarBuilder};

#[cfg(feature = "tabs")]
pub use tabs::{TabsContainer, Tab, TabsContainerFlags, TabsContainerBuilder, TabBuilder};

#[cfg(feature = "tree-view")]
pub use treeview::{TreeView, TreeViewBuilder, TreeItem, TreeInsert, TreeItemAction, ExpandState, TreeItemState, TreeViewFlags};

#[cfg(all(feature = "tree-view-iterator", feature = "tree-view") )]
pub use treeview_iterator::TreeViewIterator;

#[cfg(feature = "tray-notification")]
pub use tray_notification::{TrayNotificationFlags, TrayNotification, TrayNotificationBuilder};

#[cfg(feature = "message-window")]
pub use message_window::{MessageWindow, MessageWindowBuilder};

#[cfg(feature = "list-view")]
pub use list_view::{ListView, ListViewStyle, ListViewBuilder, ListViewFlags, ListViewExFlags, InsertListViewItem, ListViewItem, InsertListViewColumn, ListViewColumn};

#[cfg(all(feature="list-view", feature="image-list"))]
pub use list_view::ListViewImageListType;

#[cfg(feature = "number-select")]
pub use number_select::{NumberSelect, NumberSelectBuilder, NumberSelectFlags};

#[cfg(feature = "extern-canvas")]
pub use extern_canvas::{ExternCanvas, ExternCanvasBuilder, ExternCanvasFlags};

#[cfg(feature = "frame")]
pub use frame::{Frame, FrameBuilder, FrameFlags};

#[cfg(feature = "scroll-bar")]
pub use scroll_bar::{ScrollBar, ScrollBarBuilder, ScrollBarFlags};

pub use handle_from_control::*;
<|endoftext|>


<commit_before>use bytes::Buf;
use futures_core::stream::Stream;
use pin_project_lite::pin_project;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncBufRead, AsyncRead, ReadBuf};

pin_project! {
    /// Convert a [`Stream`] of byte chunks into an [`AsyncRead`].
    ///
    /// This type performs the inverse operation of [`ReaderStream`].
    ///
    /// # Example
    ///
    /// ```
    /// use bytes::Bytes;
    /// use tokio::io::{AsyncReadExt, Result};
    /// use tokio_util::io::StreamReader;
    /// # #[tokio::main]
    /// # async fn main() -> std::io::Result<()> {
    ///
    /// // Create a stream from an iterator.
    /// let stream = tokio_stream::iter(vec![
    ///     Result::Ok(Bytes::from_static(&[0, 1, 2, 3])),
    ///     Result::Ok(Bytes::from_static(&[4, 5, 6, 7])),
    ///     Result::Ok(Bytes::from_static(&[8, 9, 10, 11])),
    /// ]);
    ///
    /// // Convert it to an AsyncRead.
    /// let mut read = StreamReader::new(stream);
    ///
    /// // Read five bytes from the stream.
    /// let mut buf = [0; 5];
    /// read.read_exact(&mut buf).await?;
    /// assert_eq!(buf, [0, 1, 2, 3, 4]);
    ///
    /// // Read the rest of the current chunk.
    /// assert_eq!(read.read(&mut buf).await?, 3);
    /// assert_eq!(&buf[..3], [5, 6, 7]);
    ///
    /// // Read the next chunk.
    /// assert_eq!(read.read(&mut buf).await?, 4);
    /// assert_eq!(&buf[..4], [8, 9, 10, 11]);
    ///
    /// // We have now reached the end.
    /// assert_eq!(read.read(&mut buf).await?, 0);
    ///
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// If the stream produces errors which are not [std::io::Error],
    /// the errors can be converted using [`StreamExt`] to map each
    /// element.
    ///
    /// ```
    /// use bytes::Bytes;
    /// use tokio::io::AsyncReadExt;
    /// use tokio_util::io::StreamReader;
    /// use tokio_stream::StreamExt;
    /// # #[tokio::main]
    /// # async fn main() -> std::io::Result<()> {
    ///
    /// // Create a stream from an iterator, including an error.
    /// let stream = tokio_stream::iter(vec![
    ///     Result::Ok(Bytes::from_static(&[0, 1, 2, 3])),
    ///     Result::Ok(Bytes::from_static(&[4, 5, 6, 7])),
    ///     Result::Err("Something bad happened!")
    /// ]);
    ///
    /// // Use StreamExt to map the stream and error to a std::io::Error
    /// let stream = stream.map(|result| result.map_err(|err| {
    ///     std::io::Error::new(std::io::ErrorKind::Other, err)
    /// }));
    ///
    /// // Convert it to an AsyncRead.
    /// let mut read = StreamReader::new(stream);
    ///
    /// // Read five bytes from the stream.
    /// let mut buf = [0; 5];
    /// read.read_exact(&mut buf).await?;
    /// assert_eq!(buf, [0, 1, 2, 3, 4]);
    ///
    /// // Read the rest of the current chunk.
    /// assert_eq!(read.read(&mut buf).await?, 3);
    /// assert_eq!(&buf[..3], [5, 6, 7]);
    ///
    /// // Reading the next chunk will produce an error
    /// let error = read.read(&mut buf).await.unwrap_err();
    /// assert_eq!(error.kind(), std::io::ErrorKind::Other);
    /// assert_eq!(error.into_inner().unwrap().to_string(), "Something bad happened!");
    ///
    /// // We have now reached the end.
    /// assert_eq!(read.read(&mut buf).await?, 0);
    ///
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// [`AsyncRead`]: tokio::io::AsyncRead
    /// [`Stream`]: futures_core::Stream
    /// [`ReaderStream`]: crate::io::ReaderStream
    /// [`StreamExt`]: tokio_stream::StreamExt
    #[derive(Debug)]
    pub struct StreamReader<S, B> {
        #[pin]
        inner: S,
        chunk: Option<B>,
    }
}

impl<S, B, E> StreamReader<S, B>
where
    S: Stream<Item = Result<B, E>>,
    B: Buf,
    E: Into<std::io::Error>,
{
    /// Convert a stream of byte chunks into an [`AsyncRead`](tokio::io::AsyncRead).
    ///
    /// The item should be a [`Result`] with the ok variant being something that
    /// implements the [`Buf`] trait (e.g. `Vec<u8>` or `Bytes`). The error
    /// should be convertible into an [io error].
    ///
    /// [`Result`]: std::result::Result
    /// [`Buf`]: bytes::Buf
    /// [io error]: std::io::Error
    pub fn new(stream: S) -> Self {
        Self {
            inner: stream,
            chunk: None,
        }
    }

    /// Do we have a chunk and is it non-empty?
    fn has_chunk(&self) -> bool {
        if let Some(ref chunk) = self.chunk {
            chunk.remaining() > 0
        } else {
            false
        }
    }

    /// Consumes this `StreamReader`, returning a Tuple consisting
    /// of the underlying stream and an Option of the internal buffer,
    /// which is Some in case the buffer contains elements.
    pub fn into_inner_with_chunk(self) -> (S, Option<B>) {
        if self.has_chunk() {
            (self.inner, self.chunk)
        } else {
            (self.inner, None)
        }
    }
}

impl<S, B> StreamReader<S, B> {
    /// Gets a reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_ref(&self) -> &S {
        &self.inner
    }

    /// Gets a mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.inner
    }

    /// Gets a pinned mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut S> {
        self.project().inner
    }

    /// Consumes this `BufWriter`, returning the underlying stream.
    ///
    /// Note that any leftover data in the internal buffer is lost.
    /// If you additionally want access to the internal buffer use
    /// [`into_inner_with_chunk`].
    ///
    /// [`into_inner_with_chunk`]: crate::io::StreamReader::into_inner_with_chunk
    pub fn into_inner(self) -> S {
        self.inner
    }
}

impl<S, B, E> AsyncRead for StreamReader<S, B>
where
    S: Stream<Item = Result<B, E>>,
    B: Buf,
    E: Into<std::io::Error>,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        if buf.remaining() == 0 {
            return Poll::Ready(Ok(()));
        }

        let inner_buf = match self.as_mut().poll_fill_buf(cx) {
            Poll::Ready(Ok(buf)) => buf,
            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
            Poll::Pending => return Poll::Pending,
        };
        let len = std::cmp::min(inner_buf.len(), buf.remaining());
        buf.put_slice(&inner_buf[..len]);

        self.consume(len);
        Poll::Ready(Ok(()))
    }
}

impl<S, B, E> AsyncBufRead for StreamReader<S, B>
where
    S: Stream<Item = Result<B, E>>,
    B: Buf,
    E: Into<std::io::Error>,
{
    fn poll_fill_buf(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        loop {
            if self.as_mut().has_chunk() {
                // This unwrap is very sad, but it can't be avoided.
                let buf = self.project().chunk.as_ref().unwrap().chunk();
                return Poll::Ready(Ok(buf));
            } else {
                match self.as_mut().project().inner.poll_next(cx) {
                    Poll::Ready(Some(Ok(chunk))) => {
                        // Go around the loop in case the chunk is empty.
                        *self.as_mut().project().chunk = Some(chunk);
                    }
                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(err.into())),
                    Poll::Ready(None) => return Poll::Ready(Ok(&[])),
                    Poll::Pending => return Poll::Pending,
                }
            }
        }
    }
    fn consume(self: Pin<&mut Self>, amt: usize) {
        if amt > 0 {
            self.project()
                .chunk
                .as_mut()
                .expect("No chunk present")
                .advance(amt);
        }
    }
}
<commit_msg>io: add lines example for `StreamReader` (#5145)<commit_after>use bytes::Buf;
use futures_core::stream::Stream;
use std::io;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncBufRead, AsyncRead, ReadBuf};

/// Convert a [`Stream`] of byte chunks into an [`AsyncRead`].
///
/// This type performs the inverse operation of [`ReaderStream`].
///
/// This type also implements the [`AsyncBufRead`] trait, so you can use it
/// to read a `Stream` of byte chunks line-by-line. See the examples below.
///
/// # Example
///
/// ```
/// use bytes::Bytes;
/// use tokio::io::{AsyncReadExt, Result};
/// use tokio_util::io::StreamReader;
/// # #[tokio::main(flavor = "current_thread")]
/// # async fn main() -> std::io::Result<()> {
///
/// // Create a stream from an iterator.
/// let stream = tokio_stream::iter(vec![
///     Result::Ok(Bytes::from_static(&[0, 1, 2, 3])),
///     Result::Ok(Bytes::from_static(&[4, 5, 6, 7])),
///     Result::Ok(Bytes::from_static(&[8, 9, 10, 11])),
/// ]);
///
/// // Convert it to an AsyncRead.
/// let mut read = StreamReader::new(stream);
///
/// // Read five bytes from the stream.
/// let mut buf = [0; 5];
/// read.read_exact(&mut buf).await?;
/// assert_eq!(buf, [0, 1, 2, 3, 4]);
///
/// // Read the rest of the current chunk.
/// assert_eq!(read.read(&mut buf).await?, 3);
/// assert_eq!(&buf[..3], [5, 6, 7]);
///
/// // Read the next chunk.
/// assert_eq!(read.read(&mut buf).await?, 4);
/// assert_eq!(&buf[..4], [8, 9, 10, 11]);
///
/// // We have now reached the end.
/// assert_eq!(read.read(&mut buf).await?, 0);
///
/// # Ok(())
/// # }
/// ```
///
/// If the stream produces errors which are not [`std::io::Error`],
/// the errors can be converted using [`StreamExt`] to map each
/// element.
///
/// ```
/// use bytes::Bytes;
/// use tokio::io::AsyncReadExt;
/// use tokio_util::io::StreamReader;
/// use tokio_stream::StreamExt;
/// # #[tokio::main(flavor = "current_thread")]
/// # async fn main() -> std::io::Result<()> {
///
/// // Create a stream from an iterator, including an error.
/// let stream = tokio_stream::iter(vec![
///     Result::Ok(Bytes::from_static(&[0, 1, 2, 3])),
///     Result::Ok(Bytes::from_static(&[4, 5, 6, 7])),
///     Result::Err("Something bad happened!")
/// ]);
///
/// // Use StreamExt to map the stream and error to a std::io::Error
/// let stream = stream.map(|result| result.map_err(|err| {
///     std::io::Error::new(std::io::ErrorKind::Other, err)
/// }));
///
/// // Convert it to an AsyncRead.
/// let mut read = StreamReader::new(stream);
///
/// // Read five bytes from the stream.
/// let mut buf = [0; 5];
/// read.read_exact(&mut buf).await?;
/// assert_eq!(buf, [0, 1, 2, 3, 4]);
///
/// // Read the rest of the current chunk.
/// assert_eq!(read.read(&mut buf).await?, 3);
/// assert_eq!(&buf[..3], [5, 6, 7]);
///
/// // Reading the next chunk will produce an error
/// let error = read.read(&mut buf).await.unwrap_err();
/// assert_eq!(error.kind(), std::io::ErrorKind::Other);
/// assert_eq!(error.into_inner().unwrap().to_string(), "Something bad happened!");
///
/// // We have now reached the end.
/// assert_eq!(read.read(&mut buf).await?, 0);
///
/// # Ok(())
/// # }
/// ```
///
/// Using the [`AsyncBufRead`] impl, you can read a `Stream` of byte chunks
/// line-by-line. Note that you will usually also need to convert the error
/// type when doing this. See the second example for an explanation of how
/// to do this.
///
/// ```
/// use tokio::io::{Result, AsyncBufReadExt};
/// use tokio_util::io::StreamReader;
/// # #[tokio::main(flavor = "current_thread")]
/// # async fn main() -> std::io::Result<()> {
///
/// // Create a stream of byte chunks.
/// let stream = tokio_stream::iter(vec![
///     Result::Ok(b"The first line.\n".as_slice()),
///     Result::Ok(b"The second line.".as_slice()),
///     Result::Ok(b"\nThe third".as_slice()),
///     Result::Ok(b" line.\nThe fourth line.\nThe fifth line.\n".as_slice()),
/// ]);
///
/// // Convert it to an AsyncRead.
/// let mut read = StreamReader::new(stream);
///
/// // Loop through the lines from the `StreamReader`.
/// let mut line = String::new();
/// let mut lines = Vec::new();
/// loop {
///     line.clear();
///     let len = read.read_line(&mut line).await?;
///     if len == 0 { break; }
///     lines.push(line.clone());
/// }
///
/// // Verify that we got the lines we expected.
/// assert_eq!(
///     lines,
///     vec![
///         "The first line.\n",
///         "The second line.\n",
///         "The third line.\n",
///         "The fourth line.\n",
///         "The fifth line.\n",
///     ]
/// );
/// # Ok(())
/// # }
/// ```
///
/// [`AsyncRead`]: tokio::io::AsyncRead
/// [`AsyncBufRead`]: tokio::io::AsyncBufRead
/// [`Stream`]: futures_core::Stream
/// [`ReaderStream`]: crate::io::ReaderStream
/// [`StreamExt`]: https://docs.rs/tokio-stream/latest/tokio_stream/trait.StreamExt.html
#[derive(Debug)]
pub struct StreamReader<S, B> {
    // This field is pinned.
    inner: S,
    // This field is not pinned.
    chunk: Option<B>,
}

impl<S, B, E> StreamReader<S, B>
where
    S: Stream<Item = Result<B, E>>,
    B: Buf,
    E: Into<std::io::Error>,
{
    /// Convert a stream of byte chunks into an [`AsyncRead`](tokio::io::AsyncRead).
    ///
    /// The item should be a [`Result`] with the ok variant being something that
    /// implements the [`Buf`] trait (e.g. `Vec<u8>` or `Bytes`). The error
    /// should be convertible into an [io error].
    ///
    /// [`Result`]: std::result::Result
    /// [`Buf`]: bytes::Buf
    /// [io error]: std::io::Error
    pub fn new(stream: S) -> Self {
        Self {
            inner: stream,
            chunk: None,
        }
    }

    /// Do we have a chunk and is it non-empty?
    fn has_chunk(&self) -> bool {
        if let Some(ref chunk) = self.chunk {
            chunk.remaining() > 0
        } else {
            false
        }
    }

    /// Consumes this `StreamReader`, returning a Tuple consisting
    /// of the underlying stream and an Option of the internal buffer,
    /// which is Some in case the buffer contains elements.
    pub fn into_inner_with_chunk(self) -> (S, Option<B>) {
        if self.has_chunk() {
            (self.inner, self.chunk)
        } else {
            (self.inner, None)
        }
    }
}

impl<S, B> StreamReader<S, B> {
    /// Gets a reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_ref(&self) -> &S {
        &self.inner
    }

    /// Gets a mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_mut(&mut self) -> &mut S {
        &mut self.inner
    }

    /// Gets a pinned mutable reference to the underlying stream.
    ///
    /// It is inadvisable to directly read from the underlying stream.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut S> {
        self.project().inner
    }

    /// Consumes this `BufWriter`, returning the underlying stream.
    ///
    /// Note that any leftover data in the internal buffer is lost.
    /// If you additionally want access to the internal buffer use
    /// [`into_inner_with_chunk`].
    ///
    /// [`into_inner_with_chunk`]: crate::io::StreamReader::into_inner_with_chunk
    pub fn into_inner(self) -> S {
        self.inner
    }
}

impl<S, B, E> AsyncRead for StreamReader<S, B>
where
    S: Stream<Item = Result<B, E>>,
    B: Buf,
    E: Into<std::io::Error>,
{
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        if buf.remaining() == 0 {
            return Poll::Ready(Ok(()));
        }

        let inner_buf = match self.as_mut().poll_fill_buf(cx) {
            Poll::Ready(Ok(buf)) => buf,
            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
            Poll::Pending => return Poll::Pending,
        };
        let len = std::cmp::min(inner_buf.len(), buf.remaining());
        buf.put_slice(&inner_buf[..len]);

        self.consume(len);
        Poll::Ready(Ok(()))
    }
}

impl<S, B, E> AsyncBufRead for StreamReader<S, B>
where
    S: Stream<Item = Result<B, E>>,
    B: Buf,
    E: Into<std::io::Error>,
{
    fn poll_fill_buf(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        loop {
            if self.as_mut().has_chunk() {
                // This unwrap is very sad, but it can't be avoided.
                let buf = self.project().chunk.as_ref().unwrap().chunk();
                return Poll::Ready(Ok(buf));
            } else {
                match self.as_mut().project().inner.poll_next(cx) {
                    Poll::Ready(Some(Ok(chunk))) => {
                        // Go around the loop in case the chunk is empty.
                        *self.as_mut().project().chunk = Some(chunk);
                    }
                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(err.into())),
                    Poll::Ready(None) => return Poll::Ready(Ok(&[])),
                    Poll::Pending => return Poll::Pending,
                }
            }
        }
    }
    fn consume(self: Pin<&mut Self>, amt: usize) {
        if amt > 0 {
            self.project()
                .chunk
                .as_mut()
                .expect("No chunk present")
                .advance(amt);
        }
    }
}

// The code below is a manual expansion of the code that pin-project-lite would
// generate. This is done because pin-project-lite fails by hitting the recusion
// limit on this struct. (Every line of documentation is handled recursively by
// the macro.)

impl<S: Unpin, B> Unpin for StreamReader<S, B> {}

struct StreamReaderProject<'a, S, B> {
    inner: Pin<&'a mut S>,
    chunk: &'a mut Option<B>,
}

impl<S, B> StreamReader<S, B> {
    #[inline]
    fn project(self: Pin<&mut Self>) -> StreamReaderProject<'_, S, B> {
        // SAFETY: We define that only `inner` should be pinned when `Self` is
        // and have an appropriate `impl Unpin` for this.
        let me = unsafe { Pin::into_inner_unchecked(self) };
        StreamReaderProject {
            inner: unsafe { Pin::new_unchecked(&mut me.inner) },
            chunk: &mut me.chunk,
        }
    }
}
<|endoftext|>


<commit_before>use crate::error::Error;
use crate::module::{AddrDetails, GlobalSpec, HeapSpec, Module, ModuleInternal, TableElement};
use libc::c_void;
use libloading::Library;
use lucet_module::{
    FunctionHandle, FunctionIndex, FunctionPointer, FunctionSpec, ModuleData, ModuleFeatures,
    ModuleSignature, PublicKey, SerializedModule, Signature, VersionInfo, LUCET_MODULE_SYM,
};
use std::ffi::CStr;
use std::mem::MaybeUninit;
use std::path::Path;
use std::slice;
use std::slice::from_raw_parts;
use std::sync::Arc;

use raw_cpuid::CpuId;

fn check_feature_support(module_features: &ModuleFeatures) -> Result<(), Error> {
    let cpuid = CpuId::new();

    fn missing_feature(feature: &str) -> Error {
        Error::Unsupported(format!(
            "Module requires feature host does not support: {}",
            feature
        ))
    }

    let info = cpuid
        .get_feature_info()
        .ok_or_else(|| Error::Unsupported("Unable to obtain host CPU feature info!".to_string()))?;

    if module_features.sse3 && !info.has_sse3() {
        return Err(missing_feature("SSE3"));
    }
    if module_features.ssse3 && !info.has_ssse3() {
        return Err(missing_feature("SSS3"));
    }
    if module_features.sse41 && !info.has_sse41() {
        return Err(missing_feature("SSE4.1"));
    }
    if module_features.sse42 && !info.has_sse42() {
        return Err(missing_feature("SSE4.2"));
    }
    if module_features.avx && !info.has_avx() {
        return Err(missing_feature("AVX"));
    }
    if module_features.popcnt && !info.has_popcnt() {
        return Err(missing_feature("POPCNT"));
    }

    let info = cpuid.get_extended_feature_info().ok_or_else(|| {
        Error::Unsupported("Unable to obtain host CPU extended feature info!".to_string())
    })?;

    if module_features.bmi1 && !info.has_bmi1() {
        return Err(missing_feature("BMI1"));
    }

    if module_features.bmi2 && !info.has_bmi2() {
        return Err(missing_feature("BMI2"));
    }

    let info = cpuid.get_extended_function_info().ok_or_else(|| {
        Error::Unsupported("Unable to obtain host CPU extended function info!".to_string())
    })?;

    if module_features.lzcnt && !info.has_lzcnt() {
        return Err(missing_feature("LZCNT"));
    }

    // Features are fine, we're compatible!
    Ok(())
}

/// A Lucet module backed by a dynamically-loaded shared object.
pub struct DlModule {
    lib: Library,

    /// Base address of the dynamically-loaded module
    fbase: *const c_void,

    /// Metadata decoded from inside the module
    module: lucet_module::Module<'static>,
}

// for the one raw pointer only
unsafe impl Send for DlModule {}
unsafe impl Sync for DlModule {}

impl DlModule {
    /// Create a module, loading code from a shared object on the filesystem.
    pub fn load<P: AsRef<Path>>(so_path: P) -> Result<Arc<Self>, Error> {
        Self::load_and_maybe_verify(so_path, None)
    }

    /// Create a module, loading code from a shared object on the filesystem
    /// and verifying it using a public key if one has been supplied.
    pub fn load_and_verify<P: AsRef<Path>>(so_path: P, pk: PublicKey) -> Result<Arc<Self>, Error> {
        Self::load_and_maybe_verify(so_path, Some(pk))
    }

    fn load_and_maybe_verify<P: AsRef<Path>>(
        so_path: P,
        pk: Option<PublicKey>,
    ) -> Result<Arc<Self>, Error> {
        // Load the dynamic library. The undefined symbols corresponding to the lucet_syscall_
        // functions will be provided by the current executable.  We trust our wasm->dylib compiler
        // to make sure these function calls are the way the dylib can touch memory outside of its
        // stack and heap.
        let abs_so_path = so_path.as_ref().canonicalize().map_err(Error::DlError)?;
        let lib = Library::new(abs_so_path.as_os_str()).map_err(Error::DlError)?;

        let serialized_module_ptr = unsafe {
            lib.get::<*const SerializedModule>(LUCET_MODULE_SYM.as_bytes())
                .map_err(|e| {
                    lucet_incorrect_module!("error loading required symbol `lucet_module`: {}", e)
                })?
        };

        let serialized_module: &SerializedModule =
            unsafe { serialized_module_ptr.as_ref().unwrap() };

        let module_version = serialized_module.version.clone();

        let runtime_version =
            VersionInfo::current(include_str!(concat!(env!("OUT_DIR"), "/commit_hash")).as_bytes());

        if !module_version.valid() {
            return Err(lucet_incorrect_module!("reserved bit is not set. This module is likely too old for this lucet-runtime to load."));
        } else if !runtime_version.compatible_with(&module_version) {
            return Err(lucet_incorrect_module!(
                "version mismatch. module has version {}, while this runtime is version {}",
                module_version,
                runtime_version,
            ));
        }

        // Deserialize the slice into ModuleData, which will hold refs into the loaded
        // shared object file in `module_data_slice`. Both of these get a 'static lifetime because
        // Rust doesn't have a safe way to describe that their lifetime matches the containing
        // struct (and the dll).
        //
        // The exposed lifetime of ModuleData will be the same as the lifetime of the
        // dynamically loaded library. This makes the interface safe.
        let module_data_slice: &'static [u8] = unsafe {
            slice::from_raw_parts(
                serialized_module.module_data_ptr as *const u8,
                serialized_module.module_data_len as usize,
            )
        };
        let module_data = ModuleData::deserialize(module_data_slice)?;

        check_feature_support(module_data.features())?;

        // If a public key has been provided, verify the module signature
        // The TOCTOU issue is unavoidable without reimplenting `dlopen(3)`
        if let Some(pk) = pk {
            ModuleSignature::verify(so_path, &pk, &module_data)?;
        }

        let fbase = if let Some(dli) =
            dladdr(serialized_module as *const SerializedModule as *const c_void)
        {
            dli.dli_fbase
        } else {
            std::ptr::null()
        };

        if serialized_module.tables_len > std::u32::MAX as u64 {
            lucet_incorrect_module!("table segment too long: {}", serialized_module.tables_len);
        }
        let tables: &'static [&'static [TableElement]] = unsafe {
            from_raw_parts(
                serialized_module.tables_ptr as *const &[TableElement],
                serialized_module.tables_len as usize,
            )
        };

        let function_manifest = if serialized_module.function_manifest_ptr != 0 {
            unsafe {
                from_raw_parts(
                    serialized_module.function_manifest_ptr as *const FunctionSpec,
                    serialized_module.function_manifest_len as usize,
                )
            }
        } else {
            &[]
        };

        Ok(Arc::new(DlModule {
            lib,
            fbase,
            module: lucet_module::Module {
                version: module_version,
                module_data,
                tables,
                function_manifest,
            },
        }))
    }
}

impl Module for DlModule {}

impl ModuleInternal for DlModule {
    fn is_instruction_count_instrumented(&self) -> bool {
        self.module.module_data.features().instruction_count
    }

    fn heap_spec(&self) -> Option<&HeapSpec> {
        self.module.module_data.heap_spec()
    }

    fn globals(&self) -> &[GlobalSpec<'_>] {
        self.module.module_data.globals_spec()
    }

    fn get_sparse_page_data(&self, page: usize) -> Option<&[u8]> {
        if let Some(ref sparse_data) = self.module.module_data.sparse_data() {
            *sparse_data.get_page(page)
        } else {
            None
        }
    }

    fn sparse_page_data_len(&self) -> usize {
        self.module
            .module_data
            .sparse_data()
            .map(|d| d.len())
            .unwrap_or(0)
    }

    fn table_elements(&self) -> Result<&[TableElement], Error> {
        match self.module.tables.get(0) {
            Some(table) => Ok(table),
            None => Err(lucet_incorrect_module!("table 0 is not present")),
        }
    }

    fn get_export_func(&self, sym: &str) -> Result<FunctionHandle, Error> {
        self.module
            .module_data
            .get_export_func_id(sym)
            .ok_or_else(|| Error::SymbolNotFound(sym.to_string()))
            .map(|id| {
                let ptr = self.function_manifest()[id.as_u32() as usize].ptr();
                FunctionHandle { ptr, id }
            })
    }

    fn get_func_from_idx(&self, table_id: u32, func_id: u32) -> Result<FunctionHandle, Error> {
        if table_id != 0 {
            return Err(Error::FuncNotFound(table_id, func_id));
        }
        let table = self.table_elements()?;
        let func = table
            .get(func_id as usize)
            .map(|element| element.function_pointer())
            .ok_or(Error::FuncNotFound(table_id, func_id))?;

        Ok(self.function_handle_from_ptr(func))
    }

    fn get_start_func(&self) -> Result<Option<FunctionHandle>, Error> {
        // `guest_start` is a pointer to the function the module designates as the start function,
        // since we can't have multiple symbols pointing to the same function and guest code might
        // call it in the normal course of execution
        if let Ok(start_func) = unsafe { self.lib.get::<*const extern "C" fn()>(b"guest_start") } {
            if start_func.is_null() {
                lucet_incorrect_module!("`guest_start` is defined but null");
            }
            Ok(Some(self.function_handle_from_ptr(
                FunctionPointer::from_usize(unsafe { **start_func } as usize),
            )))
        } else {
            Ok(None)
        }
    }

    fn function_manifest(&self) -> &[FunctionSpec] {
        self.module.function_manifest
    }

    fn addr_details(&self, addr: *const c_void) -> Result<Option<AddrDetails>, Error> {
        if let Some(dli) = dladdr(addr) {
            let file_name = if dli.dli_fname.is_null() {
                None
            } else {
                Some(unsafe { CStr::from_ptr(dli.dli_fname).to_owned().into_string()? })
            };
            let sym_name = if dli.dli_sname.is_null() {
                None
            } else {
                Some(unsafe { CStr::from_ptr(dli.dli_sname).to_owned().into_string()? })
            };
            Ok(Some(AddrDetails {
                in_module_code: dli.dli_fbase as *const c_void == self.fbase,
                file_name,
                sym_name,
            }))
        } else {
            Ok(None)
        }
    }

    fn get_signature(&self, fn_id: FunctionIndex) -> &Signature {
        self.module.module_data.get_signature(fn_id)
    }
}

// TODO: PR to nix or libloading?
// TODO: possibly not safe to use without grabbing the mutex within libloading::Library?
fn dladdr(addr: *const c_void) -> Option<libc::Dl_info> {
    let mut info = MaybeUninit::<libc::Dl_info>::uninit();
    let res = unsafe { libc::dladdr(addr, info.as_mut_ptr()) };
    if res != 0 {
        Some(unsafe { info.assume_init() })
    } else {
        None
    }
}
<commit_msg>Check extended CPU features only when necessary (#426)<commit_after>use crate::error::Error;
use crate::module::{AddrDetails, GlobalSpec, HeapSpec, Module, ModuleInternal, TableElement};
use libc::c_void;
use libloading::Library;
use lucet_module::{
    FunctionHandle, FunctionIndex, FunctionPointer, FunctionSpec, ModuleData, ModuleFeatures,
    ModuleSignature, PublicKey, SerializedModule, Signature, VersionInfo, LUCET_MODULE_SYM,
};
use std::ffi::CStr;
use std::mem::MaybeUninit;
use std::path::Path;
use std::slice;
use std::slice::from_raw_parts;
use std::sync::Arc;

use raw_cpuid::CpuId;

fn check_feature_support(module_features: &ModuleFeatures) -> Result<(), Error> {
    let cpuid = CpuId::new();

    fn missing_feature(feature: &str) -> Error {
        Error::Unsupported(format!(
            "Module requires feature host does not support: {}",
            feature
        ))
    }

    let info = cpuid
        .get_feature_info()
        .ok_or_else(|| Error::Unsupported("Unable to obtain host CPU feature info!".to_string()))?;

    if module_features.sse3 && !info.has_sse3() {
        return Err(missing_feature("SSE3"));
    }
    if module_features.ssse3 && !info.has_ssse3() {
        return Err(missing_feature("SSS3"));
    }
    if module_features.sse41 && !info.has_sse41() {
        return Err(missing_feature("SSE4.1"));
    }
    if module_features.sse42 && !info.has_sse42() {
        return Err(missing_feature("SSE4.2"));
    }
    if module_features.avx && !info.has_avx() {
        return Err(missing_feature("AVX"));
    }
    if module_features.popcnt && !info.has_popcnt() {
        return Err(missing_feature("POPCNT"));
    }

    if module_features.bmi1 || module_features.bmi2 {
        let info = cpuid.get_extended_feature_info().ok_or_else(|| {
            Error::Unsupported("Unable to obtain host CPU extended feature info!".to_string())
        })?;

        if module_features.bmi1 && !info.has_bmi1() {
            return Err(missing_feature("BMI1"));
        }

        if module_features.bmi2 && !info.has_bmi2() {
            return Err(missing_feature("BMI2"));
        }
    }

    if module_features.lzcnt {
        let info = cpuid.get_extended_function_info().ok_or_else(|| {
            Error::Unsupported("Unable to obtain host CPU extended function info!".to_string())
        })?;

        if module_features.lzcnt && !info.has_lzcnt() {
            return Err(missing_feature("LZCNT"));
        }
    }

    // Features are fine, we're compatible!
    Ok(())
}

/// A Lucet module backed by a dynamically-loaded shared object.
pub struct DlModule {
    lib: Library,

    /// Base address of the dynamically-loaded module
    fbase: *const c_void,

    /// Metadata decoded from inside the module
    module: lucet_module::Module<'static>,
}

// for the one raw pointer only
unsafe impl Send for DlModule {}
unsafe impl Sync for DlModule {}

impl DlModule {
    /// Create a module, loading code from a shared object on the filesystem.
    pub fn load<P: AsRef<Path>>(so_path: P) -> Result<Arc<Self>, Error> {
        Self::load_and_maybe_verify(so_path, None)
    }

    /// Create a module, loading code from a shared object on the filesystem
    /// and verifying it using a public key if one has been supplied.
    pub fn load_and_verify<P: AsRef<Path>>(so_path: P, pk: PublicKey) -> Result<Arc<Self>, Error> {
        Self::load_and_maybe_verify(so_path, Some(pk))
    }

    fn load_and_maybe_verify<P: AsRef<Path>>(
        so_path: P,
        pk: Option<PublicKey>,
    ) -> Result<Arc<Self>, Error> {
        // Load the dynamic library. The undefined symbols corresponding to the lucet_syscall_
        // functions will be provided by the current executable.  We trust our wasm->dylib compiler
        // to make sure these function calls are the way the dylib can touch memory outside of its
        // stack and heap.
        let abs_so_path = so_path.as_ref().canonicalize().map_err(Error::DlError)?;
        let lib = Library::new(abs_so_path.as_os_str()).map_err(Error::DlError)?;

        let serialized_module_ptr = unsafe {
            lib.get::<*const SerializedModule>(LUCET_MODULE_SYM.as_bytes())
                .map_err(|e| {
                    lucet_incorrect_module!("error loading required symbol `lucet_module`: {}", e)
                })?
        };

        let serialized_module: &SerializedModule =
            unsafe { serialized_module_ptr.as_ref().unwrap() };

        let module_version = serialized_module.version.clone();

        let runtime_version =
            VersionInfo::current(include_str!(concat!(env!("OUT_DIR"), "/commit_hash")).as_bytes());

        if !module_version.valid() {
            return Err(lucet_incorrect_module!("reserved bit is not set. This module is likely too old for this lucet-runtime to load."));
        } else if !runtime_version.compatible_with(&module_version) {
            return Err(lucet_incorrect_module!(
                "version mismatch. module has version {}, while this runtime is version {}",
                module_version,
                runtime_version,
            ));
        }

        // Deserialize the slice into ModuleData, which will hold refs into the loaded
        // shared object file in `module_data_slice`. Both of these get a 'static lifetime because
        // Rust doesn't have a safe way to describe that their lifetime matches the containing
        // struct (and the dll).
        //
        // The exposed lifetime of ModuleData will be the same as the lifetime of the
        // dynamically loaded library. This makes the interface safe.
        let module_data_slice: &'static [u8] = unsafe {
            slice::from_raw_parts(
                serialized_module.module_data_ptr as *const u8,
                serialized_module.module_data_len as usize,
            )
        };
        let module_data = ModuleData::deserialize(module_data_slice)?;

        check_feature_support(module_data.features())?;

        // If a public key has been provided, verify the module signature
        // The TOCTOU issue is unavoidable without reimplenting `dlopen(3)`
        if let Some(pk) = pk {
            ModuleSignature::verify(so_path, &pk, &module_data)?;
        }

        let fbase = if let Some(dli) =
            dladdr(serialized_module as *const SerializedModule as *const c_void)
        {
            dli.dli_fbase
        } else {
            std::ptr::null()
        };

        if serialized_module.tables_len > std::u32::MAX as u64 {
            lucet_incorrect_module!("table segment too long: {}", serialized_module.tables_len);
        }
        let tables: &'static [&'static [TableElement]] = unsafe {
            from_raw_parts(
                serialized_module.tables_ptr as *const &[TableElement],
                serialized_module.tables_len as usize,
            )
        };

        let function_manifest = if serialized_module.function_manifest_ptr != 0 {
            unsafe {
                from_raw_parts(
                    serialized_module.function_manifest_ptr as *const FunctionSpec,
                    serialized_module.function_manifest_len as usize,
                )
            }
        } else {
            &[]
        };

        Ok(Arc::new(DlModule {
            lib,
            fbase,
            module: lucet_module::Module {
                version: module_version,
                module_data,
                tables,
                function_manifest,
            },
        }))
    }
}

impl Module for DlModule {}

impl ModuleInternal for DlModule {
    fn is_instruction_count_instrumented(&self) -> bool {
        self.module.module_data.features().instruction_count
    }

    fn heap_spec(&self) -> Option<&HeapSpec> {
        self.module.module_data.heap_spec()
    }

    fn globals(&self) -> &[GlobalSpec<'_>] {
        self.module.module_data.globals_spec()
    }

    fn get_sparse_page_data(&self, page: usize) -> Option<&[u8]> {
        if let Some(ref sparse_data) = self.module.module_data.sparse_data() {
            *sparse_data.get_page(page)
        } else {
            None
        }
    }

    fn sparse_page_data_len(&self) -> usize {
        self.module
            .module_data
            .sparse_data()
            .map(|d| d.len())
            .unwrap_or(0)
    }

    fn table_elements(&self) -> Result<&[TableElement], Error> {
        match self.module.tables.get(0) {
            Some(table) => Ok(table),
            None => Err(lucet_incorrect_module!("table 0 is not present")),
        }
    }

    fn get_export_func(&self, sym: &str) -> Result<FunctionHandle, Error> {
        self.module
            .module_data
            .get_export_func_id(sym)
            .ok_or_else(|| Error::SymbolNotFound(sym.to_string()))
            .map(|id| {
                let ptr = self.function_manifest()[id.as_u32() as usize].ptr();
                FunctionHandle { ptr, id }
            })
    }

    fn get_func_from_idx(&self, table_id: u32, func_id: u32) -> Result<FunctionHandle, Error> {
        if table_id != 0 {
            return Err(Error::FuncNotFound(table_id, func_id));
        }
        let table = self.table_elements()?;
        let func = table
            .get(func_id as usize)
            .map(|element| element.function_pointer())
            .ok_or(Error::FuncNotFound(table_id, func_id))?;

        Ok(self.function_handle_from_ptr(func))
    }

    fn get_start_func(&self) -> Result<Option<FunctionHandle>, Error> {
        // `guest_start` is a pointer to the function the module designates as the start function,
        // since we can't have multiple symbols pointing to the same function and guest code might
        // call it in the normal course of execution
        if let Ok(start_func) = unsafe { self.lib.get::<*const extern "C" fn()>(b"guest_start") } {
            if start_func.is_null() {
                lucet_incorrect_module!("`guest_start` is defined but null");
            }
            Ok(Some(self.function_handle_from_ptr(
                FunctionPointer::from_usize(unsafe { **start_func } as usize),
            )))
        } else {
            Ok(None)
        }
    }

    fn function_manifest(&self) -> &[FunctionSpec] {
        self.module.function_manifest
    }

    fn addr_details(&self, addr: *const c_void) -> Result<Option<AddrDetails>, Error> {
        if let Some(dli) = dladdr(addr) {
            let file_name = if dli.dli_fname.is_null() {
                None
            } else {
                Some(unsafe { CStr::from_ptr(dli.dli_fname).to_owned().into_string()? })
            };
            let sym_name = if dli.dli_sname.is_null() {
                None
            } else {
                Some(unsafe { CStr::from_ptr(dli.dli_sname).to_owned().into_string()? })
            };
            Ok(Some(AddrDetails {
                in_module_code: dli.dli_fbase as *const c_void == self.fbase,
                file_name,
                sym_name,
            }))
        } else {
            Ok(None)
        }
    }

    fn get_signature(&self, fn_id: FunctionIndex) -> &Signature {
        self.module.module_data.get_signature(fn_id)
    }
}

// TODO: PR to nix or libloading?
// TODO: possibly not safe to use without grabbing the mutex within libloading::Library?
fn dladdr(addr: *const c_void) -> Option<libc::Dl_info> {
    let mut info = MaybeUninit::<libc::Dl_info>::uninit();
    let res = unsafe { libc::dladdr(addr, info.as_mut_ptr()) };
    if res != 0 {
        Some(unsafe { info.assume_init() })
    } else {
        None
    }
}
<|endoftext|>


<commit_before>use common::util::*;

static UTIL_NAME: &'static str = "basename";
fn new_ucmd() -> UCommand {
    TestScenario::new(UTIL_NAME).ucmd()
}

fn expect_successful_stdout(input: Vec<&str>, expected: &str) {
    let results = new_ucmd()
        .args(&input).run();
    assert_empty_stderr!(results);
    assert!(results.success);
    assert_eq!(expected, results.stdout.trim_right());
}

#[test]
fn test_directory() {
    let dir = "/root/alpha/beta/gamma/delta/epsilon/omega/";
    expect_successful_stdout(vec![dir], "omega");
}

#[test]
fn test_file() {
    let file = "/etc/passwd";
    expect_successful_stdout(vec![file], "passwd");
}

#[test]
fn test_remove_suffix() {
    let path = "/usr/local/bin/reallylongexecutable.exe";
    expect_successful_stdout(vec![path, ".exe"], "reallylongexecutable");
}

#[test]
fn test_dont_remove_suffix() {
    let path = "/foo/bar/baz";
    expect_successful_stdout(vec![path, "baz"], "baz");
}

fn expect_error(input: Vec<&str>, expected_stdout: &str) {
    let results = new_ucmd()
        .args(&input).run();
    assert!(!results.success);
    assert!(results.stderr.len() > 0);
    assert_eq!(expected_stdout, results.stdout.trim_right());
}

#[cfg_attr(not(feature="test_unimplemented"),ignore)]
#[test]
fn test_multiple_param() {
    for multiple_param in vec!["-a", "--multiple"] {
        let path = "/foo/bar/baz";
        expect_successful_stdout(vec![multiple_param, path, path], "baz\nbaz");
    }
}

#[cfg_attr(not(feature="test_unimplemented"),ignore)]
#[test]
fn test_suffix_param() {
    for suffix_param in vec!["-s", "--suffix"] {
        let path = "/foo/bar/baz.exe";
        let suffix = ".exe";
        expect_successful_stdout(
            vec![suffix_param, suffix, path, path],
            "baz\nbaz"
        );
    }
}

#[cfg_attr(not(feature="test_unimplemented"),ignore)]
#[test]
fn test_zero_param() {
    for zero_param in vec!["-z", "--zero"] {
        let path = "/foo/bar/baz";
        expect_successful_stdout(vec![zero_param, "-a", path, path], "baz\0baz\0");
    }
}

#[test]
fn test_invalid_option() {
    let path = "/foo/bar/baz";
    expect_error(vec!["-q", path], "");
}

#[test]
fn test_no_args() {
    expect_error(vec![], "");
}

#[test]
fn test_too_many_args() {
    expect_error(vec!["a", "b", "c"], "");
}
<commit_msg>tests/basename: use chaining for brevity and consistency<commit_after>use common::util::*;

static UTIL_NAME: &'static str = "basename";
fn new_ucmd() -> UCommand {
    TestScenario::new(UTIL_NAME).ucmd()
}

#[test]
fn test_directory() {
    new_ucmd().args(&["/root/alpha/beta/gamma/delta/epsilon/omega/"])
        .succeeds().stdout_only("omega");
}

#[test]
fn test_file() {
    new_ucmd().args(&["/etc/passwd"]).succeeds().stdout_only("passwd");
}

#[test]
fn test_remove_suffix() {
    new_ucmd().args(&["/usr/local/bin/reallylongexecutable.exe", ".exe"])
        .succeeds().stdout_only("reallylongexecutable");
}

#[test]
fn test_dont_remove_suffix() {
    new_ucmd().args(&["/foo/bar/baz", "baz"]).succeeds().stdout_only( "baz");
}

#[cfg_attr(not(feature="test_unimplemented"),ignore)]
#[test]
fn test_multiple_param() {
    for multiple_param in vec!["-a", "--multiple"] {
        let path = "/foo/bar/baz";
        new_ucmd().args(&[multiple_param, path, path])
            .succeeds().stdout_only("baz\nbaz");
    }
}

#[cfg_attr(not(feature="test_unimplemented"),ignore)]
#[test]
fn test_suffix_param() {
    for suffix_param in vec!["-s", "--suffix"] {
        let path = "/foo/bar/baz.exe";
        new_ucmd()
            .args(&[suffix_param, ".exe", path, path])
            .succeeds().stdout_only("baz\nbaz");
    }
}

#[cfg_attr(not(feature="test_unimplemented"),ignore)]
#[test]
fn test_zero_param() {
    for zero_param in vec!["-z", "--zero"] {
        let path = "/foo/bar/baz";
        new_ucmd().args(&[zero_param, "-a", path, path])
            .succeeds().stdout_only("baz\0baz\0");
    }
}

fn expect_error(input: Vec<&str>) {
    assert!(new_ucmd().args(&input)
                .fails().no_stdout().stderr.len() > 0);
}

#[test]
fn test_invalid_option() {
    let path = "/foo/bar/baz";
    expect_error(vec!["-q", path]);
}

#[test]
fn test_no_args() {
    expect_error(vec![]);
}

#[test]
fn test_too_many_args() {
    expect_error(vec!["a", "b", "c"]);
}
<|endoftext|>