<commit_before>use log::{debug};
use goblin::{Object};
use failure::{Error};
use byteorder::{ByteOrder, LittleEndian};

use super::super::super::arch::{RVA};
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct ImportsAnalyzer {}

impl ImportsAnalyzer {
    pub fn new() -> ImportsAnalyzer {
        ImportsAnalyzer {}
    }
}

pub struct ImageImportDescriptor {
    pub original_first_thunk: RVA,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name: RVA,
    pub first_thunk: RVA,
}

impl ImageImportDescriptor {
    pub fn is_empty(&self) -> bool {
        self.original_first_thunk == RVA(0x0) &&
            self.time_date_stamp == 0x0 &&
            self.forwarder_chain == 0x0 &&
            self.name == RVA(0x0) &&
            self.first_thunk == RVA(0x0)
    }
}

impl std::fmt::Debug for ImageImportDescriptor {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_DESCRIPTOR(FT: {} OFT: {} name: {})",
               self.first_thunk,
               self.original_first_thunk,
               self.name
        )
    }
}

pub fn read_image_import_descriptor(ws: &Workspace, rva: RVA) -> Result<ImageImportDescriptor, Error> {
    let buf = ws.read_bytes(rva, 5 * 4)?;

    // these fields are all u32, even on 64-bit
    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    Ok(ImageImportDescriptor {
        original_first_thunk: RVA::from(entries[0] as i64),
        time_date_stamp: entries[1],
        forwarder_chain: entries[2],
        name: RVA::from(entries[3] as i64),
        first_thunk: RVA::from(entries[4] as i64),
    })
}

pub enum ImageThunkData {
    Function(RVA),
    Ordinal(u32),
}

pub fn read_image_thunk_data(ws: &Workspace, rva: RVA) -> Result<ImageThunkData, Error> {
    // see: https://reverseengineering.stackexchange.com/a/13387/17194
    let thunk = ws.read_rva(rva)?;
    let v: i64 = thunk.into();
    if v & (1 << (ws.loader.get_arch().get_pointer_size() * 8 - 1)) > 0x0 {
        // MSB is set, this is an ordinal
        Ok(ImageThunkData::Ordinal((v & 0xFFFF) as u32))
    } else {
        Ok(ImageThunkData::Function(thunk))
    }
}

pub struct ImageImportByName {
    pub hint: u16,
    pub name: String,
}

impl std::fmt::Debug for ImageImportByName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_BY_NAME(hint: {} name: {})",
               self.hint,
               self.name
        )
    }
}

pub fn read_image_import_by_name(ws: &Workspace, rva: RVA) -> Result<ImageImportByName, Error> {
    Ok(ImageImportByName {
        hint: ws.read_u16(rva)?,
        name: ws.read_utf8(rva + RVA::from(2))?,
    })
}

pub fn read_best_thunk_data(ws: &Workspace, original_first_thunk: RVA, first_thunk: RVA) -> Result<ImageThunkData, Error> {
    // the Original First Thunk (OFT) remains constant, and points to the IMAGE_IMPORT_BY_NAME.
    // FT and OFT are parallel arrays.

    // the First Thunk (FT) is the pointer that will be overwritten upon load.
    // entries here may not point to the IMAGE_IMPORT_BY_NAME.
    //
    // in practice, using this array works better, as some OFT entries may be empty.
    // see: be24e9d47cfe588a8ced0ac3e453d390 hotfix2.exe
    //
    // however, this doesn't work if the PE has been dumped from memory
    // (and the FTs fixed up).
    match read_image_thunk_data(ws, first_thunk) {
        Ok(ImageThunkData::Function(rva)) => {
            if rva == RVA(0x0) {
                // end of array, this is a valid entry, return it.
                Ok(ImageThunkData::Function(rva))
            } else if ws.probe(rva, 1, Permissions::R) {
                // seems to be a valid address, return it.
                Ok(ImageThunkData::Function(rva))
            } else {
                // invalid address, so let's try from the OFT array
                // this might happen if we're dealing with a PE that was dumped from memory.
                // the FT might have been overwritten, so we need to use the OFT, instead.
                read_image_thunk_data(ws, original_first_thunk)
            }
        },
        Ok(ImageThunkData::Ordinal(ord)) => {
            // ordinal is a valid entry, return it.
            Ok(ImageThunkData::Ordinal(ord))
        }
        Err(_) => {
            // fall back to the OFT array
            read_image_thunk_data(ws, original_first_thunk)
        }
    }
}

impl Analyzer for ImportsAnalyzer {
    fn get_name(&self) -> String {
        "PE imports analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::ImportsAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = ImportsAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    /// assert_eq!(ws.get_symbol(RVA(0x77448)).unwrap(), "api-ms-win-core-appcompat-l1-1-1.dll!BaseReadAppCompatDataForProcess");
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let import_directory = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let import_directory = match opt_header.data_directories.get_import_table() {
                Some(import_directory) => import_directory,
                _ => return Ok(()),
            };

            RVA::from(import_directory.virtual_address as i64)
        };

        debug!("found import directory: {}", import_directory);

        let mut symbols: Vec<(RVA, String)> = vec![];

        let psize: usize = ws.loader.get_arch().get_pointer_size() as usize;
        for i in 0..std::usize::MAX {
            //
            //  0x0                    0x14
            //  +-------------------------+           0x0          ptrsize
            //  | IMAGE_IMPORT_DESCRIPTOR | --------> +------------------+
            //  +-------------------------+ \         | IMAGE_THUNK_DATA | -------> +----------------------+
            //  | ...                     |  |        +------------------+          | IMAGE_IMPORT_BY_NAME |
            //  +-------------------------+  |        | ...              |          +----------------------+
            //  | 00 00 00 00 00 00 000   |  |        +------------------+          | hint u16             |
            //  +-------------------------+  |        | 00 00 00 00 00   |          | name ascii           |
            //                               |        +------------------+          +----------------------+
            //                               |
            //                               +> dll-name (ascii)
            let import_descriptor_rva = import_directory + RVA::from(i * 0x14);
            let import_descriptor = read_image_import_descriptor(ws, import_descriptor_rva)?;
            if import_descriptor.is_empty() {
                break;
            }

            let dll_name = ws.read_utf8(import_descriptor.name)?;
            debug!("found {:?} -> {}", import_descriptor, dll_name);

            for j in 0..std::usize::MAX {
                // the Original First Thunk (OFT) remains constant, and points to the IMAGE_IMPORT_BY_NAME.
                // FT and OFT are parallel arrays.
                let original_first_thunk = import_descriptor.original_first_thunk + RVA::from(j * psize);

                // the First Thunk (FT) is the pointer that will be overwritten upon load.
                // entries here may not point to the IMAGE_IMPORT_BY_NAME.
                let first_thunk = import_descriptor.first_thunk + RVA::from(j * psize);

                match read_best_thunk_data(ws, original_first_thunk, first_thunk)? {
                    ImageThunkData::Function(rva) => {
                        if rva == RVA(0x0) {
                            // end of array
                            break;
                        } else {
                            let imp = read_image_import_by_name(ws, rva)?;
                            debug!("found {:?}", imp);
                            symbols.push((first_thunk, format!("{}!{}", dll_name, imp.name)))
                        }
                    },
                    ImageThunkData::Ordinal(ord) => {
                        symbols.push((first_thunk, format!("{}!#{}", dll_name, ord)))
                    },
                };
            }
        }

        for (rva, name) in symbols.iter() {
            debug!("found import: {} -> {}", rva, name);
            ws.make_symbol(*rva, name)?;
            ws.analyze()?;
        }

        Ok(())
    }
}
<commit_msg>imports: correct detect ordinals<commit_after>use log::{debug};
use goblin::{Object};
use failure::{Error};
use byteorder::{ByteOrder, LittleEndian};

use super::super::super::arch::{RVA};
use super::super::super::loader::{Permissions};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

pub struct ImportsAnalyzer {}

impl ImportsAnalyzer {
    pub fn new() -> ImportsAnalyzer {
        ImportsAnalyzer {}
    }
}

pub struct ImageImportDescriptor {
    pub original_first_thunk: RVA,
    pub time_date_stamp: u32,
    pub forwarder_chain: u32,
    pub name: RVA,
    pub first_thunk: RVA,
}

impl ImageImportDescriptor {
    pub fn is_empty(&self) -> bool {
        self.original_first_thunk == RVA(0x0) &&
            self.time_date_stamp == 0x0 &&
            self.forwarder_chain == 0x0 &&
            self.name == RVA(0x0) &&
            self.first_thunk == RVA(0x0)
    }
}

impl std::fmt::Debug for ImageImportDescriptor {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_DESCRIPTOR(FT: {} OFT: {} name: {})",
               self.first_thunk,
               self.original_first_thunk,
               self.name
        )
    }
}

pub fn read_image_import_descriptor(ws: &Workspace, rva: RVA) -> Result<ImageImportDescriptor, Error> {
    let buf = ws.read_bytes(rva, 5 * 4)?;

    // these fields are all u32, even on 64-bit
    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    Ok(ImageImportDescriptor {
        original_first_thunk: RVA::from(entries[0] as i64),
        time_date_stamp: entries[1],
        forwarder_chain: entries[2],
        name: RVA::from(entries[3] as i64),
        first_thunk: RVA::from(entries[4] as i64),
    })
}

pub enum ImageThunkData {
    Function(RVA),
    Ordinal(u32),
}

pub fn read_image_thunk_data(ws: &Workspace, rva: RVA) -> Result<ImageThunkData, Error> {
    // see: https://reverseengineering.stackexchange.com/a/13387/17194
    let thunk = ws.read_rva(rva)?;
    let v: u64 = thunk.into();

    let ordinal_mask: u64 = (1 as u64) << ((ws.loader.get_arch().get_pointer_size() * 8) - 1);
    if v & ordinal_mask > 0x0 {
        // MSB is set, this is an ordinal
        Ok(ImageThunkData::Ordinal((v & 0xFFFF) as u32))
    } else {
        Ok(ImageThunkData::Function(thunk))
    }
}

pub struct ImageImportByName {
    pub hint: u16,
    pub name: String,
}

impl std::fmt::Debug for ImageImportByName {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "IMAGE_IMPORT_BY_NAME(hint: {} name: {})",
               self.hint,
               self.name
        )
    }
}

pub fn read_image_import_by_name(ws: &Workspace, rva: RVA) -> Result<ImageImportByName, Error> {
    Ok(ImageImportByName {
        hint: ws.read_u16(rva)?,
        name: ws.read_utf8(rva + RVA::from(2))?,
    })
}

pub fn read_best_thunk_data(ws: &Workspace, original_first_thunk: RVA, first_thunk: RVA) -> Result<ImageThunkData, Error> {
    // the Original First Thunk (OFT) remains constant, and points to the IMAGE_IMPORT_BY_NAME.
    // FT and OFT are parallel arrays.

    // the First Thunk (FT) is the pointer that will be overwritten upon load.
    // entries here may not point to the IMAGE_IMPORT_BY_NAME.
    //
    // in practice, using this array works better, as some OFT entries may be empty.
    // see: be24e9d47cfe588a8ced0ac3e453d390 hotfix2.exe
    //
    // however, this doesn't work if the PE has been dumped from memory
    // (and the FTs fixed up).
    match read_image_thunk_data(ws, first_thunk) {
        Ok(ImageThunkData::Function(rva)) => {
            if rva == RVA(0x0) {
                // end of array, this is a valid entry, return it.
                Ok(ImageThunkData::Function(rva))
            } else if ws.probe(rva, 1, Permissions::R) {
                // seems to be a valid address, return it.
                Ok(ImageThunkData::Function(rva))
            } else {
                // invalid address, so let's try from the OFT array
                // this might happen if we're dealing with a PE that was dumped from memory.
                // the FT might have been overwritten, so we need to use the OFT, instead.
                read_image_thunk_data(ws, original_first_thunk)
            }
        },
        Ok(ImageThunkData::Ordinal(ord)) => {
            // ordinal is a valid entry, return it.
            Ok(ImageThunkData::Ordinal(ord))
        }
        Err(_) => {
            // fall back to the OFT array
            read_image_thunk_data(ws, original_first_thunk)
        }
    }
}

impl Analyzer for ImportsAnalyzer {
    fn get_name(&self) -> String {
        "PE imports analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::ImportsAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = ImportsAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    /// assert_eq!(ws.get_symbol(RVA(0x77448)).unwrap(), "api-ms-win-core-appcompat-l1-1-1.dll!BaseReadAppCompatDataForProcess");
    /// ```
    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let import_directory = {
            let pe = match Object::parse(&ws.buf) {
                Ok(Object::PE(pe)) => pe,
                _ => panic!("can't analyze unexpected format"),
            };

            let opt_header = match pe.header.optional_header {
                Some(opt_header) => opt_header,
                _ => return Ok(()),
            };

            let import_directory = match opt_header.data_directories.get_import_table() {
                Some(import_directory) => import_directory,
                _ => return Ok(()),
            };

            RVA::from(import_directory.virtual_address as i64)
        };

        debug!("found import directory: {}", import_directory);

        let mut symbols: Vec<(RVA, String)> = vec![];

        let psize: usize = ws.loader.get_arch().get_pointer_size() as usize;
        for i in 0..std::usize::MAX {
            //
            //  0x0                    0x14
            //  +-------------------------+           0x0          ptrsize
            //  | IMAGE_IMPORT_DESCRIPTOR | --------> +------------------+
            //  +-------------------------+ \         | IMAGE_THUNK_DATA | -------> +----------------------+
            //  | ...                     |  |        +------------------+          | IMAGE_IMPORT_BY_NAME |
            //  +-------------------------+  |        | ...              |          +----------------------+
            //  | 00 00 00 00 00 00 000   |  |        +------------------+          | hint u16             |
            //  +-------------------------+  |        | 00 00 00 00 00   |          | name ascii           |
            //                               |        +------------------+          +----------------------+
            //                               |
            //                               +> dll-name (ascii)
            let import_descriptor_rva = import_directory + RVA::from(i * 0x14);
            let import_descriptor = read_image_import_descriptor(ws, import_descriptor_rva)?;
            if import_descriptor.is_empty() {
                break;
            }

            let dll_name = ws.read_utf8(import_descriptor.name)?;
            debug!("found {:?} -> {}", import_descriptor, dll_name);

            for j in 0..std::usize::MAX {
                // the Original First Thunk (OFT) remains constant, and points to the IMAGE_IMPORT_BY_NAME.
                // FT and OFT are parallel arrays.
                let original_first_thunk = import_descriptor.original_first_thunk + RVA::from(j * psize);

                // the First Thunk (FT) is the pointer that will be overwritten upon load.
                // entries here may not point to the IMAGE_IMPORT_BY_NAME.
                let first_thunk = import_descriptor.first_thunk + RVA::from(j * psize);

                match read_best_thunk_data(ws, original_first_thunk, first_thunk)? {
                    ImageThunkData::Function(rva) => {
                        if rva == RVA(0x0) {
                            // end of array
                            break;
                        } else {
                            let imp = read_image_import_by_name(ws, rva)?;
                            debug!("found {:?}", imp);
                            symbols.push((first_thunk, format!("{}!{}", dll_name, imp.name)))
                        }
                    },
                    ImageThunkData::Ordinal(ord) => {
                        symbols.push((first_thunk, format!("{}!#{}", dll_name, ord)))
                    },
                };
            }
        }

        for (rva, name) in symbols.iter() {
            debug!("found import: {} -> {}", rva, name);
            ws.make_symbol(*rva, name)?;
            ws.analyze()?;
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>use winapi::um::winuser::{WS_VISIBLE, WS_DISABLED, SS_WORDELLIPSIS};
use crate::win32::window_helper as wh;
use crate::win32::base_helper::check_hwnd;
use crate::{Font, NwgError, HTextAlign, VTextAlign, RawEventHandler, unbind_raw_event_handler};
use super::{ControlBase, ControlHandle};
use std::cell::RefCell;

const NOT_BOUND: &'static str = "Label is not yet bound to a winapi object";
const BAD_HANDLE: &'static str = "INTERNAL ERROR: Label handle is not HWND!";

bitflags! {
    /**
        The label flags

        * NONE:     No flags. Equivalent to a invisible blank label.
        * VISIBLE:  The label is immediatly visible after creation
        * DISABLED: The label cannot be interacted with by the user. It also has a grayed out look.
    */
    pub struct LabelFlags: u32 {
        const NONE = 0;
        const VISIBLE = WS_VISIBLE;
        const DISABLED = WS_DISABLED;

        /// Truncate the label if the text is too long. A label with this style CANNOT have multiple lines.
        const ELIPSIS = SS_WORDELLIPSIS;
    }
}

/**
A label is a single line of static text. Use `\r\n` to split the text on multiple lines.

Label is not behind any features.

**Builder parameters:**
  * `parent`:           **Required.** The label parent container.
  * `text`:             The label text.
  * `size`:             The label size.
  * `position`:         The label position.
  * `enabled`:          If the label is enabled. A disabled label won't trigger events
  * `flags`:            A combination of the LabelFlags values.
  * `font`:             The font used for the label text
  * `background_color`: The background color of the label
  * `h_align`:          The horizontal aligment of the label

**Control events:**
  * `OnLabelClick`: When the user click the label
  * `OnLabelDoubleClick`: When the user double click a label
  * `MousePress(_)`: Generic mouse press events on the label
  * `OnMouseMove`: Generic mouse mouse event
  * `OnMouseWheel`: Generic mouse wheel event

** Example **

```rust
use native_windows_gui as nwg;
fn build_label(label: &mut nwg::Label, window: &nwg::Window, font: &nwg::Font) {
    nwg::Label::builder()
        .text("Hello")
        .font(Some(font))
        .parent(window)
        .build(label);
}
```

*/
#[derive(Default)]
pub struct Label {
    pub handle: ControlHandle,
    handler0: RefCell<Option<RawEventHandler>>,
    handler1: RefCell<Option<RawEventHandler>>,
}

impl Label {

    pub fn builder<'a>() -> LabelBuilder<'a> {
        LabelBuilder {
            text: "A label",
            size: (130, 25),
            position: (0, 0),
            flags: None,
            font: None,
            parent: None,
            h_align: HTextAlign::Left,
            v_align: VTextAlign::Top,
            background_color: None
        }
    }

    /// Return the font of the control
    pub fn font(&self) -> Option<Font> {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);

        let font_handle = wh::get_window_font(handle);
        if font_handle.is_null() {
            None
        } else {
            Some(Font { handle: font_handle })
        }
    }

    /// Set the font of the control
    pub fn set_font(&self, font: Option<&Font>) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_font(handle, font.map(|f| f.handle), true); }
    }

    /// Return true if the control currently has the keyboard focus
    pub fn focus(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_focus(handle) }
    }

    /// Set the keyboard focus on the button.
    pub fn set_focus(&self) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_focus(handle); }
    }

    /// Return true if the control user can interact with the control, return false otherwise
    pub fn enabled(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_enabled(handle) }
    }

    /// Enable or disable the control
    pub fn set_enabled(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_enabled(handle, v) }
    }

    /// Return true if the control is visible to the user. Will return true even if the
    /// control is outside of the parent client view (ex: at the position (10000, 10000))
    pub fn visible(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_visibility(handle) }
    }

    /// Show or hide the control to the user
    pub fn set_visible(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_visibility(handle, v) }
    }

    /// Return the size of the label in the parent window
    pub fn size(&self) -> (u32, u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_size(handle) }
    }

    /// Set the size of the label in the parent window
    pub fn set_size(&self, x: u32, y: u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_size(handle, x, y, false) }
    }

    /// Return the position of the label in the parent window
    pub fn position(&self) -> (i32, i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_position(handle) }
    }

    /// Set the position of the label in the parent window
    pub fn set_position(&self, x: i32, y: i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_position(handle, x, y) }
    }

    /// Return the label text
    pub fn text(&self) -> String {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_text(handle) }
    }

    /// Set the label text
    pub fn set_text<'a>(&self, v: &'a str) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_text(handle, v) }
    }

    /// Winapi class name used during control creation
    pub fn class_name(&self) -> &'static str {
        "STATIC"
    }

    /// Winapi base flags used during window creation
    pub fn flags(&self) -> u32 {
        use winapi::um::winuser::{SS_NOPREFIX, SS_LEFT};

        WS_VISIBLE | SS_NOPREFIX | SS_LEFT
    }

    /// Winapi flags required by the control
    pub fn forced_flags(&self) -> u32 {
        use winapi::um::winuser::{SS_NOTIFY, WS_CHILD};

        WS_CHILD | SS_NOTIFY
    }

    /// Center the text vertically.
    fn hook_non_client_size(&self, bg: Option<[u8; 3]>, v_align: VTextAlign) {
        use crate::bind_raw_event_handler_inner;
        use winapi::shared::windef::{HWND, HGDIOBJ, RECT, HBRUSH, POINT};
        use winapi::shared::{basetsd::UINT_PTR, minwindef::LRESULT};
        use winapi::um::winuser::{WM_CTLCOLORSTATIC, WM_NCCALCSIZE, WM_NCPAINT, WM_SIZE, DT_CALCRECT, DT_LEFT, NCCALCSIZE_PARAMS, COLOR_WINDOW};
        use winapi::um::winuser::{SWP_NOOWNERZORDER, SWP_NOSIZE, SWP_NOMOVE, SWP_FRAMECHANGED};
        use winapi::um::winuser::{GetDC, DrawTextW, ReleaseDC, GetClientRect, GetWindowRect, FillRect, ScreenToClient, SetWindowPos};
        use winapi::um::wingdi::{SelectObject, CreateSolidBrush, RGB};
        use std::{mem, ptr};

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let parent_handle = ControlHandle::Hwnd(wh::get_window_parent(handle));

        let brush = match bg {
            Some(c) => unsafe { CreateSolidBrush(RGB(c[0], c[1], c[2])) },
            None => COLOR_WINDOW as HBRUSH
        };

        match v_align {
            VTextAlign::Top => {},
            VTextAlign::Center => {},
            VTextAlign::Bottom => {},
        };

        unsafe {

        if bg.is_some() {
            let handler0 = bind_raw_event_handler_inner(&parent_handle, handle as UINT_PTR, move |_hwnd, msg, _w, l| {
                match msg {
                    WM_CTLCOLORSTATIC => {
                        let child = l as HWND;
                        if child == handle {
                            return Some(brush as LRESULT);
                        }
                    },
                    _ => {}
                }

                None
            });

            *self.handler0.borrow_mut() = Some(handler0.unwrap());
        }

        let handler1 = bind_raw_event_handler_inner(&self.handle, 0, move |hwnd, msg, w, l| {
            match msg {
                WM_NCCALCSIZE  => {
                    if w == 0 { return None }

                    // Calculate client area height needed for a font
                    let font_handle = wh::get_window_font(hwnd);
                    let mut r: RECT = mem::zeroed();
                    let dc = GetDC(hwnd);

                    let old = SelectObject(dc, font_handle as HGDIOBJ);
                    let calc: [u16;2] = [75, 121];
                    DrawTextW(dc, calc.as_ptr(), 2, &mut r, DT_CALCRECT | DT_LEFT);

                    let client_height = r.bottom;

                    SelectObject(dc, old);
                    ReleaseDC(hwnd, dc);

                    // Calculate NC area to center text.
                    let mut client: RECT = mem::zeroed();
                    let mut window: RECT = mem::zeroed();
                    GetClientRect(hwnd, &mut client);
                    GetWindowRect(hwnd, &mut window);

                    let window_height = window.bottom - window.top;
                    let center = ((window_height - client_height) / 2) - 1;

                    // Save the info
                    let info_ptr: *mut NCCALCSIZE_PARAMS = l as *mut NCCALCSIZE_PARAMS;
                    let info = &mut *info_ptr;

                    info.rgrc[0].top += center;
                    info.rgrc[0].bottom -= center;
                },
                WM_NCPAINT  => {
                    let mut window: RECT = mem::zeroed();
                    let mut client: RECT = mem::zeroed();
                    GetWindowRect(hwnd, &mut window);
                    GetClientRect(hwnd, &mut client);

                    let mut pt1 = POINT {x: window.left, y: window.top};
                    ScreenToClient(hwnd, &mut pt1);

                    let mut pt2 = POINT {x: window.right, y: window.bottom};
                    ScreenToClient(hwnd, &mut pt2);

                    let top = RECT {
                        left: 0,
                        top: pt1.y,
                        right: client.right,
                        bottom: client.top
                    };

                    let bottom = RECT {
                        left: 0,
                        top: client.bottom,
                        right: client.right,
                        bottom: pt2.y
                    };

                    let dc = GetDC(hwnd);
                    FillRect(dc, &top, brush);
                    FillRect(dc, &bottom, brush);
                    ReleaseDC(hwnd, dc);
                },
                WM_SIZE => {
                    SetWindowPos(hwnd, ptr::null_mut(), 0, 0, 0, 0, SWP_NOOWNERZORDER | SWP_NOSIZE | SWP_NOMOVE | SWP_FRAMECHANGED);
                },
                _ => {}
            }

            None
        });

        *self.handler1.borrow_mut() = Some(handler1.unwrap());

        }
    }

}

impl PartialEq for Label {
    fn eq(&self, other: &Self) -> bool {
        self.handle == other.handle
    }
}

impl Drop for Label {
    fn drop(&mut self) {
        let handler = self.handler0.borrow();
        if let Some(h) = handler.as_ref() {
            drop(unbind_raw_event_handler(h));
        }

        let handler = self.handler1.borrow();
        if let Some(h) = handler.as_ref() {
            drop(unbind_raw_event_handler(h));
        }

        self.handle.destroy();
    }
}

pub struct LabelBuilder<'a> {
    text: &'a str,
    size: (i32, i32),
    position: (i32, i32),
    background_color: Option<[u8; 3]>,
    flags: Option<LabelFlags>,
    font: Option<&'a Font>,
    h_align: HTextAlign,
    v_align: VTextAlign,
    parent: Option<ControlHandle>
}

impl<'a> LabelBuilder<'a> {

    pub fn flags(mut self, flags: LabelFlags) -> LabelBuilder<'a> {
        self.flags = Some(flags);
        self
    }

    pub fn text(mut self, text: &'a str) -> LabelBuilder<'a> {
        self.text = text;
        self
    }

    pub fn size(mut self, size: (i32, i32)) -> LabelBuilder<'a> {
        self.size = size;
        self
    }

    pub fn position(mut self, pos: (i32, i32)) -> LabelBuilder<'a> {
        self.position = pos;
        self
    }

    pub fn font(mut self, font: Option<&'a Font>) -> LabelBuilder<'a> {
        self.font = font;
        self
    }

    pub fn background_color(mut self, color: Option<[u8;3]>) -> LabelBuilder<'a> {
        self.background_color = color;
        self
    }

    pub fn h_align(mut self, align: HTextAlign) -> LabelBuilder<'a> {
        self.h_align = align;
        self
    }

    pub fn parent<C: Into<ControlHandle>>(mut self, p: C) -> LabelBuilder<'a> {
        self.parent = Some(p.into());
        self
    }

    pub fn build(self, out: &mut Label) -> Result<(), NwgError> {
        use winapi::um::winuser::{SS_LEFT, SS_RIGHT, SS_CENTER};

        let mut flags = self.flags.map(|f| f.bits()).unwrap_or(out.flags());

        match self.h_align {
            HTextAlign::Left => { flags |= SS_LEFT; },
            HTextAlign::Right => { flags |= SS_RIGHT; },
            HTextAlign::Center => { flags |= SS_CENTER; },
        }

        let parent = match self.parent {
            Some(p) => Ok(p),
            None => Err(NwgError::no_parent("Label"))
        }?;

        // Drop the old object
        *out = Label::default();

        out.handle = ControlBase::build_hwnd()
            .class_name(out.class_name())
            .forced_flags(out.forced_flags())
            .flags(flags)
            .size(self.size)
            .position(self.position)
            .text(self.text)
            .parent(Some(parent))
            .build()?;

        if self.font.is_some() {
            out.set_font(self.font);
        } else {
            out.set_font(Font::global_default().as_ref());
        }

        out.hook_non_client_size(self.background_color, self.v_align);

        Ok(())
    }

}

<commit_msg>v text align<commit_after>use winapi::um::winuser::{WS_VISIBLE, WS_DISABLED, SS_WORDELLIPSIS};
use crate::win32::window_helper as wh;
use crate::win32::base_helper::check_hwnd;
use crate::{Font, NwgError, HTextAlign, VTextAlign, RawEventHandler, unbind_raw_event_handler};
use super::{ControlBase, ControlHandle};
use std::cell::RefCell;

const NOT_BOUND: &'static str = "Label is not yet bound to a winapi object";
const BAD_HANDLE: &'static str = "INTERNAL ERROR: Label handle is not HWND!";

bitflags! {
    /**
        The label flags

        * NONE:     No flags. Equivalent to a invisible blank label.
        * VISIBLE:  The label is immediatly visible after creation
        * DISABLED: The label cannot be interacted with by the user. It also has a grayed out look.
    */
    pub struct LabelFlags: u32 {
        const NONE = 0;
        const VISIBLE = WS_VISIBLE;
        const DISABLED = WS_DISABLED;

        /// Truncate the label if the text is too long. A label with this style CANNOT have multiple lines.
        const ELIPSIS = SS_WORDELLIPSIS;
    }
}

/**
A label is a single line of static text. Use `\r\n` to split the text on multiple lines.

Label is not behind any features.

**Builder parameters:**
  * `parent`:           **Required.** The label parent container.
  * `text`:             The label text.
  * `size`:             The label size.
  * `position`:         The label position.
  * `enabled`:          If the label is enabled. A disabled label won't trigger events
  * `flags`:            A combination of the LabelFlags values.
  * `font`:             The font used for the label text
  * `background_color`: The background color of the label
  * `h_align`:          The horizontal aligment of the label

**Control events:**
  * `OnLabelClick`: When the user click the label
  * `OnLabelDoubleClick`: When the user double click a label
  * `MousePress(_)`: Generic mouse press events on the label
  * `OnMouseMove`: Generic mouse mouse event
  * `OnMouseWheel`: Generic mouse wheel event

** Example **

```rust
use native_windows_gui as nwg;
fn build_label(label: &mut nwg::Label, window: &nwg::Window, font: &nwg::Font) {
    nwg::Label::builder()
        .text("Hello")
        .font(Some(font))
        .parent(window)
        .build(label);
}
```

*/
#[derive(Default)]
pub struct Label {
    pub handle: ControlHandle,
    handler0: RefCell<Option<RawEventHandler>>,
    handler1: RefCell<Option<RawEventHandler>>,
}

impl Label {

    pub fn builder<'a>() -> LabelBuilder<'a> {
        LabelBuilder {
            text: "A label",
            size: (130, 25),
            position: (0, 0),
            flags: None,
            font: None,
            parent: None,
            h_align: HTextAlign::Left,
            v_align: VTextAlign::Top,
            background_color: None
        }
    }

    /// Return the font of the control
    pub fn font(&self) -> Option<Font> {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);

        let font_handle = wh::get_window_font(handle);
        if font_handle.is_null() {
            None
        } else {
            Some(Font { handle: font_handle })
        }
    }

    /// Set the font of the control
    pub fn set_font(&self, font: Option<&Font>) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_font(handle, font.map(|f| f.handle), true); }
    }

    /// Return true if the control currently has the keyboard focus
    pub fn focus(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_focus(handle) }
    }

    /// Set the keyboard focus on the button.
    pub fn set_focus(&self) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_focus(handle); }
    }

    /// Return true if the control user can interact with the control, return false otherwise
    pub fn enabled(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_enabled(handle) }
    }

    /// Enable or disable the control
    pub fn set_enabled(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_enabled(handle, v) }
    }

    /// Return true if the control is visible to the user. Will return true even if the
    /// control is outside of the parent client view (ex: at the position (10000, 10000))
    pub fn visible(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_visibility(handle) }
    }

    /// Show or hide the control to the user
    pub fn set_visible(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_visibility(handle, v) }
    }

    /// Return the size of the label in the parent window
    pub fn size(&self) -> (u32, u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_size(handle) }
    }

    /// Set the size of the label in the parent window
    pub fn set_size(&self, x: u32, y: u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_size(handle, x, y, false) }
    }

    /// Return the position of the label in the parent window
    pub fn position(&self) -> (i32, i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_position(handle) }
    }

    /// Set the position of the label in the parent window
    pub fn set_position(&self, x: i32, y: i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_position(handle, x, y) }
    }

    /// Return the label text
    pub fn text(&self) -> String {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_text(handle) }
    }

    /// Set the label text
    pub fn set_text<'a>(&self, v: &'a str) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_text(handle, v) }
    }

    /// Winapi class name used during control creation
    pub fn class_name(&self) -> &'static str {
        "STATIC"
    }

    /// Winapi base flags used during window creation
    pub fn flags(&self) -> u32 {
        use winapi::um::winuser::{SS_NOPREFIX, SS_LEFT};

        WS_VISIBLE | SS_NOPREFIX | SS_LEFT
    }

    /// Winapi flags required by the control
    pub fn forced_flags(&self) -> u32 {
        use winapi::um::winuser::{SS_NOTIFY, WS_CHILD};

        WS_CHILD | SS_NOTIFY
    }

    /// Center the text vertically.
    fn hook_non_client_size(&self, bg: Option<[u8; 3]>, v_align: VTextAlign) {
        use crate::bind_raw_event_handler_inner;
        use winapi::shared::windef::{HWND, HGDIOBJ, RECT, HBRUSH, POINT};
        use winapi::shared::{basetsd::UINT_PTR, minwindef::LRESULT};
        use winapi::um::winuser::{WM_CTLCOLORSTATIC, WM_NCCALCSIZE, WM_NCPAINT, WM_SIZE, DT_CALCRECT, DT_LEFT, NCCALCSIZE_PARAMS, COLOR_WINDOW};
        use winapi::um::winuser::{SWP_NOOWNERZORDER, SWP_NOSIZE, SWP_NOMOVE, SWP_FRAMECHANGED};
        use winapi::um::winuser::{GetDC, DrawTextW, ReleaseDC, GetClientRect, GetWindowRect, FillRect, ScreenToClient, SetWindowPos};
        use winapi::um::wingdi::{SelectObject, CreateSolidBrush, RGB};
        use std::{mem, ptr};

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let parent_handle = ControlHandle::Hwnd(wh::get_window_parent(handle));

        let brush = match bg {
            Some(c) => unsafe { CreateSolidBrush(RGB(c[0], c[1], c[2])) },
            None => COLOR_WINDOW as HBRUSH
        };

        unsafe {

        if bg.is_some() {
            let handler0 = bind_raw_event_handler_inner(&parent_handle, handle as UINT_PTR, move |_hwnd, msg, _w, l| {
                match msg {
                    WM_CTLCOLORSTATIC => {
                        let child = l as HWND;
                        if child == handle {
                            return Some(brush as LRESULT);
                        }
                    },
                    _ => {}
                }

                None
            });

            *self.handler0.borrow_mut() = Some(handler0.unwrap());
        }

        let handler1 = bind_raw_event_handler_inner(&self.handle, 0, move |hwnd, msg, w, l| {
            match msg {
                WM_NCCALCSIZE  => {
                    if w == 0 { return None }

                    // Calculate client area height needed for a font
                    let font_handle = wh::get_window_font(hwnd);
                    let mut r: RECT = mem::zeroed();
                    let dc = GetDC(hwnd);

                    let old = SelectObject(dc, font_handle as HGDIOBJ);
                    let calc: [u16;2] = [75, 121];
                    DrawTextW(dc, calc.as_ptr(), 2, &mut r, DT_CALCRECT | DT_LEFT);

                    let client_height = r.bottom;

                    SelectObject(dc, old);
                    ReleaseDC(hwnd, dc);

                    // Calculate NC area to center text.
                    let mut client: RECT = mem::zeroed();
                    let mut window: RECT = mem::zeroed();
                    GetClientRect(hwnd, &mut client);
                    GetWindowRect(hwnd, &mut window);

                    let window_height = window.bottom - window.top;
                    match v_align {
                        VTextAlign::Top => {
                            // The winapi default, nothing to be done
                        },
                        VTextAlign::Center => {
                            let center = ((window_height - client_height) / 2) - 1;

                            // Save the info
                            let info_ptr: *mut NCCALCSIZE_PARAMS = l as *mut NCCALCSIZE_PARAMS;
                            let info = &mut *info_ptr;

                            info.rgrc[0].top += center;
                            info.rgrc[0].bottom -= center;
                        },
                        VTextAlign::Bottom => {
                            let info_ptr: *mut NCCALCSIZE_PARAMS = l as *mut NCCALCSIZE_PARAMS;
                            let info = &mut *info_ptr;
                            info.rgrc[0].top += window_height - client_height;
                        },
                    }
                },
                WM_NCPAINT  => {
                    let mut window: RECT = mem::zeroed();
                    let mut client: RECT = mem::zeroed();
                    GetWindowRect(hwnd, &mut window);
                    GetClientRect(hwnd, &mut client);

                    let mut pt1 = POINT {x: window.left, y: window.top};
                    ScreenToClient(hwnd, &mut pt1);

                    let mut pt2 = POINT {x: window.right, y: window.bottom};
                    ScreenToClient(hwnd, &mut pt2);

                    let top = RECT {
                        left: 0,
                        top: pt1.y,
                        right: client.right,
                        bottom: client.top
                    };

                    let bottom = RECT {
                        left: 0,
                        top: client.bottom,
                        right: client.right,
                        bottom: pt2.y
                    };

                    let dc = GetDC(hwnd);
                    FillRect(dc, &top, brush);
                    FillRect(dc, &bottom, brush);
                    ReleaseDC(hwnd, dc);
                },
                WM_SIZE => {
                    SetWindowPos(hwnd, ptr::null_mut(), 0, 0, 0, 0, SWP_NOOWNERZORDER | SWP_NOSIZE | SWP_NOMOVE | SWP_FRAMECHANGED);
                },
                _ => {}
            }

            None
        });

        *self.handler1.borrow_mut() = Some(handler1.unwrap());

        }
    }

}

impl PartialEq for Label {
    fn eq(&self, other: &Self) -> bool {
        self.handle == other.handle
    }
}

impl Drop for Label {
    fn drop(&mut self) {
        let handler = self.handler0.borrow();
        if let Some(h) = handler.as_ref() {
            drop(unbind_raw_event_handler(h));
        }

        let handler = self.handler1.borrow();
        if let Some(h) = handler.as_ref() {
            drop(unbind_raw_event_handler(h));
        }

        self.handle.destroy();
    }
}

pub struct LabelBuilder<'a> {
    text: &'a str,
    size: (i32, i32),
    position: (i32, i32),
    background_color: Option<[u8; 3]>,
    flags: Option<LabelFlags>,
    font: Option<&'a Font>,
    h_align: HTextAlign,
    v_align: VTextAlign,
    parent: Option<ControlHandle>
}

impl<'a> LabelBuilder<'a> {

    pub fn flags(mut self, flags: LabelFlags) -> LabelBuilder<'a> {
        self.flags = Some(flags);
        self
    }

    pub fn text(mut self, text: &'a str) -> LabelBuilder<'a> {
        self.text = text;
        self
    }

    pub fn size(mut self, size: (i32, i32)) -> LabelBuilder<'a> {
        self.size = size;
        self
    }

    pub fn position(mut self, pos: (i32, i32)) -> LabelBuilder<'a> {
        self.position = pos;
        self
    }

    pub fn font(mut self, font: Option<&'a Font>) -> LabelBuilder<'a> {
        self.font = font;
        self
    }

    pub fn background_color(mut self, color: Option<[u8;3]>) -> LabelBuilder<'a> {
        self.background_color = color;
        self
    }

    pub fn h_align(mut self, align: HTextAlign) -> LabelBuilder<'a> {
        self.h_align = align;
        self
    }

    pub fn v_align(mut self, align: VTextAlign) -> LabelBuilder<'a> {
        self.v_align = align;
        self
    }

    pub fn parent<C: Into<ControlHandle>>(mut self, p: C) -> LabelBuilder<'a> {
        self.parent = Some(p.into());
        self
    }

    pub fn build(self, out: &mut Label) -> Result<(), NwgError> {
        use winapi::um::winuser::{SS_LEFT, SS_RIGHT, SS_CENTER};

        let mut flags = self.flags.map(|f| f.bits()).unwrap_or(out.flags());

        match self.h_align {
            HTextAlign::Left => { flags |= SS_LEFT; },
            HTextAlign::Right => { flags |= SS_RIGHT; },
            HTextAlign::Center => { flags |= SS_CENTER; },
        }

        let parent = match self.parent {
            Some(p) => Ok(p),
            None => Err(NwgError::no_parent("Label"))
        }?;

        // Drop the old object
        *out = Label::default();

        out.handle = ControlBase::build_hwnd()
            .class_name(out.class_name())
            .forced_flags(out.forced_flags())
            .flags(flags)
            .size(self.size)
            .position(self.position)
            .text(self.text)
            .parent(Some(parent))
            .build()?;

        if self.font.is_some() {
            out.set_font(self.font);
        } else {
            out.set_font(Font::global_default().as_ref());
        }

        out.hook_non_client_size(self.background_color, self.v_align);

        Ok(())
    }

}

<|endoftext|>


<commit_before>//! Flycheck provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.

#![warn(rust_2018_idioms, unused_lifetimes, semicolon_in_expressions_from_macros)]

use std::{
    fmt, io,
    process::{ChildStderr, ChildStdout, Command, Stdio},
    time::Duration,
};

use crossbeam_channel::{never, select, unbounded, Receiver, Sender};
use paths::AbsPathBuf;
use rustc_hash::FxHashMap;
use serde::Deserialize;
use stdx::{process::streaming_output, JodChild};

pub use cargo_metadata::diagnostic::{
    Applicability, Diagnostic, DiagnosticCode, DiagnosticLevel, DiagnosticSpan,
    DiagnosticSpanMacroExpansion,
};

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FlycheckConfig {
    CargoCommand {
        command: String,
        target_triple: Option<String>,
        all_targets: bool,
        no_default_features: bool,
        all_features: bool,
        features: Vec<String>,
        extra_args: Vec<String>,
        extra_env: FxHashMap<String, String>,
    },
    CustomCommand {
        command: String,
        args: Vec<String>,
        extra_env: FxHashMap<String, String>,
    },
}

impl fmt::Display for FlycheckConfig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            FlycheckConfig::CargoCommand { command, .. } => write!(f, "cargo {}", command),
            FlycheckConfig::CustomCommand { command, args, .. } => {
                write!(f, "{} {}", command, args.join(" "))
            }
        }
    }
}

/// Flycheck wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct FlycheckHandle {
    // XXX: drop order is significant
    sender: Sender<Restart>,
    _thread: jod_thread::JoinHandle,
    id: usize,
}

impl FlycheckHandle {
    pub fn spawn(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckHandle {
        let actor = FlycheckActor::new(id, sender, config, workspace_root);
        let (sender, receiver) = unbounded::<Restart>();
        let thread = jod_thread::Builder::new()
            .name("Flycheck".to_owned())
            .spawn(move || actor.run(receiver))
            .expect("failed to spawn thread");
        FlycheckHandle { id, sender, _thread: thread }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn restart(&self) {
        self.sender.send(Restart::Yes).unwrap();
    }

    /// Stop this cargo check worker.
    pub fn cancel(&self) {
        self.sender.send(Restart::No).unwrap();
    }

    pub fn id(&self) -> usize {
        self.id
    }
}

pub enum Message {
    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { id: usize, workspace_root: AbsPathBuf, diagnostic: Diagnostic },

    /// Request check progress notification to client
    Progress {
        /// Flycheck instance ID
        id: usize,
        progress: Progress,
    },
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Message::AddDiagnostic { id, workspace_root, diagnostic } => f
                .debug_struct("AddDiagnostic")
                .field("id", id)
                .field("workspace_root", workspace_root)
                .field("diagnostic_code", &diagnostic.code.as_ref().map(|it| &it.code))
                .finish(),
            Message::Progress { id, progress } => {
                f.debug_struct("Progress").field("id", id).field("progress", progress).finish()
            }
        }
    }
}

#[derive(Debug)]
pub enum Progress {
    DidStart,
    DidCheckCrate(String),
    DidFinish(io::Result<()>),
    DidCancel,
    DidFailToRestart(String),
}

enum Restart {
    Yes,
    No,
}

struct FlycheckActor {
    id: usize,
    sender: Box<dyn Fn(Message) + Send>,
    config: FlycheckConfig,
    workspace_root: AbsPathBuf,
    /// CargoHandle exists to wrap around the communication needed to be able to
    /// run `cargo check` without blocking. Currently the Rust standard library
    /// doesn't provide a way to read sub-process output without blocking, so we
    /// have to wrap sub-processes output handling in a thread and pass messages
    /// back over a channel.
    cargo_handle: Option<CargoHandle>,
}

enum Event {
    Restart(Restart),
    CheckEvent(Option<CargoMessage>),
}

impl FlycheckActor {
    fn new(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckActor {
        tracing::info!(%id, ?workspace_root, "Spawning flycheck");
        FlycheckActor { id, sender, config, workspace_root, cargo_handle: None }
    }
    fn progress(&self, progress: Progress) {
        self.send(Message::Progress { id: self.id, progress });
    }
    fn next_event(&self, inbox: &Receiver<Restart>) -> Option<Event> {
        let check_chan = self.cargo_handle.as_ref().map(|cargo| &cargo.receiver);
        select! {
            recv(inbox) -> msg => msg.ok().map(Event::Restart),
            recv(check_chan.unwrap_or(&never())) -> msg => Some(Event::CheckEvent(msg.ok())),
        }
    }
    fn run(mut self, inbox: Receiver<Restart>) {
        while let Some(event) = self.next_event(&inbox) {
            match event {
                Event::Restart(Restart::No) => {
                    self.cancel_check_process();
                }
                Event::Restart(Restart::Yes) => {
                    // Cancel the previously spawned process
                    self.cancel_check_process();
                    while let Ok(_) = inbox.recv_timeout(Duration::from_millis(50)) {}

                    let command = self.check_command();
                    tracing::debug!(?command, "will restart flycheck");
                    match CargoHandle::spawn(command) {
                        Ok(cargo_handle) => {
                            tracing::debug!(
                                command = ?self.check_command(),
                                "did  restart flycheck"
                            );
                            self.cargo_handle = Some(cargo_handle);
                            self.progress(Progress::DidStart);
                        }
                        Err(error) => {
                            self.progress(Progress::DidFailToRestart(format!(
                                "Failed to run the following command: {:?} error={}",
                                self.check_command(),
                                error
                            )));
                        }
                    }
                }
                Event::CheckEvent(None) => {
                    tracing::debug!(flycheck_id = self.id, "flycheck finished");

                    // Watcher finished
                    let cargo_handle = self.cargo_handle.take().unwrap();
                    let res = cargo_handle.join();
                    if res.is_err() {
                        tracing::error!(
                            "Flycheck failed to run the following command: {:?}",
                            self.check_command()
                        );
                    }
                    self.progress(Progress::DidFinish(res));
                }
                Event::CheckEvent(Some(message)) => match message {
                    CargoMessage::CompilerArtifact(msg) => {
                        self.progress(Progress::DidCheckCrate(msg.target.name));
                    }

                    CargoMessage::Diagnostic(msg) => {
                        self.send(Message::AddDiagnostic {
                            id: self.id,
                            workspace_root: self.workspace_root.clone(),
                            diagnostic: msg,
                        });
                    }
                },
            }
        }
        // If we rerun the thread, we need to discard the previous check results first
        self.cancel_check_process();
    }

    fn cancel_check_process(&mut self) {
        if let Some(cargo_handle) = self.cargo_handle.take() {
            tracing::debug!(
                command = ?self.check_command(),
                "did  cancel flycheck"
            );
            cargo_handle.cancel();
            self.progress(Progress::DidCancel);
        }
    }

    fn check_command(&self) -> Command {
        let mut cmd = match &self.config {
            FlycheckConfig::CargoCommand {
                command,
                target_triple,
                no_default_features,
                all_targets,
                all_features,
                extra_args,
                features,
                extra_env,
            } => {
                let mut cmd = Command::new(toolchain::cargo());
                cmd.arg(command);
                cmd.current_dir(&self.workspace_root);
                cmd.args(&["--workspace", "--message-format=json", "--manifest-path"])
                    .arg(self.workspace_root.join("Cargo.toml").as_os_str());

                if let Some(target) = target_triple {
                    cmd.args(&["--target", target.as_str()]);
                }
                if *all_targets {
                    cmd.arg("--all-targets");
                }
                if *all_features {
                    cmd.arg("--all-features");
                } else {
                    if *no_default_features {
                        cmd.arg("--no-default-features");
                    }
                    if !features.is_empty() {
                        cmd.arg("--features");
                        cmd.arg(features.join(" "));
                    }
                }
                cmd.args(extra_args);
                cmd.envs(extra_env);
                cmd
            }
            FlycheckConfig::CustomCommand { command, args, extra_env } => {
                let mut cmd = Command::new(command);
                cmd.args(args);
                cmd.envs(extra_env);
                cmd
            }
        };
        cmd.current_dir(&self.workspace_root);
        cmd
    }

    fn send(&self, check_task: Message) {
        (self.sender)(check_task);
    }
}

/// A handle to a cargo process used for fly-checking.
struct CargoHandle {
    /// The handle to the actual cargo process. As we cannot cancel directly from with
    /// a read syscall dropping and therefor terminating the process is our best option.
    child: JodChild,
    thread: jod_thread::JoinHandle<io::Result<(bool, String)>>,
    receiver: Receiver<CargoMessage>,
}

impl CargoHandle {
    fn spawn(mut command: Command) -> std::io::Result<CargoHandle> {
        command.stdout(Stdio::piped()).stderr(Stdio::piped()).stdin(Stdio::null());
        let mut child = JodChild::spawn(command)?;

        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();

        let (sender, receiver) = unbounded();
        let actor = CargoActor::new(sender, stdout, stderr);
        let thread = jod_thread::Builder::new()
            .name("CargoHandle".to_owned())
            .spawn(move || actor.run())
            .expect("failed to spawn thread");
        Ok(CargoHandle { child, thread, receiver })
    }

    fn cancel(mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }

    fn join(mut self) -> io::Result<()> {
        let _ = self.child.kill();
        let exit_status = self.child.wait()?;
        let (read_at_least_one_message, error) = self.thread.join()?;
        if read_at_least_one_message || exit_status.success() {
            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::Other, format!(
                "Cargo watcher failed, the command produced no valid metadata (exit code: {:?}):\n{}",
                exit_status, error
            )))
        }
    }
}

struct CargoActor {
    sender: Sender<CargoMessage>,
    stdout: ChildStdout,
    stderr: ChildStderr,
}

impl CargoActor {
    fn new(sender: Sender<CargoMessage>, stdout: ChildStdout, stderr: ChildStderr) -> CargoActor {
        CargoActor { sender, stdout, stderr }
    }

    fn run(self) -> io::Result<(bool, String)> {
        // We manually read a line at a time, instead of using serde's
        // stream deserializers, because the deserializer cannot recover
        // from an error, resulting in it getting stuck, because we try to
        // be resilient against failures.
        //
        // Because cargo only outputs one JSON object per line, we can
        // simply skip a line if it doesn't parse, which just ignores any
        // erroneous output.

        let mut error = String::new();
        let mut read_at_least_one_message = false;
        let output = streaming_output(
            self.stdout,
            self.stderr,
            &mut |line| {
                read_at_least_one_message = true;

                // Try to deserialize a message from Cargo or Rustc.
                let mut deserializer = serde_json::Deserializer::from_str(line);
                deserializer.disable_recursion_limit();
                if let Ok(message) = JsonMessage::deserialize(&mut deserializer) {
                    match message {
                        // Skip certain kinds of messages to only spend time on what's useful
                        JsonMessage::Cargo(message) => match message {
                            cargo_metadata::Message::CompilerArtifact(artifact)
                                if !artifact.fresh =>
                            {
                                self.sender.send(CargoMessage::CompilerArtifact(artifact)).unwrap();
                            }
                            cargo_metadata::Message::CompilerMessage(msg) => {
                                self.sender.send(CargoMessage::Diagnostic(msg.message)).unwrap();
                            }
                            _ => (),
                        },
                        JsonMessage::Rustc(message) => {
                            self.sender.send(CargoMessage::Diagnostic(message)).unwrap();
                        }
                    }
                }
            },
            &mut |line| {
                error.push_str(line);
                error.push('\n');
            },
        );
        match output {
            Ok(_) => Ok((read_at_least_one_message, error)),
            Err(e) => Err(io::Error::new(e.kind(), format!("{:?}: {}", e, error))),
        }
    }
}

enum CargoMessage {
    CompilerArtifact(cargo_metadata::Artifact),
    Diagnostic(Diagnostic),
}

#[derive(Deserialize)]
#[serde(untagged)]
enum JsonMessage {
    Cargo(cargo_metadata::Message),
    Rustc(Diagnostic),
}
<commit_msg>Prioritize restart messages in flycheck<commit_after>//! Flycheck provides the functionality needed to run `cargo check` or
//! another compatible command (f.x. clippy) in a background thread and provide
//! LSP diagnostics based on the output of the command.

#![warn(rust_2018_idioms, unused_lifetimes, semicolon_in_expressions_from_macros)]

use std::{
    fmt, io,
    process::{ChildStderr, ChildStdout, Command, Stdio},
    time::Duration,
};

use crossbeam_channel::{never, select, unbounded, Receiver, Sender};
use paths::AbsPathBuf;
use rustc_hash::FxHashMap;
use serde::Deserialize;
use stdx::{process::streaming_output, JodChild};

pub use cargo_metadata::diagnostic::{
    Applicability, Diagnostic, DiagnosticCode, DiagnosticLevel, DiagnosticSpan,
    DiagnosticSpanMacroExpansion,
};

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FlycheckConfig {
    CargoCommand {
        command: String,
        target_triple: Option<String>,
        all_targets: bool,
        no_default_features: bool,
        all_features: bool,
        features: Vec<String>,
        extra_args: Vec<String>,
        extra_env: FxHashMap<String, String>,
    },
    CustomCommand {
        command: String,
        args: Vec<String>,
        extra_env: FxHashMap<String, String>,
    },
}

impl fmt::Display for FlycheckConfig {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            FlycheckConfig::CargoCommand { command, .. } => write!(f, "cargo {}", command),
            FlycheckConfig::CustomCommand { command, args, .. } => {
                write!(f, "{} {}", command, args.join(" "))
            }
        }
    }
}

/// Flycheck wraps the shared state and communication machinery used for
/// running `cargo check` (or other compatible command) and providing
/// diagnostics based on the output.
/// The spawned thread is shut down when this struct is dropped.
#[derive(Debug)]
pub struct FlycheckHandle {
    // XXX: drop order is significant
    sender: Sender<Restart>,
    _thread: jod_thread::JoinHandle,
    id: usize,
}

impl FlycheckHandle {
    pub fn spawn(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckHandle {
        let actor = FlycheckActor::new(id, sender, config, workspace_root);
        let (sender, receiver) = unbounded::<Restart>();
        let thread = jod_thread::Builder::new()
            .name("Flycheck".to_owned())
            .spawn(move || actor.run(receiver))
            .expect("failed to spawn thread");
        FlycheckHandle { id, sender, _thread: thread }
    }

    /// Schedule a re-start of the cargo check worker.
    pub fn restart(&self) {
        self.sender.send(Restart::Yes).unwrap();
    }

    /// Stop this cargo check worker.
    pub fn cancel(&self) {
        self.sender.send(Restart::No).unwrap();
    }

    pub fn id(&self) -> usize {
        self.id
    }
}

pub enum Message {
    /// Request adding a diagnostic with fixes included to a file
    AddDiagnostic { id: usize, workspace_root: AbsPathBuf, diagnostic: Diagnostic },

    /// Request check progress notification to client
    Progress {
        /// Flycheck instance ID
        id: usize,
        progress: Progress,
    },
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Message::AddDiagnostic { id, workspace_root, diagnostic } => f
                .debug_struct("AddDiagnostic")
                .field("id", id)
                .field("workspace_root", workspace_root)
                .field("diagnostic_code", &diagnostic.code.as_ref().map(|it| &it.code))
                .finish(),
            Message::Progress { id, progress } => {
                f.debug_struct("Progress").field("id", id).field("progress", progress).finish()
            }
        }
    }
}

#[derive(Debug)]
pub enum Progress {
    DidStart,
    DidCheckCrate(String),
    DidFinish(io::Result<()>),
    DidCancel,
    DidFailToRestart(String),
}

enum Restart {
    Yes,
    No,
}

struct FlycheckActor {
    id: usize,
    sender: Box<dyn Fn(Message) + Send>,
    config: FlycheckConfig,
    workspace_root: AbsPathBuf,
    /// CargoHandle exists to wrap around the communication needed to be able to
    /// run `cargo check` without blocking. Currently the Rust standard library
    /// doesn't provide a way to read sub-process output without blocking, so we
    /// have to wrap sub-processes output handling in a thread and pass messages
    /// back over a channel.
    cargo_handle: Option<CargoHandle>,
}

enum Event {
    Restart(Restart),
    CheckEvent(Option<CargoMessage>),
}

impl FlycheckActor {
    fn new(
        id: usize,
        sender: Box<dyn Fn(Message) + Send>,
        config: FlycheckConfig,
        workspace_root: AbsPathBuf,
    ) -> FlycheckActor {
        tracing::info!(%id, ?workspace_root, "Spawning flycheck");
        FlycheckActor { id, sender, config, workspace_root, cargo_handle: None }
    }
    fn progress(&self, progress: Progress) {
        self.send(Message::Progress { id: self.id, progress });
    }
    fn next_event(&self, inbox: &Receiver<Restart>) -> Option<Event> {
        let check_chan = self.cargo_handle.as_ref().map(|cargo| &cargo.receiver);
        if let Ok(msg) = inbox.try_recv() {
            // give restarts a preference so check outputs don't block a restart or stop
            return Some(Event::Restart(msg));
        }
        select! {
            recv(inbox) -> msg => msg.ok().map(Event::Restart),
            recv(check_chan.unwrap_or(&never())) -> msg => Some(Event::CheckEvent(msg.ok())),
        }
    }
    fn run(mut self, inbox: Receiver<Restart>) {
        'event: while let Some(event) = self.next_event(&inbox) {
            match event {
                Event::Restart(Restart::No) => {
                    self.cancel_check_process();
                }
                Event::Restart(Restart::Yes) => {
                    // Cancel the previously spawned process
                    self.cancel_check_process();
                    while let Ok(restart) = inbox.recv_timeout(Duration::from_millis(50)) {
                        // restart chained with a stop, so just cancel
                        if let Restart::No = restart {
                            continue 'event;
                        }
                    }

                    let command = self.check_command();
                    tracing::debug!(?command, "will restart flycheck");
                    match CargoHandle::spawn(command) {
                        Ok(cargo_handle) => {
                            tracing::debug!(
                                command = ?self.check_command(),
                                "did  restart flycheck"
                            );
                            self.cargo_handle = Some(cargo_handle);
                            self.progress(Progress::DidStart);
                        }
                        Err(error) => {
                            self.progress(Progress::DidFailToRestart(format!(
                                "Failed to run the following command: {:?} error={}",
                                self.check_command(),
                                error
                            )));
                        }
                    }
                }
                Event::CheckEvent(None) => {
                    tracing::debug!(flycheck_id = self.id, "flycheck finished");

                    // Watcher finished
                    let cargo_handle = self.cargo_handle.take().unwrap();
                    let res = cargo_handle.join();
                    if res.is_err() {
                        tracing::error!(
                            "Flycheck failed to run the following command: {:?}",
                            self.check_command()
                        );
                    }
                    self.progress(Progress::DidFinish(res));
                }
                Event::CheckEvent(Some(message)) => match message {
                    CargoMessage::CompilerArtifact(msg) => {
                        self.progress(Progress::DidCheckCrate(msg.target.name));
                    }

                    CargoMessage::Diagnostic(msg) => {
                        self.send(Message::AddDiagnostic {
                            id: self.id,
                            workspace_root: self.workspace_root.clone(),
                            diagnostic: msg,
                        });
                    }
                },
            }
        }
        // If we rerun the thread, we need to discard the previous check results first
        self.cancel_check_process();
    }

    fn cancel_check_process(&mut self) {
        if let Some(cargo_handle) = self.cargo_handle.take() {
            tracing::debug!(
                command = ?self.check_command(),
                "did  cancel flycheck"
            );
            cargo_handle.cancel();
            self.progress(Progress::DidCancel);
        }
    }

    fn check_command(&self) -> Command {
        let mut cmd = match &self.config {
            FlycheckConfig::CargoCommand {
                command,
                target_triple,
                no_default_features,
                all_targets,
                all_features,
                extra_args,
                features,
                extra_env,
            } => {
                let mut cmd = Command::new(toolchain::cargo());
                cmd.arg(command);
                cmd.current_dir(&self.workspace_root);
                cmd.args(&["--workspace", "--message-format=json", "--manifest-path"])
                    .arg(self.workspace_root.join("Cargo.toml").as_os_str());

                if let Some(target) = target_triple {
                    cmd.args(&["--target", target.as_str()]);
                }
                if *all_targets {
                    cmd.arg("--all-targets");
                }
                if *all_features {
                    cmd.arg("--all-features");
                } else {
                    if *no_default_features {
                        cmd.arg("--no-default-features");
                    }
                    if !features.is_empty() {
                        cmd.arg("--features");
                        cmd.arg(features.join(" "));
                    }
                }
                cmd.args(extra_args);
                cmd.envs(extra_env);
                cmd
            }
            FlycheckConfig::CustomCommand { command, args, extra_env } => {
                let mut cmd = Command::new(command);
                cmd.args(args);
                cmd.envs(extra_env);
                cmd
            }
        };
        cmd.current_dir(&self.workspace_root);
        cmd
    }

    fn send(&self, check_task: Message) {
        (self.sender)(check_task);
    }
}

/// A handle to a cargo process used for fly-checking.
struct CargoHandle {
    /// The handle to the actual cargo process. As we cannot cancel directly from with
    /// a read syscall dropping and therefor terminating the process is our best option.
    child: JodChild,
    thread: jod_thread::JoinHandle<io::Result<(bool, String)>>,
    receiver: Receiver<CargoMessage>,
}

impl CargoHandle {
    fn spawn(mut command: Command) -> std::io::Result<CargoHandle> {
        command.stdout(Stdio::piped()).stderr(Stdio::piped()).stdin(Stdio::null());
        let mut child = JodChild::spawn(command)?;

        let stdout = child.stdout.take().unwrap();
        let stderr = child.stderr.take().unwrap();

        let (sender, receiver) = unbounded();
        let actor = CargoActor::new(sender, stdout, stderr);
        let thread = jod_thread::Builder::new()
            .name("CargoHandle".to_owned())
            .spawn(move || actor.run())
            .expect("failed to spawn thread");
        Ok(CargoHandle { child, thread, receiver })
    }

    fn cancel(mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }

    fn join(mut self) -> io::Result<()> {
        let _ = self.child.kill();
        let exit_status = self.child.wait()?;
        let (read_at_least_one_message, error) = self.thread.join()?;
        if read_at_least_one_message || exit_status.success() {
            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::Other, format!(
                "Cargo watcher failed, the command produced no valid metadata (exit code: {:?}):\n{}",
                exit_status, error
            )))
        }
    }
}

struct CargoActor {
    sender: Sender<CargoMessage>,
    stdout: ChildStdout,
    stderr: ChildStderr,
}

impl CargoActor {
    fn new(sender: Sender<CargoMessage>, stdout: ChildStdout, stderr: ChildStderr) -> CargoActor {
        CargoActor { sender, stdout, stderr }
    }

    fn run(self) -> io::Result<(bool, String)> {
        // We manually read a line at a time, instead of using serde's
        // stream deserializers, because the deserializer cannot recover
        // from an error, resulting in it getting stuck, because we try to
        // be resilient against failures.
        //
        // Because cargo only outputs one JSON object per line, we can
        // simply skip a line if it doesn't parse, which just ignores any
        // erroneous output.

        let mut error = String::new();
        let mut read_at_least_one_message = false;
        let output = streaming_output(
            self.stdout,
            self.stderr,
            &mut |line| {
                read_at_least_one_message = true;

                // Try to deserialize a message from Cargo or Rustc.
                let mut deserializer = serde_json::Deserializer::from_str(line);
                deserializer.disable_recursion_limit();
                if let Ok(message) = JsonMessage::deserialize(&mut deserializer) {
                    match message {
                        // Skip certain kinds of messages to only spend time on what's useful
                        JsonMessage::Cargo(message) => match message {
                            cargo_metadata::Message::CompilerArtifact(artifact)
                                if !artifact.fresh =>
                            {
                                self.sender.send(CargoMessage::CompilerArtifact(artifact)).unwrap();
                            }
                            cargo_metadata::Message::CompilerMessage(msg) => {
                                self.sender.send(CargoMessage::Diagnostic(msg.message)).unwrap();
                            }
                            _ => (),
                        },
                        JsonMessage::Rustc(message) => {
                            self.sender.send(CargoMessage::Diagnostic(message)).unwrap();
                        }
                    }
                }
            },
            &mut |line| {
                error.push_str(line);
                error.push('\n');
            },
        );
        match output {
            Ok(_) => Ok((read_at_least_one_message, error)),
            Err(e) => Err(io::Error::new(e.kind(), format!("{:?}: {}", e, error))),
        }
    }
}

enum CargoMessage {
    CompilerArtifact(cargo_metadata::Artifact),
    Diagnostic(Diagnostic),
}

#[derive(Deserialize)]
#[serde(untagged)]
enum JsonMessage {
    Cargo(cargo_metadata::Message),
    Rustc(Diagnostic),
}
<|endoftext|>


<commit_before>use std::error;
use std::io::{self, Read};
use std::fmt::{self, Display, Formatter};
use std::str::{Utf8Error, from_utf8};

use Marker;
use super::{read_marker, read_data_u8, read_data_u16, read_data_u32, Error, ValueReadError};

#[derive(Debug)]
pub enum DecodeStringError<'a> {
    InvalidMarkerRead(Error),
    InvalidDataRead(Error),
    TypeMismatch(Marker),
    /// The given buffer is not large enough to accumulate the specified amount of bytes.
    BufferSizeTooSmall(u32),
    InvalidUtf8(&'a [u8], Utf8Error),
}

impl<'a> error::Error for DecodeStringError<'a> {
    fn description(&self) -> &str {
        "error while decoding string"
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            DecodeStringError::InvalidMarkerRead(ref err) |
            DecodeStringError::InvalidDataRead(ref err) => Some(err),
            DecodeStringError::TypeMismatch(..) |
            DecodeStringError::BufferSizeTooSmall(..) => None,
            DecodeStringError::InvalidUtf8(.., ref err) => Some(err),
        }
    }
}

impl<'a> Display for DecodeStringError<'a> {
    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
        error::Error::description(self).fmt(f)
    }
}

impl<'a> From<ValueReadError> for DecodeStringError<'a> {
    fn from(err: ValueReadError) -> DecodeStringError<'a> {
        match err {
            ValueReadError::InvalidMarkerRead(err) => DecodeStringError::InvalidMarkerRead(err),
            ValueReadError::InvalidDataRead(err) => DecodeStringError::InvalidDataRead(err),
            ValueReadError::TypeMismatch(marker) => DecodeStringError::TypeMismatch(marker),
        }
    }
}

/// Attempts to read up to 9 bytes from the given reader and to decode them as a string `u32` size
/// value.
///
/// According to the MessagePack specification, the string format family stores an byte array in 1,
/// 2, 3, or 5 bytes of extra bytes in addition to the size of the byte array.
///
/// # Errors
///
/// This function will return `ValueReadError` on any I/O error while reading either the marker or
/// the data, except the EINTR, which is handled internally.
///
/// It also returns `ValueReadError::TypeMismatch` if the actual type is not equal with the
/// expected one, indicating you with the actual type.
///
/// # Note
///
/// This function will silently retry on every EINTR received from the underlying `Read` until
/// successful read.
pub fn read_str_len<R: Read>(rd: &mut R) -> Result<u32, ValueReadError> {
    match try!(read_marker(rd)) {
        Marker::FixStr(size) => Ok(size as u32),
        Marker::Str8 => Ok(try!(read_data_u8(rd)) as u32),
        Marker::Str16 => Ok(try!(read_data_u16(rd)) as u32),
        Marker::Str32 => Ok(try!(read_data_u32(rd))),
        marker => Err(ValueReadError::TypeMismatch(marker)),
    }
}

/// Attempts to read a string data from the given reader and copy it to the buffer provided.
///
/// On success returns a borrowed string type, allowing to view the copyed bytes as properly utf-8
/// string.
/// According to the spec, the string's data must to be encoded using utf-8.
///
/// # Errors
///
/// Returns `Err` in the following cases:
///
///  - if any IO error (including unexpected EOF) occurs, while reading an `rd`, except the EINTR,
///    which is handled internally.
///  - if the `out` buffer size is not large enough to keep all the data copyed.
///  - if the data is not utf-8, with a description as to why the provided data is not utf-8 and
///    with a size of bytes actually copyed to be able to get them from `out`.
///
/// # Examples
/// ```
/// use rmp::decode::read_str;
///
/// let buf = [0xaa, 0x6c, 0x65, 0x20, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65];
/// let mut out = [0u8; 16];
///
/// assert_eq!("le message", read_str(&mut &buf[..], &mut &mut out[..]).unwrap());
/// ```
///
/// # Unstable
///
/// This function is **unstable**, because it needs review.
///
/// # Note
///
/// This function will silently retry on every EINTR received from the underlying `Read` until
/// successful read.
// TODO: Stabilize. Mark error values for each error case (in docs).
pub fn read_str<'r, R>(rd: &mut R, mut buf: &'r mut [u8]) -> Result<&'r str, DecodeStringError<'r>>
    where R: Read
{
    let len = try!(read_str_len(rd));
    let ulen = len as usize;

    if buf.len() < ulen {
        return Err(DecodeStringError::BufferSizeTooSmall(len));
    }

    read_str_data(rd, len, &mut buf[0..ulen])
}

pub fn read_str_data<'r, R>(rd: &mut R,
                            len: u32,
                            buf: &'r mut [u8])
                            -> Result<&'r str, DecodeStringError<'r>>
    where R: Read
{
    debug_assert_eq!(len as usize, buf.len());

    // Trying to copy exact `len` bytes.
    match rd.read_exact(buf) {
        Ok(()) => {
            match from_utf8(buf) {
                Ok(decoded) => Ok(decoded),
                Err(err) => Err(DecodeStringError::InvalidUtf8(buf, err)),
            }
        }
        Err(err) => Err(DecodeStringError::InvalidDataRead(From::from(err))),
    }
}

/// Attempts to read and decode a string value from the reader, returning a borrowed slice from it.
///
// TODO: it is better to return &str; may panic on len mismatch; extend documentation.
// TODO: Also it's possible to implement all borrowing functions for all `BufRead` implementors.
// TODO: It's not necessary to use cursor, use slices instead.
// TODO: Candidate to be removed/replaced.
pub fn read_str_ref(rd: &[u8]) -> Result<&[u8], DecodeStringError> {
    let mut cur = io::Cursor::new(rd);
    let len = try!(read_str_len(&mut cur));
    let start = cur.position() as usize;
    Ok(&rd[start..start + len as usize])
}
<commit_msg>fix: fix compilation on rust 1.13.0<commit_after>use std::error;
use std::io::{self, Read};
use std::fmt::{self, Display, Formatter};
use std::str::{Utf8Error, from_utf8};

use Marker;
use super::{read_marker, read_data_u8, read_data_u16, read_data_u32, Error, ValueReadError};

#[derive(Debug)]
pub enum DecodeStringError<'a> {
    InvalidMarkerRead(Error),
    InvalidDataRead(Error),
    TypeMismatch(Marker),
    /// The given buffer is not large enough to accumulate the specified amount of bytes.
    BufferSizeTooSmall(u32),
    InvalidUtf8(&'a [u8], Utf8Error),
}

impl<'a> error::Error for DecodeStringError<'a> {
    fn description(&self) -> &str {
        "error while decoding string"
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            DecodeStringError::InvalidMarkerRead(ref err) |
            DecodeStringError::InvalidDataRead(ref err) => Some(err),
            DecodeStringError::TypeMismatch(..) |
            DecodeStringError::BufferSizeTooSmall(..) => None,
            DecodeStringError::InvalidUtf8(_, ref err) => Some(err),
        }
    }
}

impl<'a> Display for DecodeStringError<'a> {
    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
        error::Error::description(self).fmt(f)
    }
}

impl<'a> From<ValueReadError> for DecodeStringError<'a> {
    fn from(err: ValueReadError) -> DecodeStringError<'a> {
        match err {
            ValueReadError::InvalidMarkerRead(err) => DecodeStringError::InvalidMarkerRead(err),
            ValueReadError::InvalidDataRead(err) => DecodeStringError::InvalidDataRead(err),
            ValueReadError::TypeMismatch(marker) => DecodeStringError::TypeMismatch(marker),
        }
    }
}

/// Attempts to read up to 9 bytes from the given reader and to decode them as a string `u32` size
/// value.
///
/// According to the MessagePack specification, the string format family stores an byte array in 1,
/// 2, 3, or 5 bytes of extra bytes in addition to the size of the byte array.
///
/// # Errors
///
/// This function will return `ValueReadError` on any I/O error while reading either the marker or
/// the data, except the EINTR, which is handled internally.
///
/// It also returns `ValueReadError::TypeMismatch` if the actual type is not equal with the
/// expected one, indicating you with the actual type.
///
/// # Note
///
/// This function will silently retry on every EINTR received from the underlying `Read` until
/// successful read.
pub fn read_str_len<R: Read>(rd: &mut R) -> Result<u32, ValueReadError> {
    match try!(read_marker(rd)) {
        Marker::FixStr(size) => Ok(size as u32),
        Marker::Str8 => Ok(try!(read_data_u8(rd)) as u32),
        Marker::Str16 => Ok(try!(read_data_u16(rd)) as u32),
        Marker::Str32 => Ok(try!(read_data_u32(rd))),
        marker => Err(ValueReadError::TypeMismatch(marker)),
    }
}

/// Attempts to read a string data from the given reader and copy it to the buffer provided.
///
/// On success returns a borrowed string type, allowing to view the copyed bytes as properly utf-8
/// string.
/// According to the spec, the string's data must to be encoded using utf-8.
///
/// # Errors
///
/// Returns `Err` in the following cases:
///
///  - if any IO error (including unexpected EOF) occurs, while reading an `rd`, except the EINTR,
///    which is handled internally.
///  - if the `out` buffer size is not large enough to keep all the data copyed.
///  - if the data is not utf-8, with a description as to why the provided data is not utf-8 and
///    with a size of bytes actually copyed to be able to get them from `out`.
///
/// # Examples
/// ```
/// use rmp::decode::read_str;
///
/// let buf = [0xaa, 0x6c, 0x65, 0x20, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65];
/// let mut out = [0u8; 16];
///
/// assert_eq!("le message", read_str(&mut &buf[..], &mut &mut out[..]).unwrap());
/// ```
///
/// # Unstable
///
/// This function is **unstable**, because it needs review.
///
/// # Note
///
/// This function will silently retry on every EINTR received from the underlying `Read` until
/// successful read.
// TODO: Stabilize. Mark error values for each error case (in docs).
pub fn read_str<'r, R>(rd: &mut R, mut buf: &'r mut [u8]) -> Result<&'r str, DecodeStringError<'r>>
    where R: Read
{
    let len = try!(read_str_len(rd));
    let ulen = len as usize;

    if buf.len() < ulen {
        return Err(DecodeStringError::BufferSizeTooSmall(len));
    }

    read_str_data(rd, len, &mut buf[0..ulen])
}

pub fn read_str_data<'r, R>(rd: &mut R,
                            len: u32,
                            buf: &'r mut [u8])
                            -> Result<&'r str, DecodeStringError<'r>>
    where R: Read
{
    debug_assert_eq!(len as usize, buf.len());

    // Trying to copy exact `len` bytes.
    match rd.read_exact(buf) {
        Ok(()) => {
            match from_utf8(buf) {
                Ok(decoded) => Ok(decoded),
                Err(err) => Err(DecodeStringError::InvalidUtf8(buf, err)),
            }
        }
        Err(err) => Err(DecodeStringError::InvalidDataRead(From::from(err))),
    }
}

/// Attempts to read and decode a string value from the reader, returning a borrowed slice from it.
///
// TODO: it is better to return &str; may panic on len mismatch; extend documentation.
// TODO: Also it's possible to implement all borrowing functions for all `BufRead` implementors.
// TODO: It's not necessary to use cursor, use slices instead.
// TODO: Candidate to be removed/replaced.
pub fn read_str_ref(rd: &[u8]) -> Result<&[u8], DecodeStringError> {
    let mut cur = io::Cursor::new(rd);
    let len = try!(read_str_len(&mut cur));
    let start = cur.position() as usize;
    Ok(&rd[start..start + len as usize])
}
<|endoftext|>