<commit_before>//! Backend using the pure-rust crossplatform crossterm library.
//!
//! Requires the `crossterm-backend` feature.

#![cfg(feature = "crossterm")]

use std::{
    cell::{Cell, RefCell, RefMut},
    fs::File,
    io::{self, BufWriter, Stdout, Write},
    time::Duration,
};

use crossterm::{
    cursor::{Hide, MoveTo, Show},
    event::{
        poll, read, DisableMouseCapture, EnableMouseCapture, Event as CEvent,
        KeyCode, KeyEvent as CKeyEvent, KeyModifiers,
        MouseButton as CMouseButton, MouseEvent as CMouseEvent,
    },
    execute, queue,
    style::{
        Attribute, Color, Print, SetAttribute, SetBackgroundColor,
        SetForegroundColor,
    },
    terminal::{
        self, disable_raw_mode, enable_raw_mode, Clear, ClearType,
        EnterAlternateScreen, LeaveAlternateScreen,
    },
};

use crate::{
    backend,
    event::{Event, Key, MouseButton, MouseEvent},
    theme,
    vec::Vec2,
};

/// Backend using crossterm
pub struct Backend {
    current_style: Cell<theme::ColorPair>,
    last_button: Option<MouseButton>,

    #[cfg(windows)]
    stdout: RefCell<BufWriter<Stdout>>,
    #[cfg(unix)]
    stdout: RefCell<BufWriter<File>>,
}

impl From<CMouseButton> for MouseButton {
    fn from(button: CMouseButton) -> Self {
        match button {
            CMouseButton::Left => MouseButton::Left,
            CMouseButton::Right => MouseButton::Right,
            CMouseButton::Middle => MouseButton::Middle,
        }
    }
}

impl From<KeyCode> for Key {
    fn from(code: KeyCode) -> Self {
        match code {
            KeyCode::Esc => Key::Esc,
            KeyCode::Backspace => Key::Backspace,
            KeyCode::Left => Key::Left,
            KeyCode::Right => Key::Right,
            KeyCode::Up => Key::Up,
            KeyCode::Down => Key::Down,
            KeyCode::Home => Key::Home,
            KeyCode::End => Key::End,
            KeyCode::PageUp => Key::PageUp,
            KeyCode::PageDown => Key::PageDown,
            KeyCode::Delete => Key::Del,
            KeyCode::Insert => Key::Ins,
            KeyCode::Enter => Key::Enter,
            KeyCode::Tab => Key::Tab,
            KeyCode::F(n) => Key::from_f(n),
            KeyCode::BackTab => Key::Tab, /* not supported */
            KeyCode::Char(_) => Key::Tab, /* is handled at `Event` level, use tab as default */
            KeyCode::Null => Key::Tab, /* is handled at `Event` level, use tab as default */
        }
    }
}

impl From<CKeyEvent> for Event {
    fn from(event: CKeyEvent) -> Self {
        const CTRL_ALT: KeyModifiers = KeyModifiers::from_bits_truncate(
            KeyModifiers::CONTROL.bits() | KeyModifiers::ALT.bits(),
        );
        const CTRL_SHIFT: KeyModifiers = KeyModifiers::from_bits_truncate(
            KeyModifiers::CONTROL.bits() | KeyModifiers::SHIFT.bits(),
        );
        const ALT_SHIFT: KeyModifiers = KeyModifiers::from_bits_truncate(
            KeyModifiers::ALT.bits() | KeyModifiers::SHIFT.bits(),
        );

        match event {
            // Handle Char + modifier.
            CKeyEvent {
                modifiers: KeyModifiers::CONTROL,
                code: KeyCode::Char('c'),
            } => Event::Exit,
            CKeyEvent {
                modifiers: KeyModifiers::CONTROL,
                code: KeyCode::Char(c),
            } => Event::CtrlChar(c),
            CKeyEvent {
                modifiers: KeyModifiers::ALT,
                code: KeyCode::Char(c),
            } => Event::AltChar(c),
            CKeyEvent {
                modifiers: KeyModifiers::SHIFT,
                code: KeyCode::Char(c),
            } => Event::Char(c),

            // Handle key + multiple modifiers
            CKeyEvent {
                modifiers: CTRL_ALT,
                code,
            } => Event::CtrlAlt(Key::from(code)),
            CKeyEvent {
                modifiers: CTRL_SHIFT,
                code,
            } => Event::CtrlShift(Key::from(code)),
            CKeyEvent {
                modifiers: ALT_SHIFT,
                code,
            } => Event::AltShift(Key::from(code)),

            // Handle key + single modifier
            CKeyEvent {
                modifiers: KeyModifiers::CONTROL,
                code,
            } => Event::Ctrl(Key::from(code)),
            CKeyEvent {
                modifiers: KeyModifiers::ALT,
                code,
            } => Event::Alt(Key::from(code)),
            CKeyEvent {
                modifiers: KeyModifiers::SHIFT,
                code,
            } => Event::Shift(Key::from(code)),

            CKeyEvent {
                code: KeyCode::Char(c),
                ..
            } => Event::Char(c),

            // Explicitly handle 'backtab' since crossterm does not sent SHIFT alongside the back tab key.
            CKeyEvent {
                code: KeyCode::BackTab,
                ..
            } => Event::Shift(Key::Tab),

            // All other keys.
            CKeyEvent { code, .. } => Event::Key(Key::from(code)),
        }
    }
}

impl From<theme::Color> for Color {
    fn from(base_color: theme::Color) -> Self {
        match base_color {
            theme::Color::Dark(theme::BaseColor::Black) => Color::Black,
            theme::Color::Dark(theme::BaseColor::Red) => Color::DarkRed,
            theme::Color::Dark(theme::BaseColor::Green) => Color::DarkGreen,
            theme::Color::Dark(theme::BaseColor::Yellow) => Color::DarkYellow,
            theme::Color::Dark(theme::BaseColor::Blue) => Color::DarkBlue,
            theme::Color::Dark(theme::BaseColor::Magenta) => {
                Color::DarkMagenta
            }
            theme::Color::Dark(theme::BaseColor::Cyan) => Color::DarkCyan,
            theme::Color::Dark(theme::BaseColor::White) => Color::Grey,
            theme::Color::Light(theme::BaseColor::Black) => Color::DarkGrey,
            theme::Color::Light(theme::BaseColor::Red) => Color::Red,
            theme::Color::Light(theme::BaseColor::Green) => Color::Green,
            theme::Color::Light(theme::BaseColor::Yellow) => Color::Yellow,
            theme::Color::Light(theme::BaseColor::Blue) => Color::Blue,
            theme::Color::Light(theme::BaseColor::Magenta) => Color::Magenta,
            theme::Color::Light(theme::BaseColor::Cyan) => Color::Cyan,
            theme::Color::Light(theme::BaseColor::White) => Color::White,
            theme::Color::Rgb(r, g, b) => Color::Rgb { r, g, b },
            theme::Color::RgbLowRes(r, g, b) => {
                debug_assert!(r <= 5,
                              "Red color fragment (r = {}) is out of bound. Make sure r ≤ 5.",
                              r);
                debug_assert!(g <= 5,
                              "Green color fragment (g = {}) is out of bound. Make sure g ≤ 5.",
                              g);
                debug_assert!(b <= 5,
                              "Blue color fragment (b = {}) is out of bound. Make sure b ≤ 5.",
                              b);

                Color::AnsiValue(16 + 36 * r + 6 * g + b)
            }
            theme::Color::TerminalDefault => Color::Reset,
        }
    }
}

impl Backend {
    /// Creates a new crossterm backend.
    pub fn init() -> Result<Box<dyn backend::Backend>, crossterm::ErrorKind>
    where
        Self: Sized,
    {
        enable_raw_mode()?;

        execute!(
            io::stdout(),
            EnterAlternateScreen,
            EnableMouseCapture,
            Hide
        )?;

        #[cfg(unix)]
        let stdout = RefCell::new(BufWriter::new(File::create("/dev/tty")?));
        #[cfg(windows)]
        let stdout = RefCell::new(BufWriter::new(io::stdout()));

        Ok(Box::new(Backend {
            current_style: Cell::new(theme::ColorPair::from_256colors(0, 0)),
            last_button: None,
            stdout,
        }))
    }

    fn apply_colors(&self, colors: theme::ColorPair) {
        queue!(
            self.stdout_mut(),
            SetForegroundColor(Color::from(colors.front)),
            SetBackgroundColor(Color::from(colors.back))
        )
        .unwrap();
    }

    #[cfg(unix)]
    fn stdout_mut(&self) -> RefMut<BufWriter<File>> {
        self.stdout.borrow_mut()
    }

    #[cfg(windows)]
    fn stdout_mut(&self) -> RefMut<BufWriter<Stdout>> {
        self.stdout.borrow_mut()
    }

    fn set_attr(&self, attr: Attribute) {
        queue!(self.stdout_mut(), SetAttribute(attr)).unwrap();
    }

    fn map_key(&mut self, event: CEvent) -> Event {
        match event {
            CEvent::Key(key_event) => Event::from(key_event),
            CEvent::Mouse(mouse_event) => {
                let position;
                let event;

                match mouse_event {
                    CMouseEvent::Down(button, x, y, _) => {
                        let button = MouseButton::from(button);
                        self.last_button = Some(button);
                        event = MouseEvent::Press(button);
                        position = (x, y).into();
                    }
                    CMouseEvent::Up(_, x, y, _) => {
                        event = MouseEvent::Release(self.last_button.unwrap());
                        position = (x, y).into();
                    }
                    CMouseEvent::Drag(_, x, y, _) => {
                        event = MouseEvent::Hold(self.last_button.unwrap());
                        position = (x, y).into();
                    }
                    CMouseEvent::ScrollDown(x, y, _) => {
                        event = MouseEvent::WheelDown;
                        position = (x, y).into();
                    }
                    CMouseEvent::ScrollUp(x, y, _) => {
                        event = MouseEvent::WheelUp;
                        position = (x, y).into();
                    }
                };

                Event::Mouse {
                    event,
                    position,
                    offset: Vec2::zero(),
                }
            }
            CEvent::Resize(_, _) => Event::WindowResize,
        }
    }
}

impl backend::Backend for Backend {
    fn poll_event(&mut self) -> Option<Event> {
        match poll(Duration::from_millis(1)) {
            Ok(true) => match read() {
                Ok(event) => Some(self.map_key(event)),
                Err(e) => panic!("{:?}", e),
            },
            _ => None,
        }
    }

    fn finish(&mut self) {
        // We have to execute the show cursor command at the `stdout`.
        execute!(
            io::stdout(),
            LeaveAlternateScreen,
            DisableMouseCapture,
            Show
        )
        .expect("Can not disable mouse capture or show cursor.");

        disable_raw_mode().unwrap();
    }

    fn refresh(&mut self) {
        self.stdout_mut().flush().unwrap();
    }

    fn has_colors(&self) -> bool {
        // TODO: color support detection?
        true
    }

    fn screen_size(&self) -> Vec2 {
        let size = terminal::size().unwrap_or((1, 1));
        Vec2::from(size)
    }

    fn print_at(&self, pos: Vec2, text: &str) {
        queue!(
            self.stdout_mut(),
            MoveTo(pos.x as u16, pos.y as u16),
            Print(text)
        )
        .unwrap();
    }

    fn print_at_rep(&self, pos: Vec2, repetitions: usize, text: &str) {
        if repetitions > 0 {
            let mut out = self.stdout_mut();

            queue!(out, MoveTo(pos.x as u16, pos.y as u16)).unwrap();

            out.write_all(text.as_bytes()).unwrap();

            let mut dupes_left = repetitions - 1;
            while dupes_left > 0 {
                out.write_all(text.as_bytes()).unwrap();
                dupes_left -= 1;
            }
        }
    }

    fn clear(&self, color: theme::Color) {
        self.apply_colors(theme::ColorPair {
            front: color,
            back: color,
        });

        queue!(self.stdout_mut(), Clear(ClearType::All)).unwrap();
    }

    fn set_color(&self, color: theme::ColorPair) -> theme::ColorPair {
        let current_style = self.current_style.get();

        if current_style != color {
            self.apply_colors(color);
            self.current_style.set(color);
        }

        current_style
    }

    fn set_effect(&self, effect: theme::Effect) {
        match effect {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => self.set_attr(Attribute::Reverse),
            theme::Effect::Bold => self.set_attr(Attribute::Bold),
            theme::Effect::Italic => self.set_attr(Attribute::Italic),
            theme::Effect::Strikethrough => {
                self.set_attr(Attribute::CrossedOut)
            }
            theme::Effect::Underline => self.set_attr(Attribute::Underlined),
        }
    }

    fn unset_effect(&self, effect: theme::Effect) {
        match effect {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => self.set_attr(Attribute::NoReverse),
            theme::Effect::Bold => self.set_attr(Attribute::NormalIntensity),
            theme::Effect::Italic => self.set_attr(Attribute::NoItalic),
            theme::Effect::Strikethrough => {
                self.set_attr(Attribute::NotCrossedOut)
            }
            theme::Effect::Underline => self.set_attr(Attribute::NoUnderline),
        }
    }

    fn name(&self) -> &str {
        "crossterm"
    }
}
<commit_msg>Crossterm: define platform-dependant Stdout type<commit_after>//! Backend using the pure-rust crossplatform crossterm library.
//!
//! Requires the `crossterm-backend` feature.

#![cfg(feature = "crossterm")]

use std::{
    cell::{Cell, RefCell, RefMut},
    fs::File,
    io::{self, BufWriter, Write},
    time::Duration,
};

use crossterm::{
    cursor::{Hide, MoveTo, Show},
    event::{
        poll, read, DisableMouseCapture, EnableMouseCapture, Event as CEvent,
        KeyCode, KeyEvent as CKeyEvent, KeyModifiers,
        MouseButton as CMouseButton, MouseEvent as CMouseEvent,
    },
    execute, queue,
    style::{
        Attribute, Color, Print, SetAttribute, SetBackgroundColor,
        SetForegroundColor,
    },
    terminal::{
        self, disable_raw_mode, enable_raw_mode, Clear, ClearType,
        EnterAlternateScreen, LeaveAlternateScreen,
    },
};

use crate::{
    backend,
    event::{Event, Key, MouseButton, MouseEvent},
    theme,
    vec::Vec2,
};

#[cfg(windows)]
type Stdout = io::Stdout;

#[cfg(unix)]
type Stdout = File;

/// Backend using crossterm
pub struct Backend {
    current_style: Cell<theme::ColorPair>,
    last_button: Option<MouseButton>,

    stdout: RefCell<BufWriter<Stdout>>,
}

impl From<CMouseButton> for MouseButton {
    fn from(button: CMouseButton) -> Self {
        match button {
            CMouseButton::Left => MouseButton::Left,
            CMouseButton::Right => MouseButton::Right,
            CMouseButton::Middle => MouseButton::Middle,
        }
    }
}

impl From<KeyCode> for Key {
    fn from(code: KeyCode) -> Self {
        match code {
            KeyCode::Esc => Key::Esc,
            KeyCode::Backspace => Key::Backspace,
            KeyCode::Left => Key::Left,
            KeyCode::Right => Key::Right,
            KeyCode::Up => Key::Up,
            KeyCode::Down => Key::Down,
            KeyCode::Home => Key::Home,
            KeyCode::End => Key::End,
            KeyCode::PageUp => Key::PageUp,
            KeyCode::PageDown => Key::PageDown,
            KeyCode::Delete => Key::Del,
            KeyCode::Insert => Key::Ins,
            KeyCode::Enter => Key::Enter,
            KeyCode::Tab => Key::Tab,
            KeyCode::F(n) => Key::from_f(n),
            KeyCode::BackTab => Key::Tab, /* not supported */
            KeyCode::Char(_) => Key::Tab, /* is handled at `Event` level, use tab as default */
            KeyCode::Null => Key::Tab, /* is handled at `Event` level, use tab as default */
        }
    }
}

impl From<CKeyEvent> for Event {
    fn from(event: CKeyEvent) -> Self {
        const CTRL_ALT: KeyModifiers = KeyModifiers::from_bits_truncate(
            KeyModifiers::CONTROL.bits() | KeyModifiers::ALT.bits(),
        );
        const CTRL_SHIFT: KeyModifiers = KeyModifiers::from_bits_truncate(
            KeyModifiers::CONTROL.bits() | KeyModifiers::SHIFT.bits(),
        );
        const ALT_SHIFT: KeyModifiers = KeyModifiers::from_bits_truncate(
            KeyModifiers::ALT.bits() | KeyModifiers::SHIFT.bits(),
        );

        match event {
            // Handle Char + modifier.
            CKeyEvent {
                modifiers: KeyModifiers::CONTROL,
                code: KeyCode::Char('c'),
            } => Event::Exit,
            CKeyEvent {
                modifiers: KeyModifiers::CONTROL,
                code: KeyCode::Char(c),
            } => Event::CtrlChar(c),
            CKeyEvent {
                modifiers: KeyModifiers::ALT,
                code: KeyCode::Char(c),
            } => Event::AltChar(c),
            CKeyEvent {
                modifiers: KeyModifiers::SHIFT,
                code: KeyCode::Char(c),
            } => Event::Char(c),

            // Handle key + multiple modifiers
            CKeyEvent {
                modifiers: CTRL_ALT,
                code,
            } => Event::CtrlAlt(Key::from(code)),
            CKeyEvent {
                modifiers: CTRL_SHIFT,
                code,
            } => Event::CtrlShift(Key::from(code)),
            CKeyEvent {
                modifiers: ALT_SHIFT,
                code,
            } => Event::AltShift(Key::from(code)),

            // Handle key + single modifier
            CKeyEvent {
                modifiers: KeyModifiers::CONTROL,
                code,
            } => Event::Ctrl(Key::from(code)),
            CKeyEvent {
                modifiers: KeyModifiers::ALT,
                code,
            } => Event::Alt(Key::from(code)),
            CKeyEvent {
                modifiers: KeyModifiers::SHIFT,
                code,
            } => Event::Shift(Key::from(code)),

            CKeyEvent {
                code: KeyCode::Char(c),
                ..
            } => Event::Char(c),

            // Explicitly handle 'backtab' since crossterm does not sent SHIFT alongside the back tab key.
            CKeyEvent {
                code: KeyCode::BackTab,
                ..
            } => Event::Shift(Key::Tab),

            // All other keys.
            CKeyEvent { code, .. } => Event::Key(Key::from(code)),
        }
    }
}

impl From<theme::Color> for Color {
    fn from(base_color: theme::Color) -> Self {
        match base_color {
            theme::Color::Dark(theme::BaseColor::Black) => Color::Black,
            theme::Color::Dark(theme::BaseColor::Red) => Color::DarkRed,
            theme::Color::Dark(theme::BaseColor::Green) => Color::DarkGreen,
            theme::Color::Dark(theme::BaseColor::Yellow) => Color::DarkYellow,
            theme::Color::Dark(theme::BaseColor::Blue) => Color::DarkBlue,
            theme::Color::Dark(theme::BaseColor::Magenta) => {
                Color::DarkMagenta
            }
            theme::Color::Dark(theme::BaseColor::Cyan) => Color::DarkCyan,
            theme::Color::Dark(theme::BaseColor::White) => Color::Grey,
            theme::Color::Light(theme::BaseColor::Black) => Color::DarkGrey,
            theme::Color::Light(theme::BaseColor::Red) => Color::Red,
            theme::Color::Light(theme::BaseColor::Green) => Color::Green,
            theme::Color::Light(theme::BaseColor::Yellow) => Color::Yellow,
            theme::Color::Light(theme::BaseColor::Blue) => Color::Blue,
            theme::Color::Light(theme::BaseColor::Magenta) => Color::Magenta,
            theme::Color::Light(theme::BaseColor::Cyan) => Color::Cyan,
            theme::Color::Light(theme::BaseColor::White) => Color::White,
            theme::Color::Rgb(r, g, b) => Color::Rgb { r, g, b },
            theme::Color::RgbLowRes(r, g, b) => {
                debug_assert!(r <= 5,
                              "Red color fragment (r = {}) is out of bound. Make sure r ≤ 5.",
                              r);
                debug_assert!(g <= 5,
                              "Green color fragment (g = {}) is out of bound. Make sure g ≤ 5.",
                              g);
                debug_assert!(b <= 5,
                              "Blue color fragment (b = {}) is out of bound. Make sure b ≤ 5.",
                              b);

                Color::AnsiValue(16 + 36 * r + 6 * g + b)
            }
            theme::Color::TerminalDefault => Color::Reset,
        }
    }
}

impl Backend {
    /// Creates a new crossterm backend.
    pub fn init() -> Result<Box<dyn backend::Backend>, crossterm::ErrorKind>
    where
        Self: Sized,
    {
        enable_raw_mode()?;

        execute!(
            io::stdout(),
            EnterAlternateScreen,
            EnableMouseCapture,
            Hide
        )?;

        #[cfg(unix)]
        let stdout = RefCell::new(BufWriter::new(File::create("/dev/tty")?));
        #[cfg(windows)]
        let stdout = RefCell::new(BufWriter::new(io::stdout()));

        Ok(Box::new(Backend {
            current_style: Cell::new(theme::ColorPair::from_256colors(0, 0)),
            last_button: None,
            stdout,
        }))
    }

    fn apply_colors(&self, colors: theme::ColorPair) {
        queue!(
            self.stdout_mut(),
            SetForegroundColor(Color::from(colors.front)),
            SetBackgroundColor(Color::from(colors.back))
        )
        .unwrap();
    }

    fn stdout_mut(&self) -> RefMut<BufWriter<Stdout>> {
        self.stdout.borrow_mut()
    }

    fn set_attr(&self, attr: Attribute) {
        queue!(self.stdout_mut(), SetAttribute(attr)).unwrap();
    }

    fn map_key(&mut self, event: CEvent) -> Event {
        match event {
            CEvent::Key(key_event) => Event::from(key_event),
            CEvent::Mouse(mouse_event) => {
                let position;
                let event;

                match mouse_event {
                    CMouseEvent::Down(button, x, y, _) => {
                        let button = MouseButton::from(button);
                        self.last_button = Some(button);
                        event = MouseEvent::Press(button);
                        position = (x, y).into();
                    }
                    CMouseEvent::Up(_, x, y, _) => {
                        event = MouseEvent::Release(self.last_button.unwrap());
                        position = (x, y).into();
                    }
                    CMouseEvent::Drag(_, x, y, _) => {
                        event = MouseEvent::Hold(self.last_button.unwrap());
                        position = (x, y).into();
                    }
                    CMouseEvent::ScrollDown(x, y, _) => {
                        event = MouseEvent::WheelDown;
                        position = (x, y).into();
                    }
                    CMouseEvent::ScrollUp(x, y, _) => {
                        event = MouseEvent::WheelUp;
                        position = (x, y).into();
                    }
                };

                Event::Mouse {
                    event,
                    position,
                    offset: Vec2::zero(),
                }
            }
            CEvent::Resize(_, _) => Event::WindowResize,
        }
    }
}

impl backend::Backend for Backend {
    fn poll_event(&mut self) -> Option<Event> {
        match poll(Duration::from_millis(1)) {
            Ok(true) => match read() {
                Ok(event) => Some(self.map_key(event)),
                Err(e) => panic!("{:?}", e),
            },
            _ => None,
        }
    }

    fn finish(&mut self) {
        // We have to execute the show cursor command at the `stdout`.
        execute!(
            io::stdout(),
            LeaveAlternateScreen,
            DisableMouseCapture,
            Show
        )
        .expect("Can not disable mouse capture or show cursor.");

        disable_raw_mode().unwrap();
    }

    fn refresh(&mut self) {
        self.stdout_mut().flush().unwrap();
    }

    fn has_colors(&self) -> bool {
        // TODO: color support detection?
        true
    }

    fn screen_size(&self) -> Vec2 {
        let size = terminal::size().unwrap_or((1, 1));
        Vec2::from(size)
    }

    fn print_at(&self, pos: Vec2, text: &str) {
        queue!(
            self.stdout_mut(),
            MoveTo(pos.x as u16, pos.y as u16),
            Print(text)
        )
        .unwrap();
    }

    fn print_at_rep(&self, pos: Vec2, repetitions: usize, text: &str) {
        if repetitions > 0 {
            let mut out = self.stdout_mut();

            queue!(out, MoveTo(pos.x as u16, pos.y as u16)).unwrap();

            out.write_all(text.as_bytes()).unwrap();

            let mut dupes_left = repetitions - 1;
            while dupes_left > 0 {
                out.write_all(text.as_bytes()).unwrap();
                dupes_left -= 1;
            }
        }
    }

    fn clear(&self, color: theme::Color) {
        self.apply_colors(theme::ColorPair {
            front: color,
            back: color,
        });

        queue!(self.stdout_mut(), Clear(ClearType::All)).unwrap();
    }

    fn set_color(&self, color: theme::ColorPair) -> theme::ColorPair {
        let current_style = self.current_style.get();

        if current_style != color {
            self.apply_colors(color);
            self.current_style.set(color);
        }

        current_style
    }

    fn set_effect(&self, effect: theme::Effect) {
        match effect {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => self.set_attr(Attribute::Reverse),
            theme::Effect::Bold => self.set_attr(Attribute::Bold),
            theme::Effect::Italic => self.set_attr(Attribute::Italic),
            theme::Effect::Strikethrough => {
                self.set_attr(Attribute::CrossedOut)
            }
            theme::Effect::Underline => self.set_attr(Attribute::Underlined),
        }
    }

    fn unset_effect(&self, effect: theme::Effect) {
        match effect {
            theme::Effect::Simple => (),
            theme::Effect::Reverse => self.set_attr(Attribute::NoReverse),
            theme::Effect::Bold => self.set_attr(Attribute::NormalIntensity),
            theme::Effect::Italic => self.set_attr(Attribute::NoItalic),
            theme::Effect::Strikethrough => {
                self.set_attr(Attribute::NotCrossedOut)
            }
            theme::Effect::Underline => self.set_attr(Attribute::NoUnderline),
        }
    }

    fn name(&self) -> &str {
        "crossterm"
    }
}
<|endoftext|>


<commit_before>use crate::alloc::{host_page_size, instance_heap_offset, Alloc, Limits, Slot};
use crate::embed_ctx::CtxMap;
use crate::error::Error;
use crate::instance::{new_instance_handle, Instance, InstanceHandle};
use crate::module::Module;
use crate::region::{Region, RegionCreate, RegionInternal};
#[cfg(not(target_os = "linux"))]
use libc::memset;
use libc::{c_void, SIGSTKSZ};
use nix::sys::mman::{madvise, mmap, munmap, MapFlags, MmapAdvise, ProtFlags};
use std::ptr;
use std::sync::{Arc, Mutex, Weak};

/// A [`Region`](trait.Region.html) backed by `mmap`.
pub struct MmapRegion {
    capacity: usize,
    freelist: Mutex<Vec<Slot>>,
    limits: Limits,
}

impl Region for MmapRegion {}

impl RegionInternal for MmapRegion {
    fn new_instance_with(
        &self,
        module: Arc<dyn Module>,
        embed_ctx: CtxMap,
    ) -> Result<InstanceHandle, Error> {
        let slot = self
            .freelist
            .lock()
            .unwrap()
            .pop()
            .ok_or(Error::RegionFull(self.capacity))?;

        if slot.heap as usize % host_page_size() != 0 {
            lucet_bail!("heap is not page-aligned; this is a bug");
        }

        let limits = &slot.limits;
        module.validate_runtime_spec(limits)?;

        for (ptr, len) in [
            // make the stack read/writable
            (slot.stack, limits.stack_size),
            // make the globals read/writable
            (slot.globals, limits.globals_size),
            // make the sigstack read/writable
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting r/w {:p}[{:x}]", *ptr, len);
            unsafe { mprotect(*ptr, *len, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? };
        }

        // note: the initial heap will be made read/writable when `new_instance_handle` calls `reset`

        let inst_ptr = slot.start as *mut Instance;

        // upgrade the slot's weak region pointer so the region can't get dropped while the instance
        // exists
        let region = slot
            .region
            .upgrade()
            // if this precondition isn't met, something is deeply wrong as some other region's slot
            // ended up in our freelist
            .expect("backing region of slot (`self`) exists");

        let alloc = Alloc {
            heap_accessible_size: 0, // the `reset` call in `new_instance_handle` will set this
            heap_inaccessible_size: slot.limits.heap_address_space_size,
            slot: Some(slot),
            region,
        };

        let inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx)?;

        Ok(inst)
    }

    fn drop_alloc(&self, alloc: &mut Alloc) {
        let slot = alloc
            .slot
            .take()
            .expect("alloc didn't have a slot during drop; dropped twice?");

        if slot.heap as usize % host_page_size() != 0 {
            panic!("heap is not page-aligned");
        }

        // clear and disable access to the heap, stack, globals, and sigstack
        for (ptr, len) in [
            (slot.heap, slot.limits.heap_address_space_size),
            (slot.stack, slot.limits.stack_size),
            (slot.globals, slot.limits.globals_size),
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting none {:p}[{:x}]", *ptr, len);
            unsafe {
                mprotect(*ptr, *len, ProtFlags::PROT_NONE).expect("mprotect succeeds during drop");
                madvise(*ptr, *len, MmapAdvise::MADV_DONTNEED)
                    .expect("madvise succeeds during drop");
            }
        }

        self.freelist.lock().unwrap().push(slot);
    }

    fn expand_heap(&self, slot: &Slot, start: u32, len: u32) -> Result<(), Error> {
        unsafe {
            mprotect(
                (slot.heap as usize + start as usize) as *mut c_void,
                len as usize,
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?;
        }
        Ok(())
    }

    fn reset_heap(&self, alloc: &mut Alloc, module: &dyn Module) -> Result<(), Error> {
        let heap = alloc.slot().heap;

        if alloc.heap_accessible_size > 0 {
            // zero the whole heap, if any of it is currently accessible
            let heap_size = alloc.slot().limits.heap_address_space_size;

            unsafe {
                // `mprotect()` and `madvise()` are sufficient to zero a page on Linux,
                // but not necessarily on all POSIX operating systems, and on macOS in particular.
                #[cfg(not(target_os = "linux"))]
                {
                    mprotect(
                        heap,
                        alloc.heap_accessible_size,
                        ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                    )?;
                    memset(heap, 0, alloc.heap_accessible_size);
                }
                mprotect(heap, heap_size, ProtFlags::PROT_NONE)?;
                madvise(heap, heap_size, MmapAdvise::MADV_DONTNEED)?;
            }
        }

        let initial_size = module
            .heap_spec()
            .map(|h| h.initial_size as usize)
            .unwrap_or(0);

        // reset the heap to the initial size, and mprotect those pages appropriately
        if initial_size > 0 {
            unsafe {
                mprotect(
                    heap,
                    initial_size,
                    ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                )?
            };
        }
        alloc.heap_accessible_size = initial_size;
        alloc.heap_inaccessible_size = alloc.slot().limits.heap_address_space_size - initial_size;

        // Initialize the heap using the module sparse page data. There cannot be more pages in the
        // sparse page data than will fit in the initial heap size.
        //
        // Pages with a corresponding Some entry in the sparse page data are initialized with
        // the contents of that data.
        //
        // Any pages which don't have an entry in the sparse page data, either because their entry
        // is None, or because the sparse data has fewer pages than the initial heap, are zeroed.
        let heap = unsafe { alloc.heap_mut() };
        let initial_pages =
            initial_size
                .checked_div(host_page_size())
                .ok_or(lucet_incorrect_module!(
                    "initial heap size {} is not divisible by host page size ({})",
                    initial_size,
                    host_page_size()
                ))?;
        for page_num in 0..initial_pages {
            let page_base = page_num * host_page_size();
            if heap.len() < page_base {
                return Err(lucet_incorrect_module!(
                    "sparse page data length exceeded initial heap size"
                ));
            }
            if let Some(contents) = module.get_sparse_page_data(page_num) {
                // otherwise copy in the page data
                heap[page_base..page_base + host_page_size()].copy_from_slice(contents);
            }
        }

        Ok(())
    }

    fn as_dyn_internal(&self) -> &dyn RegionInternal {
        self
    }
}

impl Drop for MmapRegion {
    fn drop(&mut self) {
        for slot in self.freelist.get_mut().unwrap().drain(0..) {
            Self::free_slot(slot);
        }
    }
}

impl RegionCreate for MmapRegion {
    const TYPE_NAME: &'static str = "MmapRegion";

    fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        MmapRegion::create(instance_capacity, limits)
    }
}

impl MmapRegion {
    /// Create a new `MmapRegion` that can support a given number instances, each subject to the
    /// same runtime limits.
    ///
    /// The region is returned in an `Arc`, because any instances created from it carry a reference
    /// back to the region.
    pub fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        assert!(
            SIGSTKSZ % host_page_size() == 0,
            "signal stack size is a multiple of host page size"
        );
        limits.validate()?;

        let region = Arc::new(MmapRegion {
            capacity: instance_capacity,
            freelist: Mutex::new(Vec::with_capacity(instance_capacity)),
            limits: limits.clone(),
        });
        {
            let mut freelist = region.freelist.lock().unwrap();
            for _ in 0..instance_capacity {
                freelist.push(MmapRegion::create_slot(&region)?);
            }
        }

        Ok(region)
    }

    fn create_slot(region: &Arc<MmapRegion>) -> Result<Slot, Error> {
        // get the chunk of virtual memory that the `Slot` will manage
        let mem = unsafe {
            mmap(
                ptr::null_mut(),
                region.limits.total_memory_size(),
                ProtFlags::PROT_NONE,
                MapFlags::MAP_ANON | MapFlags::MAP_PRIVATE,
                0,
                0,
            )?
        };

        // set the first part of the memory to read/write so that the `Instance` can be stored there
        // TODO: post slot refactor, is this necessary/desirable?
        unsafe {
            mprotect(
                mem,
                instance_heap_offset(),
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?
        };

        // lay out the other sections in memory
        let heap = mem as usize + instance_heap_offset();
        let stack = heap + region.limits.heap_address_space_size;
        let globals = stack + region.limits.stack_size + host_page_size();
        let sigstack = globals + host_page_size();

        Ok(Slot {
            start: mem,
            heap: heap as *mut c_void,
            stack: stack as *mut c_void,
            globals: globals as *mut c_void,
            sigstack: sigstack as *mut c_void,
            limits: region.limits.clone(),
            region: Arc::downgrade(region) as Weak<dyn RegionInternal>,
        })
    }

    fn free_slot(slot: Slot) {
        // eprintln!(
        //     "unmapping {:p}[{:x}]",
        //     slot.start,
        //     slot.limits.total_memory_size()
        // );
        let res = unsafe { munmap(slot.start, slot.limits.total_memory_size()) };
        res.expect("munmap succeeded");
    }
}

// TODO: remove this once `nix` PR https://github.com/nix-rust/nix/pull/991 is merged
unsafe fn mprotect(addr: *mut c_void, length: libc::size_t, prot: ProtFlags) -> nix::Result<()> {
    nix::errno::Errno::result(libc::mprotect(addr, length, prot.bits())).map(drop)
}
<commit_msg>Add another explicit page zeroing when MADV_DONTNEED is not enough<commit_after>use crate::alloc::{host_page_size, instance_heap_offset, Alloc, Limits, Slot};
use crate::embed_ctx::CtxMap;
use crate::error::Error;
use crate::instance::{new_instance_handle, Instance, InstanceHandle};
use crate::module::Module;
use crate::region::{Region, RegionCreate, RegionInternal};
#[cfg(not(target_os = "linux"))]
use libc::memset;
use libc::{c_void, SIGSTKSZ};
use nix::sys::mman::{madvise, mmap, munmap, MapFlags, MmapAdvise, ProtFlags};
use std::ptr;
use std::sync::{Arc, Mutex, Weak};

/// A [`Region`](trait.Region.html) backed by `mmap`.
pub struct MmapRegion {
    capacity: usize,
    freelist: Mutex<Vec<Slot>>,
    limits: Limits,
}

impl Region for MmapRegion {}

impl RegionInternal for MmapRegion {
    fn new_instance_with(
        &self,
        module: Arc<dyn Module>,
        embed_ctx: CtxMap,
    ) -> Result<InstanceHandle, Error> {
        let slot = self
            .freelist
            .lock()
            .unwrap()
            .pop()
            .ok_or(Error::RegionFull(self.capacity))?;

        if slot.heap as usize % host_page_size() != 0 {
            lucet_bail!("heap is not page-aligned; this is a bug");
        }

        let limits = &slot.limits;
        module.validate_runtime_spec(limits)?;

        for (ptr, len) in [
            // make the stack read/writable
            (slot.stack, limits.stack_size),
            // make the globals read/writable
            (slot.globals, limits.globals_size),
            // make the sigstack read/writable
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting r/w {:p}[{:x}]", *ptr, len);
            unsafe { mprotect(*ptr, *len, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)? };
        }

        // note: the initial heap will be made read/writable when `new_instance_handle` calls `reset`

        let inst_ptr = slot.start as *mut Instance;

        // upgrade the slot's weak region pointer so the region can't get dropped while the instance
        // exists
        let region = slot
            .region
            .upgrade()
            // if this precondition isn't met, something is deeply wrong as some other region's slot
            // ended up in our freelist
            .expect("backing region of slot (`self`) exists");

        let alloc = Alloc {
            heap_accessible_size: 0, // the `reset` call in `new_instance_handle` will set this
            heap_inaccessible_size: slot.limits.heap_address_space_size,
            slot: Some(slot),
            region,
        };

        let inst = new_instance_handle(inst_ptr, module, alloc, embed_ctx)?;

        Ok(inst)
    }

    fn drop_alloc(&self, alloc: &mut Alloc) {
        let slot = alloc
            .slot
            .take()
            .expect("alloc didn't have a slot during drop; dropped twice?");

        if slot.heap as usize % host_page_size() != 0 {
            panic!("heap is not page-aligned");
        }

        // clear and disable access to the heap, stack, globals, and sigstack
        for (ptr, len) in [
            (slot.heap, slot.limits.heap_address_space_size),
            (slot.stack, slot.limits.stack_size),
            (slot.globals, slot.limits.globals_size),
            (slot.sigstack, SIGSTKSZ),
        ]
        .into_iter()
        {
            // eprintln!("setting none {:p}[{:x}]", *ptr, len);
            unsafe {
                // MADV_DONTNEED is not guaranteed to clear pages on non-Linux systems
                #[cfg(not(target_os = "linux"))]
                {
                    mprotect(*ptr, *len, ProtFlags::PROT_READ | ProtFlags::PROT_WRITE)
                        .expect("mprotect succeeds during drop");
                    memset(*ptr, 0, *len);
                }
                mprotect(*ptr, *len, ProtFlags::PROT_NONE).expect("mprotect succeeds during drop");
                madvise(*ptr, *len, MmapAdvise::MADV_DONTNEED)
                    .expect("madvise succeeds during drop");
            }
        }

        self.freelist.lock().unwrap().push(slot);
    }

    fn expand_heap(&self, slot: &Slot, start: u32, len: u32) -> Result<(), Error> {
        unsafe {
            mprotect(
                (slot.heap as usize + start as usize) as *mut c_void,
                len as usize,
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?;
        }
        Ok(())
    }

    fn reset_heap(&self, alloc: &mut Alloc, module: &dyn Module) -> Result<(), Error> {
        let heap = alloc.slot().heap;

        if alloc.heap_accessible_size > 0 {
            // zero the whole heap, if any of it is currently accessible
            let heap_size = alloc.slot().limits.heap_address_space_size;

            unsafe {
                // `mprotect()` and `madvise()` are sufficient to zero a page on Linux,
                // but not necessarily on all POSIX operating systems, and on macOS in particular.
                #[cfg(not(target_os = "linux"))]
                {
                    mprotect(
                        heap,
                        alloc.heap_accessible_size,
                        ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                    )?;
                    memset(heap, 0, alloc.heap_accessible_size);
                }
                mprotect(heap, heap_size, ProtFlags::PROT_NONE)?;
                madvise(heap, heap_size, MmapAdvise::MADV_DONTNEED)?;
            }
        }

        let initial_size = module
            .heap_spec()
            .map(|h| h.initial_size as usize)
            .unwrap_or(0);

        // reset the heap to the initial size, and mprotect those pages appropriately
        if initial_size > 0 {
            unsafe {
                mprotect(
                    heap,
                    initial_size,
                    ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                )?
            };
        }
        alloc.heap_accessible_size = initial_size;
        alloc.heap_inaccessible_size = alloc.slot().limits.heap_address_space_size - initial_size;

        // Initialize the heap using the module sparse page data. There cannot be more pages in the
        // sparse page data than will fit in the initial heap size.
        //
        // Pages with a corresponding Some entry in the sparse page data are initialized with
        // the contents of that data.
        //
        // Any pages which don't have an entry in the sparse page data, either because their entry
        // is None, or because the sparse data has fewer pages than the initial heap, are zeroed.
        let heap = unsafe { alloc.heap_mut() };
        let initial_pages =
            initial_size
                .checked_div(host_page_size())
                .ok_or(lucet_incorrect_module!(
                    "initial heap size {} is not divisible by host page size ({})",
                    initial_size,
                    host_page_size()
                ))?;
        for page_num in 0..initial_pages {
            let page_base = page_num * host_page_size();
            if heap.len() < page_base {
                return Err(lucet_incorrect_module!(
                    "sparse page data length exceeded initial heap size"
                ));
            }
            if let Some(contents) = module.get_sparse_page_data(page_num) {
                // otherwise copy in the page data
                heap[page_base..page_base + host_page_size()].copy_from_slice(contents);
            }
        }

        Ok(())
    }

    fn as_dyn_internal(&self) -> &dyn RegionInternal {
        self
    }
}

impl Drop for MmapRegion {
    fn drop(&mut self) {
        for slot in self.freelist.get_mut().unwrap().drain(0..) {
            Self::free_slot(slot);
        }
    }
}

impl RegionCreate for MmapRegion {
    const TYPE_NAME: &'static str = "MmapRegion";

    fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        MmapRegion::create(instance_capacity, limits)
    }
}

impl MmapRegion {
    /// Create a new `MmapRegion` that can support a given number instances, each subject to the
    /// same runtime limits.
    ///
    /// The region is returned in an `Arc`, because any instances created from it carry a reference
    /// back to the region.
    pub fn create(instance_capacity: usize, limits: &Limits) -> Result<Arc<Self>, Error> {
        assert!(
            SIGSTKSZ % host_page_size() == 0,
            "signal stack size is a multiple of host page size"
        );
        limits.validate()?;

        let region = Arc::new(MmapRegion {
            capacity: instance_capacity,
            freelist: Mutex::new(Vec::with_capacity(instance_capacity)),
            limits: limits.clone(),
        });
        {
            let mut freelist = region.freelist.lock().unwrap();
            for _ in 0..instance_capacity {
                freelist.push(MmapRegion::create_slot(&region)?);
            }
        }

        Ok(region)
    }

    fn create_slot(region: &Arc<MmapRegion>) -> Result<Slot, Error> {
        // get the chunk of virtual memory that the `Slot` will manage
        let mem = unsafe {
            mmap(
                ptr::null_mut(),
                region.limits.total_memory_size(),
                ProtFlags::PROT_NONE,
                MapFlags::MAP_ANON | MapFlags::MAP_PRIVATE,
                0,
                0,
            )?
        };

        // set the first part of the memory to read/write so that the `Instance` can be stored there
        // TODO: post slot refactor, is this necessary/desirable?
        unsafe {
            mprotect(
                mem,
                instance_heap_offset(),
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            )?
        };

        // lay out the other sections in memory
        let heap = mem as usize + instance_heap_offset();
        let stack = heap + region.limits.heap_address_space_size;
        let globals = stack + region.limits.stack_size + host_page_size();
        let sigstack = globals + host_page_size();

        Ok(Slot {
            start: mem,
            heap: heap as *mut c_void,
            stack: stack as *mut c_void,
            globals: globals as *mut c_void,
            sigstack: sigstack as *mut c_void,
            limits: region.limits.clone(),
            region: Arc::downgrade(region) as Weak<dyn RegionInternal>,
        })
    }

    fn free_slot(slot: Slot) {
        // eprintln!(
        //     "unmapping {:p}[{:x}]",
        //     slot.start,
        //     slot.limits.total_memory_size()
        // );
        let res = unsafe { munmap(slot.start, slot.limits.total_memory_size()) };
        res.expect("munmap succeeded");
    }
}

// TODO: remove this once `nix` PR https://github.com/nix-rust/nix/pull/991 is merged
unsafe fn mprotect(addr: *mut c_void, length: libc::size_t, prot: ProtFlags) -> nix::Result<()> {
    nix::errno::Errno::result(libc::mprotect(addr, length, prot.bits())).map(drop)
}
<|endoftext|>


<commit_before>//! Module for reading and writing Mobile-Originated (MO) SBD messages.
//!
//! Though messages technically come in two flavors, mobile originated and mobile terminated, we
//! only handle mobile originated messages in this library.

use {Error, Result};

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use chrono::{DateTime, Duration, TimeZone, Utc};
use information_element::{InformationElement, InformationElementType};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Cursor, Read, Write};
use std::path::Path;
use std::str;

const INFORMATION_ELEMENT_HEADER_LENGTH: u16 = 3;
const MOBILE_ORIGINATED_HEADER_LENGTH: u16 = 28;
const ASCII_ZERO: u8 = 48;

/// The protocol number of an SBD message.
///
/// At this point, this can *only* be one.
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
struct ProtocolRevisionNumber(u8);

impl ProtocolRevisionNumber {
    /// Returns true if this is a valid protocol revision number.
    pub fn valid(&self) -> bool {
        self.0 == 1
    }
}

/// The modem IMEI identifier.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct Imei([u8; 15]);

impl Default for Imei {
    fn default() -> Imei {
        Imei([ASCII_ZERO; 15])
    }
}

impl Imei {
    /// Returns this IMEI number as a `str`.
    ///
    /// # Panics
    ///
    /// Panics if the string is not valid utf8.
    fn as_str(&self) -> &str {
        str::from_utf8(&self.0).unwrap()
    }
}

/// The status of a mobile-originated session.
///
/// The descriptions for these codes are taken directly from the `DirectIP` documentation.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize)]
pub enum SessionStatus {
    /// The SBD session completed successfully.
    Ok = 0,
    /// The MO message transfer, if any, was successful. The MT message queued at the GSS is too
    /// large to be transferred within a single SBD session.
    OkMobileTerminatedTooLarge = 1,
    /// The MO message transfer, if any, was successful. The reported location was determined to be
    /// of unacceptable quality. This value is only applicable to IMEIs using SBD protocol revision
    /// 1.
    OkLocationUnacceptableQuality = 2,
    /// The SBD session timed out before session completion.
    Timeout = 10,
    /// The MO message being transferred by the IMEI is too large to be transerred within a single
    /// SBD session.
    MobileOriginatedTooLarge = 12,
    /// An RF link loss ocurred during the SBD session.
    RFLinkLoss = 13,
    /// An IMEI protocol anomaly occurred during SBD session.
    IMEIProtocolAnomaly = 14,
    /// The IMEI is prohibited from accessing the GSS.
    Prohibited = 15,
    /// Unknown session status code.
    Unknown,
}

impl Default for SessionStatus {
    fn default() -> SessionStatus {
        SessionStatus::Unknown
    }
}

impl From<u8> for SessionStatus {
    fn from(n: u8) -> Self {
        match n {
            0 => SessionStatus::Ok,
            1 => SessionStatus::OkMobileTerminatedTooLarge,
            2 => SessionStatus::OkLocationUnacceptableQuality,
            10 => SessionStatus::Timeout,
            12 => SessionStatus::MobileOriginatedTooLarge,
            13 => SessionStatus::RFLinkLoss,
            14 => SessionStatus::IMEIProtocolAnomaly,
            15 => SessionStatus::Prohibited,
            _ => SessionStatus::Unknown,
        }
    }
}

/// A mobile-origined Iridium SBD message.
///
/// `Message`s can be ordered by time of session.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Message {
    protocol_revision_number: ProtocolRevisionNumber,
    cdr_reference: u32,
    imei: Imei,
    session_status: SessionStatus,
    momsn: u16,
    mtmsn: u16,
    time_of_session: DateTime<Utc>,
    payload: Vec<u8>,
}

impl Default for Message {
    fn default() -> Message {
        Message {
            protocol_revision_number: ProtocolRevisionNumber(1),
            cdr_reference: Default::default(),
            imei: Default::default(),
            session_status: Default::default(),
            momsn: Default::default(),
            mtmsn: Default::default(),
            time_of_session: Utc.ymd(1970, 1, 1).and_hms(0, 0, 0),
            payload: Vec::new(),
        }
    }
}

impl PartialOrd for Message {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Message {
    fn cmp(&self, other: &Self) -> Ordering {
        self.time_of_session.cmp(&other.time_of_session)
    }
}

impl Message {
    /// Reads in a message from a file path.
    ///
    /// # Examples
    ///
    /// ```
    /// use sbd::mo::Message;
    /// let message = Message::from_path("data/0-mo.sbd").unwrap();
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Message> {
        let file = File::open(path)?;
        Message::read_from(file)
    }

    /// Reads in a message from an object that implements `Read`.
    ///
    /// Per the specification, oversized and undersized messages will result in an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::File;
    /// use sbd::mo::Message;
    /// let mut file = File::open("data/0-mo.sbd").unwrap();
    /// let message = Message::read_from(file).unwrap();
    /// ```
    pub fn read_from<R: Read>(mut readable: R) -> Result<Message> {
        let mut message: Message = Default::default();

        message.protocol_revision_number = ProtocolRevisionNumber(readable.read_u8()?);
        if !message.protocol_revision_number.valid() {
            return Err(Error::InvalidProtocolRevisionNumber(
                message.protocol_revision_number.0,
            ));
        }
        let overall_message_length = readable.read_u16::<BigEndian>()?;

        let mut information_elements: HashMap<InformationElementType, InformationElement> =
            HashMap::new();
        let mut bytes_read = 0u16;
        loop {
            let ie = match InformationElement::read_from(&mut readable) {
                Ok(ie) => ie,
                Err(e) => return Err(e),
            };
            bytes_read += ie.len();
            information_elements.insert(ie.id(), ie);
            if bytes_read >= overall_message_length {
                break;
            }
        }

        if readable.take(1).read_to_end(&mut Vec::new())? != 0 {
            return Err(Error::Oversized);
        }

        let header =
            match information_elements.remove(&InformationElementType::MobileOriginatedHeader) {
                Some(ie) => ie,
                None => return Err(Error::MissingMobileOriginatedHeader),
            };
        let mut cursor = &mut Cursor::new(header.contents_ref());
        message.cdr_reference = cursor.read_u32::<BigEndian>()?;
        let bytes_read = cursor.take(message.imei.0.len() as u64).read(
            &mut message.imei.0,
        )?;
        if bytes_read != message.imei.0.len() {
            return Err(Error::InvalidImei);
        }
        message.session_status = SessionStatus::from(cursor.read_u8()?);
        message.momsn = cursor.read_u16::<BigEndian>()?;
        message.mtmsn = cursor.read_u16::<BigEndian>()?;
        message.time_of_session = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0) +
            Duration::seconds(cursor.read_u32::<BigEndian>()? as i64);

        let payload =
            match information_elements.remove(&InformationElementType::MobileOriginatedPayload) {
                Some(ie) => ie,
                None => return Err(Error::MissingMobileOriginatedPayload),
            };
        message.payload = payload.into_contents();

        if !information_elements.is_empty() {
            return Err(Error::UnhandledInformationElements(information_elements));
        }

        Ok(message)
    }

    /// Write this message back to a object that can `Write`.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use sbd::mo::Message;
    /// let message = Message::from_path("data/0-mo.sbd").unwrap();
    /// let mut cursor = Cursor::new(Vec::new());
    /// message.write_to(&mut cursor);
    /// ```
    pub fn write_to<W: Write>(&self, w: &mut W) -> Result<()> {
        w.write_u8(self.protocol_revision_number.0)?;
        w.write_u16::<BigEndian>(self.overall_message_length())?;
        w.write_u8(
            InformationElementType::MobileOriginatedHeader as u8,
        )?;
        w.write_u16::<BigEndian>(MOBILE_ORIGINATED_HEADER_LENGTH)?;
        w.write_u32::<BigEndian>(self.cdr_reference)?;
        w.write_all(&self.imei.0)?;
        w.write_u8(self.session_status as u8)?;
        w.write_u16::<BigEndian>(self.momsn)?;
        w.write_u16::<BigEndian>(self.mtmsn)?;
        w.write_u32::<BigEndian>(
            self.time_of_session.timestamp() as u32,
        )?;
        w.write_u8(
            InformationElementType::MobileOriginatedPayload as u8,
        )?;
        // TODO can we check to make sure the payload is of appropriate size?
        w.write_u16::<BigEndian>(self.payload.len() as u16)?;
        w.write_all(&self.payload[..])?;
        Ok(())
    }

    /// Returns this message's protocol revision number.
    pub fn protocol_revision_number(&self) -> u8 {
        self.protocol_revision_number.0
    }
    /// Returns this message's IMEI number as a string.
    pub fn imei(&self) -> &str {
        self.imei.as_str()
    }
    /// Returns this message's cdr reference number (also called auto id).
    pub fn cdr_reference(&self) -> u32 {
        self.cdr_reference
    }
    /// Returns this message's session status.
    pub fn session_status(&self) -> SessionStatus {
        self.session_status
    }
    /// Returns this message's mobile originated message number.
    pub fn momsn(&self) -> u16 {
        self.momsn
    }
    /// Returns this message's mobile terminated message number.
    pub fn mtmsn(&self) -> u16 {
        self.mtmsn
    }
    /// Returns this message's time of session.
    pub fn time_of_session(&self) -> DateTime<Utc> {
        self.time_of_session
    }
    /// Returns a reference to this message's payload.
    pub fn payload_ref(&self) -> &Vec<u8> {
        &self.payload
    }
    /// Returns this message's payload as a str.
    pub fn payload_str(&self) -> Result<&str> {
        str::from_utf8(&self.payload).map_err(Error::from)
    }

    /// Returns the overall message length, as is contained in the message's header.
    ///
    /// The whole message is actually three bytes longer than this value, thanks to the message
    /// header itself.
    fn overall_message_length(&self) -> u16 {
        INFORMATION_ELEMENT_HEADER_LENGTH + MOBILE_ORIGINATED_HEADER_LENGTH +
            INFORMATION_ELEMENT_HEADER_LENGTH + self.payload.len() as u16
    }
}

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

    use chrono::{TimeZone, Utc};

    use std::fs::File;
    use std::io::{Cursor, Read};
    use std::str;

    #[test]
    fn from_path() {
        Message::from_path("data/0-mo.sbd").unwrap();
    }

    #[test]
    fn from_read() {
        let file = File::open("data/0-mo.sbd").unwrap();
        Message::read_from(file).unwrap();
    }

    #[test]
    fn from_path_that_doesnt_exist() {
        assert!(Message::from_path("notafile.sbd").is_err());
    }

    #[test]
    fn from_path_that_is_not_an_sbd_message() {
        assert!(Message::from_path("data/1-invalid.sbd").is_err());
    }

    #[test]
    fn undersized() {
        let file = File::open("data/0-mo.sbd").unwrap();
        let readable = file.take(58);
        assert!(Message::read_from(readable).is_err());
    }

    #[test]
    fn oversized() {
        let file = File::open("data/0-mo.sbd").unwrap();
        let readable = file.chain(Cursor::new(vec![0]));
        assert!(Message::read_from(readable).is_err());
    }

    #[test]
    fn values() {
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        assert_eq!(1, message.protocol_revision_number());
        assert_eq!(1894516585, message.cdr_reference());
        assert_eq!("300234063904190", message.imei());
        assert_eq!(SessionStatus::Ok, message.session_status());
        assert_eq!(75, message.momsn());
        assert_eq!(0, message.mtmsn());
        assert_eq!(
            Utc.ymd(2015, 7, 9).and_hms(18, 15, 8),
            message.time_of_session()
        );
        assert_eq!(
            "test message from pete",
            str::from_utf8(message.payload_ref()).unwrap()
        );
    }

    #[test]
    fn write() {
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        let mut cursor = Cursor::new(Vec::new());
        message.write_to(&mut cursor).unwrap();
        cursor.set_position(0);
        let message2 = Message::read_from(cursor).unwrap();
        assert_eq!(message, message2);
    }

    #[test]
    fn write_default() {
        let message: Message = Default::default();
        let mut cursor = Cursor::new(Vec::new());
        message.write_to(&mut cursor).unwrap();
        cursor.set_position(0);
        let message2 = Message::read_from(cursor).unwrap();
        assert_eq!(message, message2);
    }

    #[test]
    fn order() {
        let default: Message = Default::default();
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        assert!(message > default);
    }

    #[test]
    fn payload_str() {
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        assert_eq!("test message from pete", message.payload_str().unwrap());
    }
}
<commit_msg>Fix mutable warning<commit_after>//! Module for reading and writing Mobile-Originated (MO) SBD messages.
//!
//! Though messages technically come in two flavors, mobile originated and mobile terminated, we
//! only handle mobile originated messages in this library.

use {Error, Result};

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use chrono::{DateTime, Duration, TimeZone, Utc};
use information_element::{InformationElement, InformationElementType};
use std::cmp::Ordering;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Cursor, Read, Write};
use std::path::Path;
use std::str;

const INFORMATION_ELEMENT_HEADER_LENGTH: u16 = 3;
const MOBILE_ORIGINATED_HEADER_LENGTH: u16 = 28;
const ASCII_ZERO: u8 = 48;

/// The protocol number of an SBD message.
///
/// At this point, this can *only* be one.
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
struct ProtocolRevisionNumber(u8);

impl ProtocolRevisionNumber {
    /// Returns true if this is a valid protocol revision number.
    pub fn valid(&self) -> bool {
        self.0 == 1
    }
}

/// The modem IMEI identifier.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct Imei([u8; 15]);

impl Default for Imei {
    fn default() -> Imei {
        Imei([ASCII_ZERO; 15])
    }
}

impl Imei {
    /// Returns this IMEI number as a `str`.
    ///
    /// # Panics
    ///
    /// Panics if the string is not valid utf8.
    fn as_str(&self) -> &str {
        str::from_utf8(&self.0).unwrap()
    }
}

/// The status of a mobile-originated session.
///
/// The descriptions for these codes are taken directly from the `DirectIP` documentation.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize)]
pub enum SessionStatus {
    /// The SBD session completed successfully.
    Ok = 0,
    /// The MO message transfer, if any, was successful. The MT message queued at the GSS is too
    /// large to be transferred within a single SBD session.
    OkMobileTerminatedTooLarge = 1,
    /// The MO message transfer, if any, was successful. The reported location was determined to be
    /// of unacceptable quality. This value is only applicable to IMEIs using SBD protocol revision
    /// 1.
    OkLocationUnacceptableQuality = 2,
    /// The SBD session timed out before session completion.
    Timeout = 10,
    /// The MO message being transferred by the IMEI is too large to be transerred within a single
    /// SBD session.
    MobileOriginatedTooLarge = 12,
    /// An RF link loss ocurred during the SBD session.
    RFLinkLoss = 13,
    /// An IMEI protocol anomaly occurred during SBD session.
    IMEIProtocolAnomaly = 14,
    /// The IMEI is prohibited from accessing the GSS.
    Prohibited = 15,
    /// Unknown session status code.
    Unknown,
}

impl Default for SessionStatus {
    fn default() -> SessionStatus {
        SessionStatus::Unknown
    }
}

impl From<u8> for SessionStatus {
    fn from(n: u8) -> Self {
        match n {
            0 => SessionStatus::Ok,
            1 => SessionStatus::OkMobileTerminatedTooLarge,
            2 => SessionStatus::OkLocationUnacceptableQuality,
            10 => SessionStatus::Timeout,
            12 => SessionStatus::MobileOriginatedTooLarge,
            13 => SessionStatus::RFLinkLoss,
            14 => SessionStatus::IMEIProtocolAnomaly,
            15 => SessionStatus::Prohibited,
            _ => SessionStatus::Unknown,
        }
    }
}

/// A mobile-origined Iridium SBD message.
///
/// `Message`s can be ordered by time of session.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Message {
    protocol_revision_number: ProtocolRevisionNumber,
    cdr_reference: u32,
    imei: Imei,
    session_status: SessionStatus,
    momsn: u16,
    mtmsn: u16,
    time_of_session: DateTime<Utc>,
    payload: Vec<u8>,
}

impl Default for Message {
    fn default() -> Message {
        Message {
            protocol_revision_number: ProtocolRevisionNumber(1),
            cdr_reference: Default::default(),
            imei: Default::default(),
            session_status: Default::default(),
            momsn: Default::default(),
            mtmsn: Default::default(),
            time_of_session: Utc.ymd(1970, 1, 1).and_hms(0, 0, 0),
            payload: Vec::new(),
        }
    }
}

impl PartialOrd for Message {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Message {
    fn cmp(&self, other: &Self) -> Ordering {
        self.time_of_session.cmp(&other.time_of_session)
    }
}

impl Message {
    /// Reads in a message from a file path.
    ///
    /// # Examples
    ///
    /// ```
    /// use sbd::mo::Message;
    /// let message = Message::from_path("data/0-mo.sbd").unwrap();
    /// ```
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Message> {
        let file = File::open(path)?;
        Message::read_from(file)
    }

    /// Reads in a message from an object that implements `Read`.
    ///
    /// Per the specification, oversized and undersized messages will result in an error.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::fs::File;
    /// use sbd::mo::Message;
    /// let mut file = File::open("data/0-mo.sbd").unwrap();
    /// let message = Message::read_from(file).unwrap();
    /// ```
    pub fn read_from<R: Read>(mut readable: R) -> Result<Message> {
        let mut message: Message = Default::default();

        message.protocol_revision_number = ProtocolRevisionNumber(readable.read_u8()?);
        if !message.protocol_revision_number.valid() {
            return Err(Error::InvalidProtocolRevisionNumber(
                message.protocol_revision_number.0,
            ));
        }
        let overall_message_length = readable.read_u16::<BigEndian>()?;

        let mut information_elements: HashMap<InformationElementType, InformationElement> =
            HashMap::new();
        let mut bytes_read = 0u16;
        loop {
            let ie = match InformationElement::read_from(&mut readable) {
                Ok(ie) => ie,
                Err(e) => return Err(e),
            };
            bytes_read += ie.len();
            information_elements.insert(ie.id(), ie);
            if bytes_read >= overall_message_length {
                break;
            }
        }

        if readable.take(1).read_to_end(&mut Vec::new())? != 0 {
            return Err(Error::Oversized);
        }

        let header =
            match information_elements.remove(&InformationElementType::MobileOriginatedHeader) {
                Some(ie) => ie,
                None => return Err(Error::MissingMobileOriginatedHeader),
            };
        let cursor = &mut Cursor::new(header.contents_ref());
        message.cdr_reference = cursor.read_u32::<BigEndian>()?;
        let bytes_read = cursor.take(message.imei.0.len() as u64).read(
            &mut message.imei.0,
        )?;
        if bytes_read != message.imei.0.len() {
            return Err(Error::InvalidImei);
        }
        message.session_status = SessionStatus::from(cursor.read_u8()?);
        message.momsn = cursor.read_u16::<BigEndian>()?;
        message.mtmsn = cursor.read_u16::<BigEndian>()?;
        message.time_of_session = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0) +
            Duration::seconds(cursor.read_u32::<BigEndian>()? as i64);

        let payload =
            match information_elements.remove(&InformationElementType::MobileOriginatedPayload) {
                Some(ie) => ie,
                None => return Err(Error::MissingMobileOriginatedPayload),
            };
        message.payload = payload.into_contents();

        if !information_elements.is_empty() {
            return Err(Error::UnhandledInformationElements(information_elements));
        }

        Ok(message)
    }

    /// Write this message back to a object that can `Write`.
    ///
    /// # Examples
    ///
    /// ```
    /// use std::io::Cursor;
    /// use sbd::mo::Message;
    /// let message = Message::from_path("data/0-mo.sbd").unwrap();
    /// let mut cursor = Cursor::new(Vec::new());
    /// message.write_to(&mut cursor);
    /// ```
    pub fn write_to<W: Write>(&self, w: &mut W) -> Result<()> {
        w.write_u8(self.protocol_revision_number.0)?;
        w.write_u16::<BigEndian>(self.overall_message_length())?;
        w.write_u8(
            InformationElementType::MobileOriginatedHeader as u8,
        )?;
        w.write_u16::<BigEndian>(MOBILE_ORIGINATED_HEADER_LENGTH)?;
        w.write_u32::<BigEndian>(self.cdr_reference)?;
        w.write_all(&self.imei.0)?;
        w.write_u8(self.session_status as u8)?;
        w.write_u16::<BigEndian>(self.momsn)?;
        w.write_u16::<BigEndian>(self.mtmsn)?;
        w.write_u32::<BigEndian>(
            self.time_of_session.timestamp() as u32,
        )?;
        w.write_u8(
            InformationElementType::MobileOriginatedPayload as u8,
        )?;
        // TODO can we check to make sure the payload is of appropriate size?
        w.write_u16::<BigEndian>(self.payload.len() as u16)?;
        w.write_all(&self.payload[..])?;
        Ok(())
    }

    /// Returns this message's protocol revision number.
    pub fn protocol_revision_number(&self) -> u8 {
        self.protocol_revision_number.0
    }
    /// Returns this message's IMEI number as a string.
    pub fn imei(&self) -> &str {
        self.imei.as_str()
    }
    /// Returns this message's cdr reference number (also called auto id).
    pub fn cdr_reference(&self) -> u32 {
        self.cdr_reference
    }
    /// Returns this message's session status.
    pub fn session_status(&self) -> SessionStatus {
        self.session_status
    }
    /// Returns this message's mobile originated message number.
    pub fn momsn(&self) -> u16 {
        self.momsn
    }
    /// Returns this message's mobile terminated message number.
    pub fn mtmsn(&self) -> u16 {
        self.mtmsn
    }
    /// Returns this message's time of session.
    pub fn time_of_session(&self) -> DateTime<Utc> {
        self.time_of_session
    }
    /// Returns a reference to this message's payload.
    pub fn payload_ref(&self) -> &Vec<u8> {
        &self.payload
    }
    /// Returns this message's payload as a str.
    pub fn payload_str(&self) -> Result<&str> {
        str::from_utf8(&self.payload).map_err(Error::from)
    }

    /// Returns the overall message length, as is contained in the message's header.
    ///
    /// The whole message is actually three bytes longer than this value, thanks to the message
    /// header itself.
    fn overall_message_length(&self) -> u16 {
        INFORMATION_ELEMENT_HEADER_LENGTH + MOBILE_ORIGINATED_HEADER_LENGTH +
            INFORMATION_ELEMENT_HEADER_LENGTH + self.payload.len() as u16
    }
}

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

    use chrono::{TimeZone, Utc};

    use std::fs::File;
    use std::io::{Cursor, Read};
    use std::str;

    #[test]
    fn from_path() {
        Message::from_path("data/0-mo.sbd").unwrap();
    }

    #[test]
    fn from_read() {
        let file = File::open("data/0-mo.sbd").unwrap();
        Message::read_from(file).unwrap();
    }

    #[test]
    fn from_path_that_doesnt_exist() {
        assert!(Message::from_path("notafile.sbd").is_err());
    }

    #[test]
    fn from_path_that_is_not_an_sbd_message() {
        assert!(Message::from_path("data/1-invalid.sbd").is_err());
    }

    #[test]
    fn undersized() {
        let file = File::open("data/0-mo.sbd").unwrap();
        let readable = file.take(58);
        assert!(Message::read_from(readable).is_err());
    }

    #[test]
    fn oversized() {
        let file = File::open("data/0-mo.sbd").unwrap();
        let readable = file.chain(Cursor::new(vec![0]));
        assert!(Message::read_from(readable).is_err());
    }

    #[test]
    fn values() {
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        assert_eq!(1, message.protocol_revision_number());
        assert_eq!(1894516585, message.cdr_reference());
        assert_eq!("300234063904190", message.imei());
        assert_eq!(SessionStatus::Ok, message.session_status());
        assert_eq!(75, message.momsn());
        assert_eq!(0, message.mtmsn());
        assert_eq!(
            Utc.ymd(2015, 7, 9).and_hms(18, 15, 8),
            message.time_of_session()
        );
        assert_eq!(
            "test message from pete",
            str::from_utf8(message.payload_ref()).unwrap()
        );
    }

    #[test]
    fn write() {
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        let mut cursor = Cursor::new(Vec::new());
        message.write_to(&mut cursor).unwrap();
        cursor.set_position(0);
        let message2 = Message::read_from(cursor).unwrap();
        assert_eq!(message, message2);
    }

    #[test]
    fn write_default() {
        let message: Message = Default::default();
        let mut cursor = Cursor::new(Vec::new());
        message.write_to(&mut cursor).unwrap();
        cursor.set_position(0);
        let message2 = Message::read_from(cursor).unwrap();
        assert_eq!(message, message2);
    }

    #[test]
    fn order() {
        let default: Message = Default::default();
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        assert!(message > default);
    }

    #[test]
    fn payload_str() {
        let message = Message::from_path("data/0-mo.sbd").unwrap();
        assert_eq!("test message from pete", message.payload_str().unwrap());
    }
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate std;
use self::std::prelude::v1::*;
use self::std::env;
use self::std::io::Error as IoError;
use self::libc::{c_void, size_t};
use core::ptr;
use stack;
use valgrind;

#[allow(raw_pointer_derive)]
#[derive(Debug)]
pub struct Stack {
  ptr: *mut u8,
  len: usize,
  valgrind_id: valgrind::stack_id_t
}

pub struct StackSource;

impl stack::StackSource for StackSource {
  type Output = Stack;
  fn get_stack(size: usize) -> Stack {
    Stack::new(size)
  }
}

impl stack::Stack for Stack {
  fn top(&mut self) -> *mut u8 {
    unsafe {
      self.ptr.offset(self.len as isize)
    }
  }

  fn limit(&self) -> *const u8 {
    unsafe {
      self.ptr.offset(env::page_size() as isize)
    }
  }
}

impl Stack {
  fn new(size: usize) -> Stack {
    let page_size = env::page_size();

    // round the page size up,
    // using the fact that it is a power of two
    let len = (size + page_size - 1) & !(page_size - 1);

    const STACK_PROT: libc::c_int = libc::PROT_READ | libc::PROT_WRITE;
    const STACK_FLAGS: libc::c_int = libc::MAP_STACK
                                   | libc::MAP_PRIVATE
                                   | libc::MAP_ANON;

    let stack = unsafe {
      let ptr = libc::mmap(ptr::null_mut(), len as size_t,
                           STACK_PROT, STACK_FLAGS, -1, 0);

      if ptr == libc::MAP_FAILED {
        panic!("mmap for stack of size {} failed: {:?}",
               len, IoError::last_os_error())
      }

      let valgrind_id =
        valgrind::stack_register(ptr.offset(len as isize) as *const _,
                                 ptr as *const _);

      Stack { ptr: ptr as *mut u8, len: len, valgrind_id: valgrind_id }
    };

    stack.protect();

    stack
  }

  fn protect(&self) {
    unsafe {
      if libc::mprotect(self.ptr as *mut c_void,
                        env::page_size() as libc::size_t,
                        libc::PROT_NONE) != 0 {
        panic!("mprotect for guard page of stack {:p} failed: {:?}",
               self.ptr, IoError::last_os_error());
      }
    }
  }
}

impl Drop for Stack {
  fn drop(&mut self) {
    unsafe {
      valgrind::stack_deregister(self.valgrind_id);
      if libc::munmap(self.ptr as *mut c_void, self.len as size_t) != 0 {
        panic!("munmap for stack {:p} of size {} failed: {:?}",
               self.ptr, self.len, IoError::last_os_error())
      }
    }
  }
}
<commit_msg>pass os::Stack addresses to Valgrind in the right order<commit_after>extern crate libc;
extern crate std;
use self::std::prelude::v1::*;
use self::std::env;
use self::std::io::Error as IoError;
use self::libc::{c_void, size_t};
use core::ptr;
use stack;
use valgrind;

#[allow(raw_pointer_derive)]
#[derive(Debug)]
pub struct Stack {
  ptr: *mut u8,
  len: usize,
  valgrind_id: valgrind::stack_id_t
}

pub struct StackSource;

impl stack::StackSource for StackSource {
  type Output = Stack;
  fn get_stack(size: usize) -> Stack {
    Stack::new(size)
  }
}

impl stack::Stack for Stack {
  fn top(&mut self) -> *mut u8 {
    unsafe {
      self.ptr.offset(self.len as isize)
    }
  }

  fn limit(&self) -> *const u8 {
    unsafe {
      self.ptr.offset(env::page_size() as isize)
    }
  }
}

impl Stack {
  fn new(size: usize) -> Stack {
    let page_size = env::page_size();

    // round the page size up,
    // using the fact that it is a power of two
    let len = (size + page_size - 1) & !(page_size - 1);

    const STACK_PROT: libc::c_int = libc::PROT_READ | libc::PROT_WRITE;
    const STACK_FLAGS: libc::c_int = libc::MAP_STACK
                                   | libc::MAP_PRIVATE
                                   | libc::MAP_ANON;

    let stack = unsafe {
      let ptr = libc::mmap(ptr::null_mut(), len as size_t,
                           STACK_PROT, STACK_FLAGS, -1, 0);

      if ptr == libc::MAP_FAILED {
        panic!("mmap for stack of size {} failed: {:?}",
               len, IoError::last_os_error())
      }

      let valgrind_id =
        valgrind::stack_register(ptr as *const _,
                                 ptr.offset(len as isize) as *const _);

      Stack { ptr: ptr as *mut u8, len: len, valgrind_id: valgrind_id }
    };

    stack.protect();

    stack
  }

  fn protect(&self) {
    unsafe {
      if libc::mprotect(self.ptr as *mut c_void,
                        env::page_size() as libc::size_t,
                        libc::PROT_NONE) != 0 {
        panic!("mprotect for guard page of stack {:p} failed: {:?}",
               self.ptr, IoError::last_os_error());
      }
    }
  }
}

impl Drop for Stack {
  fn drop(&mut self) {
    unsafe {
      valgrind::stack_deregister(self.valgrind_id);
      if libc::munmap(self.ptr as *mut c_void, self.len as size_t) != 0 {
        panic!("munmap for stack {:p} of size {} failed: {:?}",
               self.ptr, self.len, IoError::last_os_error())
      }
    }
  }
}
<|endoftext|>


<commit_before>use std::ffi::{CStr,CString};
use std::ptr;
use std::mem;
use std::str;

use libc;
use z3_sys;

pub struct Z3 {
    ctx: z3_sys::Z3_context
}
pub struct Z3Ast<'a> {
    ast: z3_sys::Z3_ast,
    z3: &'a Z3
}
pub struct Z3Sort<'a> {
    sort: z3_sys::Z3_sort,
    z3: &'a Z3
}
pub struct Z3Model<'a> {
    model: *mut z3_sys::Z3_model,
    res: i32,
    z3: &'a Z3
}

macro_rules! impl_mk_1 {
    ($name:ident, $fun:ident) => (
        pub fn $name<'a>(&'a self, a: &Z3Ast) -> Z3Ast<'a> {
            unsafe {
                let ast = z3_sys::$fun(self.ctx, a.ast);
                Z3Ast {
                    ast: ast,
                    z3: &self
                }
            }
        }
    )
}
macro_rules! impl_mk_2 {
    ($name:ident, $fun:ident) => (
        pub fn $name<'a>(&'a self, t1: &Z3Ast, t2: &Z3Ast) -> Z3Ast<'a> {
            unsafe {
                let ast = z3_sys::$fun(self.ctx, t1.ast, t2.ast);
                Z3Ast {
                    ast: ast,
                    z3: &self
                }
            }
        }
    )
}

impl Z3 {
    pub fn new() -> Z3 {
        unsafe {
            let cfg = z3_sys::Z3_mk_config();
            let ctx = z3_sys::Z3_mk_context(cfg);
            Z3 {
                ctx: ctx
            }
        }
    }
    #[inline]
    pub fn ctx(&self) -> z3_sys::Z3_context {
        self.ctx
    }
    pub unsafe fn get_model_str(&self, m: z3_sys::Z3_model) -> &str {
        let cstr = z3_sys::Z3_model_to_string(self.ctx, m);
        let slice = CStr::from_ptr(cstr);
        let buf: &[u8] = slice.to_bytes();
        let str_slice: &str = str::from_utf8(buf).unwrap();
        str_slice
    }

    pub fn check_and_get_model<'a>(&'a self, f: &Z3Ast) -> Z3Model<'a> {
        unsafe {
            let mut m: *mut z3_sys::Z3_model =
                libc::malloc(mem::size_of::<usize>() as libc::size_t)
                as *mut z3_sys::Z3_model;

            z3_sys::Z3_push(self.ctx);
            let not_f = z3_sys::Z3_mk_not(self.ctx, f.ast);
            z3_sys::Z3_assert_cnstr(self.ctx, not_f);

            let res = z3_sys::Z3_check_and_get_model(self.ctx, m);

            z3_sys::Z3_pop(self.ctx, 1);

            Z3Model {
                model: m,
                res: res,
                z3: &self
            }
        }
    }

    pub fn prove(&self, f: &Z3Ast, is_valid: bool) {
        unsafe {
            z3_sys::Z3_push(self.ctx);
            let not_f = z3_sys::Z3_mk_not(self.ctx, f.ast);
            z3_sys::Z3_assert_cnstr(self.ctx, not_f);

            let mut m: *mut z3_sys::Z3_model =
                libc::malloc(mem::size_of::<usize>() as libc::size_t)
                as *mut z3_sys::Z3_model;

            match z3_sys::Z3_check_and_get_model(self.ctx, m) {
                z3_sys::Z3_L_FALSE => {
                    /* proved */
                    println!("valid\n");
                    if !is_valid {
                        //exitf("unexpected result");
                    }
                },
                z3_sys::Z3_L_UNDEF => {
                    /* Z3 failed to prove/disprove f. */
                    println!("unknown\n");
                    if *m != ptr::null_mut() {
                        /* m should be viewed as a potential counterexample. */
                        println!("potential counterexample:\n{}\n",
                                 self.get_model_str(*m));
                    }
                    if is_valid {
                        //exitf("unexpected result");
                    }
                },
                z3_sys::Z3_L_TRUE => {
                    /* disproved */
                    println!("invalid\n");
                    if *m != ptr::null_mut() {
                        /* the model returned by Z3 is a counterexample */
                        println!("counterexample:\n{}\n",
                                 self.get_model_str(*m));
                    }
                    if is_valid {
                        //exitf("unexpected result");
                    }
                },
                _ => panic!("not supported")
            }
            if *m != ptr::null_mut() {
                z3_sys::Z3_del_model(self.ctx, *m);
            }
            z3_sys::Z3_pop(self.ctx(), 1);
        }
    }
    pub fn bv_sort<'a>(&'a self, width: u32) -> Z3Sort<'a> {
        unsafe {
            Z3Sort {
                sort: z3_sys::Z3_mk_bv_sort(self.ctx, width),
                z3: &self
            }
        }
    }

    impl_mk_2!(div, Z3_mk_div);
    impl_mk_2!(od, Z3_mk_mod);
    impl_mk_2!(re, Z3_mk_rem);
    impl_mk_2!(power, Z3_mk_power);
    impl_mk_2!(lt, Z3_mk_lt);
    impl_mk_2!(le, Z3_mk_le);
    impl_mk_2!(gt, Z3_mk_gt);
    impl_mk_2!(ge, Z3_mk_ge);
    impl_mk_2!(bvand, Z3_mk_bvand);
    impl_mk_2!(bvor, Z3_mk_bvor);
    impl_mk_2!(bvxor, Z3_mk_bvxor);
    impl_mk_2!(bvnand, Z3_mk_bvnand);
    impl_mk_2!(bvnor, Z3_mk_bvnor);
    impl_mk_2!(bvxnor, Z3_mk_bvxnor);
    impl_mk_2!(bvadd, Z3_mk_bvadd);
    impl_mk_2!(bvsub, Z3_mk_bvsub);
    impl_mk_2!(bvmul, Z3_mk_bvmul);
    impl_mk_2!(bvudiv, Z3_mk_bvudiv);
    impl_mk_2!(bvsdiv, Z3_mk_bvsdiv);
    impl_mk_2!(bvure, Z3_mk_bvurem);
    impl_mk_2!(bvsre, Z3_mk_bvsrem);
    impl_mk_2!(bvsmod, Z3_mk_bvsmod);
    impl_mk_2!(bvult, Z3_mk_bvult);
    impl_mk_2!(bvslt, Z3_mk_bvslt);
    impl_mk_2!(bvule, Z3_mk_bvule);
    impl_mk_2!(bvsle, Z3_mk_bvsle);
    impl_mk_2!(bvuge, Z3_mk_bvuge);
    impl_mk_2!(bvsge, Z3_mk_bvsge);
    impl_mk_2!(bvugt, Z3_mk_bvugt);
    impl_mk_2!(bvsgt, Z3_mk_bvsgt);
    impl_mk_2!(concat, Z3_mk_concat);
    impl_mk_2!(bvshl, Z3_mk_bvshl);
    impl_mk_2!(bvlshr, Z3_mk_bvlshr);
    impl_mk_2!(bvashr, Z3_mk_bvashr);
    impl_mk_2!(ext_rotate_left, Z3_mk_ext_rotate_left);
    impl_mk_2!(ext_rotate_right, Z3_mk_ext_rotate_right);
    impl_mk_2!(bvadd_no_underflow, Z3_mk_bvadd_no_underflow);
    impl_mk_2!(bvsub_no_overflow, Z3_mk_bvsub_no_overflow);
    impl_mk_2!(bvsdiv_no_overflow, Z3_mk_bvsdiv_no_overflow);
    impl_mk_2!(bvmul_no_underflow, Z3_mk_bvmul_no_underflow);

    impl_mk_1!(not, Z3_mk_not);
    impl_mk_1!(int2real, Z3_mk_int2real);
    impl_mk_1!(real2int, Z3_mk_real2int);
    impl_mk_1!(is_int, Z3_mk_is_int);
    impl_mk_1!(bvnot, Z3_mk_bvnot);
    impl_mk_1!(bvredand, Z3_mk_bvredand);
    impl_mk_1!(bvredor, Z3_mk_bvredor);
    impl_mk_1!(bvneg_no_overflow, Z3_mk_bvneg_no_overflow);
    impl_mk_1!(bvneg, Z3_mk_bvneg);

    pub fn eq<'a>(&'a self, t1: &Z3Ast, t2: &Z3Ast) -> Z3Ast<'a> {
        unsafe {
            let ast = z3_sys::Z3_mk_eq(self.ctx, t1.ast, t2.ast);
            Z3Ast {
                ast: ast,
                z3: &self
            }
        }
    }
}

impl<'a> Z3Sort<'a> {
    pub fn mk_const_i(&'a self, i: i32) -> Z3Ast<'a> {
        unsafe {
            let sym = z3_sys::Z3_mk_int_symbol(self.z3.ctx(), i);
            Z3Ast {
                ast: z3_sys::Z3_mk_const(self.z3.ctx(), sym, self.sort),
                z3: self.z3
            }
        }
    }
    pub fn mk_const_str(&'a self, s: &str) -> Z3Ast<'a> {
        unsafe {
            let cs = CString::new(s).unwrap();
            let sym = z3_sys::Z3_mk_string_symbol(self.z3.ctx(), cs.as_ptr());
            Z3Ast {
                ast: z3_sys::Z3_mk_const(self.z3.ctx(), sym, self.sort),
                z3: self.z3
            }
        }
    }
}

impl <'a> Z3Ast<'a> {
    pub fn get_u64(&self) -> Option<u64> {
        unsafe {
            let mut res: u64 = 0;
            if z3_sys::Z3_get_numeral_uint64(self.z3.ctx(), self.ast, &mut res)
                == z3_sys::Z3_L_TRUE
            {
                Some(res)
            } else {
                None
            }
        }
    }
}

impl <'a> Z3Model<'a> {
    pub fn eval(&'a self, ast: &Z3Ast) -> Option<Z3Ast<'a>> {
        if self.is_model() {
            unsafe {
                let mut res_ast: z3_sys::Z3_ast = ptr::null_mut();
                let res = z3_sys::Z3_model_eval(self.z3.ctx(), *self.model,
                                            ast.ast, z3_sys::Z3_L_TRUE,
                                            &mut res_ast);
                Some(Z3Ast {
                    ast: res_ast,
                    z3: self.z3
                })
            }
        } else {
            None
        }
    }
    pub fn is_model(&self) -> bool {
        match self.res {
            z3_sys::Z3_L_UNDEF | z3_sys::Z3_L_TRUE => true,
            z3_sys::Z3_L_FALSE => false,
            _ => panic!("not supported")
        }
    }

    pub fn get_str(&self) -> &str {
        if self.is_model() {
                unsafe {
                    let cstr = z3_sys::Z3_model_to_string(self.z3.ctx(), *self.model);
                    let slice = CStr::from_ptr(cstr);
                    let buf: &[u8] = slice.to_bytes();
                    let str_slice: &str = str::from_utf8(buf).unwrap();
                    str_slice
                }
        } else {
            "no model"
        }
    }
}

impl Drop for Z3 {
    fn drop(&mut self) {
        unsafe {
            z3_sys::Z3_del_context(self.ctx);
        }
    }
}

#[test]
fn it_works() {
    let z3 = Z3::new();
    let bv_sorts = z3.bv_sort(32);
    let a = bv_sorts.mk_const_i(1);
    let b = bv_sorts.mk_const_i(2);
    let c = bv_sorts.mk_const_str("c");

    let d = z3.bvand(&a, &b);
    let e = z3.bvor(&a, &b);

    let gt = z3.bvugt(&d, &e);
    let eq = z3.eq(&d, &e);
    let neq = z3.not(&gt);

    let h = z3.prove(&eq, true);

    let model = z3.check_and_get_model(&eq);
    println!("model: {}", model.get_str());

    let new_a = model.eval(&a).unwrap();
    let val_a = new_a.get_u64().unwrap();
    println!("val_a: 0x{:x}", val_a);
}
<commit_msg>making ctx() unsafe because it could escape lifetimes<commit_after>use std::ffi::{CStr,CString};
use std::ptr;
use std::mem;
use std::str;

use libc;
use z3_sys;

pub struct Z3 {
    ctx: z3_sys::Z3_context
}
pub struct Z3Ast<'a> {
    ast: z3_sys::Z3_ast,
    z3: &'a Z3
}
pub struct Z3Sort<'a> {
    sort: z3_sys::Z3_sort,
    z3: &'a Z3
}
pub struct Z3Model<'a> {
    model: *mut z3_sys::Z3_model,
    res: i32,
    z3: &'a Z3
}

macro_rules! impl_mk_1 {
    ($name:ident, $fun:ident) => (
        pub fn $name<'a>(&'a self, a: &Z3Ast) -> Z3Ast<'a> {
            unsafe {
                let ast = z3_sys::$fun(self.ctx, a.ast);
                Z3Ast {
                    ast: ast,
                    z3: &self
                }
            }
        }
    )
}
macro_rules! impl_mk_2 {
    ($name:ident, $fun:ident) => (
        pub fn $name<'a>(&'a self, t1: &Z3Ast, t2: &Z3Ast) -> Z3Ast<'a> {
            unsafe {
                let ast = z3_sys::$fun(self.ctx, t1.ast, t2.ast);
                Z3Ast {
                    ast: ast,
                    z3: &self
                }
            }
        }
    )
}

impl Z3 {
    pub fn new() -> Z3 {
        unsafe {
            let cfg = z3_sys::Z3_mk_config();
            let ctx = z3_sys::Z3_mk_context(cfg);
            Z3 {
                ctx: ctx
            }
        }
    }
    #[inline]
    pub unsafe fn ctx(&self) -> z3_sys::Z3_context {
        self.ctx
    }
    pub unsafe fn get_model_str(&self, m: z3_sys::Z3_model) -> &str {
        let cstr = z3_sys::Z3_model_to_string(self.ctx, m);
        let slice = CStr::from_ptr(cstr);
        let buf: &[u8] = slice.to_bytes();
        let str_slice: &str = str::from_utf8(buf).unwrap();
        str_slice
    }

    pub fn check_and_get_model<'a>(&'a self, f: &Z3Ast) -> Z3Model<'a> {
        unsafe {
            let mut m: *mut z3_sys::Z3_model =
                libc::malloc(mem::size_of::<usize>() as libc::size_t)
                as *mut z3_sys::Z3_model;

            z3_sys::Z3_push(self.ctx);
            let not_f = z3_sys::Z3_mk_not(self.ctx, f.ast);
            z3_sys::Z3_assert_cnstr(self.ctx, not_f);

            let res = z3_sys::Z3_check_and_get_model(self.ctx, m);

            z3_sys::Z3_pop(self.ctx, 1);

            Z3Model {
                model: m,
                res: res,
                z3: &self
            }
        }
    }

    pub fn prove(&self, f: &Z3Ast, is_valid: bool) {
        unsafe {
            z3_sys::Z3_push(self.ctx);
            let not_f = z3_sys::Z3_mk_not(self.ctx, f.ast);
            z3_sys::Z3_assert_cnstr(self.ctx, not_f);

            let mut m: *mut z3_sys::Z3_model =
                libc::malloc(mem::size_of::<usize>() as libc::size_t)
                as *mut z3_sys::Z3_model;

            match z3_sys::Z3_check_and_get_model(self.ctx, m) {
                z3_sys::Z3_L_FALSE => {
                    /* proved */
                    println!("valid\n");
                    if !is_valid {
                        //exitf("unexpected result");
                    }
                },
                z3_sys::Z3_L_UNDEF => {
                    /* Z3 failed to prove/disprove f. */
                    println!("unknown\n");
                    if *m != ptr::null_mut() {
                        /* m should be viewed as a potential counterexample. */
                        println!("potential counterexample:\n{}\n",
                                 self.get_model_str(*m));
                    }
                    if is_valid {
                        //exitf("unexpected result");
                    }
                },
                z3_sys::Z3_L_TRUE => {
                    /* disproved */
                    println!("invalid\n");
                    if *m != ptr::null_mut() {
                        /* the model returned by Z3 is a counterexample */
                        println!("counterexample:\n{}\n",
                                 self.get_model_str(*m));
                    }
                    if is_valid {
                        //exitf("unexpected result");
                    }
                },
                _ => panic!("not supported")
            }
            if *m != ptr::null_mut() {
                z3_sys::Z3_del_model(self.ctx, *m);
            }
            z3_sys::Z3_pop(self.ctx(), 1);
        }
    }
    pub fn bv_sort<'a>(&'a self, width: u32) -> Z3Sort<'a> {
        unsafe {
            Z3Sort {
                sort: z3_sys::Z3_mk_bv_sort(self.ctx, width),
                z3: &self
            }
        }
    }

    impl_mk_2!(div, Z3_mk_div);
    impl_mk_2!(od, Z3_mk_mod);
    impl_mk_2!(re, Z3_mk_rem);
    impl_mk_2!(power, Z3_mk_power);
    impl_mk_2!(lt, Z3_mk_lt);
    impl_mk_2!(le, Z3_mk_le);
    impl_mk_2!(gt, Z3_mk_gt);
    impl_mk_2!(ge, Z3_mk_ge);
    impl_mk_2!(bvand, Z3_mk_bvand);
    impl_mk_2!(bvor, Z3_mk_bvor);
    impl_mk_2!(bvxor, Z3_mk_bvxor);
    impl_mk_2!(bvnand, Z3_mk_bvnand);
    impl_mk_2!(bvnor, Z3_mk_bvnor);
    impl_mk_2!(bvxnor, Z3_mk_bvxnor);
    impl_mk_2!(bvadd, Z3_mk_bvadd);
    impl_mk_2!(bvsub, Z3_mk_bvsub);
    impl_mk_2!(bvmul, Z3_mk_bvmul);
    impl_mk_2!(bvudiv, Z3_mk_bvudiv);
    impl_mk_2!(bvsdiv, Z3_mk_bvsdiv);
    impl_mk_2!(bvure, Z3_mk_bvurem);
    impl_mk_2!(bvsre, Z3_mk_bvsrem);
    impl_mk_2!(bvsmod, Z3_mk_bvsmod);
    impl_mk_2!(bvult, Z3_mk_bvult);
    impl_mk_2!(bvslt, Z3_mk_bvslt);
    impl_mk_2!(bvule, Z3_mk_bvule);
    impl_mk_2!(bvsle, Z3_mk_bvsle);
    impl_mk_2!(bvuge, Z3_mk_bvuge);
    impl_mk_2!(bvsge, Z3_mk_bvsge);
    impl_mk_2!(bvugt, Z3_mk_bvugt);
    impl_mk_2!(bvsgt, Z3_mk_bvsgt);
    impl_mk_2!(concat, Z3_mk_concat);
    impl_mk_2!(bvshl, Z3_mk_bvshl);
    impl_mk_2!(bvlshr, Z3_mk_bvlshr);
    impl_mk_2!(bvashr, Z3_mk_bvashr);
    impl_mk_2!(ext_rotate_left, Z3_mk_ext_rotate_left);
    impl_mk_2!(ext_rotate_right, Z3_mk_ext_rotate_right);
    impl_mk_2!(bvadd_no_underflow, Z3_mk_bvadd_no_underflow);
    impl_mk_2!(bvsub_no_overflow, Z3_mk_bvsub_no_overflow);
    impl_mk_2!(bvsdiv_no_overflow, Z3_mk_bvsdiv_no_overflow);
    impl_mk_2!(bvmul_no_underflow, Z3_mk_bvmul_no_underflow);

    impl_mk_1!(not, Z3_mk_not);
    impl_mk_1!(int2real, Z3_mk_int2real);
    impl_mk_1!(real2int, Z3_mk_real2int);
    impl_mk_1!(is_int, Z3_mk_is_int);
    impl_mk_1!(bvnot, Z3_mk_bvnot);
    impl_mk_1!(bvredand, Z3_mk_bvredand);
    impl_mk_1!(bvredor, Z3_mk_bvredor);
    impl_mk_1!(bvneg_no_overflow, Z3_mk_bvneg_no_overflow);
    impl_mk_1!(bvneg, Z3_mk_bvneg);

    pub fn eq<'a>(&'a self, t1: &Z3Ast, t2: &Z3Ast) -> Z3Ast<'a> {
        unsafe {
            let ast = z3_sys::Z3_mk_eq(self.ctx, t1.ast, t2.ast);
            Z3Ast {
                ast: ast,
                z3: &self
            }
        }
    }
}

impl<'a> Z3Sort<'a> {
    pub fn mk_const_i(&'a self, i: i32) -> Z3Ast<'a> {
        unsafe {
            let sym = z3_sys::Z3_mk_int_symbol(self.z3.ctx(), i);
            Z3Ast {
                ast: z3_sys::Z3_mk_const(self.z3.ctx(), sym, self.sort),
                z3: self.z3
            }
        }
    }
    pub fn mk_const_str(&'a self, s: &str) -> Z3Ast<'a> {
        unsafe {
            let cs = CString::new(s).unwrap();
            let sym = z3_sys::Z3_mk_string_symbol(self.z3.ctx(), cs.as_ptr());
            Z3Ast {
                ast: z3_sys::Z3_mk_const(self.z3.ctx(), sym, self.sort),
                z3: self.z3
            }
        }
    }
}

impl <'a> Z3Ast<'a> {
    pub fn get_u64(&self) -> Option<u64> {
        unsafe {
            let mut res: u64 = 0;
            if z3_sys::Z3_get_numeral_uint64(self.z3.ctx(), self.ast, &mut res)
                == z3_sys::Z3_L_TRUE
            {
                Some(res)
            } else {
                None
            }
        }
    }
}

impl <'a> Z3Model<'a> {
    pub fn eval(&'a self, ast: &Z3Ast) -> Option<Z3Ast<'a>> {
        if self.is_model() {
            unsafe {
                let mut res_ast: z3_sys::Z3_ast = ptr::null_mut();
                if z3_sys::Z3_model_eval(self.z3.ctx(), *self.model,
                                         ast.ast, z3_sys::Z3_L_TRUE,
                                         &mut res_ast) == z3_sys::Z3_L_TRUE {
                    Some(Z3Ast {
                        ast: res_ast,
                        z3: self.z3
                    })
                } else {
                    None
                }
            }
        } else {
            None
        }
    }
    pub fn is_model(&self) -> bool {
        match self.res {
            z3_sys::Z3_L_UNDEF | z3_sys::Z3_L_TRUE => true,
            z3_sys::Z3_L_FALSE => false,
            _ => panic!("not supported")
        }
    }

    pub fn get_str(&self) -> &str {
        if self.is_model() {
                unsafe {
                    let cstr = z3_sys::Z3_model_to_string(self.z3.ctx(), *self.model);
                    let slice = CStr::from_ptr(cstr);
                    let buf: &[u8] = slice.to_bytes();
                    let str_slice: &str = str::from_utf8(buf).unwrap();
                    str_slice
                }
        } else {
            "no model"
        }
    }
}

impl Drop for Z3 {
    fn drop(&mut self) {
        unsafe {
            z3_sys::Z3_del_context(self.ctx);
        }
    }
}

#[test]
fn it_works() {
    let z3 = Z3::new();
    let bv_sorts = z3.bv_sort(32);
    let a = bv_sorts.mk_const_i(1);
    let b = bv_sorts.mk_const_i(2);
    let c = bv_sorts.mk_const_str("c");

    let d = z3.bvand(&a, &b);
    let e = z3.bvor(&a, &b);

    let gt = z3.bvugt(&d, &e);
    let eq = z3.eq(&d, &e);
    let neq = z3.not(&gt);

    let h = z3.prove(&eq, true);

    let model = z3.check_and_get_model(&eq);
    println!("model: {}", model.get_str());

    let new_a = model.eval(&a).unwrap();
    let val_a = new_a.get_u64().unwrap();
    println!("val_a: 0x{:x}", val_a);
}
<|endoftext|>